diff --git a/.circleci/config.yml b/.circleci/config.yml index c3a34a97b3b..790cd6c7010 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -88,6 +88,36 @@ commands: rm -f /tmp/uv-install.sh echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$BASH_ENV" export PATH="$HOME/.local/bin:$PATH" + install_rust: + description: "Install pinned rustup (1.28.2) and Rust toolchain (1.97.1) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself." + steps: + - run: + name: Install Rust (rustup 1.28.2, toolchain 1.97.1) + command: | + case "$(uname -m)" in + x86_64) + RUSTUP_TRIPLE=x86_64-unknown-linux-gnu + RUSTUP_SHA256=20a06e644b0d9bd2fbdbfd52d42540bdde820ea7df86e92e533c073da0cdd43c + ;; + aarch64) + RUSTUP_TRIPLE=aarch64-unknown-linux-gnu + RUSTUP_SHA256=e3853c5a252fca15252d07cb23a1bdd9377a8c6f3efa01531109281ae47f841c + ;; + *) + echo "install_rust: unsupported architecture $(uname -m)" >&2 + exit 1 + ;; + esac + curl -sSLf -o /tmp/rustup-init \ + "https://static.rust-lang.org/rustup/archive/1.28.2/${RUSTUP_TRIPLE}/rustup-init" + echo "${RUSTUP_SHA256} /tmp/rustup-init" | sha256sum -c - + chmod +x /tmp/rustup-init + /tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.97.1 + rm -f /tmp/rustup-init + echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> "$BASH_ENV" + export PATH="$HOME/.cargo/bin:$PATH" + rustc --version + cargo --version start_postgres: description: "Start a postgres-db container on port 5432 and wait until it accepts connections." parameters: @@ -163,6 +193,26 @@ commands: done echo "fake OpenAI endpoint did not become ready" >&2 exit 1 + start_cost_center_service: + description: "Start the stand-in cost center validation service (tests/store_model_in_db_tests/cost_center_service.py) on host port 9414 and wait until healthy. The proxy's team-metadata validator (team_metadata_validator_e2e.py, impl 'http') reaches it via TEAM_METADATA_VALIDATION_SERVICE_URL=http://host.docker.internal:9414/validate. Run after uv deps are synced." + steps: + - run: + name: Start cost center validation service + background: true + command: | + uv run --no-sync python tests/store_model_in_db_tests/cost_center_service.py --host 0.0.0.0 --port 9414 + - run: + name: Wait for cost center validation service + command: | + for i in $(seq 1 30); do + if curl -sf http://localhost:9414/health >/dev/null 2>&1; then + echo "cost center validation service is up" + exit 0 + fi + sleep 1 + done + echo "cost center validation service did not become ready" >&2 + exit 1 setup_litellm_enterprise_pip: steps: - run: @@ -178,6 +228,7 @@ commands: - checkout - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -292,6 +343,7 @@ jobs: - checkout - setup_google_dns - install_uv + - install_rust - run: name: Build the wheel environment: @@ -324,6 +376,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -397,6 +450,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -471,6 +525,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -522,6 +577,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -588,6 +644,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -628,6 +685,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -669,6 +727,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -702,6 +761,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -752,6 +812,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -803,6 +864,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -836,6 +898,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -882,6 +945,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -928,6 +992,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -970,6 +1035,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1016,6 +1082,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1063,6 +1130,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -1103,6 +1171,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1148,6 +1217,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1192,6 +1262,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1224,6 +1295,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1267,6 +1339,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1311,6 +1384,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1355,6 +1429,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1386,6 +1461,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1432,6 +1508,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1477,6 +1554,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1527,6 +1605,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1551,6 +1630,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1577,6 +1657,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1678,6 +1759,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1773,6 +1855,7 @@ jobs: at: ~/project - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1861,6 +1944,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1944,6 +2028,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2076,6 +2161,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2162,6 +2248,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2258,12 +2345,14 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | uv sync --frozen --all-groups --all-extras --python 3.12 - start_postgres - start_fake_openai_endpoint + - start_cost_center_service - attach_workspace: at: ~/project - run: @@ -2283,11 +2372,13 @@ jobs: -e STORE_MODEL_IN_DB="True" \ -e LITELLM_MASTER_KEY="sk-1234" \ -e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \ + -e TEAM_METADATA_VALIDATION_SERVICE_URL=http://host.docker.internal:9414/validate \ -e LITELLM_LICENSE=$LITELLM_LICENSE \ -e LITELLM_LOG=ERROR \ --add-host host.docker.internal:host-gateway \ --name my-app \ -v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \ + -v $(pwd)/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py:/app/team_metadata_validator_e2e.py \ litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 @@ -2333,6 +2424,7 @@ jobs: - setup_google_dns # Remove Docker CLI installation since it's already available in machine executor - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2414,6 +2506,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2553,6 +2646,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2743,6 +2837,7 @@ jobs: category: client - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -2885,6 +2980,7 @@ jobs: category: client - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index acdf97cb386..3180cea2568 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 29684 + "limit": 29682 }, "reportArgumentType": { "limit": 2645 @@ -24,7 +24,7 @@ "limit": 42 }, "reportExplicitAny": { - "limit": 9442 + "limit": 9440 }, "reportFunctionMemberAccess": { "limit": 11 diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 8f9cd74f171..4eee525f6a4 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -23,6 +23,7 @@ from typing import Any, cast # Import all the data structures that define what can be lazy-loaded # These are just lists of names and maps of where to find them from ._lazy_imports_registry import ( + # Import maps _BEDROCK_TYPES_IMPORT_MAP, _CACHING_IMPORT_MAP, _COST_CALCULATOR_IMPORT_MAP, @@ -33,12 +34,11 @@ from ._lazy_imports_registry import ( _TOKEN_COUNTER_IMPORT_MAP, _TYPES_IMPORT_MAP, _TYPES_UTILS_IMPORT_MAP, - # Import maps _UTILS_IMPORT_MAP, _UTILS_MODULE_IMPORT_MAP, + # Name tuples BEDROCK_TYPES_NAMES, CACHING_NAMES, - # Name tuples COST_CALCULATOR_NAMES, DOTPROMPT_NAMES, HTTP_HANDLER_NAMES, diff --git a/litellm/batch_completion/main.py b/litellm/batch_completion/main.py index 792be3ff7ad..fb892789b15 100644 --- a/litellm/batch_completion/main.py +++ b/litellm/batch_completion/main.py @@ -249,7 +249,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: {e!s}") + print_verbose(f"batch_completion_models_all_responses: model request failed: {e}") continue return responses diff --git a/litellm/batches/main.py b/litellm/batches/main.py index b27939be8bf..3a057d41744 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -182,7 +182,7 @@ def create_batch( ) except Exception as e: verbose_logger.exception( - f"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - {e!s}" + f"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - {e}" ) _is_async = kwargs.pop("acreate_batch", False) is True @@ -890,7 +890,7 @@ def cancel_batch( ) except Exception as e: verbose_logger.exception( - f"litellm.batches.main.py::cancel_batch() - Error inferring custom_llm_provider - {e!s}" + f"litellm.batches.main.py::cancel_batch() - Error inferring custom_llm_provider - {e}" ) optional_params = GenericLiteLLMParams(**kwargs) litellm_params = get_litellm_params( diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index f69c2fa3b58..9542be0999a 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -353,13 +353,13 @@ class Cache: if param in combined_kwargs: param_value: str | None = self._get_param_value(param, kwargs) if param_value is not None: - cache_key += f"{param!s}: {param_value!s}" + cache_key += f"{param}: {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 if kwargs[param] is None: continue # ignore None params param_value = kwargs[param] - cache_key += f"{param!s}: {param_value!s}" + cache_key += f"{param}: {param_value}" if is_semantic_cache: cache_key += self._get_semantic_cache_tenant_scope(kwargs) @@ -676,7 +676,7 @@ class Cache: 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: {e!s}") + verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {e}") async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs): """ @@ -695,7 +695,7 @@ class Cache: else: await self.cache.async_set_cache(cache_key, cached_data, **kwargs) except Exception as e: - verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {e!s}") + verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {e}") def _convert_to_cached_embedding( self, @@ -874,7 +874,7 @@ class Cache: else: 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: {e!s}") + verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {e}") def should_use_cache(self, **kwargs): """ diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 5b56789e8db..b641c600a0e 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -147,7 +147,7 @@ class DualCache(BaseCache): return result except Exception as e: - verbose_logger.error(f"LiteLLM Cache: Excepton async add_cache: {e!s}") + verbose_logger.error(f"LiteLLM Cache: Excepton async add_cache: {e}") raise e def get_cache( @@ -347,7 +347,7 @@ 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: {e!s}") + verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {e}") # async_batch_set_cache async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs): @@ -366,7 +366,7 @@ class DualCache(BaseCache): cache_list=cache_list, ttl=kwargs.pop("ttl", None), **kwargs ) except Exception as e: - verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {e!s}") + verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {e}") async def async_increment_cache( self, diff --git a/litellm/caching/evicted_client_closer.py b/litellm/caching/evicted_client_closer.py new file mode 100644 index 00000000000..bca9656b252 --- /dev/null +++ b/litellm/caching/evicted_client_closer.py @@ -0,0 +1,276 @@ +""" +Deferred close of HTTP/SDK clients that the LLM client cache has evicted. + +Eviction only drops the cache's reference to a client. Every OpenAI/Azure SDK +client is a reference cycle (each resource namespace holds the client back), so +an evicted client and its pooled TCP connections survive until a generational +collection runs, which under load is thousands of requests later. + +Closing at eviction time is not an option: a request that was handed the client +just before it was evicted is still using it, and closing it underneath that +request raises ``RuntimeError: Cannot send a request, as the client has been +closed.`` + +So an evicted client is closed once two conditions hold. A grace window must +have passed since its eviction, which covers a request that holds the client +but is momentarily not on the wire, and the client must report no connection in +flight. The second condition is what keeps the first honest: a request may run +for ``litellm.request_timeout`` seconds, 6000 by default, and a streaming +response is bounded only by how long the upstream keeps sending, so no deadline +on its own can promise that a request has finished. + +Only clients litellm itself created are closed; a client the caller supplied is +left alone because litellm does not own its lifecycle. + +A client that closes synchronously is closed from wherever the cache is next +used. One whose close is a coroutine needs the event loop it was evicted on, so +it waits for a call from that loop rather than having work scheduled onto a loop +it does not belong to. Queued clients are therefore bucketed by what it takes to +close them, and each bucket is ordered by deadline, so a reap walks the entries +that are due rather than the whole queue. + +The queue holds its clients weakly, so waiting out a grace window never keeps +alive anything the collector would have reclaimed first. +""" + +import asyncio +import inspect +import threading +import time +import weakref +from collections import deque +from collections.abc import Awaitable, Callable, Iterator +from dataclasses import dataclass, replace + +from litellm.constants import ( + EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS, + EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING, +) + +_CLOSABLE_ANYWHERE = "closable-anywhere" +_CLOSABLE_ON_ANY_LOOP = "closable-on-any-loop" + +_BucketKey = str | int + + +@dataclass(frozen=True, slots=True) +class _PendingClose: + """A queued close. + + The client is held weakly, so queueing one never keeps alive anything the + collector would otherwise have reclaimed first. + + ``needs_loop`` is set for a client whose close is a coroutine; those can only + be closed from the event loop they were evicted on, recorded in ``loop_id``. + A client that closes synchronously carries neither constraint. + """ + + client_ref: "weakref.ref[object]" + loop_id: int | None + needs_loop: bool + close_after: float + + +def _bucket_key(pending: _PendingClose) -> _BucketKey: + """Which reaps can close this entry: any at all, any running a loop, or one loop's.""" + if not pending.needs_loop: + return _CLOSABLE_ANYWHERE + if pending.loop_id is None: + return _CLOSABLE_ON_ANY_LOOP + return pending.loop_id + + +def _running_loop_id() -> int | None: + try: + return id(asyncio.get_running_loop()) + except RuntimeError: + return None + + +def _close_function(client: object) -> Callable[[], object] | None: + close_fn: Callable[[], object] | None = getattr(client, "aclose", None) or getattr(client, "close", None) + return close_fn + + +def _transport_of(client: object) -> object: + """The httpx transport behind an SDK wrapper, a litellm handler, or a bare client.""" + for holder in (getattr(client, "_client", None), getattr(client, "client", None), client): + transport: object = getattr(holder, "_transport", None) + if transport is not None: + return transport + return None + + +def _connection_is_idle(connection: object) -> bool: + """A pooled connection is idle unless it is servicing a request.""" + is_idle: object = getattr(connection, "is_idle", None) + return bool(is_idle()) if callable(is_idle) else True + + +def _pool_has_busy_connection(transport: object) -> bool | None: + """Whether the httpcore pool behind the transport is servicing a request. + + ``None`` when there is no such pool, so the caller can ask the other backend. + """ + pooled: object = getattr(getattr(transport, "_pool", None), "connections", None) + if not isinstance(pooled, (list, tuple)): + return None + return any( + not _connection_is_idle(connection) # pyright: ignore[reportUnknownArgumentType] # untyped pool list + for connection in pooled # pyright: ignore[reportUnknownVariableType] # untyped pool list + ) + + +def _has_connection_in_flight(client: object) -> bool: + """Whether the client is servicing a request right now. + + Both connection backends litellm uses already account for the connections + they have handed out, so this reads the client's own lease accounting rather + than inferring it from elapsed time: httpcore reports a non-idle connection + for the whole of a response including a stream, and aiohttp holds the + connection in ``_acquired`` over the same span. + + A client that cannot answer is reported as idle, which leaves the grace + window as the only guard, exactly as it was before this check existed. + """ + try: + transport = _transport_of(client) + pooled_busy = _pool_has_busy_connection(transport) + if pooled_busy is not None: + return pooled_busy + session: object = getattr(transport, "client", None) + return bool(getattr(getattr(session, "connector", None), "_acquired", None)) + except Exception: # noqa: BLE001 - a client that cannot report its state is treated as idle + return False + + +async def _close_quietly(closing: Awaitable[object]) -> None: + try: + await closing + except Exception: # noqa: BLE001 - a discarded client's close must never surface to callers + pass + + +class EvictedClientCloser: + """Closes evicted, litellm-owned clients once they are idle and out of grace.""" + + def __init__( + self, + grace_seconds: float = EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS, + max_pending: int = EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self._grace_seconds = grace_seconds + self._max_pending = max_pending + self._clock = clock + self._owned: weakref.WeakSet[object] = weakref.WeakSet() + self._buckets: dict[_BucketKey, deque[_PendingClose]] = {} # mutable-ok: deadline-ordered queues + self._pending_count = 0 + self._queue_lock = threading.Lock() # the cache is reachable from every worker thread's loop + self._close_tasks: set[asyncio.Task[None]] = set() # mutable-ok: strong refs to running closes + + def mark_owned(self, client: object) -> None: + """Record that litellm created this client, so it may be closed on eviction.""" + try: + self._owned.add(client) + except TypeError: + pass # values that cannot be weak-referenced are never litellm clients + + def _is_owned(self, client: object) -> bool: + try: + return client in self._owned + except TypeError: + return False # unhashable values are never litellm clients + + def schedule(self, client: object) -> None: + """Queue an evicted client for closing once it is idle and out of grace. + + Past ``max_pending`` the client is left to the collector instead, so a + workload that churns the cache cannot grow this queue without bound. + Every queued entry comes due within one grace window, so the capacity it + occupies is returned within that window rather than held. + """ + if client is None or not self._is_owned(client): + return + close_fn = _close_function(client) + if close_fn is None: + return + if self._pending_count >= self._max_pending: + return + self._enqueue( + _PendingClose( + client_ref=weakref.ref(client), + loop_id=_running_loop_id(), + needs_loop=inspect.iscoroutinefunction(close_fn), + close_after=self._clock() + self._grace_seconds, + ) + ) + + def reap(self) -> None: + """Close every queued client that is due, idle, and closable from here. + + Called from the cache's read path, so the empty-queue exit comes first and + the work done past it is proportional to what is due, not to the queue. + """ + if not self._pending_count: + return + now = self._clock() + for pending in self._take_due(_running_loop_id(), now): + client = pending.client_ref() + if client is None: + continue + if _has_connection_in_flight(client): + self._enqueue(replace(pending, close_after=now + self._grace_seconds)) + continue + self._close(client) + + @property + def pending_count(self) -> int: + return self._pending_count + + def _enqueue(self, pending: _PendingClose) -> None: + """Append to the entry's bucket, dropping any dead entries it queues behind. + + Deadlines only ever move forward, so appending keeps each bucket ordered + by deadline, and entries whose client the collector already took sit at + the front rather than having to be searched for. + """ + with self._queue_lock: + bucket = self._buckets.setdefault(_bucket_key(pending), deque()) # mutable-ok: FIFO by design + while bucket and bucket[0].client_ref() is None: + bucket.popleft() + self._pending_count -= 1 + bucket.append(pending) + self._pending_count += 1 + + def _take_due(self, loop_id: int | None, now: float) -> tuple[_PendingClose, ...]: + buckets = (_CLOSABLE_ANYWHERE,) if loop_id is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_id) + with self._queue_lock: + return tuple(pending for key in buckets for pending in self._drain_locked(key, now)) + + def _drain_locked(self, key: _BucketKey, now: float) -> Iterator[_PendingClose]: + bucket = self._buckets.get(key) + if bucket is None: + return + while bucket and bucket[0].close_after <= now: + self._pending_count -= 1 + yield bucket.popleft() + if not bucket: + del self._buckets[key] + + def _close(self, client: object) -> None: + close_fn = _close_function(client) + if close_fn is None: + return + try: + closing = close_fn() + except Exception: # noqa: BLE001 - a discarded client's close must never surface to callers + return + if not inspect.isawaitable(closing): + return + task = asyncio.get_running_loop().create_task(_close_quietly(closing)) + self._close_tasks.add(task) + task.add_done_callback(self._close_tasks.discard) + + +default_evicted_client_closer = EvictedClientCloser() diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py index c2274713bb9..7eae8ee3749 100644 --- a/litellm/caching/llm_caching_handler.py +++ b/litellm/caching/llm_caching_handler.py @@ -4,21 +4,44 @@ Add the event loop to the cache key, to prevent event loop closed errors. import asyncio +from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer from .in_memory_cache import InMemoryCache class LLMClientCache(InMemoryCache): """Cache for LLM HTTP clients (OpenAI, Azure, httpx, etc.). - IMPORTANT: This cache intentionally does NOT close clients on eviction. - Evicted clients may still be in use by in-flight requests. Closing them - eagerly causes ``RuntimeError: Cannot send a request, as the client has - been closed.`` errors in production after the TTL (1 hour) expires. + An evicted client is never closed on the spot: a request handed the client + just before eviction is still using it, and closing it there raises + ``RuntimeError: Cannot send a request, as the client has been closed.`` - Clients that are no longer referenced will be garbage-collected normally. - For explicit shutdown cleanup, use ``close_litellm_async_clients()``. + Nor can eviction be left to rely on garbage collection. The SDK clients are + reference cycles, so an evicted client and its open TCP connections survive + until a generational collection runs. Instead a client litellm created is + handed to ``EvictedClientCloser``, which closes it once a grace window has + passed. Clients the caller supplied are left untouched. """ + def __init__( + self, + max_size_in_memory: int | None = 200, + default_ttl: int | None = 600, + max_size_per_item: int | None = 1024, + evicted_client_closer: EvictedClientCloser | None = None, + ): + super().__init__( + max_size_in_memory=max_size_in_memory, + default_ttl=default_ttl, + max_size_per_item=max_size_per_item, + ) + self.evicted_client_closer = evicted_client_closer or default_evicted_client_closer + + def _remove_key(self, key: str) -> None: + evicted: object = self.cache_dict.get(key) + super()._remove_key(key) + self.evicted_client_closer.schedule(evicted) + self.evicted_client_closer.reap() + def update_cache_key_with_event_loop(self, key): """ Add the event loop to the cache key, to prevent event loop closed errors. @@ -31,16 +54,22 @@ class LLMClientCache(InMemoryCache): except RuntimeError: # handle no current running event loop return key - def set_cache(self, key, value, **kwargs): + def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs): + """``litellm_owned_client`` marks a client litellm built, so it may be closed once evicted.""" + if litellm_owned_client: + self.evicted_client_closer.mark_owned(value) key = self.update_cache_key_with_event_loop(key) return super().set_cache(key, value, **kwargs) - async def async_set_cache(self, key, value, **kwargs): + async def async_set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs): + if litellm_owned_client: + self.evicted_client_closer.mark_owned(value) key = self.update_cache_key_with_event_loop(key) return await super().async_set_cache(key, value, **kwargs) def get_cache(self, key, **kwargs): key = self.update_cache_key_with_event_loop(key) + self.evicted_client_closer.reap() return super().get_cache(key, **kwargs) diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 6e36dfbc096..98fd9cfd1d2 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -178,7 +178,7 @@ class QdrantSemanticCache(BaseCache): if response.status_code not in (200, 201): print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {response.text}") except Exception as exc: - print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {exc!s}") + print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {exc}") def _payload_matches_cache_key(self, payload: dict, key: str) -> bool: # Pre-isolation points stored only prompt + response with no cache-key diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 1b0aa778f4e..e3c0e3616f0 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -346,7 +346,7 @@ class RedisCache(BaseCache): verbose_logger.debug("Ignoring async redis ping. No running event loop.") else: verbose_logger.error( - f"Error connecting to Async Redis client - {e!s}", + f"Error connecting to Async Redis client - {e}", extra={"error": str(e)}, ) self._handle_async_ping_error(e) @@ -483,7 +483,7 @@ 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 : {e!s}") + print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {e}") def increment_cache(self, key, value: int, ttl: float | None = None, **kwargs) -> int: _redis_client = self.redis_client @@ -1139,7 +1139,7 @@ class RedisCache(BaseCache): return decoded_results except Exception as e: - verbose_logger.error(f"Error occurred in batch get cache - {e!s}") + verbose_logger.error(f"Error occurred in batch get cache - {e}") return key_value_dict @_redis_circuit_breaker_guard @@ -1185,7 +1185,7 @@ class RedisCache(BaseCache): event_metadata={"key": key}, ) ) - print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {e!s}") + print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {e}") _record_swallowed_redis_failure(self._circuit_breaker, e) @_redis_circuit_breaker_guard @@ -1257,7 +1257,7 @@ class RedisCache(BaseCache): parent_otel_span=parent_otel_span, ) ) - verbose_logger.error(f"Error occurred in async batch get cache - {e!s}") + verbose_logger.error(f"Error occurred in async batch get cache - {e}") _record_swallowed_redis_failure(self._circuit_breaker, e) return key_value_dict @@ -1292,7 +1292,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 : {e!s}") + verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {e}") raise e async def ping(self) -> bool: @@ -1326,7 +1326,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 : {e!s}") + verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {e}") raise e @_redis_circuit_breaker_guard @@ -1388,10 +1388,10 @@ class RedisCache(BaseCache): else: return {"status": "failed", "message": "Redis ping returned False"} except Exception as e: - verbose_logger.error(f"Redis connection test failed: {e!s}") + verbose_logger.error(f"Redis connection test failed: {e}") return { "status": "failed", - "message": f"Redis connection failed: {e!s}", + "message": f"Redis connection failed: {e}", "error": str(e), } @@ -1565,7 +1565,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 : {e!s}") + verbose_logger.error(f"LiteLLM Redis Cache RPUSH: - Got exception from REDIS : {e}") raise e async def _pipeline_rpush_helper( @@ -1711,7 +1711,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 : {e!s}") + verbose_logger.error(f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {e}") raise e async def _pipeline_lpop_helper( diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index 1e4c4684f48..127a5c3bd29 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -100,9 +100,9 @@ class RedisClusterCache(RedisCache): except Exception as e: from litellm._logging import verbose_logger - verbose_logger.error(f"Redis Cluster connection test failed: {e!s}") + verbose_logger.error(f"Redis Cluster connection test failed: {e}") return { "status": "failed", - "message": f"Redis Cluster connection failed: {e!s}", + "message": f"Redis Cluster connection failed: {e}", "error": str(e), } diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index b2d8efa1dba..f55274d446d 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -364,7 +364,7 @@ class RedisSemanticCache(BaseCache): try: cached_response = ast.literal_eval(cached_response) except (ValueError, SyntaxError) as e: - print_verbose(f"Error parsing cached response: {e!s}") + print_verbose(f"Error parsing cached response: {e}") return None return cached_response @@ -403,7 +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: {e!s}") + print_verbose(f"Error setting {value_str or value} in the Redis semantic cache: {e}") def get_cache(self, key: str, **kwargs) -> Any: """ @@ -468,7 +468,7 @@ class RedisSemanticCache(BaseCache): return self._get_cache_logic(cached_response=cached_response) except Exception as e: - print_verbose(f"Error retrieving from Redis semantic cache: {e!s}") + print_verbose(f"Error retrieving from Redis semantic cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> list[float]: @@ -505,8 +505,8 @@ class RedisSemanticCache(BaseCache): ) return embedding_response["data"][0]["embedding"] except Exception as e: - print_verbose(f"Error generating async embedding: {e!s}") - raise ValueError(f"Failed to generate embedding: {e!s}") from e + print_verbose(f"Error generating async embedding: {e}") + raise ValueError(f"Failed to generate embedding: {e}") from e async def async_set_cache(self, key: str, value: Any, **kwargs) -> None: """ @@ -546,7 +546,7 @@ class RedisSemanticCache(BaseCache): **store_kwargs, ) except Exception as e: - print_verbose(f"Error in async_set_cache: {e!s}") + print_verbose(f"Error in async_set_cache: {e}") async def async_get_cache(self, key: str, **kwargs) -> Any: """ @@ -612,7 +612,7 @@ class RedisSemanticCache(BaseCache): return self._get_cache_logic(cached_response=cached_response) except Exception as e: - print_verbose(f"Error in async_get_cache: {e!s}") + print_verbose(f"Error in async_get_cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 async def _index_info(self) -> dict[str, Any]: @@ -639,4 +639,4 @@ class RedisSemanticCache(BaseCache): tasks.append(self.async_set_cache(val[0], val[1], **kwargs)) await asyncio.gather(*tasks) except Exception as e: - print_verbose(f"Error in async_set_cache_pipeline: {e!s}") + print_verbose(f"Error in async_set_cache_pipeline: {e}") diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py index 86e687c0009..e01bb430987 100644 --- a/litellm/caching/valkey_semantic_cache.py +++ b/litellm/caching/valkey_semantic_cache.py @@ -249,7 +249,7 @@ class ValkeySemanticCache(RedisSemanticCache): if ttl is not None: self.sync_client.expire(doc_key, ttl) except Exception as e: - print_verbose(f"Error in Valkey semantic-cache set_cache: {e!s}") + print_verbose(f"Error in Valkey semantic-cache set_cache: {e}") def get_cache(self, key: str, **kwargs: Any) -> Any: print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}") @@ -268,7 +268,7 @@ class ValkeySemanticCache(RedisSemanticCache): ) return self._resolve_hit(self._first_hit(search_result), key, **kwargs) except Exception as e: - print_verbose(f"Error in Valkey semantic-cache get_cache: {e!s}") + print_verbose(f"Error in Valkey semantic-cache get_cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None: @@ -288,7 +288,7 @@ class ValkeySemanticCache(RedisSemanticCache): if ttl is not None: await self.async_client.expire(doc_key, ttl) except Exception as e: - print_verbose(f"Error in async Valkey semantic-cache set_cache: {e!s}") + print_verbose(f"Error in async Valkey semantic-cache set_cache: {e}") async def async_get_cache(self, key: str, **kwargs: Any) -> Any: print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}") @@ -307,14 +307,14 @@ class ValkeySemanticCache(RedisSemanticCache): ) return self._resolve_hit(self._first_hit(search_result), key, **kwargs) except Exception as e: - print_verbose(f"Error in async Valkey semantic-cache get_cache: {e!s}") + print_verbose(f"Error in async Valkey semantic-cache get_cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 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]) except Exception as e: - print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e!s}") + print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e}") async def _index_info(self) -> dict: return await self.async_client.ft(self.index_name).info() diff --git a/litellm/constants.py b/litellm/constants.py index d46f62af000..06421e6ed6a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -197,6 +197,16 @@ RUNWAYML_POLLING_TIMEOUT = int(os.getenv("RUNWAYML_POLLING_TIMEOUT", 600)) # 10 ########## Networking constants ############################################################## _DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client for 1 hour +# The earliest an evicted, litellm-created client may be closed. A request handed the +# client just before eviction is still using it, so nothing is closed inside this window; +# past it, the client is closed once it reports no connection in flight. +EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS = 900 + +# How many evicted clients may be queued for closing at once. Past this, an evicted client +# is left to the collector rather than letting a cache-churning workload grow the queue +# without bound. Each queued entry is ~100 bytes and comes due within one grace window. +EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING = 10_000 + # 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)) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index f10a9e327d6..f04a9d61d4a 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -715,7 +715,7 @@ def _get_provider_for_cost_calc( _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) except Exception as e: verbose_logger.debug( - f"litellm.cost_calculator.py::_get_provider_for_cost_calc() - Error inferring custom_llm_provider - {e!s}" + f"litellm.cost_calculator.py::_get_provider_for_cost_calc() - Error inferring custom_llm_provider - {e}" ) return None @@ -1092,7 +1092,7 @@ def _store_cost_breakdown_in_logging_obj( ) except Exception as breakdown_error: - verbose_logger.debug(f"Error storing cost breakdown: {breakdown_error!s}") + verbose_logger.debug(f"Error storing cost breakdown: {breakdown_error}") # Don't fail the main cost calculation if breakdown storage fails @@ -1315,7 +1315,7 @@ def completion_cost( ) # strip the llm provider from the model name -> for image gen cost calculation except Exception as e: verbose_logger.debug( - f"litellm.cost_calculator.py::completion_cost() - Error inferring custom_llm_provider - {e!s}" + f"litellm.cost_calculator.py::completion_cost() - Error inferring custom_llm_provider - {e}" ) if CostCalculatorUtils._call_type_has_image_response(call_type) and isinstance( completion_response, ImageResponse @@ -1662,7 +1662,7 @@ def completion_cost( return _final_cost except Exception as e: verbose_logger.debug( - f"litellm.cost_calculator.py::completion_cost() - Error calculating cost for model={model} - {e!s}" + f"litellm.cost_calculator.py::completion_cost() - Error calculating cost for model={model} - {e}" ) if idx == len(potential_model_names) - 1: raise e diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 0d85c795c7b..c4a64e0ad9b 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -1140,7 +1140,7 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore if self.max_retries: _message += f", LiteLLM Max Retries: {self.max_retries}" if self.original_exception: - _message += f" Original exception: {type(self.original_exception).__name__}: {self.original_exception!s}" + _message += f" Original exception: {type(self.original_exception).__name__}: {self.original_exception}" return _message def __repr__(self): diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 72248c4448d..8815c38192b 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -515,7 +515,7 @@ class MCPClient: _log( f"MCP client list_tools failed - " f"Error Type: {error_type}, " - f"Error: {e!s}, " + f"Error: {e}, " f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) @@ -536,7 +536,7 @@ class MCPClient: def error_tool_result(exc: Exception) -> MCPCallToolResult: """The error result ``call_tool`` returns when it swallows a failure (no re-execution).""" return MCPCallToolResult( - content=[TextContent(type="text", text=f"{type(exc).__name__}: {exc!s}")], + content=[TextContent(type="text", text=f"{type(exc).__name__}: {exc}")], isError=True, ) @@ -601,7 +601,7 @@ class MCPClient: _log( f"MCP client call_tool failed - " f"Error Type: {error_type}, " - f"Error: {e!s}, " + f"Error: {e}, " f"Tool: {call_tool_request_params.name}, " f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" @@ -640,7 +640,7 @@ class MCPClient: verbose_logger.error( f"MCP client list_prompts failed - " f"Error Type: {error_type}, " - f"Error: {e!s}, " + f"Error: {e}, " f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) @@ -681,7 +681,7 @@ class MCPClient: verbose_logger.error( f"MCP client get_prompt failed - " f"Error Type: {error_type}, " - f"Error: {e!s}, " + f"Error: {e}, " f"Prompt: {get_prompt_request_params.name}, " f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" @@ -717,7 +717,7 @@ class MCPClient: verbose_logger.error( f"MCP client list_resources failed - " f"Error Type: {error_type}, " - f"Error: {e!s}, " + f"Error: {e}, " f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) @@ -753,7 +753,7 @@ class MCPClient: verbose_logger.error( f"MCP client list_resource_templates failed - " f"Error Type: {error_type}, " - f"Error: {e!s}, " + f"Error: {e}, " f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) @@ -791,7 +791,7 @@ class MCPClient: verbose_logger.error( f"MCP client read_resource failed - " f"Error Type: {error_type}, " - f"Error: {e!s}, " + f"Error: {e}, " f"Url: {url}, " f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index 573f0633af5..5236e207cc5 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -98,7 +98,7 @@ class GenerateContentToCompletionHandler: return generate_content_response except Exception as e: - raise ValueError(f"Error calling litellm.acompletion for generate_content: {e!s}") + raise ValueError(f"Error calling litellm.acompletion for generate_content: {e}") @staticmethod def generate_content_handler( @@ -159,4 +159,4 @@ class GenerateContentToCompletionHandler: return generate_content_response except Exception as e: - raise ValueError(f"Error calling litellm.completion for generate_content: {e!s}") + raise ValueError(f"Error calling litellm.completion for generate_content: {e}") diff --git a/litellm/integrations/SlackAlerting/batching_handler.py b/litellm/integrations/SlackAlerting/batching_handler.py index e5a60640ee2..da905b606a5 100644 --- a/litellm/integrations/SlackAlerting/batching_handler.py +++ b/litellm/integrations/SlackAlerting/batching_handler.py @@ -70,6 +70,6 @@ async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item, count) if response.status_code != 200: 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: {e!s}") + verbose_proxy_logger.debug(f"Error sending slack alert: {e}") finally: _print_alerting_payload_warning(payload, slackAlertingInstance=slackAlertingInstance) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 4378b2f754e..114924e7359 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -1467,7 +1467,7 @@ Model Info: try: await self._flush_digest_buckets() except Exception as e: - verbose_proxy_logger.debug(f"Error flushing digest buckets: {e!s}") + verbose_proxy_logger.debug(f"Error flushing digest buckets: {e}") await self.flush_queue() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -1502,7 +1502,7 @@ Model Info: ) except Exception as e: verbose_proxy_logger.error( - f"[Non-Blocking Error] Slack Alerting: Got error in logging LLM deployment latency: {e!s}" + f"[Non-Blocking Error] Slack Alerting: Got error in logging LLM deployment latency: {e}" ) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): @@ -1522,7 +1522,7 @@ Model Info: ) ) except Exception as e: - verbose_logger.debug(f"Exception raises -{e!s}") + verbose_logger.debug(f"Exception raises -{e}") if isinstance(kwargs.get("exception", ""), APIError): if "outage_alerts" in self.alert_types: diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index 9d743659135..86e861afb8a 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -169,7 +169,7 @@ class ArizeLogger(OpenTelemetry): except Exception as e: return { "status": "unhealthy", - "error_message": f"Arize health check failed: {e!s}", + "error_message": f"Arize health check failed: {e}", } def construct_dynamic_otel_headers( diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index f0200b75c43..e0ed0cd7cf3 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -203,7 +203,7 @@ class AzureSentinelLogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception(f"Azure Sentinel Layer Error - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Azure Sentinel Layer Error - {e}\n{traceback.format_exc()}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """ @@ -233,7 +233,7 @@ class AzureSentinelLogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception(f"Azure Sentinel Layer Error - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Azure Sentinel Layer Error - {e}\n{traceback.format_exc()}") async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None: """ @@ -256,7 +256,7 @@ class AzureSentinelLogger(CustomBatchLogger): await self.async_send_audit_batch() except Exception as e: - verbose_logger.exception(f"Azure Sentinel Audit Log Layer Error - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Azure Sentinel Audit Log Layer Error - {e}\n{traceback.format_exc()}") async def async_send_batch(self): """ @@ -323,7 +323,7 @@ class AzureSentinelLogger(CustomBatchLogger): ) except Exception as e: - verbose_logger.exception(f"Azure Sentinel Error sending batch API - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Azure Sentinel Error sending batch API - {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 bbd6e9698bb..d2dd3d37dc7 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -53,9 +53,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): self.log_queue: list[StandardLoggingPayload] = [] super().__init__(**kwargs, flush_lock=self.flush_lock) except Exception as e: - verbose_logger.exception( - f"AzureBlobStorageLogger: Got exception on init AzureBlobStorageLogger client {e!s}" - ) + verbose_logger.exception(f"AzureBlobStorageLogger: Got exception on init AzureBlobStorageLogger client {e}") raise e async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -79,7 +77,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): self.log_queue.append(standard_logging_payload) except Exception as e: - verbose_logger.exception(f"AzureBlobStorageLogger Layer Error - {e!s}") + verbose_logger.exception(f"AzureBlobStorageLogger Layer Error - {e}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """ @@ -101,7 +99,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): self.log_queue.append(standard_logging_payload) except Exception as e: - verbose_logger.exception(f"AzureBlobStorageLogger Layer Error - {e!s}") + verbose_logger.exception(f"AzureBlobStorageLogger Layer Error - {e}") async def async_send_batch(self): """ @@ -124,7 +122,7 @@ 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 - {e!s}") + verbose_logger.exception(f"AzureBlobStorageLogger Error sending batch API - {e}") async def async_upload_payload_to_azure_blob_storage(self, payload: StandardLoggingPayload): """ @@ -153,7 +151,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): 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: {e!s}") + verbose_logger.exception(f"Error uploading to Azure Blob Storage: {e}") raise e async def _create_file(self, client: AsyncHTTPHandler, base_url: str): @@ -169,7 +167,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): response.raise_for_status() verbose_logger.debug("Successfully created file resource") except Exception as e: - verbose_logger.exception(f"Error creating file resource: {e!s}") + verbose_logger.exception(f"Error creating file resource: {e}") raise async def _append_data(self, client: AsyncHTTPHandler, base_url: str, json_payload: str): @@ -189,7 +187,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): response.raise_for_status() verbose_logger.debug("Successfully appended data") except Exception as e: - verbose_logger.exception(f"Error appending data: {e!s}") + verbose_logger.exception(f"Error appending data: {e}") raise async def _flush_data(self, client: AsyncHTTPHandler, base_url: str, position: int): @@ -205,7 +203,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): response.raise_for_status() verbose_logger.debug("Successfully flushed data") except Exception as e: - verbose_logger.exception(f"Error flushing data: {e!s}") + verbose_logger.exception(f"Error flushing data: {e}") raise ####### Helper methods to managing Authentication to Azure Storage ####### @@ -345,4 +343,4 @@ class AzureBlobStorageLogger(CustomBatchLogger): verbose_logger.debug(f"Successfully uploaded and wrote to {today}/{file_name}") except Exception as e: - verbose_logger.exception(f"Error occurred: {e!s}") + verbose_logger.exception(f"Error occurred: {e}") diff --git a/litellm/integrations/cloudzero/cloudzero.py b/litellm/integrations/cloudzero/cloudzero.py index e6faf4a6a62..52b41f74fce 100644 --- a/litellm/integrations/cloudzero/cloudzero.py +++ b/litellm/integrations/cloudzero/cloudzero.py @@ -153,7 +153,7 @@ class CloudZeroLogger(CustomLogger): 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: {e!s}") + verbose_logger.error(f"CloudZero Logger: Error exporting usage data: {e}") raise async def dry_run_export_usage_data(self, limit: int | None = 10000): @@ -244,8 +244,8 @@ class CloudZeroLogger(CustomLogger): } except Exception as e: - verbose_logger.error(f"CloudZero Logger: Error in dry run export: {e!s}") - verbose_logger.error(f"CloudZero Dry Run Error: {e!s}") + verbose_logger.error(f"CloudZero Logger: Error in dry run export: {e}") + verbose_logger.error(f"CloudZero Dry Run Error: {e}") raise def _display_cbf_data_on_screen(self, cbf_data): diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 16fb99517ae..2d0f81af98b 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -98,4 +98,4 @@ class LiteLLMDatabase: # This prevents schema mismatch errors when data types vary across rows return pl.DataFrame(db_response, infer_schema_length=None) except Exception as e: - raise Exception(f"Error retrieving usage data: {e!s}") + raise Exception(f"Error retrieving usage data: {e}") diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 971d53ffec4..9915224ba09 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -927,7 +927,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}: {e!s}") + verbose_logger.debug(f"Error in handle_callback_failure for {callback_name}: {e}") async def _strip_base64_from_messages( self, diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 047d69c9c9c..fa14e1fa459 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -171,7 +171,7 @@ class DataDogLogger( batch_size=_resolve_dd_batch_size(), ) except Exception as e: - verbose_logger.exception(f"Datadog: Got exception on init Datadog client {e!s}") + verbose_logger.exception(f"Datadog: Got exception on init Datadog client {e}") raise e def _get_datadog_params(self) -> dict: @@ -257,7 +257,7 @@ class DataDogLogger( await self._log_async_event(kwargs, response_obj, start_time, end_time) except Exception as e: - verbose_logger.exception(f"Datadog Layer Error - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Datadog Layer Error - {e}\n{traceback.format_exc()}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: @@ -265,7 +265,7 @@ class DataDogLogger( await self._log_async_event(kwargs, response_obj, start_time, end_time) except Exception as e: - verbose_logger.exception(f"Datadog Layer Error - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Datadog Layer Error - {e}\n{traceback.format_exc()}") async def async_post_call_failure_hook( self, @@ -340,7 +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 - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Datadog: async_post_call_failure_hook - {e}\n{traceback.format_exc()}") return None async def async_send_batch(self): @@ -380,7 +380,7 @@ class DataDogLogger( except Exception as e: self.log_queue = batch_to_send + self.log_queue - verbose_logger.exception(f"Datadog Error sending batch API - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Datadog Error sending batch API - {e}\n{traceback.format_exc()}") async def _send_with_413_split(self, batch: list) -> list: """ @@ -411,7 +411,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 - {e!s}") + verbose_logger.exception(f"Datadog Error sending batch API - {e}") return self._undelivered(chunk, pending) if response.status_code == 413: @@ -515,7 +515,7 @@ class DataDogLogger( ) except Exception as e: - verbose_logger.exception(f"Datadog Layer Error - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Datadog Layer Error - {e}\n{traceback.format_exc()}") async def _log_async_event(self, kwargs, response_obj, start_time, end_time): dd_payload = self.create_datadog_logging_payload( diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py index 7b22f4658f2..da45f94f02b 100644 --- a/litellm/integrations/datadog/datadog_cost_management.py +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -84,7 +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: {e!s}") + verbose_logger.exception(f"Datadog Cost Management: Error in async_log_success_event: {e}") async def async_send_batch(self): if not self.log_queue: @@ -104,7 +104,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): 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: {e!s}") + verbose_logger.exception(f"Datadog Cost Management: Error in async_send_batch: {e}") def _aggregate_costs(self, logs: list[StandardLoggingPayload]) -> list[DatadogFOCUSCostEntry]: """ diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index e10071cb083..02e1affd361 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -89,7 +89,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): kwargs.update(dict_datadog_llm_obs_params) CustomBatchLogger.__init__(self, **kwargs, flush_lock=self.flush_lock) except Exception as e: - verbose_logger.exception(f"DataDogLLMObs: Error initializing - {e!s}") + verbose_logger.exception(f"DataDogLLMObs: Error initializing - {e}") raise e def _configure_dd_agent(self, dd_agent_host: str): @@ -145,7 +145,7 @@ 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 - {e!s}") + verbose_logger.exception(f"DataDogLLMObs: Error logging success event - {e}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: @@ -157,7 +157,7 @@ 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 - {e!s}") + verbose_logger.exception(f"DataDogLLMObs: Error logging failure event - {e}") async def async_send_batch(self): try: @@ -214,7 +214,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): except httpx.HTTPStatusError as e: verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {e.response.text}") except Exception as e: - verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {e!s}") + verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {e}") def create_llm_obs_payload(self, kwargs: dict, start_time: datetime, end_time: datetime) -> LLMObsPayload: standard_logging_payload: StandardLoggingPayload | None = kwargs.get("standard_logging_object") @@ -707,7 +707,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): 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}: {e!s}") + verbose_logger.debug(f"DataDogLLMObs: Error processing tool call {idx}: {e}") continue return kv_pairs @@ -747,6 +747,6 @@ class DataDogLLMObsLogger(CustomBatchLogger): tool_call_metadata[f"output_{key}"] = value except Exception as e: - verbose_logger.debug(f"DataDogLLMObs: Error extracting tool call metadata: {e!s}") + verbose_logger.debug(f"DataDogLLMObs: Error extracting tool call metadata: {e}") return tool_call_metadata diff --git a/litellm/integrations/datadog/datadog_metrics.py b/litellm/integrations/datadog/datadog_metrics.py index 3fbd0f917dc..9fb86bfb125 100644 --- a/litellm/integrations/datadog/datadog_metrics.py +++ b/litellm/integrations/datadog/datadog_metrics.py @@ -180,7 +180,7 @@ class DatadogMetricsLogger(CustomBatchLogger): await self.flush_queue() except Exception as e: - verbose_logger.exception(f"Datadog Metrics: Error in async_log_success_event: {e!s}") + verbose_logger.exception(f"Datadog Metrics: Error in async_log_success_event: {e}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: @@ -202,7 +202,7 @@ class DatadogMetricsLogger(CustomBatchLogger): await self.flush_queue() except Exception as e: - verbose_logger.exception(f"Datadog Metrics: Error in async_log_failure_event: {e!s}") + verbose_logger.exception(f"Datadog Metrics: Error in async_log_failure_event: {e}") async def async_send_batch(self): if not self.log_queue: @@ -214,7 +214,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: {e!s}") + verbose_logger.exception(f"Datadog Metrics: Error in async_send_batch: {e}") raise async def _upload_to_datadog(self, payload: DatadogMetricsPayload): diff --git a/litellm/integrations/dynamodb.py b/litellm/integrations/dynamodb.py index 5826a06b0ec..a41130cbab1 100644 --- a/litellm/integrations/dynamodb.py +++ b/litellm/integrations/dynamodb.py @@ -70,7 +70,7 @@ class DyanmoDBLogger: # Assuming log_data is a dictionary with log information response = table.put_item(Item=payload) - print_verbose(f"Response from DynamoDB:{response!s}") + print_verbose(f"Response from DynamoDB:{response}") print_verbose(f"DynamoDB Layer Logging - final response object: {response_obj}") return response diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index 0180af51992..f7870a6c0f8 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -128,7 +128,7 @@ class GalileoObserve(CustomLogger): except Exception as e: return IntegrationHealthCheckStatus( status="unhealthy", - error_message=f"Galileo health check failed: {e!s}", + error_message=f"Galileo health check failed: {e}", ) async def async_set_galileo_headers(self) -> None: diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 552e078cb60..b5b3d4e81a3 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -76,7 +76,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): 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: {e!s}") + verbose_logger.exception(f"GCS Bucket logging error: {e}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: @@ -95,7 +95,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): 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: {e!s}") + verbose_logger.exception(f"GCS Bucket logging error: {e}") def _drain_queue_batch(self) -> list[GCSLogQueueItem]: """ @@ -218,7 +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: {e!s}") + verbose_logger.exception(f"GCS Bucket error logging batch payload to GCS bucket: {e}") return (success_count, error_count) async def _send_individual_logs(self, items: list[GCSLogQueueItem]) -> None: @@ -255,7 +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: {e!s}") + verbose_logger.exception(f"GCS Bucket error logging individual payload to GCS bucket: {e}") async def async_send_batch(self): """ @@ -336,7 +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}: {e!s}") + verbose_logger.debug(f"Failed to fetch payload for date {date_str}: {e}") continue return None diff --git a/litellm/integrations/gcs_pubsub/pub_sub.py b/litellm/integrations/gcs_pubsub/pub_sub.py index 6ade70ab6d6..b43e7626b77 100644 --- a/litellm/integrations/gcs_pubsub/pub_sub.py +++ b/litellm/integrations/gcs_pubsub/pub_sub.py @@ -132,7 +132,7 @@ class GcsPubSubLogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception(f"PubSub Layer Error - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"PubSub Layer Error - {e}\n{traceback.format_exc()}") async def async_send_batch(self): """ @@ -148,7 +148,7 @@ class GcsPubSubLogger(CustomBatchLogger): await self.publish_message(message) except Exception as e: - verbose_logger.exception(f"PubSub Error sending batch - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"PubSub Error sending batch - {e}\n{traceback.format_exc()}") finally: self.log_queue.clear() diff --git a/litellm/integrations/generic_api/generic_api_callback.py b/litellm/integrations/generic_api/generic_api_callback.py index a524755540e..c7f2661a5ad 100644 --- a/litellm/integrations/generic_api/generic_api_callback.py +++ b/litellm/integrations/generic_api/generic_api_callback.py @@ -42,7 +42,7 @@ def load_compatible_callbacks() -> dict: 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: {e!s}") + verbose_logger.warning(f"Error loading generic_api_compatible_callbacks.json: {e}") return {} @@ -214,7 +214,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: {e!s}") + verbose_logger.warning(f"Error parsing headers from environment variables: {e}") # 2. Update with litellm generic headers if available if litellm.generic_logger_headers: @@ -308,7 +308,7 @@ class GenericAPILogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception(f"Generic API Logger Error - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Generic API Logger Error - {e}\n{traceback.format_exc()}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """ @@ -339,7 +339,7 @@ class GenericAPILogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception(f"Generic API Logger Error - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Generic API Logger Error - {e}\n{traceback.format_exc()}") async def async_send_batch(self): """ @@ -395,7 +395,7 @@ class GenericAPILogger(CustomBatchLogger): ) except Exception as e: - verbose_logger.exception(f"Generic API Logger Error sending batch - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Generic API Logger Error sending batch - {e}\n{traceback.format_exc()}") finally: self.log_queue.clear() diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 3fb50e07b01..2dab1874c01 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -330,7 +330,7 @@ class LangFuseLogger: return {"trace_id": trace_id, "generation_id": generation_id} except Exception as e: - verbose_logger.exception(f"Langfuse Layer Error(): Exception occured - {e!s}") + verbose_logger.exception(f"Langfuse Layer Error(): Exception occured - {e}") return {"trace_id": None, "generation_id": None} def _get_langfuse_input_output_content( diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index 9a5ee49bd0d..56383b45a8c 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -317,7 +317,7 @@ 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: {e!s}") + verbose_logger.exception(f"Langfuse Layer Error - Exception occurred while logging success event: {e}") self.handle_callback_failure(callback_name="langfuse") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): @@ -347,5 +347,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: {e!s}") + verbose_logger.exception(f"Langfuse Layer Error - Exception occurred while logging failure event: {e}") self.handle_callback_failure(callback_name="langfuse") diff --git a/litellm/integrations/logfire_logger.py b/litellm/integrations/logfire_logger.py index 78735c47e5b..c94fb832ccc 100644 --- a/litellm/integrations/logfire_logger.py +++ b/litellm/integrations/logfire_logger.py @@ -35,7 +35,7 @@ class LogfireLogger: if logfire.DEFAULT_LOGFIRE_INSTANCE.config.send_to_logfire: logfire.configure(token=os.getenv("LOGFIRE_TOKEN")) except Exception as e: - print_verbose(f"Got exception on init logfire client {e!s}") + print_verbose(f"Got exception on init logfire client {e}") raise e def _get_span_config(self, payload) -> SpanConfig: @@ -159,4 +159,4 @@ class LogfireLogger: print_verbose(f"Logfire Layer Logging - final response object: {response_obj}") except Exception as e: - verbose_logger.debug(f"Logfire Layer Error - {e!s}\n{traceback.format_exc()}") + verbose_logger.debug(f"Logfire Layer Error - {e}\n{traceback.format_exc()}") diff --git a/litellm/integrations/opik/opik.py b/litellm/integrations/opik/opik.py index deb325286e9..e4d40a1af8f 100644 --- a/litellm/integrations/opik/opik.py +++ b/litellm/integrations/opik/opik.py @@ -81,7 +81,7 @@ class OpikLogger(CustomBatchLogger): self.flush_lock: asyncio.Lock | None = asyncio.Lock() except Exception as e: verbose_logger.exception( - f"OpikLogger - Asynchronous processing not initialized as we are not running in an async context {e!s}" + f"OpikLogger - Asynchronous processing not initialized as we are not running in an async context {e}" ) self.flush_lock = None @@ -161,7 +161,7 @@ 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 - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"OpikLogger failed to log success event - {e}\n{traceback.format_exc()}") def _sync_send(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None: try: @@ -174,7 +174,7 @@ class OpikLogger(CustomBatchLogger): if response.status_code != 204: 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 - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"OpikLogger failed to send batch - {e}\n{traceback.format_exc()}") def log_success_event( self, @@ -245,7 +245,7 @@ class OpikLogger(CustomBatchLogger): batch={"spans": [span_payload.__dict__]}, ) except Exception as e: - verbose_logger.exception(f"OpikLogger failed to log success event - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"OpikLogger failed to log success event - {e}\n{traceback.format_exc()}") async def _submit_batch(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None: try: @@ -261,7 +261,7 @@ class OpikLogger(CustomBatchLogger): else: 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 - {e!s}") + verbose_logger.exception(f"OpikLogger failed to send batch - {e}") def _create_opik_headers(self) -> dict[str, str]: headers: dict[str, str] = {} diff --git a/litellm/integrations/posthog.py b/litellm/integrations/posthog.py index b61eeb8198f..216edc44d3f 100644 --- a/litellm/integrations/posthog.py +++ b/litellm/integrations/posthog.py @@ -72,7 +72,7 @@ class PostHogLogger(CustomBatchLogger): 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 {e!s}") + verbose_logger.exception(f"PostHog: Got exception on init PostHog client {e}") raise e def log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -107,7 +107,7 @@ class PostHogLogger(CustomBatchLogger): verbose_logger.debug("PostHog: Sync event successfully sent") except Exception as e: - verbose_logger.exception(f"PostHog Sync Layer Error - {e!s}") + verbose_logger.exception(f"PostHog Sync Layer Error - {e}") async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: @@ -115,7 +115,7 @@ class PostHogLogger(CustomBatchLogger): 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 - {e!s}") + verbose_logger.exception(f"PostHog Layer Error - {e}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: @@ -123,7 +123,7 @@ class PostHogLogger(CustomBatchLogger): 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 - {e!s}") + verbose_logger.exception(f"PostHog Layer Error - {e}") 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 @@ -367,7 +367,7 @@ class PostHogLogger(CustomBatchLogger): else: 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 - {e!s}") + verbose_logger.exception(f"PostHog Error sending batch API - {e}") def _ensure_async_setup(self): if not self._async_initialized: @@ -377,7 +377,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: {e!s}") + verbose_logger.error(f"PostHog: Failed to initialize async components: {e}") raise def _extract_metadata(self, kwargs: dict[str, Any]) -> dict[str, Any]: @@ -445,4 +445,4 @@ class PostHogLogger(CustomBatchLogger): self.log_queue.clear() except Exception as e: - verbose_logger.error(f"PostHog: Error flushing events on exit: {e!s}") + verbose_logger.error(f"PostHog: Error flushing events on exit: {e}") diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index c84a6c34f1f..b7705a40e0c 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -683,7 +683,7 @@ class PrometheusLogger(CustomLogger): ) except Exception as e: - print_verbose(f"Got exception on init prometheus client {e!s}") + print_verbose(f"Got exception on init prometheus client {e}") raise e def _parse_prometheus_config(self) -> dict[str, list[str]]: @@ -2132,7 +2132,7 @@ class PrometheusLogger(CustomLogger): response_cost=0, ) except Exception as e: - verbose_logger.exception(f"prometheus Layer Error(): Exception occured - {e!s}") + verbose_logger.exception(f"prometheus Layer Error(): Exception occured - {e}") def _extract_status_code( self, @@ -2383,7 +2383,7 @@ class PrometheusLogger(CustomLogger): ) except Exception as e: - verbose_logger.exception(f"prometheus Layer Error(): Exception occured - {e!s}") + verbose_logger.exception(f"prometheus Layer Error(): Exception occured - {e}") async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): """ @@ -2608,7 +2608,7 @@ class PrometheusLogger(CustomLogger): ) except Exception as e: - verbose_logger.debug(f"Prometheus Error: set_llm_deployment_failure_metrics. Exception occured - {e!s}") + verbose_logger.debug(f"Prometheus Error: set_llm_deployment_failure_metrics. Exception occured - {e}") def _set_deployment_tpm_rpm_limit_metrics( self, @@ -2722,9 +2722,7 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_requests_metric.labels(**_labels).set(remaining_requests) except Exception as e: - verbose_logger.exception( - f"Prometheus Error: _async_set_router_remaining_metrics. Exception occured - {e!s}" - ) + verbose_logger.exception(f"Prometheus Error: _async_set_router_remaining_metrics. Exception occured - {e}") def set_llm_deployment_success_metrics( self, @@ -2867,7 +2865,7 @@ class PrometheusLogger(CustomLogger): self.litellm_deployment_latency_per_output_token.labels(**_labels).observe(latency_per_token) except Exception as e: - verbose_logger.exception(f"Prometheus Error: set_llm_deployment_success_metrics. Exception occured - {e!s}") + verbose_logger.exception(f"Prometheus Error: set_llm_deployment_success_metrics. Exception occured - {e}") return def _record_guardrail_metrics( @@ -2912,7 +2910,7 @@ class PrometheusLogger(CustomLogger): hook_type=hook_type, ).inc() except Exception as e: - verbose_logger.debug(f"Error recording guardrail metrics: {e!s}") + verbose_logger.debug(f"Error recording guardrail metrics: {e}") ######################################## # Managed Batch Metric Recording Methods @@ -3315,7 +3313,7 @@ class PrometheusLogger(CustomLogger): await set_metrics_function(data) except Exception as e: - verbose_logger.exception(f"Error initializing {data_type} budget metrics: {e!s}") + verbose_logger.exception(f"Error initializing {data_type} budget metrics: {e}") async def _initialize_team_budget_metrics(self): """ @@ -3506,7 +3504,7 @@ class PrometheusLogger(CustomLogger): self.litellm_teams_count_metric.set(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: {e!s}") + verbose_logger.exception(f"Error initializing user/team count metrics: {e}") async def _set_key_list_budget_metrics(self, keys: list[str | UserAPIKeyAuth]): """Helper function to set budget metrics for a list of keys""" @@ -3597,7 +3595,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: {e!s}") + verbose_logger.debug(f"[Non-Blocking] Prometheus: Error getting team info: {e}") return team_object if team_info: @@ -3695,7 +3693,7 @@ class PrometheusLogger(CustomLogger): include_budget_table=True, ) except Exception as e: - verbose_logger.debug(f"[Non-Blocking] Prometheus: Error getting org info: {e!s}") + verbose_logger.debug(f"[Non-Blocking] Prometheus: Error getting org info: {e}") return if org_info is None: @@ -3852,7 +3850,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: {e!s}") + verbose_logger.debug(f"[Non-Blocking] Prometheus: Error getting key info: {e}") return user_api_key_dict @@ -3917,7 +3915,7 @@ class PrometheusLogger(CustomLogger): check_db_only=False, ) except Exception as e: - verbose_logger.debug(f"[Non-Blocking] Prometheus: Error getting user info: {e!s}") + verbose_logger.debug(f"[Non-Blocking] Prometheus: Error getting user info: {e}") return user_object if user_info: diff --git a/litellm/integrations/prometheus_services.py b/litellm/integrations/prometheus_services.py index f07606a3192..002d61265a4 100644 --- a/litellm/integrations/prometheus_services.py +++ b/litellm/integrations/prometheus_services.py @@ -82,7 +82,7 @@ class PrometheusServicesLogger: self.mock_testing_failure_calls = 0 except Exception as e: - print_verbose(f"Got exception on init prometheus client {e!s}") + print_verbose(f"Got exception on init prometheus client {e}") raise e def _get_service_metrics_initialize(self, service: ServiceTypes) -> list[ServiceMetrics]: diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 2e49da45ce9..9942776bc00 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -1,12 +1,14 @@ -"""Rubrik LiteLLM Plugin for tool blocking and batch logging.""" +"""Rubrik LiteLLM Plugin for prompt/response moderation and batch logging.""" import asyncio import os import random import time -import urllib.parse import uuid from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Literal, Optional import httpx @@ -18,6 +20,10 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.litellm_core_utils.core_helpers import safe_deep_copy +from litellm.litellm_core_utils.model_param_helper import ModelParamHelper +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -28,22 +34,25 @@ from litellm.types.utils import ( Function, GenericGuardrailAPIInputs, StandardLoggingPayload, + StandardLoggingUserAPIKeyMetadata, ) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObj, ) + from litellm.proxy._types import UserAPIKeyAuth -_ENDPOINT_ANTHROPIC_MESSAGES = "/v1/messages" -_WEBHOOK_PATH_TOOL_BLOCKING = "/v1/after_completion/openai/v1" +_WEBHOOK_PATH_RESPONSE_MODERATION = "/v1/after_completion/openai/v1" +_WEBHOOK_PATH_PROMPT_MODERATION = "/v1/before_prompt/openai/v1" _WEBHOOK_PATH_LOGGING_BATCH = "/v1/litellm/batch" _MAX_QUEUE_SIZE = 10_000 _DROP_WARNING_INTERVAL_SECONDS = 60.0 +_EMPTY_MAPPING: Mapping[str, Any] = MappingProxyType({}) class _MalformedToolBlockingResponseError(Exception): - """Raised when the tool blocking service returns a structurally invalid + """Raised when the response moderation service returns a structurally invalid response (e.g. empty ``choices``). Distinct from transient network/HTTP errors so callers can surface a @@ -52,11 +61,15 @@ class _MalformedToolBlockingResponseError(Exception): """ -class RubrikLogger(CustomGuardrail, CustomBatchLogger): - @classmethod - def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: - return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] +@dataclass +class BlockedResponseResult: + """Returned by _extract_response_block when the response was blocked + (response text replaced, or at least one tool call removed).""" + explanation: str + + +class RubrikLogger(CustomGuardrail, CustomBatchLogger): def __init__( self, api_key: str | None = None, @@ -67,21 +80,82 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): kwargs.setdefault("guardrail_name", "rubrik") # `initialize_guardrail` always passes these kwargs explicitly, with # value `None` when the user omits `mode` / `default_on` from the - # guardrail config. Coerce None (omitted) to the desired default - # while preserving any explicit value the caller did set -- - # in particular `default_on=False` if the user wants the guardrail - # off by default. + # guardrail config. Follow the standard litellm convention: omitted + # resolves to False (off by default, user must opt in explicitly). kwargs["event_hook"] = kwargs.get("event_hook") or GuardrailEventHooks.post_call if kwargs.get("default_on") is None: - kwargs["default_on"] = True - kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) + kwargs["default_on"] = False super().__init__( flush_lock=self.flush_lock, + supported_event_hooks=list(self.get_supported_event_hooks()), **kwargs, ) verbose_logger.debug("initializing rubrik logger") + # Defining ``apply_guardrail`` routes streaming responses through + # litellm's ``unified_guardrail.async_post_call_streaming_iterator_hook``. + # By default that hook samples intermediate chunks + # (``streaming_sampling_rate``, default 5) and also moderates at + # end-of-stream, so a streamed response costs ~ceil(N/5)+1 Rubrik + # webhook round-trips. litellm reads this attribute via + # ``getattr(guardrail, "streaming_end_of_stream_only", False)``; when + # True it yields chunks unprocessed and only moderates the fully + # assembled response once at end of stream. + self.streaming_end_of_stream_only = True + + # ``streaming_end_of_stream_only`` is detect-only: it releases every + # chunk to the client *before* moderating, so a block can only append a + # trailing message -- the original content has already been delivered. + # ``streaming_buffer_until_moderated`` (litellm >= BerriAI/litellm#31389) + # withholds all chunks until end-of-stream moderation passes, then + # releases the original response (clean) or only the block message + # (blocked). On older litellm this attribute is ignored and we fall + # back to the detect-only behavior above. + self.streaming_buffer_until_moderated = True + + self._parse_sampling_rate() + + 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.") + + self._parse_batch_size() + + # Cap the in-memory retry queue so a Rubrik webhook outage cannot let + # authenticated traffic accumulate prompt/response payloads until the + # proxy runs out of memory. Once the cap is reached, oldest events are + # dropped to make room for fresh ones (drop-oldest backpressure). + self.max_queue_size = _MAX_QUEUE_SIZE + self._dropped_since_warning = 0 + self._last_drop_warning_time = 0.0 + + _webhook_url = api_base or os.getenv("RUBRIK_WEBHOOK_URL") + if not _webhook_url: + raise ValueError("Rubrik webhook URL not configured. Set RUBRIK_WEBHOOK_URL or pass api_base.") + + _webhook_url = _webhook_url.rstrip("/").removesuffix("/v1") + self._setup_clients(_webhook_url) + + self._headers: Mapping[str, str] = MappingProxyType( + {"Content-Type": "application/json", "Authorization": f"Bearer {self.key}"} + if self.key + else {"Content-Type": "application/json"} + ) + + self._periodic_flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task() + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: + """Return the guardrail event hooks this integration supports. + + Prompt moderation (``pre_call``) evaluates the user's message before + the LLM is called. Response moderation (``post_call``) evaluates the + assistant's reply and tool calls after the LLM returns. + """ + return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] + + def _parse_sampling_rate(self) -> None: self.sampling_rate = 1.0 rbrk_sampling_rate = os.getenv("RUBRIK_SAMPLING_RATE") if rbrk_sampling_rate is not None: @@ -93,80 +167,54 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): except ValueError: 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.") + def _parse_batch_size(self) -> None: _batch_size = os.getenv("RUBRIK_BATCH_SIZE") - if _batch_size: try: - self.batch_size = int(_batch_size) + parsed_size = int(_batch_size) + if parsed_size <= 0: + verbose_logger.warning(f"RUBRIK_BATCH_SIZE={_batch_size!r} must be > 0, using default") + else: + self.batch_size = parsed_size except ValueError: 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 - # proxy runs out of memory. Once the cap is reached, oldest events are - # dropped to make room for fresh ones (drop-oldest backpressure). - self.max_queue_size = _MAX_QUEUE_SIZE - self._dropped_since_warning = 0 - self._last_drop_warning_time = 0.0 - - _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.") - - _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}" + def _setup_clients(self, webhook_url: str) -> None: + self.response_moderation_endpoint = f"{webhook_url}{_WEBHOOK_PATH_RESPONSE_MODERATION}" + self.prompt_moderation_endpoint = f"{webhook_url}{_WEBHOOK_PATH_PROMPT_MODERATION}" + self.logging_endpoint = f"{webhook_url}{_WEBHOOK_PATH_LOGGING_BATCH}" self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) - self.tool_blocking_client = get_async_httpx_client( + self.moderation_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback, params={"timeout": httpx.Timeout(5.0, connect=2.0)}, ) - self._headers: dict[str, str] = {"Content-Type": "application/json"} - if self.key: - self._headers["Authorization"] = f"Bearer {self.key}" - - # 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: asyncio.Task[Any] | None = self._start_periodic_flush_task() - def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None: """Start the periodic flush task only when an event loop is already running.""" try: 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." - ) return None return loop.create_task(self.periodic_flush()) def _ensure_periodic_flush_task(self) -> None: - # Synchronous helper: in asyncio's cooperative model there is no await - # between the check and assignment, so two callers cannot race here. - if self._flush_task is None or self._flush_task.done(): - self._flush_task = self._start_periodic_flush_task() + if self._periodic_flush_task is None or self._periodic_flush_task.done(): + self._periodic_flush_task = self._start_periodic_flush_task() async def aclose(self): - """Close the dedicated HTTP clients used by this logger.""" - # Cancel the periodic flush task before closing the HTTP clients so - # the loop doesn't wake up and try to POST via a closed client. - if self._flush_task is not None and not self._flush_task.done(): - self._flush_task.cancel() - try: - await self._flush_task - except (asyncio.CancelledError, Exception): - pass - self._flush_task = None - await self.tool_blocking_client.close() - await self.async_httpx_client.close() + """Cancel the periodic flush task. + + ``moderation_client`` and ``async_httpx_client`` are shared objects + from LiteLLM's global HTTP-client cache (``get_async_httpx_client`` + uses the same cache key for all instances with equal parameters). + Closing them here would close the shared connection pool for every + other logger instance; let LiteLLM manage their lifecycle instead. + """ + task = getattr(self, "_periodic_flush_task", None) + if task is not None: + task.cancel() # -- Guardrail hook -------------------------------------------------------- @@ -177,67 +225,104 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: - """Validate tool calls against the blocking service (fail-open).""" - if input_type != "response": - return inputs + """Moderate prompts (request) and responses (response); fail-open. - tool_calls = inputs.get("tool_calls") - if not tool_calls: - return inputs + - ``request``: evaluate the prompt via the before_prompt webhook and + block disallowed prompts before the model is called. + - ``response``: evaluate the assistant's response text and tool calls + via the after_completion webhook and block on a policy violation. + litellm's guardrail-translation layer normalizes Anthropic and OpenAI + requests/responses into ``inputs`` before this runs, so a single code + path covers both wire formats. The configured guardrail ``mode`` + selects which surface(s) run. + """ + if input_type == "request": + return await self._guarded( + self._moderate_prompt(inputs, request_data, logging_obj), + inputs, + "Prompt moderation", + ) + if input_type == "response": + return await self._guarded( + self._moderate_response(inputs, request_data, logging_obj), + inputs, + "Response moderation", + ) + return inputs + + @staticmethod + async def _guarded( + coro: Any, + inputs: GenericGuardrailAPIInputs, + label: str, + ) -> GenericGuardrailAPIInputs: + """Await a moderation coroutine fail-open: re-raise an intentional + block, log at critical for malformed service responses, and swallow + any other error returning ``inputs`` unchanged.""" try: - return await self._check_tool_calls(inputs, tool_calls, request_data, logging_obj) + return await coro except ModifyResponseException: raise except _MalformedToolBlockingResponseError as e: - # Distinct from transient errors: the service responded but the - # payload was structurally invalid, which usually indicates a - # misconfigured webhook or a breaking change in its response - # format. Log loudly so operators notice their tool-blocking - # policy is not actually being enforced. + # The service responded but the payload was structurally invalid, + # which usually indicates a misconfigured webhook or a breaking + # change in its response format. Log loudly so operators notice + # their moderation policy is not actually being enforced. verbose_logger.critical( - "Tool blocking service returned a malformed response: %s. " - "Tool calls are NOT being checked -- verify the webhook " - "configuration. Returning original response unchanged.", + "Response moderation service returned a malformed response: %s. " + "Requests are NOT being checked -- verify the webhook " + "configuration. Returning original inputs unchanged.", e, exc_info=True, ) return inputs except Exception as e: verbose_logger.error( - f"Tool blocking hook failed: {e}. Returning original response unchanged.", + f"{label} hook failed: {e}. Returning original inputs unchanged.", exc_info=True, ) return inputs - async def _check_tool_calls( + async def _moderate_response( self, inputs: GenericGuardrailAPIInputs, - tool_calls: Any, request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"], ) -> GenericGuardrailAPIInputs: - """Send tool calls to blocking service, raise if any are blocked.""" - message_tool_calls = self._normalize_tool_calls(tool_calls) + """Send response text + tool calls to the after_completion webhook and + raise if either the response text or any tool call is blocked.""" + tool_calls = inputs.get("tool_calls") + texts = inputs.get("texts") + if not tool_calls and not texts: + return inputs - 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 + message_tool_calls = self._normalize_tool_calls(tool_calls or ()) + sent_content = self._join_texts(texts) + + call_details = getattr(logging_obj, "model_call_details", _EMPTY_MAPPING) if logging_obj else _EMPTY_MAPPING 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" ) - response_data = self._build_tool_call_payload(message_tool_calls, request_id) - req_data = self._extract_request_data(call_details) + # The moderation payload's ``id`` becomes the tool-blocking log's + # correlation key (the S3 filename), so it must match the failure + # (response) log written for the same blocked request. Both use + # ``litellm_call_id`` -- see ``_correlation_id``. + request_id = self._correlation_id(call_details, request_data) - service_response = await self._post_to_tool_blocking_service(response_data, req_data) - blocked_explanation = self._extract_blocked_tools(service_response, message_tool_calls) + response_data = self._build_response_moderation_payload(message_tool_calls, sent_content, request_id) + req_data = self._extract_request_data(call_details, request_data) - if blocked_explanation is not None: + service_response = await self._post_to_response_moderation_endpoint(response_data, req_data) + blocked = self._extract_response_block(service_response, message_tool_calls, sent_content) + + if blocked: model = self._resolve_model(request_data, call_details) + self._stash_block_context(logging_obj, request_data) raise ModifyResponseException( - message=blocked_explanation, + message=blocked.explanation, model=model, request_data=request_data, guardrail_name=self.guardrail_name, @@ -245,43 +330,125 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return inputs - @staticmethod - def _normalize_tool_calls(tool_calls: Any) -> list[ChatCompletionMessageToolCall]: - """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" - result = [] - for tc in tool_calls: - if isinstance(tc, ChatCompletionMessageToolCall): - result.append(tc) - elif isinstance(tc, dict): - func = tc.get("function", {}) - result.append( - ChatCompletionMessageToolCall( - id=tc.get("id", ""), - type=tc.get("type", "function"), - function=Function( - name=func.get("name", ""), - arguments=func.get("arguments", ""), - ), - ) - ) - elif hasattr(tc, "id") and hasattr(tc, "function"): - result.append( - ChatCompletionMessageToolCall( - id=tc.id or "", - type=getattr(tc, "type", None) or "function", - function=tc.function, - ) - ) - else: - raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}") - return result + async def _moderate_prompt( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + logging_obj: Optional["LiteLLMLoggingObj"], + ) -> GenericGuardrailAPIInputs: + """Send the (normalized) prompt to the before_prompt webhook and raise + if the prompt is blocked.""" + messages = inputs.get("structured_messages") + if not messages: + # For non-chat request types (e.g. /v1/completions), litellm + # supplies the prompt as ``texts`` with no structured_messages. + # Synthesise a user-message so the webhook can evaluate the prompt. + texts = inputs.get("texts") + if texts: + joined = "\n".join(t for t in texts if t) + if joined: + messages = [{"role": "user", "content": joined}] + if not messages: + return inputs + + payload = self._build_prompt_moderation_payload(inputs, request_data) + service_response = await self._post_to_prompt_moderation_endpoint(payload) + refusal = self._extract_prompt_refusal(service_response) + if refusal is None: + return inputs + + model = inputs.get("model") or request_data.get("model") or "unknown" + self._stash_block_context(logging_obj, request_data) + raise ModifyResponseException( + message=refusal, + model=model, + request_data=request_data, + guardrail_name=self.guardrail_name, + ) @staticmethod - def _build_tool_call_payload( - tool_calls: list[ChatCompletionMessageToolCall], + def _stash_block_context( + logging_obj: Optional["LiteLLMLoggingObj"], + request_data: dict, + ) -> None: + """Stash signals so the deferred success-event skips this request and + ``async_post_call_failure_hook`` can build the failure payload. + + - Sets a flag on ``logging_obj.model_call_details`` so the deferred + success-event handler short-circuits. + - Stashes a reference to ``logging_obj`` on ``request_data`` under a + custom key. ``ProxyLogging.post_call_failure_hook`` pops only + ``litellm_logging_obj`` before iterating callbacks, so this key + survives. + + When ``logging_obj`` is ``None`` the success-event has no way to + observe the block (the flag has nowhere to live), so we log an error + instead of silently dropping the signal. + """ + if logging_obj is None: + verbose_logger.error( + "Rubrik: moderation block fired with logging_obj=None for " + f"litellm_call_id={request_data.get('litellm_call_id')}; " + "cannot suppress success event or attach failure payload." + ) + request_data["_rubrik_logging_obj"] = None + return + logging_obj.model_call_details["_rubrik_blocked"] = True + request_data["_rubrik_logging_obj"] = logging_obj + + @staticmethod + def _normalize_tool_calls(tool_calls: Any) -> tuple[ChatCompletionMessageToolCall, ...]: + """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" + return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls) + + @staticmethod + def _normalize_tool_call(tc: Any) -> ChatCompletionMessageToolCall: + if isinstance(tc, ChatCompletionMessageToolCall): + return tc + if isinstance(tc, dict): + func = tc.get("function") or _EMPTY_MAPPING + return ChatCompletionMessageToolCall( + id=tc.get("id", ""), + type=tc.get("type", "function"), + function=Function( + name=func.get("name", ""), + arguments=func.get("arguments", ""), + ), + ) + if hasattr(tc, "id") and hasattr(tc, "function"): + return ChatCompletionMessageToolCall( + id=tc.id or "", + type=getattr(tc, "type", None) or "function", + function=tc.function, + ) + raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}: {tc!r}") + + @staticmethod + def _join_texts(texts: Any) -> str: + """Join response text segments into the single content string the + webhook evaluates. Empty when there is no assistant text.""" + if not texts: + return "" + return "\n".join(t for t in texts if t) + + @staticmethod + def _build_response_moderation_payload( + tool_calls: Sequence[ChatCompletionMessageToolCall], + content: str, request_id: str | None, - ) -> dict[str, Any]: - """Build a full OpenAI ChatCompletion-format dict for the blocking service.""" + ) -> Mapping[str, Any]: + """Build an OpenAI ChatCompletion-format dict (assistant text + tool + calls) for the after_completion webhook. + + ``content`` is sent so the webhook can moderate the response text; + ``None`` when the assistant produced no text (tool-call-only response). + """ + message: dict[str, Any] = { + "role": "assistant", + "content": content or None, + } + if tool_calls: + message["tool_calls"] = tuple(tc.model_dump(exclude_none=True) for tc in tool_calls) return { "id": request_id or f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", @@ -290,42 +457,133 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "choices": [ { "index": 0, - "message": { - "role": "assistant", - "content": None, - "tool_calls": [tc.model_dump(exclude_none=True) for tc in tool_calls], - }, - "finish_reason": "tool_calls", + "message": message, + "finish_reason": "tool_calls" if tool_calls else "stop", } ], } @staticmethod - def _extract_request_data(call_details: dict[str, Any]) -> dict[str, Any]: - """Extract original request data from model_call_details.""" - if not call_details: - return {} - litellm_params = call_details.get("litellm_params", {}) or {} + def _flatten_messages_for_moderation(messages: Any) -> tuple[Mapping[str, Any], ...]: + """Collapse each message's content to a plain string for the webhook. + + litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape, + but a turn sent as content-parts (``[{"type": "text", ...}]``) stays a + list. The before_prompt webhook reads ``content`` as a string and drops + non-string content, so we flatten text parts here (images skipped, per + ``convert_content_list_to_str``) -- otherwise block-content prompts + would pass through unmoderated. Builds a new list; never mutates the + shared ``structured_messages``. + """ + return tuple( + { + "role": message.get("role"), + "content": "\n".join(p for p in RubrikLogger._moderation_text_parts(message) if p), + } + for message in messages or () + if isinstance(message, dict) + ) + + @staticmethod + def _moderation_text_parts(message: Mapping[str, Any]) -> tuple[str, ...]: + """Every attacker-controlled text segment of a message: its content plus + the arguments of any tool call or deprecated function call.""" + fc = message.get("function_call") + return ( + # Base text content (flattens Anthropic content-part arrays) + convert_content_list_to_str(message), # pyright: ignore[reportArgumentType] # dict[str,Any] is AllMessageValues at runtime + *( + str((tc.get("function") or _EMPTY_MAPPING).get("arguments") or "") + for tc in message.get("tool_calls") or () + if isinstance(tc, dict) + ), + str((fc.get("arguments") if isinstance(fc, dict) else None) or ""), + ) + + @staticmethod + def _build_prompt_moderation_payload( + inputs: GenericGuardrailAPIInputs, + request_data: Mapping[str, Any], + ) -> Mapping[str, Any]: + """Build the bare OpenAI request the before_prompt webhook consumes. + + Unlike the after_completion envelope, this endpoint takes a raw OpenAI + chat-completions request. ``structured_messages`` is litellm's + OpenAI-normalized view of the prompt, so this works for Anthropic + ``/v1/messages`` requests too. Optional fields are sent only when + present so the payload stays clean. + """ + payload: dict[str, Any] = { + "model": inputs.get("model") or request_data.get("model") or "", + "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), + } + tools = inputs.get("tools") + if tools is not None: + payload["tools"] = tools + user = request_data.get("user") + if user: + payload["user"] = user + # Fall back to litellm_call_id, the stable cross-provider join key the + # response/tool path uses (see _correlation_id). LiteLLM does not + # populate request_data["correlation_key"]; it carries litellm_call_id. + # The before_prompt webhook skips the *_prompt_moderation.json S3 write + # when correlation_key is empty, so without this the block fires but no + # log is ever written. An explicit correlation_key still wins. + correlation_key = request_data.get("correlation_key") or request_data.get("litellm_call_id") + if correlation_key: + payload["correlation_key"] = correlation_key + return payload + + @staticmethod + def _extract_request_data( + call_details: Mapping[str, Any], + request_data: Mapping[str, Any] | None, + ) -> Mapping[str, Any]: + """Extract original request data from model_call_details for the + response moderation service envelope. + + Includes the agent's declared ``tools`` (OpenAI-format) when available + so the webhook's hallucination evaluator can compare returned tool calls + against the declared tool list. + """ + if not call_details and not request_data: + return _EMPTY_MAPPING + call_details = call_details or _EMPTY_MAPPING + request_data = request_data or _EMPTY_MAPPING + optional_params = call_details.get("optional_params") or _EMPTY_MAPPING + + # Use ``in`` rather than truthy ``or`` so an explicit empty list + # (caller declared the agent has NO tools) is forwarded as-is. + # The response moderation service uses that signal to flag tool-call + # hallucinations -- ``or`` would mask it by falling through to + # optional_params. + if "tools" in request_data: + tools = request_data["tools"] + else: + tools = optional_params.get("tools") + + # The response moderation service consumes only messages/model/tools. + # Don't forward proxy_server_request -- in litellm >=1.83 its ``body`` + # snapshot carries a UserAPIKeyAuth instance that breaks json.dumps, + # silently fail-opening the guardrail. return { "messages": call_details.get("messages"), "model": call_details.get("model"), - "proxy_server_request": RubrikLogger._sanitize_proxy_server_request( - litellm_params.get("proxy_server_request") - ), + "tools": tools, } @staticmethod def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any: """Allowlist only routing fields (``url``, ``method``) when forwarding - ``proxy_server_request`` to the external Rubrik webhook, dropping - inbound ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw + ``proxy_server_request`` to an external webhook, dropping inbound + ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw 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} @staticmethod - def _resolve_model(request_data: dict[str, Any], call_details: dict[str, Any]) -> str: + def _resolve_model(request_data: Mapping[str, Any], call_details: Mapping[str, Any]) -> str: """Get the model name for the ModifyResponseException.""" response = request_data.get("response") if response and hasattr(response, "model"): @@ -334,8 +592,70 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Logging hooks --------------------------------------------------------- - async def _prepare_log_payload(self, kwargs: dict, event_type: str) -> StandardLoggingPayload | None: - """Shared logic for success and failure logging.""" + @staticmethod + def _correlation_id(call_details: Mapping[str, Any], request_data: Mapping[str, Any] | None = None) -> str | None: + """The id that joins a blocked request's two S3 logs by filename: the + moderation (``_blocking``) log and the failure (response) log. + + Always ``litellm_call_id``. It is assigned at request start and is + present identically in both the guardrail path (``model_call_details`` + / ``request_data``) and the failure-hook path. Unlike ``response.id`` + or ``standard_logging_object["id"]`` it is immune to the race where a + block fires before the response/logging object is populated, so the + two logs correlate for every provider (OpenAI and Anthropic alike). + """ + return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id") + + @classmethod + def _apply_correlation_id(cls, payload: dict[str, Any], source: Mapping[str, Any]) -> None: + """Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log + shares its S3 filename id with the moderation (``_blocking``) and + failure logs for the same request -- for every provider. + + ``standard_logging_object["id"]`` is the provider response id + (``response_obj.get("id", litellm_call_id)``), a ``chatcmpl-*`` value + for OpenAI, which would not correlate. ``litellm_call_id`` is assigned + at request start and is identical across all log paths. Falls back to + the existing id when ``litellm_call_id`` is somehow absent rather than + writing a null filename key. + + ``source`` may be ``model_call_details`` directly or a ``kwargs`` dict + that aliases it -- same shape either way. + """ + correlated = cls._correlation_id(source) + if correlated: + payload["id"] = correlated + + @staticmethod + def _prepend_system_prompt(payload: dict[str, Any], source: Mapping[str, Any]) -> None: + """Prepend ``source["system"]`` onto ``payload["messages"]``. + + Builds a NEW messages list rather than mutating ``payload["messages"]`` + in place. The fallback branch of ``_prepare_block_failure_payload`` + aliases ``call_details["messages"]`` directly, so an in-place + ``list.insert(0, ...)`` would mutate the shared source dict. + + No-op if no system prompt is present. Tolerates list/dict/str + message shapes; on unexpected shape, leaves payload alone. + """ + system_prompt = source.get("system") + if not system_prompt: + return + try: + system_scaffold = {"role": "system", "content": system_prompt} + messages = payload.get("messages") + if isinstance(messages, list): + payload["messages"] = (system_scaffold, *messages) + elif isinstance(messages, (dict, str)): + payload["messages"] = (system_scaffold, messages) + except Exception as e: + verbose_logger.warning( + f"Rubrik: failed to prepend system prompt: {e}", + exc_info=True, + ) + + async def _prepare_log_payload(self, kwargs: Mapping[str, Any], event_type: str) -> StandardLoggingPayload | None: + """Shared logic for success logging (sampled).""" if random.random() > self.sampling_rate: verbose_logger.debug(f"Skipping Rubrik {event_type} logging (sampling_rate={self.sampling_rate})") return None @@ -343,59 +663,17 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # Deep-copy so mutations don't affect other callbacks sharing this 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 - # differs from the original Anthropic msg-* id on the response dict. - # Normalize to litellm_call_id so that the logging and tool-blocking - # endpoints see the same request identifier. - litellm_params = kwargs.get("litellm_params", {}) or {} - proxy_request = litellm_params.get("proxy_server_request", {}) or {} - url_path = urllib.parse.urlparse(proxy_request.get("url", "")).path - if url_path.endswith(_ENDPOINT_ANTHROPIC_MESSAGES): - _litellm_call_id = kwargs.get("litellm_call_id") - if _litellm_call_id: - standard_logging_payload["id"] = _litellm_call_id # type: ignore[literal-required] - - if "system" in kwargs: - system_prompt_msg_list = kwargs["system"] - try: - if system_prompt_msg_list: - system_scaffold = { - "role": "system", - "content": system_prompt_msg_list, - } - if isinstance(standard_logging_payload["messages"], list): - standard_logging_payload["messages"].insert(0, system_scaffold) - elif isinstance(standard_logging_payload["messages"], (dict, str)): - standard_logging_payload["messages"] = [ - system_scaffold, - standard_logging_payload["messages"], - ] - except Exception as e: - verbose_logger.warning( - f"Rubrik: failed to prepend system prompt: {e}", - exc_info=True, - ) + self._apply_correlation_id(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime + self._prepend_system_prompt(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime return standard_logging_payload - async def _enqueue_log_event(self, kwargs: dict, event_type: str): - try: - self._ensure_periodic_flush_task() - payload = await self._prepare_log_payload(kwargs, event_type) - if payload is None: - return - - self.log_queue.append(payload) - self._enforce_max_queue_size() - - if len(self.log_queue) >= self.batch_size: - await self.flush_queue() - except Exception as e: - verbose_logger.error( - f"Rubrik {event_type} logging hook failed: {e}. Skipping logging for this event.", - exc_info=True, - ) + async def _append_and_maybe_flush(self, payload) -> None: + self._ensure_periodic_flush_task() + self.log_queue.append(payload) + self._enforce_max_queue_size() + if len(self.log_queue) >= self.batch_size: + await self.flush_queue() def _enforce_max_queue_size(self) -> None: overflow = len(self.log_queue) - self.max_queue_size @@ -415,18 +693,237 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): self._dropped_since_warning = 0 self._last_drop_warning_time = now + async def _enqueue_log_event(self, kwargs: Mapping[str, Any], event_type: str): + try: + payload = await self._prepare_log_payload(kwargs, event_type) + if payload is None: + return + await self._append_and_maybe_flush(payload) + except Exception as e: + verbose_logger.error( + f"Rubrik {event_type} logging hook failed: {e}. Skipping logging for this event.", + exc_info=True, + ) + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + # Blocked requests are logged via async_post_call_failure_hook; + # skip here to avoid double-logging the pre-block response. + if kwargs.get("_rubrik_blocked"): + verbose_logger.debug( + f"Rubrik: skipping success event for blocked request litellm_call_id={kwargs.get('litellm_call_id')}" + ) + return await self._enqueue_log_event(kwargs, "success") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + # Log regular LLM failures (timeouts, upstream errors, etc.) to Rubrik. + # NOTE: ``ModifyResponseException`` blocks are NOT routed here; they + # bypass ``Logging.async_failure_handler`` entirely and reach + # ``async_post_call_failure_hook`` instead. So there is no risk of + # double-logging a block through this path. await self._enqueue_log_event(kwargs, "failure") + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: "UserAPIKeyAuth", + traceback_str: str | None = None, + ) -> None: + """Log blocked requests signalled via ``ModifyResponseException`` + (prompt blocks, response/tool blocks, streaming blocks). + + Carries the stashed ``_rubrik_logging_obj``. For every other + exception we no-op; LiteLLM's standard failure plumbing handles those. + """ + if not isinstance(original_exception, ModifyResponseException): + return + + # Guard by guardrail_name so that when multiple Rubrik instances are + # registered, only the instance that raised the block handles it. + # The failure hook is called for every registered callback; without + # this check the first instance pops the stash and the originating + # instance finds None and silently skips logging. + if getattr(original_exception, "guardrail_name", None) != self.guardrail_name: + return + + logging_obj = request_data.pop("_rubrik_logging_obj", None) + if logging_obj is None: + # Legitimate when a non-Rubrik guardrail raised the block; + # problematic if Rubrik did and the stash was lost (e.g. + # ``_stash_block_context`` ran with ``logging_obj=None``). Either + # way we cannot build the payload. + verbose_logger.warning( + "Rubrik: block exception without stashed logging_obj. " + f"litellm_call_id={request_data.get('litellm_call_id')}, " + f"model={request_data.get('model')}, " + f"user_id={user_api_key_dict.user_id}, " + f"raising_guardrail=" + f"{getattr(original_exception, 'guardrail_name', None)}" + ) + return + + call_id: str | None = None + await self._build_and_enqueue_block_event(logging_obj, original_exception, call_id, user_api_key_dict) + + async def _build_and_enqueue_block_event( + self, + logging_obj: "LiteLLMLoggingObj", + exception: "ModifyResponseException", + call_id: str | None, + user_api_key_dict: "UserAPIKeyAuth", + ) -> None: + try: + call_details = logging_obj.model_call_details + # Do NOT pop "_rubrik_blocked" here. The deferred success-handler + # task may still be iterating callbacks, and popping mid-iteration + # (between two awaited callback invocations) would cause this + # plugin's success-event callback to read the flag as absent and + # log the pre-block response -- the exact bug this hook exists to + # prevent. The flag dies with model_call_details when the request + # completes; there's nothing to clean up. + call_id = call_details.get("litellm_call_id") + payload = self._prepare_block_failure_payload(logging_obj, exception, user_api_key_dict) + except (AttributeError, ImportError, KeyError, TypeError) as e: + verbose_logger.error( + f"Rubrik: failed to build blocked-tool payload for " + f"litellm_call_id={call_id}: {e}. Event will NOT be logged.", + exc_info=True, + ) + return + + try: + await self._append_and_maybe_flush(payload) + except Exception as e: + verbose_logger.error( + f"Rubrik: failed to enqueue blocked-tool event for litellm_call_id={call_id}: {e}.", + exc_info=True, + ) + + def _prepare_block_failure_payload( + self, + logging_obj: "LiteLLMLoggingObj", + exception: "ModifyResponseException", + user_api_key_dict: "UserAPIKeyAuth", + ) -> StandardLoggingPayload: + """Build a failure-style payload using the exception text as response. + + Blocked-tool events are security-relevant and **bypass sampling**: + every block is logged. + + A non-streaming block always takes the fallback, and not because of a + race: registering a post_call guardrail sets ``_defer_async_logging``, + which parks the success handler that would have written + ``standard_logging_object``, and ``_flush_deferred_async_logging`` + returns early once an exception was raised. Streaming requests never set + that flag, so a streamed block can arrive with the object already + populated; the branch below covers it and wins over the fallback. + + For prompt blocks the LLM is never called, so ``standard_logging_object`` + is never populated. The fallback therefore must carry enough fields to + pass the log processor's ``LogEntry`` schema (``BaseLogEntry`` requires + ``metadata``, ``model_id``, ``model_group``, ``model_parameters``, + ``startTime``, ``endTime``, and ``completionStartTime``). Without a + parseable payload the log processor discards the entry with a parse + error and no session is created, so prompt-moderation violations are + silently dropped even though the ``_prompt_moderation.json`` forensic + log is written correctly. + + Field sourcing for the fallback path: + - ``model`` / ``model_group``: ``call_details["model"]`` -- this is the + model-group name (e.g. "gpt-4o") set by the proxy before the guardrail + fires. The router writes ``metadata["model_group"]`` only inside + ``acompletion()``, which hasn't run yet for a prompt block. + - ``model_id``: not available before the LLM returns hidden_params; + defaults to empty string. + - caller identity: ``_caller_metadata`` off the ``user_api_key_dict`` + the failure hook is handed. The enriched litellm metadata lives under + ``call_details["litellm_params"]["metadata"]``, never at the top + level, so the previous top-level read resolved to an empty string for + every block. + - time fields: ``call_details["start_time"]`` reused for all three; + end/completion times are meaningless for a prompt block. + """ + call_details = logging_obj.model_call_details + exception_text = f"{type(exception).__name__}: {exception.message}" + + base = call_details.get("standard_logging_object") + if base is not None: + payload: dict = safe_deep_copy(base) + else: + verbose_logger.debug( + "Rubrik: standard_logging_object not yet on model_call_details " + f"for litellm_call_id={call_details.get('litellm_call_id')}; " + "using best-effort fallback payload." + ) + payload = self._build_fallback_payload(call_details, user_api_key_dict) + + payload["response"] = exception_text + + # Pin the correlation key to litellm_call_id so this failure log shares + # its S3 filename id with the moderation (``_blocking``) log for the + # same request. The copied ``standard_logging_object["id"]`` is + # ``response_obj.get("id", litellm_call_id)`` -- a provider ``chatcmpl-*`` + # value for OpenAI -- which would not correlate; overwrite it. + payload["id"] = self._correlation_id(call_details) or f"chatcmpl-{uuid.uuid4()}" + self._prepend_system_prompt(payload, call_details) + + return payload # type: ignore[return-value] + + @staticmethod + def _caller_metadata(user_api_key_dict: "UserAPIKeyAuth") -> StandardLoggingUserAPIKeyMetadata: + """Identify the caller whose request was blocked. + + Uses the same mapper the success path and the proxy spend logger use, so + a block log and a success log agree on the caller key set. + + The import is deferred because ``litellm/integrations/`` is SDK-side + while the mapper lives under ``proxy/``: ``rubrik.py`` is imported during + guardrail discovery and must not pull proxy-only dependencies into its + import chain. It is unguarded because the only dispatcher of this hook, + ``ProxyLogging.post_call_failure_hook``, already imports fastapi at + module scope, so there is no path where this hook runs and the mapper is + missing. + """ + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + return LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict) + + @classmethod + def _build_fallback_payload( + cls, + call_details: Mapping[str, Any], + user_api_key_dict: "UserAPIKeyAuth", + ) -> dict[str, Any]: + # Convert datetime to a Unix float so json.dumps can serialize it. + # httpx's json= parameter uses stdlib json.dumps with no custom encoder. + _raw_start = call_details.get("start_time") + _start = _raw_start.timestamp() if _raw_start is not None else None + return { + "id": call_details.get("litellm_call_id"), + "model": call_details.get("model") or "", + # model_group is set by the router inside acompletion(), which + # hasn't run for a prompt block; use the model name instead. + "model_group": call_details.get("model") or "", + # model_id comes from response.hidden_params -- unavailable here. + "model_id": "", + "model_parameters": ModelParamHelper.get_standard_logging_model_parameters( + call_details.get("optional_params") or _EMPTY_MAPPING # pyright: ignore[reportArgumentType] # helper only reads the mapping + ), + "startTime": _start, + "endTime": _start, + "completionStartTime": _start, + "messages": call_details.get("messages") or (), + "metadata": cls._caller_metadata(user_api_key_dict), + "status": "failure", + } + # -- Batch logging --------------------------------------------------------- async def _log_batch_to_rubrik(self, data): - # NOTE: this method intentionally re-raises on failure so the parent - # CustomBatchLogger.flush_queue keeps the unsent events in the queue - # for the next flush attempt instead of silently dropping them. + # NOTE: this method intentionally re-raises on failure so flush_queue + # can preserve the unsent events for the next flush attempt instead of + # silently dropping them. try: response = await self.async_httpx_client.post( url=self.logging_endpoint, @@ -452,10 +949,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): if not self.log_queue: return - log_queue_snapshot = list(self.log_queue) - verbose_logger.debug("Rubrik: Flushing batch of %s events", len(log_queue_snapshot)) await self._log_batch_to_rubrik( - data=log_queue_snapshot, + data=self.log_queue, ) async def flush_queue(self): @@ -463,8 +958,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): Overrides the base implementation so the same snapshot drives both the HTTP send and the queue truncation. This avoids the subtle - coupling where the base class captures `len(self.log_queue)` - separately from the snapshot taken inside `async_send_batch`, + coupling where the base class captures ``len(self.log_queue)`` + separately from the snapshot taken inside ``async_send_batch``, which could otherwise drift in a future refactor and cause duplicate deliveries to Rubrik. """ @@ -485,70 +980,141 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): del self.log_queue[: len(snapshot)] self.last_flush_time = time.time() - # -- Tool blocking service ------------------------------------------------- + # -- Webhook services ------------------------------------------------------ - async def _post_to_tool_blocking_service( + async def _post_json(self, endpoint: str, payload: Mapping[str, Any], service_name: str) -> Mapping[str, Any]: + """POST ``payload`` to a Rubrik webhook and return its dict response. + + Raises: + Exception: If the service is unavailable or returns an error. + TypeError: If the response JSON is not a dict. + """ + verbose_logger.debug(f"Sending request to {service_name}: {endpoint}") + http_response = await self.moderation_client.post( + endpoint, + json=payload, + headers=self._headers, + ) + http_response.raise_for_status() + result = http_response.json() + if not isinstance(result, dict): + raise TypeError( + f"{service_name} returned non-dict JSON " + f"({type(result).__name__}); expected OpenAI chat completion " + "shape or empty object." + ) + return result + + async def _post_to_response_moderation_endpoint( self, - response_data: dict[str, Any], - request_data: dict[str, Any], - ) -> dict[str, Any]: - """Post a payload to the tool blocking service and return the response. + response_data: Mapping[str, Any], + request_data: Mapping[str, Any], + ) -> Mapping[str, Any]: + """Post the ``{request, response}`` envelope to the after_completion + webhook and return its (possibly rewritten) response. Args: response_data: The OpenAI-formatted response payload to send. request_data: Original LLM request data to include alongside the response for additional context. Empty dict if unavailable. - - Raises: - Exception: If the service is unavailable or returns an error. """ - envelope = { - "request": request_data, - "response": response_data, - } - 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, - headers=self._headers, + envelope = {"request": request_data, "response": response_data} + return await self._post_json( + self.response_moderation_endpoint, + envelope, + "Response moderation service", ) - http_response.raise_for_status() - result: dict[str, Any] = http_response.json() - return result + + async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """Post a bare OpenAI request to the before_prompt webhook. + + Returns ``{}`` (passthrough) or a synthetic chat.completion (block). + """ + return await self._post_json(self.prompt_moderation_endpoint, payload, "Prompt moderation service") @staticmethod - def _extract_blocked_tools( - service_response: dict[str, Any], - all_tool_calls: list[ChatCompletionMessageToolCall], - ) -> str | None: - """Return the blocking explanation if any tool calls were blocked. + def _extract_prompt_refusal(service_response: Mapping[str, Any]) -> str | None: + """Return the refusal text when the prompt was blocked, else None. - Compares the service response (which contains only allowed tools) against - the full set of tool calls. Returns ``None`` if all tools are allowed, or - the explanation string (prefixed with newlines) otherwise. + The before_prompt webhook returns ``{}`` (passthrough) or a synthetic + chat.completion whose ``choices[0].message.content`` is the refusal + explanation. + """ + choices = service_response.get("choices") + if not choices: + return None + message = choices[0].get("message") or _EMPTY_MAPPING + content = message.get("content") + return content or "Request blocked by policy." + + @staticmethod + def _extract_response_block( + service_response: Mapping[str, Any], + all_tool_calls: Sequence[ChatCompletionMessageToolCall], + sent_content: str, + ) -> BlockedResponseResult | None: + """Detect whether the webhook moderated the response text or tool calls. + + The after_completion webhook rewrites the response in place with no + explicit "blocked" flag, so we infer a block by diffing what we sent + against what came back: + + - Tool block: a tool call we sent is absent from the returned (allowed) + set. + - Text block: the returned content was REPLACED wholesale (a text + violation), as opposed to having a tool-block explanation APPENDED to + the original content. We tell them apart with ``startswith``, which + mirrors the webhook's own append-vs-replace behavior. + + Returns None when nothing was moderated. A text block supersedes a tool + block (mirroring the webhook, which drops tool calls on a text block). Expects service_response in OpenAI chat completion format: {"choices": [{"message": {"tool_calls": [...], "content": "..."}}]} """ - choices = service_response.get("choices", []) + choices = service_response.get("choices") or () if not choices: - raise _MalformedToolBlockingResponseError("Tool blocking service returned empty response") + raise _MalformedToolBlockingResponseError("Response moderation service returned empty response") - message = choices[0].get("message", {}) - returned_tool_calls = message.get("tool_calls") or [] - blocking_explanation = message.get("content", "") + message = choices[0].get("message") or _EMPTY_MAPPING + returned_tool_calls = message.get("tool_calls") or () + returned_content = message.get("content") or "" - allowed_id_counts: Counter = Counter( - 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() + # Use Counter so duplicate IDs are handled correctly: if the model + # emits two calls with the same ID (one allowed, one prohibited) and + # the service returns only the allowed one, a set-based check would + # miss the block. Counter preserves multiplicity. + returned_id_counts: Counter[str] = Counter(tc["id"] for tc in returned_tool_calls if tc.get("id")) + required_id_counts: Counter[str] = Counter(tc.id for tc in all_tool_calls if tc.id) + # Cardinality check catches ID-less tool calls (not counted in + # required_id_counts because tc.id is falsy); Counter check catches + # duplicate-ID attacks where one occurrence is silently removed. + tools_blocked = len(returned_tool_calls) < len(all_tool_calls) or not all( + returned_id_counts.get(tc_id, 0) >= count for tc_id, count in required_id_counts.items() ) - if all_allowed: - return None + # The webhook either replaces content wholesale (text block) or appends + # a tool-block explanation to the original text. ``appended`` tells the + # two apart, and is reused below to recover just the explanation. A text + # block requires there to have been assistant text to block. + # Use the documented ``\n\n`` separator to distinguish a tool-block + # append from a text replacement that shares the original as a prefix. + # Without the separator, a replacement like "Hello, blocked." where the + # original was "Hello" would be classified as an append (not a text + # block) and silently pass through to the client. + appended = bool(sent_content) and returned_content.startswith(f"{sent_content}\n\n") + text_blocked = bool(sent_content) and returned_content != sent_content and not appended - explanation = blocking_explanation or "Tool call blocked by policy." - return f"\n\n{explanation}" + if text_blocked: + return BlockedResponseResult(explanation=returned_content or "Response blocked by policy.") + + if tools_blocked: + if appended: + # Recover just the appended explanation: drop the original text + # and the leading separator the webhook inserted before it. + explanation = returned_content[len(sent_content) :].lstrip("\n") + else: + explanation = returned_content + return BlockedResponseResult(explanation=explanation or "Tool call blocked by policy.") + + return None diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index 51de43e302c..c35cc88107f 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -78,7 +78,7 @@ class S3Logger: **kwargs, ) except Exception as e: - print_verbose(f"Got exception on init s3 client {e!s}") + print_verbose(f"Got exception on init s3 client {e}") raise e async def _async_log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): @@ -163,12 +163,12 @@ class S3Logger: **sse_params, ) - print_verbose(f"Response from s3:{response!s}") + print_verbose(f"Response from s3:{response}") print_verbose(f"s3 Layer Logging - final response object: {response_obj}") return response except Exception as e: - verbose_logger.exception(f"s3 Layer Error - {e!s}") + verbose_logger.exception(f"s3 Layer Error - {e}") def _validated_sse_value(name: str, value: str | None) -> str | None: diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 8c6cadd5356..44c6e42f9f0 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -125,7 +125,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): BaseAWSLLM.__init__(self) except Exception as e: - print_verbose(f"Got exception on init s3 client {e!s}") + print_verbose(f"Got exception on init s3 client {e}") raise e def _init_s3_params( @@ -284,7 +284,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): self.batch_size, ) except Exception as e: - verbose_logger.exception(f"s3 Layer Error - {e!s}") + verbose_logger.exception(f"s3 Layer Error - {e}") self.handle_callback_failure(callback_name="S3Logger") async def async_upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement): @@ -383,7 +383,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): response.raise_for_status() break except Exception as e: - verbose_logger.exception(f"Error uploading to s3: {e!s}") + verbose_logger.exception(f"Error uploading to s3: {e}") self.handle_callback_failure(callback_name="S3Logger") async def async_send_batch(self): @@ -557,7 +557,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): response.raise_for_status() break except Exception as e: - verbose_logger.exception(f"Error uploading to s3: {e!s}") + verbose_logger.exception(f"Error uploading to s3: {e}") self.handle_callback_failure(callback_name="S3Logger") async def _download_object_from_s3(self, s3_object_key: str) -> dict | None: @@ -642,7 +642,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): return response.json() except Exception as e: - verbose_logger.exception(f"Error downloading from S3: {e!s}") + verbose_logger.exception(f"Error downloading from S3: {e}") return None async def get_proxy_server_request_from_cold_storage_with_object_key( @@ -666,5 +666,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: {e!s}") + verbose_logger.exception(f"Error retrieving object {object_key} from cold storage: {e}") return None diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 18717790207..56618b62368 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -113,7 +113,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): BaseAWSLLM.__init__(self) except Exception as e: - print_verbose(f"Got exception on init sqs client {e!s}") + print_verbose(f"Got exception on init sqs client {e}") raise e def _init_sqs_params( @@ -215,7 +215,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): self.batch_size, ) except Exception as e: - verbose_logger.exception(f"sqs Layer Error - {e!s}") + verbose_logger.exception(f"sqs Layer Error - {e}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: @@ -233,7 +233,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): ) except Exception as e: - verbose_logger.exception(f"Datadog Layer Error - {e!s}\n{traceback.format_exc()}") + verbose_logger.exception(f"Datadog Layer Error - {e}\n{traceback.format_exc()}") async def async_send_batch(self) -> None: verbose_logger.debug(f"sqs logger - sending batch of {len(self.log_queue)}") @@ -305,7 +305,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): ) response.raise_for_status() except Exception as e: - verbose_logger.exception(f"Error sending to SQS: {e!s}") + verbose_logger.exception(f"Error sending to SQS: {e}") async def async_health_check(self) -> IntegrationHealthCheckStatus: """ 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 73c48f72d34..6eac7a27e73 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 @@ -146,7 +146,7 @@ class VectorStorePreCallHook(CustomLogger): return model, modified_messages, non_default_params except Exception as e: - verbose_logger.exception(f"Error in VectorStorePreCallHook: {e!s}") + verbose_logger.exception(f"Error in VectorStorePreCallHook: {e}") # Return original parameters on error return model, messages, non_default_params @@ -275,7 +275,7 @@ class VectorStorePreCallHook(CustomLogger): return response except Exception as e: - verbose_logger.exception(f"Error adding search results to response: {e!s}") + verbose_logger.exception(f"Error adding search results to response: {e}") # Don't fail the request if search results fail to be added return None @@ -322,6 +322,6 @@ class VectorStorePreCallHook(CustomLogger): return response_chunk except Exception as e: - verbose_logger.exception(f"Error adding search results to streaming chunk: {e!s}") + verbose_logger.exception(f"Error adding search results to streaming chunk: {e}") # Don't fail the request if search results fail to be added return response_chunk diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 54278afafc4..718f7b8fcd7 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -224,7 +224,7 @@ class WebSearchInterceptionLogger(CustomLogger): content.append({"type": "text", "text": search_result_text}) response: dict[str, object] = { - "id": f"msg_{uuid.uuid4()!s}", + "id": f"msg_{uuid.uuid4()}", "type": "message", "role": "assistant", "model": model, @@ -1038,8 +1038,8 @@ class WebSearchInterceptionLogger(CustomLogger): @staticmethod def _extract_search_text(result: object) -> str: if isinstance(result, Exception): - verbose_logger.error(f"WebSearchInterception: Responses search failed with error: {result!s}") - return f"Search failed: {result!s}" + verbose_logger.error(f"WebSearchInterception: Responses search failed with error: {result}") + return f"Search failed: {result}" if isinstance(result, tuple) and len(result) == 2: text_value, _ = result return text_value if isinstance(text_value, str) else str(text_value) @@ -1194,8 +1194,8 @@ class WebSearchInterceptionLogger(CustomLogger): structured_results: list[SearchResponse | None] = [] for i, result in enumerate(search_results): if isinstance(result, Exception): - verbose_logger.error(f"WebSearchInterception: Search {i} failed with error: {result!s}") - final_search_results.append(f"Search failed: {result!s}") + verbose_logger.error(f"WebSearchInterception: Search {i} failed with error: {result}") + final_search_results.append(f"Search failed: {result}") structured_results.append(None) elif isinstance(result, tuple) and len(result) == 2: text_value, structured_value = result @@ -1308,7 +1308,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) return search_result_text, result except Exception as e: - verbose_logger.error(f"WebSearchInterception: Search failed for '{query}': {e!s}") + verbose_logger.error(f"WebSearchInterception: Search failed for '{query}': {e}") raise async def _authorize_search_tool( @@ -1486,8 +1486,8 @@ 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: {result!s}") - final_search_results.append(f"Search failed: {result!s}") + verbose_logger.error(f"WebSearchInterception: Search {i} failed with error: {result}") + final_search_results.append(f"Search failed: {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)) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 47b7aa6d568..101cbae23f9 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -679,7 +679,7 @@ def _map_replicate_exception( ) raise APIError( status_code=500, - message=f"ReplicateException - {original_exception!s}", + message=f"ReplicateException - {original_exception}", llm_provider="replicate", model=model, request=httpx.Request( @@ -2459,7 +2459,7 @@ def exception_type( # type: ignore ): # deal with edge-case invalid request error bug in openai-python sdk exception_mapping_worked = True raise BadRequestError( - message=f"{exception_provider} BadRequestError : This can happen due to missing AZURE_API_VERSION: {original_exception!s}", + message=f"{exception_provider} BadRequestError : This can happen due to missing AZURE_API_VERSION: {original_exception}", model=model, llm_provider=custom_llm_provider, response=getattr(original_exception, "response", None), @@ -2478,7 +2478,7 @@ def exception_type( # type: ignore ) else: raise APIConnectionError( - message=f"{original_exception!s}\n{_redact_string(traceback.format_exc())}", + message=f"{original_exception}\n{_redact_string(traceback.format_exc())}", llm_provider=custom_llm_provider, model=model, request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), # stub the request diff --git a/litellm/litellm_core_utils/fallback_utils.py b/litellm/litellm_core_utils/fallback_utils.py index ff4a4c9c74c..4e7ce828a58 100644 --- a/litellm/litellm_core_utils/fallback_utils.py +++ b/litellm/litellm_core_utils/fallback_utils.py @@ -70,7 +70,7 @@ async def async_completion_with_fallbacks(**kwargs): ) except Exception as e: - verbose_logger.exception(f"Fallback attempt failed for model {model}: {e!s}") + verbose_logger.exception(f"Fallback attempt failed for model {model}: {e}") most_recent_exception_str = str(e) continue diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index f869909e751..32e517883b7 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -501,9 +501,9 @@ def get_llm_provider( if isinstance(e, litellm.exceptions.BadRequestError): raise e else: - error_str = f"GetLLMProvider Exception - {e!s}\n\noriginal model: {model}" + error_str = f"GetLLMProvider Exception - {e}\n\noriginal model: {model}" raise litellm.exceptions.BadRequestError( # type: ignore - message=f"GetLLMProvider Exception - {e!s}\n\noriginal model: {model}", + message=f"GetLLMProvider Exception - {e}\n\noriginal model: {model}", model=model, response=None, llm_provider="", diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 0addc7586fe..e87e3d8aca2 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -292,7 +292,7 @@ def get_model_cost_map(url: str) -> dict: str(e), ) _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = f"Remote fetch failed: {e!s}" + _cost_map_source_info.fallback_reason = f"Remote fetch failed: {e}" return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()) # Validate using cached count (cheap int comparison, no file I/O) diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index 11668acb21e..171165d01be 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -1,7 +1,7 @@ -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import Any -from litellm.types.utils import StandardCallbackDynamicParams +from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, StandardCallbackDynamicParams _CLIENT_CALLBACK_METADATA_SLOTS: tuple[str, ...] = ("litellm_metadata", "metadata") @@ -75,14 +75,32 @@ _supported_callback_params = [ "turn_off_message_logging", ] -_request_blocked_callback_params = { - "gcs_bucket_name", - "gcs_path_service_account", - "dd_api_key", - "dd_site", - "dd_agent_host", - "dd_agent_port", -} +_request_blocked_callback_params = frozenset( + { + "gcs_bucket_name", + "gcs_path_service_account", + "dd_api_key", + "dd_site", + "dd_agent_host", + "dd_agent_port", + } +) + + +def get_trusted_callback_params(kwargs: Mapping[str, Any] | None) -> tuple[tuple[str, str], ...]: + """ + Read callback params the proxy itself stamped from admin-configured team/key callback settings. + + Request-body values never reach this field: the proxy strips it from client input before + setting it, so callbacks can consume credentials and destinations here without re-validating. + + Returned as pairs rather than a mapping because the caller keeps this on the Logging object, + which the proxy deep-copies; a mappingproxy is not copyable and a dict would be mutable. + """ + trusted_vars = kwargs.get(TRUSTED_CALLBACK_VARS_FIELD) if kwargs else None + if not isinstance(trusted_vars, Mapping): + return () + return tuple((key, str(value)) for key, value in trusted_vars.items() if isinstance(key, str)) def initialize_standard_callback_dynamic_params( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index db10e18e324..66d82bd18f1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -166,6 +166,9 @@ from ..integrations.s3_v2 import S3Logger as S3V2Logger from ..integrations.supabase import Supabase from ..integrations.traceloop import TraceloopLogger from .exception_mapping_utils import _get_response_headers +from .initialize_dynamic_callback_params import ( + get_trusted_callback_params, +) from .initialize_dynamic_callback_params import ( initialize_standard_callback_dynamic_params as _initialize_standard_callback_dynamic_params, ) @@ -199,7 +202,7 @@ try: EnterpriseStandardLoggingPayloadSetup ) except Exception as e: - verbose_logger.debug(f"[Non-Blocking] Unable to import GenericAPILogger - LiteLLM Enterprise Feature - {e!s}") + verbose_logger.debug(f"[Non-Blocking] Unable to import GenericAPILogger - LiteLLM Enterprise Feature - {e}") GenericAPILogger = CustomLogger # type: ignore ResendEmailLogger = CustomLogger # type: ignore SendGridEmailLogger = CustomLogger # type: ignore @@ -362,6 +365,7 @@ class Logging(LiteLLMLoggingBaseClass): self.standard_callback_dynamic_params: StandardCallbackDynamicParams = ( self.initialize_standard_callback_dynamic_params(kwargs) ) + self._trusted_callback_vars: tuple[tuple[str, str], ...] = get_trusted_callback_params(kwargs) # Process dynamic callbacks (after standard_callback_dynamic_params is initialized, # so team-scoped credentials are available for callback initialization) @@ -459,9 +463,10 @@ class Logging(LiteLLMLoggingBaseClass): # pass only the relevant dynamic params as custom_logger_init_args. _custom_logger_init_args: dict | None = None if callback == "datadog": - _custom_logger_init_args = { - k: v for k, v in self.standard_callback_dynamic_params.items() if k.startswith("dd_") - } + # dd_* params are blocked from standard_callback_dynamic_params + # (request-level security); only the proxy-stamped team/key + # callback vars are admin-configured and trusted. + _custom_logger_init_args = {k: v for k, v in self._trusted_callback_vars if k.startswith("dd_")} callback_class = _init_custom_logger_compatible_class( callback, # type: ignore[arg-type] @@ -968,7 +973,7 @@ class Logging(LiteLLMLoggingBaseClass): error=str(e), ) _metadata["raw_request"] = f"Unable to Log \ - raw request: {e!s}" + raw request: {e}" if getattr(self, "logger_fn", None) and callable(self.logger_fn): try: self.logger_fn( @@ -976,7 +981,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( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {e!s}" + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {e}" ) self.model_call_details["api_call_start_time"] = datetime.datetime.now() @@ -1036,14 +1041,14 @@ class Logging(LiteLLMLoggingBaseClass): callback_func=callback, ) except Exception as e: - verbose_logger.exception(f"litellm.Logging.pre_call(): Exception occured - {e!s}") + verbose_logger.exception(f"litellm.Logging.pre_call(): Exception occured - {e}") verbose_logger.debug( 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(f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {e!s}") + verbose_logger.exception(f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {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) @@ -1159,7 +1164,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( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {e!s}" + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {e}" ) original_response = redact_message_input_output_from_logging( model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), @@ -1196,7 +1201,7 @@ class Logging(LiteLLMLoggingBaseClass): ) except Exception as e: verbose_logger.exception( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while post-call logging with integrations {e!s}" + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while post-call logging with integrations {e}" ) verbose_logger.debug( f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" @@ -1204,7 +1209,7 @@ class Logging(LiteLLMLoggingBaseClass): if capture_exception: # log this error to sentry for debugging capture_exception(e) except Exception as e: - verbose_logger.exception(f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {e!s}") + verbose_logger.exception(f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {e}") async def async_post_mcp_tool_call_hook( self, @@ -1244,7 +1249,7 @@ class Logging(LiteLLMLoggingBaseClass): if response is not None: response_obj = self._parse_post_mcp_call_hook_response(response=response) except Exception as e: - verbose_logger.exception(f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {e!s}") + verbose_logger.exception(f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {e}") return response_obj def _parse_post_mcp_call_hook_response(self, response: MCPPostCallResponseObject | None) -> Any: @@ -1889,7 +1894,7 @@ class Logging(LiteLLMLoggingBaseClass): return start_time, end_time, result except Exception as e: - raise Exception(f"[Non-Blocking] LiteLLM.Success_Call Error: {e!s}") + raise Exception(f"[Non-Blocking] LiteLLM.Success_Call Error: {e}") def _is_recognized_call_type_for_logging( self, @@ -2378,7 +2383,7 @@ class Logging(LiteLLMLoggingBaseClass): pass except Exception as e: verbose_logger.exception( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging {e!s}", + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging {e}", ) async def async_success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs): @@ -2694,7 +2699,7 @@ class Logging(LiteLLMLoggingBaseClass): break # Only increment once except Exception as e: - verbose_logger.debug(f"Error in _handle_callback_failure: {e!s}") + verbose_logger.debug(f"Error in _handle_callback_failure: {e}") def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): if start_time is None: @@ -2931,14 +2936,14 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: print_verbose( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging with integrations {e!s}" + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging with integrations {e}" ) 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( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging {e!s}" + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging {e}" ) async def async_failure_handler(self, exception, traceback_exception, start_time=None, end_time=None): @@ -2995,7 +3000,7 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: verbose_logger.exception( f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure \ - logging {e!s}\nCallback={callback}" + logging {e}\nCallback={callback}" ) # Track callback logging failures in Prometheus self._handle_callback_failure(callback=callback) @@ -5426,7 +5431,7 @@ def get_standard_logging_object_payload( return payload except Exception as e: - verbose_logger.exception(f"Error creating standard logging object - {e!s}") + verbose_logger.exception(f"Error creating standard logging object - {e}") 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 5bc6107dbec..face1d1b49f 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -150,7 +150,7 @@ def _generic_cost_per_character( prompt_cost = prompt_characters * custom_prompt_cost except Exception as e: verbose_logger.exception( - f"litellm.litellm_core_utils.llm_cost_calc.utils.py::cost_per_character(): Exception occured - {e!s}\nDefaulting to None" + f"litellm.litellm_core_utils.llm_cost_calc.utils.py::cost_per_character(): Exception occured - {e}\nDefaulting to None" ) prompt_cost = None @@ -165,7 +165,7 @@ def _generic_cost_per_character( completion_cost = completion_characters * custom_completion_cost except Exception as e: verbose_logger.exception( - f"litellm.litellm_core_utils.llm_cost_calc.utils.py::cost_per_character(): Exception occured - {e!s}\nDefaulting to None" + f"litellm.litellm_core_utils.llm_cost_calc.utils.py::cost_per_character(): Exception occured - {e}\nDefaulting to None" ) completion_cost = None 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 5e332f4c8d6..1982e40448d 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 @@ -53,7 +53,7 @@ def get_api_base(model: str, optional_params: dict | LiteLLM_Params) -> str | No api_key=_optional_params.api_key, ) except Exception as e: - verbose_logger.debug(f"Error occurred in getting api base - {e!s}") + verbose_logger.debug(f"Error occurred in getting api base - {e}") custom_llm_provider = None dynamic_api_base = None diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 32e2abc53b0..9340554b6d9 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -178,7 +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: {e!s}") + verbose_logger.exception(f"Error in _get_parent_otel_span_from_logging_obj: {e}") return None @@ -265,7 +265,7 @@ def _set_duration_in_model_call_details( else: 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`: {e!s}") + verbose_logger.warning(f"Error setting `llm_api_duration_ms`: {e}") def track_llm_api_timing(): @@ -321,7 +321,7 @@ def track_llm_api_timing(): ) ) except Exception as e: - verbose_logger.debug(f"Error in service logging: {e!s}") + verbose_logger.debug(f"Error in service logging: {e}") @functools.wraps(func) def sync_wrapper(*args, **kwargs): @@ -366,7 +366,7 @@ def track_llm_api_timing(): parent_otel_span=parent_otel_span, ) except Exception as e: - verbose_logger.debug(f"Error in service logging: {e!s}") + verbose_logger.debug(f"Error in service logging: {e}") # Check if the function is async or sync if inspect.iscoroutinefunction(func): diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 639c93dfb80..90c9fb05e4c 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1683,7 +1683,7 @@ def parse_tool_call_arguments( if context: error_parts.append(f"({context})") - error_message = " ".join(error_parts) + f". Error: {original_error!s}. Arguments: {arguments}" + error_message = " ".join(error_parts) + f". Error: {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 147280af1b1..8aa4f60b7c5 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -438,9 +438,7 @@ def _render_chat_template(env, chat_template: str, bos_token: str, eos_token: st return rendered_text except Exception as e: - raise Exception( - f"Error rendering template - {e!s}" - ) # don't use verbose_logger.exception, if exception is raised + raise Exception(f"Error rendering template - {e}") # don't use verbose_logger.exception, if exception is raised async def _afetch_and_extract_template( @@ -858,7 +856,7 @@ def convert_to_anthropic_image_obj(openai_image_url: str, format: str | None) -> raise except Exception as e: raise Exception( - f"""Image url not in expected format. Example Expected input - "image_url": "data:image/jpeg;base64,{{base64_image}}". Supported formats - ['image/jpeg', 'image/png', 'image/gif', 'image/webp']. Error: {e!s}""" + f"""Image url not in expected format. Example Expected input - "image_url": "data:image/jpeg;base64,{{base64_image}}". Supported formats - ['image/jpeg', 'image/png', 'image/gif', 'image/webp']. Error: {e}""" ) @@ -1361,7 +1359,7 @@ def convert_to_gemini_tool_call_invoke( ) return _parts_list except Exception as e: - raise Exception(f"Unable to convert openai tool calls={message} to gemini tool calls. Received error={e!s}") + raise Exception(f"Unable to convert openai tool calls={message} to gemini tool calls. Received error={e}") def convert_to_gemini_tool_call_result( @@ -3713,7 +3711,7 @@ def _convert_to_bedrock_tool_call_invoke( _parts_list.append(cache_point_block) return _parts_list except Exception as e: - raise Exception(f"Unable to convert openai tool calls={tool_calls} to bedrock tool calls. Received error={e!s}") + raise Exception(f"Unable to convert openai tool calls={tool_calls} to bedrock tool calls. Received error={e}") def _append_bedrock_tool_result_media_block( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index fb7d06bee93..25155068baa 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -618,7 +618,7 @@ class CustomStreamWrapper: else: return "" except Exception as e: - verbose_logger.exception(f"litellm.CustomStreamWrapper.handle_baseten_chunk(): Exception occured - {e!s}") + verbose_logger.exception(f"litellm.CustomStreamWrapper.handle_baseten_chunk(): Exception occured - {e}") return "" def handle_triton_stream(self, chunk): @@ -1179,7 +1179,7 @@ class CustomStreamWrapper: content=None, tool_calls=[ { - "id": f"call_{uuid.uuid4()!s}", + "id": f"call_{uuid.uuid4()}", "function": { "arguments": args_str, "name": function_call.name, @@ -1204,7 +1204,7 @@ class CustomStreamWrapper: ) except Exception: if chunk.candidates[0].finish_reason.name == "SAFETY": # type: ignore - raise Exception(f"The response was blocked by VertexAI. {chunk!s}") + raise Exception(f"The response was blocked by VertexAI. {chunk}") else: completion_obj["content"] = str(chunk) elif self.custom_llm_provider == "petals": @@ -1430,7 +1430,7 @@ class CustomStreamWrapper: model_response.choices[0].delta = Delta(**_json_delta) except Exception as e: verbose_logger.exception( - f"litellm.CustomStreamWrapper.chunk_creator(): Exception occured - {e!s}" + f"litellm.CustomStreamWrapper.chunk_creator(): Exception occured - {e}" ) model_response.choices[0].delta = Delta() elif self._has_any_special_delta_attributes(delta): @@ -1538,7 +1538,7 @@ class CustomStreamWrapper: except Exception as e: from litellm._logging import verbose_logger - verbose_logger.exception(f"Error in post-call streaming deployment hook: {e!s}") + verbose_logger.exception(f"Error in post-call streaming deployment hook: {e}") return chunk def _add_mcp_list_tools_to_first_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: @@ -1578,7 +1578,7 @@ class CustomStreamWrapper: except Exception as e: from litellm._logging import verbose_logger - verbose_logger.exception(f"Error adding MCP list tools to first chunk: {e!s}") + verbose_logger.exception(f"Error adding MCP list tools to first chunk: {e}") return chunk @@ -1615,7 +1615,7 @@ class CustomStreamWrapper: except Exception as e: from litellm._logging import verbose_logger - verbose_logger.exception(f"Error adding MCP metadata to final chunk: {e!s}") + verbose_logger.exception(f"Error adding MCP metadata to final chunk: {e}") return chunk diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index fbd19b43f3e..ff94965f628 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -104,7 +104,7 @@ def get_modified_max_tokens( return user_max_tokens except Exception as e: verbose_logger.debug( - f"litellm.litellm_core_utils.token_counter.py::get_modified_max_tokens() - Error while checking max token limit: {e!s}\nmodel={model}, base_model={base_model}" + f"litellm.litellm_core_utils.token_counter.py::get_modified_max_tokens() - Error while checking max token limit: {e}\nmodel={model}, base_model={base_model}" ) return user_max_tokens diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index 3de584d1d5f..967fcc354a5 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -279,7 +279,7 @@ class A2AConfig(BaseConfig): except Exception as e: raise A2AError( status_code=raw_response.status_code, - message=f"Failed to parse A2A response: {e!s}", + message=f"Failed to parse A2A response: {e}", headers=dict(raw_response.headers), ) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 40d1dbac187..51b862e79d9 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1875,7 +1875,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): except Exception as e: raise AnthropicError( status_code=400, - message=f"{e!s}\nReceived Messages={messages}", + message=f"{e}\nReceived Messages={messages}", ) # don't use verbose_logger.exception, if exception is raised ## Auto-strip advisor blocks from history if advisor tool is absent. @@ -2454,7 +2454,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): except Exception as e: response_headers = getattr(raw_response, "headers", None) raise AnthropicError( - message=f"Unable to get json response - {e!s}, Original Response: {raw_response.text}", + message=f"Unable to get json response - {e}, Original Response: {raw_response.text}", status_code=raw_response.status_code, headers=response_headers, ) diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index b1584b98456..0c3d0e931a2 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -109,14 +109,14 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): raise except httpx.HTTPStatusError as e: # HTTP errors - preserve the actual status code - verbose_logger.error(f"HTTP error in CountTokens handler: {e!s}") + verbose_logger.error(f"HTTP error in CountTokens handler: {e}") raise AnthropicError( status_code=e.response.status_code, message=e.response.text, ) except Exception as e: - verbose_logger.error(f"Error in CountTokens handler: {e!s}") + verbose_logger.error(f"Error in CountTokens handler: {e}") raise AnthropicError( status_code=500, - message=f"CountTokens processing error: {e!s}", + message=f"CountTokens processing error: {e}", ) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 373460c151d..a04bb29d7a5 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -684,7 +684,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except json.JSONDecodeError as json_error: raise AzureOpenAIError( status_code=raw_response.status_code or 500, - message=f"Failed to parse raw Azure embedding response: {json_error!s}", + message=f"Failed to parse raw Azure embedding response: {json_error}", ) from json_error if isinstance(response, str): raise AzureOpenAIError( diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index dcbd3985dfd..8db422e00ff 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -333,7 +333,7 @@ def get_azure_ad_token( verbose_logger.debug("Azure AD Token Provider could not be used.") except Exception as e: verbose_logger.error( - f"Error calling Azure AD token provider: {e!s}. Follow docs - https://docs.litellm.ai/docs/providers/azure/#azure-ad-token-refresh---defaultazurecredential" + f"Error calling Azure AD token provider: {e}. Follow docs - https://docs.litellm.ai/docs/providers/azure/#azure-ad-token-refresh---defaultazurecredential" ) raise e @@ -359,8 +359,8 @@ def get_azure_ad_token( # Re-raise TypeError directly raise except Exception as e: - verbose_logger.error(f"Error calling Azure AD token provider: {e!s}") - raise RuntimeError(f"Failed to get Azure AD token: {e!s}") from e + verbose_logger.error(f"Error calling Azure AD token provider: {e}") + raise RuntimeError(f"Failed to get Azure AD token: {e}") from e return azure_ad_token @@ -393,7 +393,7 @@ class BaseAzureLLM(BaseOpenAILLM): 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: {e!s}") + verbose_logger.debug(f"DefaultAzureCredential failed: {e}") return None def get_azure_openai_client( @@ -508,6 +508,8 @@ class BaseAzureLLM(BaseOpenAILLM): openai_client=openai_client, client_initialization_params=client_initialization_params, client_type="azure", + litellm_owned_client=client is None + and self.owns_wrapped_http_client(azure_client_params.get("http_client")), ) return openai_client @@ -580,7 +582,7 @@ class BaseAzureLLM(BaseOpenAILLM): # only show first 5 chars of api_key _api_key = _api_key[:8] + "*" * 15 verbose_logger.debug( - f"Initializing Azure OpenAI Client for {model_name}, Api Base: {api_base!s}, Api Key:{_api_key}" + f"Initializing Azure OpenAI Client for {model_name}, Api Base: {api_base}, Api Key:{_api_key}" ) azure_client_params = { "api_key": api_key, diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index b12f2203e51..7023dbca0b8 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -193,7 +193,7 @@ class AzureAIAgentsHandler: ), ) except Exception as e: - verbose_logger.warning(f"Failed to calculate token usage: {e!s}") + verbose_logger.warning(f"Failed to calculate token usage: {e}") return model_response diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 3ac04729267..65d8c0182ee 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -114,14 +114,14 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): raise except httpx.HTTPStatusError as e: # HTTP errors - preserve the actual status code - verbose_logger.error(f"HTTP error in CountTokens handler: {e!s}") + verbose_logger.error(f"HTTP error in CountTokens handler: {e}") raise AnthropicError( status_code=e.response.status_code, message=e.response.text, ) except Exception as e: - verbose_logger.error(f"Error in CountTokens handler: {e!s}") + verbose_logger.error(f"Error in CountTokens handler: {e}") raise AnthropicError( status_code=500, - message=f"CountTokens processing error: {e!s}", + message=f"CountTokens processing error: {e}", ) diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index 88a38fc1ec7..943232dc348 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -132,7 +132,7 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): ) query_vector = embedding_response.data[0]["embedding"] except Exception as e: - raise Exception(f"Failed to generate embedding for query: {e!s}") + raise Exception(f"Failed to generate embedding for query: {e}") # Azure AI Search endpoint for search index_name = vector_store_id # vector_store_id is the index name 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 7c76003de3a..33255657287 100644 --- a/litellm/llms/base_llm/files/azure_blob_storage_backend.py +++ b/litellm/llms/base_llm/files/azure_blob_storage_backend.py @@ -133,7 +133,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): return storage_url except Exception as e: - verbose_logger.exception(f"Error uploading file to Azure Blob Storage: {e!s}") + verbose_logger.exception(f"Error uploading file to Azure Blob Storage: {e}") raise async def _upload_file_with_account_key(self, file_content: bytes, full_path: str) -> str: @@ -247,7 +247,7 @@ 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: {e!s}") + verbose_logger.exception(f"Error downloading file from Azure Blob Storage: {e}") raise async def _download_file_with_account_key(self, file_path: str) -> bytes: diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 40b12e17e8a..d6626562393 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -186,7 +186,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return session_id # Generate a session ID with 33+ characters - generated_id = f"litellm-session-{uuid.uuid4()!s}" + generated_id = f"litellm-session-{uuid.uuid4()}" verbose_logger.debug(f"Generated new session ID: {generated_id}") return generated_id @@ -370,7 +370,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): total_tokens=total_tokens, ) except Exception as e: - verbose_logger.warning(f"Failed to calculate token usage: {e!s}") + verbose_logger.warning(f"Failed to calculate token usage: {e}") return None def _parse_json_response(self, response_json: dict) -> AgentCoreParsedResponse: @@ -1023,9 +1023,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return model_response except Exception as e: - verbose_logger.error(f"Error processing Bedrock AgentCore response: {e!s}") + verbose_logger.error(f"Error processing Bedrock AgentCore response: {e}") raise BedrockError( - message=f"Error processing response: {e!s}", + message=f"Error processing response: {e}", status_code=raw_response.status_code, ) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 5bd498a465e..2b34c9f2654 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -2073,7 +2073,7 @@ class AmazonConverseConfig(BaseConfig): completion_response = ConverseResponseBlock(**response.json()) # type: ignore except Exception as e: raise BedrockError( - message=f"Error converting to valid response block={e!s}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues", + message=f"Error converting to valid response block={e}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues", status_code=422, ) diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index d877ca81244..da6224ec487 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -464,9 +464,9 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): ) except Exception as e: - verbose_logger.error(f"Error processing Bedrock Invoke Agent response: {e!s}") + verbose_logger.error(f"Error processing Bedrock Invoke Agent response: {e}") raise BedrockError( - message=f"Error processing response: {e!s}", + message=f"Error processing response: {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 d069929df92..4a429b639d2 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -590,7 +590,7 @@ class AWSEventStreamDecoder: return response except Exception as e: - raise Exception(f"Received streaming error - {e!s}") + raise Exception(f"Received streaming error - {e}") def _chunk_parser(self, chunk_data: dict) -> Union[GChunk, ModelResponseStream, dict]: text = "" 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 b96756f1e4e..8de9c3de3f2 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 @@ -208,7 +208,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): completion_response = raw_response.json() except Exception as e: raise BedrockError( - message=f"Error parsing response: {raw_response.text}, error: {e!s}", + message=f"Error parsing response: {raw_response.text}, error: {e}", status_code=raw_response.status_code, ) @@ -237,7 +237,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): raise Exception("Unable to set message content") except Exception as e: raise BedrockError( - message=f"Error setting response content: {e!s}. Response: {completion_response}", + message=f"Error setting response content: {e}. Response: {completion_response}", status_code=raw_response.status_code, ) 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 0c6436030af..a54bf8d6b2b 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -356,7 +356,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): outputText = completion_response.get("results")[0].get("outputText") except Exception as e: raise BedrockError( - message=f"Error processing={raw_response.text}, Received error={e!s}", + message=f"Error processing={raw_response.text}, Received error={e}", status_code=422, ) @@ -379,7 +379,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): raise Exception() except Exception as e: raise BedrockError( - message=f"Error parsing received text={outputText}.\nError-{e!s}", + message=f"Error parsing received text={outputText}.\nError-{e}", status_code=raw_response.status_code, ) diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 8e993c6f8b2..44cc535385d 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -120,14 +120,14 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): raise except httpx.HTTPStatusError as e: # HTTP errors - preserve the actual status code - verbose_logger.error(f"HTTP error in CountTokens handler: {e!s}") + verbose_logger.error(f"HTTP error in CountTokens handler: {e}") raise BedrockError( status_code=e.response.status_code, message=e.response.text, ) except Exception as e: - verbose_logger.error(f"Error in CountTokens handler: {e!s}") + verbose_logger.error(f"Error in CountTokens handler: {e}") raise BedrockError( status_code=500, - message=f"CountTokens processing error: {e!s}", + message=f"CountTokens processing error: {e}", ) diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index 12ebc52dff3..8f590bd917c 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -130,7 +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: {e!s}") + raise ValueError(f"Failed to download file from S3: {s3_uri}. Error: {e}") # Create mock HTTP response mock_response = httpx.Response( diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 3656088cb9d..d3e61829681 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -652,7 +652,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) except Exception as e: verbose_logger.exception( - f"litellm.llms.bedrock.files.transformation.py::_transform_openai_jsonl_content_to_bedrock_jsonl_content() - Error inferring custom_llm_provider - {e!s}" + f"litellm.llms.bedrock.files.transformation.py::_transform_openai_jsonl_content_to_bedrock_jsonl_content() - Error inferring custom_llm_provider - {e}" ) # Determine provider from model name diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 17007f48fb0..a8969894dda 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -175,7 +175,7 @@ class BedrockRealtime(BaseAWSLLM): except Exception as e: verbose_proxy_logger.exception(f"Error in BedrockRealtime.async_realtime: {e}") try: - await websocket.close(code=1011, reason=_redact_string(f"Internal error: {e!s}")) + await websocket.close(code=1011, reason=_redact_string(f"Internal error: {e}")) except Exception: pass raise diff --git a/litellm/llms/black_forest_labs/image_edit/handler.py b/litellm/llms/black_forest_labs/image_edit/handler.py index 62aaa6da77a..cf0cc31283b 100644 --- a/litellm/llms/black_forest_labs/image_edit/handler.py +++ b/litellm/llms/black_forest_labs/image_edit/handler.py @@ -159,7 +159,7 @@ class BlackForestLabsImageEdit: except Exception as e: raise BlackForestLabsError( status_code=500, - message=f"Request failed: {e!s}", + message=f"Request failed: {e}", ) # Poll for result @@ -262,7 +262,7 @@ class BlackForestLabsImageEdit: except Exception as e: raise BlackForestLabsError( status_code=500, - message=f"Request failed: {e!s}", + message=f"Request failed: {e}", ) # Poll for result diff --git a/litellm/llms/black_forest_labs/image_generation/handler.py b/litellm/llms/black_forest_labs/image_generation/handler.py index af321fad580..054d28003f1 100644 --- a/litellm/llms/black_forest_labs/image_generation/handler.py +++ b/litellm/llms/black_forest_labs/image_generation/handler.py @@ -156,7 +156,7 @@ class BlackForestLabsImageGeneration: except Exception as e: raise BlackForestLabsError( status_code=500, - message=f"Request failed: {e!s}", + message=f"Request failed: {e}", ) # Poll for result @@ -262,7 +262,7 @@ class BlackForestLabsImageGeneration: except Exception as e: raise BlackForestLabsError( status_code=500, - message=f"Request failed: {e!s}", + message=f"Request failed: {e}", ) # Poll for result diff --git a/litellm/llms/clarifai/chat/transformation.py b/litellm/llms/clarifai/chat/transformation.py index 147c7986f2a..bf922893f13 100644 --- a/litellm/llms/clarifai/chat/transformation.py +++ b/litellm/llms/clarifai/chat/transformation.py @@ -106,7 +106,7 @@ class ClarifaiConfig(OpenAIGPTConfig): except Exception as e: raise OpenAIError( status_code=raw_response.status_code, - message=f"Failed to parse Clarifai response: {e!s}", + message=f"Failed to parse Clarifai response: {e}", headers=raw_response.headers, ) from e diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index 1261604e6a7..eb4b8acd71f 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -356,7 +356,7 @@ class CodestralTextCompletion: ) except Exception as e: raise TextCompletionCodestralError( - status_code=500, message=f"{e!s}" + status_code=500, message=f"{e}" ) # don't use verbose_logger.exception, if exception is raised return self.process_text_completion_response( model=model, diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 046840e6fd0..3e34b483002 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1411,6 +1411,7 @@ def get_async_httpx_client( key=_cache_key_name, value=_new_client, ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, + litellm_owned_client=True, ) return _new_client @@ -1456,5 +1457,6 @@ def _get_httpx_client(params: dict | None = None) -> HTTPHandler: key=_cache_key_name, value=_new_client, ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, + litellm_owned_client=True, ) return _new_client diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a203f0d6c8c..f7bf174f9ac 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5659,7 +5659,7 @@ class BaseLLMHTTPHandler: fingerprint=fingerprint, ) except Exception as e: - verbose_logger.exception(f"LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: {e!s}") + verbose_logger.exception(f"LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: {e}") # Check if we need to convert response to fake stream for chat completions # This happens when: @@ -5906,7 +5906,7 @@ 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: {e!s}")) + await websocket.close(code=1011, reason=_redact_string(f"Internal server error: {e}")) except RuntimeError as 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 @@ -6303,7 +6303,7 @@ 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: {e!s}")) + await websocket.close(code=1011, reason=_redact_string(f"Internal server error: {e}")) except RuntimeError as close_error: if "already completed" in str(close_error) or "websocket.close" in str(close_error): pass diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index 55722ce35d1..870c96edb65 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -130,7 +130,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): except Exception as e: raise DashScopeError( status_code=raw_response.status_code, - message=f"Failed to parse DashScope response as JSON: {e!s}", + message=f"Failed to parse DashScope response as JSON: {e}", ) logging_obj.post_call( diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 9f6f669a264..b7c9dc23762 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -630,7 +630,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): except Exception as e: response_headers = getattr(raw_response, "headers", None) raise DatabricksException( - message=f"Unable to get json response - {e!s}, Original Response: {raw_response.text}", + message=f"Unable to get json response - {e}, Original Response: {raw_response.text}", status_code=raw_response.status_code, headers=response_headers, ) diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index 2fb7cacb9bf..62e2245db99 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -245,7 +245,7 @@ class DatabricksBase: except requests.RequestException as e: raise DatabricksException( status_code=500, - message=f"OAuth M2M token request failed: {e!s}", + message=f"OAuth M2M token request failed: {e}", ) if response.status_code != 200: diff --git a/litellm/llms/deepgram/audio_transcription/transformation.py b/litellm/llms/deepgram/audio_transcription/transformation.py index 034c41c79fb..4c21f6eb3c7 100644 --- a/litellm/llms/deepgram/audio_transcription/transformation.py +++ b/litellm/llms/deepgram/audio_transcription/transformation.py @@ -122,7 +122,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): return response except Exception as e: - raise ValueError(f"Error transforming Deepgram response: {e!s}\nResponse: {raw_response.text}") + raise ValueError(f"Error transforming Deepgram response: {e}\nResponse: {raw_response.text}") def _reconstruct_diarized_transcript(self, words: list) -> str: """ diff --git a/litellm/llms/elevenlabs/audio_transcription/transformation.py b/litellm/llms/elevenlabs/audio_transcription/transformation.py index a33e221dafd..3672d080b22 100644 --- a/litellm/llms/elevenlabs/audio_transcription/transformation.py +++ b/litellm/llms/elevenlabs/audio_transcription/transformation.py @@ -144,7 +144,7 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): return response except Exception as e: - raise ValueError(f"Error transforming ElevenLabs response: {e!s}\nResponse: {raw_response.text}") + raise ValueError(f"Error transforming ElevenLabs response: {e}\nResponse: {raw_response.text}") def get_complete_url( self, diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 9fcb81e00e3..64ef731a0e2 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -542,7 +542,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): except Exception as e: response_headers = getattr(raw_response, "headers", None) raise FireworksAIException( - message=f"Unable to get json response - {e!s}, Original Response: {raw_response.text}", + message=f"Unable to get json response - {e}, Original Response: {raw_response.text}", status_code=raw_response.status_code, headers=response_headers, ) diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index 7979eeeba42..c8a79878b56 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -178,7 +178,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): raw_response_json = raw_response.json() except Exception as e: raise self.get_error_class( - error_message=f"Failed to parse response: {e!s}", + error_message=f"Failed to parse response: {e}", status_code=raw_response.status_code, headers=raw_response.headers, ) diff --git a/litellm/llms/gdc/chat/transformation.py b/litellm/llms/gdc/chat/transformation.py index 0416d246ea1..69056075a9d 100644 --- a/litellm/llms/gdc/chat/transformation.py +++ b/litellm/llms/gdc/chat/transformation.py @@ -220,7 +220,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): AttributeError, ) as e: raise litellm.utils.AuthenticationError( - message=f"Failed to load service account credentials from api_key: {e!s}", + message=f"Failed to load service account credentials from api_key: {e}", llm_provider="gdc", model=model, ) from e diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index ed82a37e47b..25f767a348e 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -155,8 +155,8 @@ class GoogleAIStudioTokenCounter: status_code=e.response.status_code, ) from e except httpx.RequestError as e: - error_msg = f"Request to Google Gen AI Studio failed: {e!s}" + error_msg = f"Request to Google Gen AI Studio failed: {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: {e!s}" + error_msg = f"Unexpected error during token counting: {e}" raise Exception(error_msg) from e diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index f91737ae613..89ac56979bb 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -190,8 +190,8 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): status_details=None, ) except Exception as e: - verbose_logger.exception(f"Error parsing file upload response: {e!s}") - raise ValueError(f"Error parsing file upload response: {e!s}") + verbose_logger.exception(f"Error parsing file upload response: {e}") + raise ValueError(f"Error parsing file upload response: {e}") def transform_retrieve_file_request( self, @@ -294,8 +294,8 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): 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: {e!s}") - raise ValueError(f"Error parsing file retrieve response: {e!s}") + verbose_logger.exception(f"Error parsing file retrieve response: {e}") + raise ValueError(f"Error parsing file retrieve response: {e}") def transform_delete_file_request( self, @@ -362,8 +362,8 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): else: raise ValueError(f"Failed to delete file: {raw_response.text}") except Exception as e: - verbose_logger.exception(f"Error parsing file delete response: {e!s}") - raise ValueError(f"Error parsing file delete response: {e!s}") + verbose_logger.exception(f"Error parsing file delete response: {e}") + raise ValueError(f"Error parsing file delete response: {e}") def transform_list_files_request( self, diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index 051c0c544f5..9a823011289 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -256,7 +256,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): except Exception as e: raise self.get_error_class( - error_message=f"Failed to parse Gemini response: {e!s}", + error_message=f"Failed to parse Gemini response: {e}", status_code=response.status_code, headers=response.headers, ) @@ -327,7 +327,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): except Exception as e: raise self.get_error_class( - error_message=f"Failed to parse Gemini create response: {e!s}", + error_message=f"Failed to parse Gemini create response: {e}", status_code=response.status_code, headers=response.headers, ) diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index f5bced63869..356d438c6b2 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -177,7 +177,7 @@ def _request_token_sync( except httpx.RequestError as e: raise GigaChatAuthError( status_code=500, - message=f"GigaChat authentication request failed: {e!s}", + message=f"GigaChat authentication request failed: {e}", ) @@ -212,7 +212,7 @@ async def _request_token_async( except httpx.RequestError as e: raise GigaChatAuthError( status_code=500, - message=f"GigaChat authentication request failed: {e!s}", + message=f"GigaChat authentication request failed: {e}", ) diff --git a/litellm/llms/github_copilot/authenticator.py b/litellm/llms/github_copilot/authenticator.py index 2cb099edfb4..180c2215212 100644 --- a/litellm/llms/github_copilot/authenticator.py +++ b/litellm/llms/github_copilot/authenticator.py @@ -68,7 +68,7 @@ class Authenticator: verbose_logger.error("Error saving access token to file") return access_token except (GetDeviceCodeError, GetAccessTokenError, RefreshAPIKeyError) as e: - verbose_logger.warning(f"Failed attempt {attempt + 1}: {e!s}") + verbose_logger.warning(f"Failed attempt {attempt + 1}: {e}") continue raise GetAccessTokenError( @@ -100,7 +100,7 @@ class Authenticator: except OSError: verbose_logger.warning("No API key file found or error opening file") except (json.JSONDecodeError, KeyError) as e: - verbose_logger.warning(f"Error reading API key from file: {e!s}") + verbose_logger.warning(f"Error reading API key from file: {e}") except APIKeyExpiredError: pass # Already logged in the try block @@ -117,14 +117,14 @@ class Authenticator: status_code=401, ) except OSError as e: - verbose_logger.error(f"Error saving API key to file: {e!s}") + verbose_logger.error(f"Error saving API key to file: {e}") raise GetAPIKeyError( - message=f"Failed to save API key: {e!s}", + message=f"Failed to save API key: {e}", status_code=500, ) except RefreshAPIKeyError as e: raise GetAPIKeyError( - message=f"Failed to refresh API key: {e!s}", + message=f"Failed to refresh API key: {e}", status_code=401, ) @@ -142,7 +142,7 @@ class Authenticator: api_endpoint = endpoints.get("api") return api_endpoint except (OSError, json.JSONDecodeError, KeyError) as e: - verbose_logger.warning(f"Error reading API endpoint from file: {e!s}") + verbose_logger.warning(f"Error reading API endpoint from file: {e}") return None def _refresh_api_key(self) -> dict[str, Any]: @@ -173,9 +173,9 @@ class Authenticator: else: 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}): {e!s}") + verbose_logger.error(f"HTTP error refreshing API key (attempt {attempt + 1}/{max_retries}): {e}") except Exception as e: - verbose_logger.error(f"Unexpected error refreshing API key: {e!s}") + verbose_logger.error(f"Unexpected error refreshing API key: {e}") raise RefreshAPIKeyError( message="Failed to refresh API key after maximum retries", @@ -245,21 +245,21 @@ class Authenticator: return resp_json except httpx.HTTPStatusError as e: - verbose_logger.error(f"HTTP error getting device code: {e!s}") + verbose_logger.error(f"HTTP error getting device code: {e}") raise GetDeviceCodeError( - message=f"Failed to get device code: {e!s}", + message=f"Failed to get device code: {e}", status_code=400, ) except json.JSONDecodeError as e: - verbose_logger.error(f"Error decoding JSON response: {e!s}") + verbose_logger.error(f"Error decoding JSON response: {e}") raise GetDeviceCodeError( - message=f"Failed to decode device code response: {e!s}", + message=f"Failed to decode device code response: {e}", status_code=400, ) except Exception as e: - verbose_logger.error(f"Unexpected error getting device code: {e!s}") + verbose_logger.error(f"Unexpected error getting device code: {e}") raise GetDeviceCodeError( - message=f"Failed to get device code: {e!s}", + message=f"Failed to get device code: {e}", status_code=400, ) @@ -304,21 +304,21 @@ class Authenticator: else: verbose_logger.warning(f"Unexpected response: {resp_json}") except httpx.HTTPStatusError as e: - verbose_logger.error(f"HTTP error polling for access token: {e!s}") + verbose_logger.error(f"HTTP error polling for access token: {e}") raise GetAccessTokenError( - message=f"Failed to get access token: {e!s}", + message=f"Failed to get access token: {e}", status_code=400, ) except json.JSONDecodeError as e: - verbose_logger.error(f"Error decoding JSON response: {e!s}") + verbose_logger.error(f"Error decoding JSON response: {e}") raise GetAccessTokenError( - message=f"Failed to decode access token response: {e!s}", + message=f"Failed to decode access token response: {e}", status_code=400, ) except Exception as e: - verbose_logger.error(f"Unexpected error polling for access token: {e!s}") + verbose_logger.error(f"Unexpected error polling for access token: {e}") raise GetAccessTokenError( - message=f"Failed to get access token: {e!s}", + message=f"Failed to get access token: {e}", status_code=400, ) diff --git a/litellm/llms/huggingface/common_utils.py b/litellm/llms/huggingface/common_utils.py index 9dbdf05d0ec..07b580e68ce 100644 --- a/litellm/llms/huggingface/common_utils.py +++ b/litellm/llms/huggingface/common_utils.py @@ -96,7 +96,7 @@ def _fetch_inference_provider_mapping(model: str) -> dict: status_code = 500 headers = {} raise HuggingFaceError( - message=f"Failed to fetch provider mapping: {e!s}", + message=f"Failed to fetch provider mapping: {e}", status_code=status_code, headers=headers, ) diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index 895bbdca656..bdaa34871cf 100644 --- a/litellm/llms/langgraph/chat/sse_iterator.py +++ b/litellm/llms/langgraph/chat/sse_iterator.py @@ -196,7 +196,7 @@ class LangGraphSSEStreamIterator: except httpx.StreamClosed: raise StopIteration except Exception as e: - verbose_logger.error(f"Error in LangGraph SSE stream: {e!s}") + verbose_logger.error(f"Error in LangGraph SSE stream: {e}") raise StopIteration async def __anext__(self) -> ModelResponseStream: @@ -224,5 +224,5 @@ class LangGraphSSEStreamIterator: except httpx.StreamClosed: raise StopAsyncIteration except Exception as e: - verbose_logger.error(f"Error in LangGraph SSE stream: {e!s}") + verbose_logger.error(f"Error in LangGraph SSE stream: {e}") raise StopAsyncIteration diff --git a/litellm/llms/langgraph/chat/transformation.py b/litellm/llms/langgraph/chat/transformation.py index a40c08738f9..2aa96ddb978 100644 --- a/litellm/llms/langgraph/chat/transformation.py +++ b/litellm/llms/langgraph/chat/transformation.py @@ -451,14 +451,14 @@ class LangGraphConfig(BaseConfig): ) setattr(model_response, "usage", usage) except Exception as e: - verbose_logger.warning(f"Failed to calculate token usage: {e!s}") + verbose_logger.warning(f"Failed to calculate token usage: {e}") return model_response except Exception as e: - verbose_logger.error(f"Error processing LangGraph response: {e!s}") + verbose_logger.error(f"Error processing LangGraph response: {e}") raise LangGraphError( - message=f"Error processing response: {e!s}", + message=f"Error processing response: {e}", status_code=raw_response.status_code, ) diff --git a/litellm/llms/litellm_proxy/skills/code_execution.py b/litellm/llms/litellm_proxy/skills/code_execution.py index c99698a5c8e..f1142a8e355 100644 --- a/litellm/llms/litellm_proxy/skills/code_execution.py +++ b/litellm/llms/litellm_proxy/skills/code_execution.py @@ -239,7 +239,7 @@ class CodeExecutionHandler: tool_result += f"\n\nError:\n{exec_result['error']}" except Exception as e: - tool_result = f"Code execution failed: {e!s}" + tool_result = f"Code execution failed: {e}" execution_results.append( { "iteration": iteration, diff --git a/litellm/llms/manus/files/transformation.py b/litellm/llms/manus/files/transformation.py index cfa6d1cc722..325f6f36814 100644 --- a/litellm/llms/manus/files/transformation.py +++ b/litellm/llms/manus/files/transformation.py @@ -279,8 +279,8 @@ class ManusFilesConfig(BaseFilesConfig): status_details=response_json.get("status_details"), ) except Exception as e: - verbose_logger.exception(f"Error parsing Manus file response: {e!s}") - raise ValueError(f"Error parsing Manus file response: {e!s}") + verbose_logger.exception(f"Error parsing Manus file response: {e}") + raise ValueError(f"Error parsing Manus file response: {e}") def transform_retrieve_file_request( self, diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 48265d095a8..8646258b3db 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -158,7 +158,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): ) query_vector = embedding_response.data[0]["embedding"] except Exception as e: - raise Exception(f"Failed to generate embedding for query: {e!s}") + raise Exception(f"Failed to generate embedding for query: {e}") # Azure AI Search endpoint for search index_name = vector_store_id # vector_store_id is the index name diff --git a/litellm/llms/minimax/text_to_speech/transformation.py b/litellm/llms/minimax/text_to_speech/transformation.py index 93845d10789..af08bb8cb5f 100644 --- a/litellm/llms/minimax/text_to_speech/transformation.py +++ b/litellm/llms/minimax/text_to_speech/transformation.py @@ -353,7 +353,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): except Exception as e: raise MinimaxException( status_code=500, - message=f"Failed to decode audio data: {e!s}", + message=f"Failed to decode audio data: {e}", headers=dict(raw_response.headers), ) @@ -378,7 +378,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): except json.JSONDecodeError as e: raise MinimaxException( status_code=500, - message=f"Failed to parse MiniMax response: {e!s}", + message=f"Failed to parse MiniMax response: {e}", headers=dict(raw_response.headers), ) except Exception as e: @@ -386,7 +386,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): raise raise MinimaxException( status_code=500, - message=f"Error processing MiniMax response: {e!s}", + message=f"Error processing MiniMax response: {e}", headers=dict(raw_response.headers), ) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index d73435dbcfc..91d12fd78ba 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -330,7 +330,7 @@ class MistralConfig(OpenAIGPTConfig): 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{existing_content!s}" + new_content = f"{reasoning_prompt}\n\n{existing_content}" messages[i] = cast(AllMessageValues, {**msg, "content": new_content}) break diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index d3ffa926c46..5db85d355c8 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -201,7 +201,7 @@ def handle_cohere_response( cohere_response = CohereChatResult(**json_response) except (TypeError, ValidationError) as e: raise OCIError( - message=f"Response cannot be casted to CohereChatResult: {e!s}", + message=f"Response cannot be casted to CohereChatResult: {e}", status_code=raw_response.status_code, ) @@ -283,7 +283,7 @@ def handle_cohere_stream_chunk( except (TypeError, ValidationError) as e: raise OCIError( status_code=500, - message=f"Chunk cannot be parsed as CohereStreamChunk: {e!s}", + message=f"Chunk cannot be parsed as CohereStreamChunk: {e}", ) if typed_chunk.index is None: diff --git a/litellm/llms/oci/chat/generic.py b/litellm/llms/oci/chat/generic.py index 354bcbed3ba..7c60b3bea65 100644 --- a/litellm/llms/oci/chat/generic.py +++ b/litellm/llms/oci/chat/generic.py @@ -309,7 +309,7 @@ def handle_generic_response( completion_response = OCICompletionResponse(**json_data) except (TypeError, ValidationError) as e: raise OCIError( - message=f"Response cannot be casted to OCICompletionResponse: {e!s}", + message=f"Response cannot be casted to OCICompletionResponse: {e}", status_code=raw_response.status_code, ) @@ -373,7 +373,7 @@ def handle_generic_stream_chunk(dict_chunk: dict) -> ModelResponseStream: except (TypeError, ValidationError) as e: raise OCIError( status_code=500, - message=f"Chunk cannot be parsed as OCIStreamChunk: {e!s}", + message=f"Chunk cannot be parsed as OCIStreamChunk: {e}", ) if typed_chunk.index is None: diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 2d441cb4515..b0fcf85e840 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -741,7 +741,7 @@ class OCIStreamWrapper(CustomStreamWrapper): except json.JSONDecodeError as e: raise OCIError( status_code=500, - message=f"Chunk cannot be parsed as JSON: {e!s}", + message=f"Chunk cannot be parsed as JSON: {e}", ) if dict_chunk.get("apiFormat") == "COHERE": diff --git a/litellm/llms/oci/common_utils.py b/litellm/llms/oci/common_utils.py index 7277972f64a..d5bacede08c 100644 --- a/litellm/llms/oci/common_utils.py +++ b/litellm/llms/oci/common_utils.py @@ -232,7 +232,7 @@ def sign_with_oci_signer( raise OCIError( status_code=500, message=( - f"Failed to sign request with provided oci_signer: {e!s}. " + f"Failed to sign request with provided oci_signer: {e}. " "The signer must implement the OCI SDK Signer interface with a " "do_request_sign(request, enforce_content_headers=True) method. " "See: https://docs.oracle.com/en-us/iaas/tools/python/latest/api/signing.html" diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index e9e60106d2d..e5afb4b87b6 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -369,7 +369,7 @@ class OllamaChatConfig(BaseConfig): content=None, tool_calls=[ { - "id": f"call_{uuid.uuid4()!s}", + "id": f"call_{uuid.uuid4()}", "function": { "name": function_call.get("name", litellm_params.get("function_name")), "arguments": json.dumps(function_call.get("arguments", function_call)), diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 0add66827f8..5823c2dad75 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -282,7 +282,7 @@ class OllamaConfig(BaseConfig): content=None, tool_calls=[ { - "id": f"call_{uuid.uuid4()!s}", + "id": f"call_{uuid.uuid4()}", "function": { "name": function_call["name"], "arguments": json.dumps(function_call["arguments"]), diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index a37e15c1f86..723b22a57b9 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -621,7 +621,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): except Exception as e: response_headers = getattr(raw_response, "headers", None) raise OpenAIError( - message=f"Unable to get json response - {e!s}, Original Response: {raw_response.text}", + message=f"Unable to get json response - {e}, Original Response: {raw_response.text}", status_code=raw_response.status_code, headers=response_headers, ) diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index e72680f387d..808998ddaf6 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -134,13 +134,33 @@ class BaseOpenAILLM: _cached_client = litellm.in_memory_llm_clients_cache.get_cache(_cache_key) return _cached_client + @staticmethod + def owns_wrapped_http_client(http_client: httpx.Client | httpx.AsyncClient | None) -> bool: + """Whether litellm may close an SDK client built around ``http_client``. + + ``_get_async_http_client`` / ``_get_sync_http_client`` hand back + ``litellm.aclient_session`` / ``litellm.client_session`` when the caller + configured one. The SDK's ``close()`` closes whatever http client it was + given, so an SDK client wrapping one of those shared sessions must never be + closed on eviction; the caller goes on using the session. ``None`` means the + SDK built its own http client, which litellm does own. + """ + if http_client is None: + return True + return http_client is not litellm.aclient_session and http_client is not litellm.client_session + @staticmethod def set_cached_openai_client( openai_client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI, client_type: Literal["openai", "azure"], client_initialization_params: dict, + litellm_owned_client: bool = False, ): - """Stores the OpenAI client in the in-memory cache for _DEFAULT_TTL_FOR_HTTPX_CLIENTS SECONDS""" + """Stores the OpenAI client in the in-memory cache for _DEFAULT_TTL_FOR_HTTPX_CLIENTS SECONDS + + ``litellm_owned_client`` says litellm built this client, so the cache may close it once it + is evicted. A client the caller supplied stays open, since litellm does not own it. + """ _cache_key = BaseOpenAILLM.get_openai_client_cache_key( client_initialization_params=client_initialization_params, client_type=client_type, @@ -149,6 +169,7 @@ class BaseOpenAILLM: key=_cache_key, value=openai_client, ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, + litellm_owned_client=litellm_owned_client, ) @staticmethod diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index e4a13f0f526..f01730a06a5 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -366,11 +366,16 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): if cached_client: if isinstance(cached_client, OpenAI) or isinstance(cached_client, AsyncOpenAI): return cached_client + http_client: httpx.Client | httpx.AsyncClient | None = ( + OpenAIChatCompletion._get_async_http_client(shared_session=shared_session) + if is_async + else OpenAIChatCompletion._get_sync_http_client() + ) if is_async: _new_client: OpenAI | AsyncOpenAI = AsyncOpenAI( api_key=api_key, base_url=api_base, - http_client=OpenAIChatCompletion._get_async_http_client(shared_session=shared_session), + http_client=http_client, timeout=timeout, max_retries=max_retries, organization=organization, @@ -379,7 +384,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): _new_client = OpenAI( api_key=api_key, base_url=api_base, - http_client=OpenAIChatCompletion._get_sync_http_client(), + http_client=http_client, timeout=timeout, max_retries=max_retries, organization=organization, @@ -390,6 +395,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): openai_client=_new_client, client_initialization_params=client_initialization_params, client_type="openai", + litellm_owned_client=self.owns_wrapped_http_client(http_client), ) return _new_client @@ -551,7 +557,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): except Exception as e: verbose_logger.exception( - f"LiteLLM.AgenticHookError: Exception in agentic completion hooks for OpenAI: {e!s}" + f"LiteLLM.AgenticHookError: Exception in agentic completion hooks for OpenAI: {e}" ) return None @@ -774,7 +780,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): # e.message except Exception as e: if print_verbose is not None: - print_verbose(f"openai.py: Received openai error - {e!s}") + print_verbose(f"openai.py: Received openai error - {e}") if ( "Conversation roles must alternate user/assistant" in str(e) or "user and assistant roles should be alternating" in str(e) @@ -1089,7 +1095,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): if response is not None and hasattr(response, "text"): raise OpenAIError( status_code=status_code, - message=f"{e!s}\n\nOriginal Response: {response.text}", # type: ignore + message=f"{e}\n\nOriginal Response: {response.text}", # type: ignore headers=error_headers, body=exception_body, ) @@ -1111,7 +1117,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): else: raise OpenAIError( status_code=500, - message=f"{e!s}", + message=f"{e}", headers=error_headers, body=exception_body, ) diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 14fa6dc9954..a9a2b476776 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -178,7 +178,7 @@ 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: {e!s}")) + await websocket.close(code=1011, reason=_redact_string(f"Internal server error: {e}")) except RuntimeError as 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 diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index b7cc3b1673a..e59a28c2d09 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -88,14 +88,14 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): except OpenAIError: raise except httpx.HTTPStatusError as e: - verbose_logger.error(f"HTTP error in CountTokens handler: {e!s}") + verbose_logger.error(f"HTTP error in CountTokens handler: {e}") raise OpenAIError( status_code=e.response.status_code, message=e.response.text, ) except (httpx.RequestError, json.JSONDecodeError, ValueError) as e: - verbose_logger.error(f"Error in CountTokens handler: {e!s}") + verbose_logger.error(f"Error in CountTokens handler: {e}") raise OpenAIError( status_code=500, - message=f"CountTokens processing error: {e!s}", + message=f"CountTokens processing error: {e}", ) diff --git a/litellm/llms/openrouter/image_edit/transformation.py b/litellm/llms/openrouter/image_edit/transformation.py index fad7d53577c..8163b92bb19 100644 --- a/litellm/llms/openrouter/image_edit/transformation.py +++ b/litellm/llms/openrouter/image_edit/transformation.py @@ -203,7 +203,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): response_json = raw_response.json() except Exception as e: raise OpenRouterException( - message=f"Error parsing OpenRouter response: {e!s}", + message=f"Error parsing OpenRouter response: {e}", status_code=raw_response.status_code, headers=raw_response.headers, ) @@ -246,7 +246,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): except Exception as e: raise OpenRouterException( - message=f"Error transforming OpenRouter image edit response: {e!s}", + message=f"Error transforming OpenRouter image edit response: {e}", status_code=500, headers={}, ) diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py index 1114bb41275..f56ca6ba89e 100644 --- a/litellm/llms/openrouter/image_generation/transformation.py +++ b/litellm/llms/openrouter/image_generation/transformation.py @@ -345,7 +345,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): response_json = raw_response.json() except Exception as e: raise OpenRouterException( - message=f"Error parsing OpenRouter response: {e!s}", + message=f"Error parsing OpenRouter response: {e}", status_code=raw_response.status_code, headers=raw_response.headers, ) @@ -394,7 +394,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): except Exception as e: raise OpenRouterException( - message=f"Error transforming OpenRouter image generation response: {e!s}", + message=f"Error transforming OpenRouter image generation response: {e}", status_code=500, headers={}, ) diff --git a/litellm/llms/predibase/chat/handler.py b/litellm/llms/predibase/chat/handler.py index 36537562638..2bf39966dd1 100644 --- a/litellm/llms/predibase/chat/handler.py +++ b/litellm/llms/predibase/chat/handler.py @@ -225,7 +225,7 @@ class PredibaseChatCompletion: if isinstance(e, exception): raise e raise PredibaseError( - status_code=500, message=f"{e!s}" + status_code=500, message=f"{e}" ) # don't use verbose_logger.exception, if exception is raised return predibase_config.transform_response( model=model, diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 868f4f9696e..406a72ffd99 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -529,7 +529,7 @@ class SagemakerLLM(BaseAWSLLM): ) raise e except Exception as e: - error_message = f"{e!s}" + error_message = f"{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=500, message=error_message) diff --git a/litellm/llms/sagemaker/embedding/transformation.py b/litellm/llms/sagemaker/embedding/transformation.py index 7221e030d97..51fad1e1c2e 100644 --- a/litellm/llms/sagemaker/embedding/transformation.py +++ b/litellm/llms/sagemaker/embedding/transformation.py @@ -97,7 +97,7 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): response_data = raw_response.json() except Exception as e: raise SagemakerError( - message=f"Failed to parse response: {e!s}", + message=f"Failed to parse response: {e}", status_code=raw_response.status_code, ) diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index e785e7ec28e..9f01a9ee506 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -270,7 +270,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): total_tokens=total_tokens, ) except Exception as e: - verbose_logger.warning(f"Failed to calculate token usage: {e!s}") + verbose_logger.warning(f"Failed to calculate token usage: {e}") return None def transform_response( @@ -335,9 +335,9 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): return model_response except Exception as e: - verbose_logger.error(f"Error processing Vertex Agent Engine response: {e!s}") + verbose_logger.error(f"Error processing Vertex Agent Engine response: {e}") raise VertexAgentEngineError( - message=f"Error processing response: {e!s}", + message=f"Error processing response: {e}", status_code=raw_response.status_code, ) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index b627444b181..81d084e7e03 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -221,7 +221,7 @@ def get_supports_system_message( supports_system_message = True except Exception as e: verbose_logger.warning( - f"Unable to identify if system message supported. Defaulting to 'False'. Received error message - {e!s}\nAdd it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json" + f"Unable to identify if system message supported. Defaulting to 'False'. Received error message - {e}\nAdd it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json" ) supports_system_message = False diff --git a/litellm/llms/vertex_ai/cost_calculator.py b/litellm/llms/vertex_ai/cost_calculator.py index 1cd4c0a9e97..a53e54e5fc2 100644 --- a/litellm/llms/vertex_ai/cost_calculator.py +++ b/litellm/llms/vertex_ai/cost_calculator.py @@ -114,7 +114,7 @@ def cost_per_character( prompt_cost = prompt_characters * model_info["input_cost_per_character"] except Exception as e: verbose_logger.debug( - f"litellm.litellm_core_utils.llm_cost_calc.google.py::cost_per_character(): Exception occured - {e!s}\nDefaulting to None" + f"litellm.litellm_core_utils.llm_cost_calc.google.py::cost_per_character(): Exception occured - {e}\nDefaulting to None" ) prompt_cost, _ = cost_per_token( model=model, @@ -152,7 +152,7 @@ def cost_per_character( completion_cost = completion_characters * model_info["output_cost_per_character"] except Exception as e: verbose_logger.debug( - f"litellm.litellm_core_utils.llm_cost_calc.google.py::cost_per_character(): Exception occured - {e!s}\nDefaulting to None" + f"litellm.litellm_core_utils.llm_cost_calc.google.py::cost_per_character(): Exception occured - {e}\nDefaulting to None" ) _, completion_cost = cost_per_token( model=model, diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index f8eea399bf6..49fdb2786e1 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -815,6 +815,6 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): "response": None, "error": { "code": "transformation_error", - "message": f"Failed to transform response: {e!s}", + "message": f"Failed to transform response: {e}", }, } diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 9d4d8a5a02e..76549d0fed2 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -356,9 +356,7 @@ def _get_gcs_object_content_type( headers["Authorization"] = f"Bearer {access_token}" except Exception as e: raise litellm.BadRequestError( - message=( - f"Unable to fetch GCS metadata with provided Vertex credentials/project. Original error: {e!s}" - ), + message=(f"Unable to fetch GCS metadata with provided Vertex credentials/project. Original error: {e}"), model=None, llm_provider="vertex_ai", ) @@ -844,7 +842,7 @@ def _gemini_convert_messages_with_history( f"{file_id or 'provided data'}, set this explicitly " f"using message[{msg_i}].content[{element_idx}].file.format " f"(or file.mime_type/content_type). " - f"Original error: {e!s}" + f"Original error: {e}" ), model=model, llm_provider="vertex_ai", 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 19c43d8000c..cadc8760601 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 @@ -2405,7 +2405,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): completion_response = GenerateContentResponseBody(**raw_response.json()) # type: ignore except Exception as e: raise VertexAIError( - message=f"Error converting to valid response block={e!s}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues", + message=f"Error converting to valid response block={e}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues", status_code=422, headers=raw_response.headers, ) @@ -2512,7 +2512,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): except Exception as e: raise VertexAIError( - message=f"Error converting to valid response block={e!s}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues", + message=f"Error converting to valid response block={e}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues", status_code=422, headers=raw_response.headers, ) 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 abad2bb73ea..d5e279ea240 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 @@ -127,7 +127,7 @@ class VertexAILlama3Config(OpenAIGPTConfig): except Exception as e: response_headers = getattr(raw_response, "headers", None) raise VertexAIError( - message=f"Unable to get json response - {e!s}, Original Response: {raw_response.text}", + message=f"Unable to get json response - {e}, Original Response: {raw_response.text}", status_code=raw_response.status_code, headers=response_headers, ) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index b3ffa1d40be..2def1acb708 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -783,7 +783,7 @@ class VertexBase: except Exception as retry_error: verbose_logger.error( f"Reauthentication retry failed for project_id: {project_id}. " - f"Original error: {error!s}. Retry error: {retry_error!s}" + f"Original error: {error}. Retry error: {retry_error}" ) # Re-raise the original error for better context raise error @@ -837,7 +837,7 @@ class VertexBase: except Exception as retry_error: verbose_logger.error( f"Async reauthentication retry failed for project_id: {project_id}. " - f"Original error: {error!s}. Retry error: {retry_error!s}" + f"Original error: {error}. Retry error: {retry_error}" ) raise error @@ -897,7 +897,7 @@ class VertexBase: _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: {e!s}" + f"Failed to load vertex credentials. Check to see if credentials containing partial/invalid information. Error: {e}" ) raise e diff --git a/litellm/llms/volcengine/embedding/transformation.py b/litellm/llms/volcengine/embedding/transformation.py index 5a0b59d411c..9923167ba31 100644 --- a/litellm/llms/volcengine/embedding/transformation.py +++ b/litellm/llms/volcengine/embedding/transformation.py @@ -162,7 +162,7 @@ class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): try: response_json = raw_response.json() except Exception as e: - raise ValueError(f"Failed to parse Volcengine response as JSON: {e!s}") + raise ValueError(f"Failed to parse Volcengine response as JSON: {e}") # Volcengine response format matches OpenAI format closely # Just need to ensure all required fields are present diff --git a/litellm/llms/watsonx/audio_transcription/transformation.py b/litellm/llms/watsonx/audio_transcription/transformation.py index 6019b2e8355..a9cd85cb674 100644 --- a/litellm/llms/watsonx/audio_transcription/transformation.py +++ b/litellm/llms/watsonx/audio_transcription/transformation.py @@ -170,7 +170,7 @@ class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTran try: raw_response_json = raw_response.json() except Exception as e: - raise ValueError(f"Error transforming response to json: {e!s}\nResponse: {raw_response.text}") + raise ValueError(f"Error transforming response to json: {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/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 6d9de3f481b..e1d5f2f3571 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -164,7 +164,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): raw_response_json = raw_response.json() except Exception as e: raise self.get_error_class( - error_message=f"Failed to parse response: {e!s}", + error_message=f"Failed to parse response: {e}", status_code=raw_response.status_code, headers=raw_response.headers, ) diff --git a/litellm/main.py b/litellm/main.py index cea9d44fb1a..731a545a267 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8303,7 +8303,7 @@ async def ahealth_check( if mode is None: return { - "error": f"error:{e!s}. Missing `mode`. Set the `mode` for the model - https://docs.litellm.ai/docs/proxy/health#embedding-models \nstacktrace: {stack_trace}", + "error": f"error:{e}. Missing `mode`. Set the `mode` for the model - https://docs.litellm.ai/docs/proxy/health#embedding-models \nstacktrace: {stack_trace}", "exception": e, } @@ -8669,7 +8669,7 @@ def stream_chunk_builder( processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) return response except Exception as e: - verbose_logger.exception(f"litellm.main.py::stream_chunk_builder() - Exception occurred - {e!s}") + verbose_logger.exception(f"litellm.main.py::stream_chunk_builder() - Exception occurred - {e}") raise litellm.APIError( status_code=500, message="Error building chunks for logging/streaming usage calculation", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 07f04136bd0..848af54cbcf 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1975,6 +1975,7 @@ "prompt_cache_min_tokens": 2048 }, "anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -2011,6 +2012,7 @@ "prompt_cache_min_tokens": 1024 }, "global.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -2047,6 +2049,7 @@ "prompt_cache_min_tokens": 1024 }, "us.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2083,6 +2086,7 @@ "prompt_cache_min_tokens": 1024 }, "eu.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2119,6 +2123,7 @@ "prompt_cache_min_tokens": 1024 }, "au.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2155,6 +2160,7 @@ "prompt_cache_min_tokens": 1024 }, "jp.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, 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 e8f39daa758..a256653f0f9 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 @@ -1008,7 +1008,7 @@ class MCPRequestHandler: limits[source.team_id] = applicable return limits or None except Exception as e: # noqa: BLE001 # throttling metadata must never fail an allowed request - verbose_logger.warning(f"Failed to resolve per-team MCP rpm limits for admitted subject: {e!s}") + verbose_logger.warning(f"Failed to resolve per-team MCP rpm limits for admitted subject: {e}") return None @staticmethod @@ -1514,9 +1514,9 @@ class MCPRequestHandler: if isinstance(e, UnloadableEntitlementError): # A ceiling we KNOW exists and cannot read. Denying is the only answer that does not # widen this caller past what an operator configured, for both caller shapes. - verbose_logger.warning(f"Denying MCP access, entitlement unreadable: {e!s}") + verbose_logger.warning(f"Denying MCP access, entitlement unreadable: {e}") else: - verbose_logger.warning(f"Failed to get allowed MCP servers: {e!s}") + verbose_logger.warning(f"Failed to get allowed MCP servers: {e}") return [] @staticmethod @@ -1649,7 +1649,7 @@ class MCPRequestHandler: # Fault isolation is per SOURCE: an unresolvable team contributes nothing (fail closed for # it alone, access only narrows) while every other source stands. Raising would collapse the # whole union to deny-all over one momentarily-unreadable row. - verbose_logger.warning(f"MCP admitted-subject source team {team_id!r} unresolvable, skipping: {e!s}") + verbose_logger.warning(f"MCP admitted-subject source team {team_id!r} unresolvable, skipping: {e}") return None if team_obj is None: return None @@ -1682,10 +1682,10 @@ class MCPRequestHandler: proxy_logging_obj=proxy_logging_obj, ) except BudgetExceededError as e: - verbose_logger.info(f"MCP admitted-subject source team {team_id!r} over budget, not a grantor: {e!s}") + verbose_logger.info(f"MCP admitted-subject source team {team_id!r} over budget, not a grantor: {e}") return None except Exception as e: # noqa: BLE001 # per-source isolation: a budget-check fault narrows, never raises - verbose_logger.warning(f"MCP budget check failed for source team {team_id!r}, skipping source: {e!s}") + verbose_logger.warning(f"MCP budget check failed for source team {team_id!r}, skipping source: {e}") return None return team_obj @@ -1738,7 +1738,7 @@ class MCPRequestHandler: billed.org_id = source.org_id return billed except Exception as e: # noqa: BLE001 # attribution must never fail an authorized call - verbose_logger.warning(f"MCP billing attribution failed for {tool_name!r}, billing the user: {e!s}") + verbose_logger.warning(f"MCP billing attribution failed for {tool_name!r}, billing the user: {e}") return auth @staticmethod @@ -1946,9 +1946,9 @@ class MCPRequestHandler: # than the None (allow-all) key auth gets for an indeterminate fault. unreadable_entitlement = isinstance(e, UnloadableEntitlementError) if unreadable_entitlement: - verbose_logger.warning(f"Denying MCP tools, entitlement unreadable: {e!s}") + verbose_logger.warning(f"Denying MCP tools, entitlement unreadable: {e}") else: - verbose_logger.warning(f"Failed to get allowed tools for server: {e!s}") + verbose_logger.warning(f"Failed to get allowed tools for server: {e}") # Fail CLOSED for a keyless admitted subject: ANY error must deny the server's tools ([]), # not collapse to allow-all (None); key/JWT auth keeps its prior allow-all-on-error. Both # keyless_source AND the marker are needed: each source resolves through an UNMARKED auth, so @@ -1999,7 +1999,7 @@ class MCPRequestHandler: raise verbose_logger.warning( f"MCP org tool ceiling unresolvable for org_id={user_api_key_auth.org_id!r}; " - f"skipping org intersect, key/team/agent restrictions stand: {e!s}" + f"skipping org intersect, key/team/agent restrictions stand: {e}" ) return allowed_tools org_tools = ( @@ -2102,7 +2102,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: {e!s}") + verbose_logger.warning(f"Failed to get key access group MCP server grants: {e}") return [] @staticmethod @@ -2180,7 +2180,7 @@ class MCPRequestHandler: all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers + toolset_servers return list(set(all_servers)) except Exception as e: - verbose_logger.warning(f"Failed to get allowed MCP servers for key: {e!s}") + verbose_logger.warning(f"Failed to get allowed MCP servers for key: {e}") return [] @staticmethod @@ -2238,7 +2238,7 @@ class MCPRequestHandler: proxy_logging_obj=proxy_logging_obj, ) except Exception as e: # noqa: BLE001 # a team-resolution blip narrows access, never raises - verbose_logger.warning(f"Failed to resolve user teams for MCP grant: {e!s}") + verbose_logger.warning(f"Failed to resolve user teams for MCP grant: {e}") return [] if user_object is None or not user_object.teams: return [] @@ -2323,7 +2323,7 @@ class MCPRequestHandler: servers = await MCPRequestHandler._team_granted_servers(team_obj, team_access_group_servers) return list(servers) except Exception as e: - verbose_logger.warning(f"Failed to get allowed MCP servers for team: {e!s}") + verbose_logger.warning(f"Failed to get allowed MCP servers for team: {e}") return [] @staticmethod @@ -2462,7 +2462,7 @@ class MCPRequestHandler: # A NAMED-but-unreadable ceiling is a stronger fact than "unresolved" and denies everywhere. if isinstance(e, UnloadableEntitlementError): raise - verbose_logger.warning(f"Failed to get allowed MCP servers for org: {e!s}") + verbose_logger.warning(f"Failed to get allowed MCP servers for org: {e}") return None @staticmethod @@ -2490,7 +2490,7 @@ class MCPRequestHandler: route="/mcp", ) except Exception as e: # noqa: BLE001 # entitlement unknown, not known-absent: no ceiling, as before this level - verbose_logger.warning(f"Failed to resolve end_user for MCP permissions: {e!s}") + verbose_logger.warning(f"Failed to resolve end_user for MCP permissions: {e}") return None if end_user_obj is None: @@ -2554,7 +2554,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: {e!s}") + verbose_logger.warning(f"Failed to get allowed MCP servers for end_user: {e}") return [] @staticmethod @@ -2637,7 +2637,7 @@ class MCPRequestHandler: ) return object_permission_id except Exception as e: # noqa: BLE001 # unknown whether entitled at all: no ceiling, as before - verbose_logger.warning(f"MCP user entitlement: link for {user_id!r} unresolved, no ceiling: {e!s}") + verbose_logger.warning(f"MCP user entitlement: link for {user_id!r} unresolved, no ceiling: {e}") return None @staticmethod @@ -2669,7 +2669,7 @@ class MCPRequestHandler: ) return list(set(direct_mcp_servers + access_group_servers + tool_perm_servers)) except Exception as e: # noqa: BLE001 # any resolution fault is an unresolved ceiling, never "no ceiling" - verbose_logger.warning(f"Failed to get allowed MCP servers for user: {e!s}") + verbose_logger.warning(f"Failed to get allowed MCP servers for user: {e}") return None @staticmethod @@ -2739,7 +2739,7 @@ class MCPRequestHandler: try: object_permissions = await MCPRequestHandler._get_user_object_permission(user_api_key_auth) except Exception as e: # noqa: BLE001 # an unresolved human entitlement must deny, not widen - verbose_logger.warning(f"MCP user tool ceiling unresolvable, denying tools on {server_id!r}: {e!s}") + verbose_logger.warning(f"MCP user tool ceiling unresolvable, denying tools on {server_id!r}: {e}") return [] if object_permissions is None or not object_permissions.mcp_tool_permissions: @@ -2785,7 +2785,7 @@ class MCPRequestHandler: ) return object_permission_id except Exception as e: # noqa: BLE001 # entitlement unknown, not known-absent: no ceiling, as before this level - verbose_logger.warning(f"Failed to resolve object_permission_id for agent {agent_id!r}: {e!s}") + verbose_logger.warning(f"Failed to resolve object_permission_id for agent {agent_id!r}: {e}") return None @staticmethod @@ -2869,7 +2869,7 @@ class MCPRequestHandler: 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: {e!s}") + verbose_logger.warning(f"Failed to get allowed MCP servers for agent: {e}") return [] @staticmethod @@ -2911,7 +2911,7 @@ class MCPRequestHandler: 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: {e!s}") + verbose_logger.warning(f"Failed to get agent tool permissions for server: {e}") return None @staticmethod @@ -2969,7 +2969,7 @@ class MCPRequestHandler: return list(server_ids) except Exception as e: - verbose_logger.warning(f"Failed to get MCP servers from access groups: {e!s}") + verbose_logger.warning(f"Failed to get MCP servers from access groups: {e}") return [] @staticmethod @@ -3029,7 +3029,7 @@ class MCPRequestHandler: return key_object_permission.mcp_access_groups or [] except Exception as e: - verbose_logger.warning(f"Failed to get MCP access groups for key: {e!s}") + verbose_logger.warning(f"Failed to get MCP access groups for key: {e}") return [] @staticmethod @@ -3077,7 +3077,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: {e!s}") + verbose_logger.warning(f"Failed to get MCP access groups for team: {e}") return [] @staticmethod diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 3c8e7d9f1ef..672396afd05 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -570,7 +570,7 @@ async def get_all_mcp_servers( decrypt_global_env_var_values(table.env_vars) return tables except Exception as e: - verbose_proxy_logger.debug(f"litellm.proxy._experimental.mcp_server.db.py::get_all_mcp_servers - {e!s}") + verbose_proxy_logger.debug(f"litellm.proxy._experimental.mcp_server.db.py::get_all_mcp_servers - {e}") return [] diff --git a/litellm/proxy/_experimental/mcp_server/elicitation_handler.py b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py index 9927afa20d0..35681a9473e 100644 --- a/litellm/proxy/_experimental/mcp_server/elicitation_handler.py +++ b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py @@ -79,7 +79,7 @@ async def handle_elicitation_request( verbose_logger.exception("MCP elicitation handler failed: %s", e) return ErrorData( code=-1, - message=f"Elicitation failed: {e!s}", + message=f"Elicitation failed: {e}", ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d8ab34a7ddb..0a6a0374d13 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1951,7 +1951,7 @@ class MCPServerManager: 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}: {e!s}") + verbose_logger.error(f"Failed to register OpenAPI tools for server {server.name}: {e}") raise e def _cleanup_server_tool_routing_artifacts(self, server: MCPServer) -> None: @@ -2326,7 +2326,7 @@ class MCPServerManager: verbose_logger.debug(f"Added MCP Server: {new_server.name}") except Exception as e: - verbose_logger.debug(f"Failed to add MCP server: {e!s}") + verbose_logger.debug(f"Failed to add MCP server: {e}") raise e async def update_server(self, mcp_server: LiteLLM_MCPServerTable): @@ -2360,7 +2360,7 @@ class MCPServerManager: verbose_logger.debug(f"Updated MCP Server: {new_server.name}") except Exception as e: - verbose_logger.debug(f"Failed to udpate MCP server: {e!s}") + verbose_logger.debug(f"Failed to udpate MCP server: {e}") raise e def get_all_mcp_server_ids(self) -> set[str]: @@ -2386,7 +2386,7 @@ class MCPServerManager: await user_api_key_cache.async_delete_cache(key=self.get_byom_submitted_servers_cache_key(user_id)) except Exception as e: # noqa: BLE001 - verbose_logger.warning(f"Failed to invalidate BYOM submitted MCP server cache: {e!s}") + verbose_logger.warning(f"Failed to invalidate BYOM submitted MCP server cache: {e}") async def _get_active_submitted_mcp_server_ids_for_user( self, user_api_key_auth: UserAPIKeyAuth | None @@ -2401,7 +2401,7 @@ class MCPServerManager: ) from litellm.proxy.proxy_server import prisma_client, user_api_key_cache except Exception as e: # noqa: BLE001 - verbose_logger.warning(f"Failed to load BYOM submitted MCP server cache dependencies: {e!s}") + verbose_logger.warning(f"Failed to load BYOM submitted MCP server cache dependencies: {e}") return [] byom_cache_key = self.get_byom_submitted_servers_cache_key(submitter_user_id) @@ -2411,7 +2411,7 @@ class MCPServerManager: if cached_submitted_server_ids is not None: submitted_server_ids = cast(list[str], cached_submitted_server_ids) except Exception as e: # noqa: BLE001 - verbose_logger.warning(f"Failed to read BYOM submitted MCP server cache: {e!s}") + verbose_logger.warning(f"Failed to read BYOM submitted MCP server cache: {e}") if submitted_server_ids is None: if prisma_client is None: @@ -2422,7 +2422,7 @@ class MCPServerManager: prisma_client, submitter_user_id ) except Exception as e: # noqa: BLE001 - verbose_logger.warning(f"Failed to read BYOM submitted MCP servers from database: {e!s}") + verbose_logger.warning(f"Failed to read BYOM submitted MCP servers from database: {e}") submitted_server_ids = [] try: await user_api_key_cache.async_set_cache( @@ -2431,7 +2431,7 @@ class MCPServerManager: ttl=60, ) except Exception as e: # noqa: BLE001 - verbose_logger.warning(f"Failed to write BYOM submitted MCP server cache: {e!s}") + verbose_logger.warning(f"Failed to write BYOM submitted MCP server cache: {e}") return [server_id for server_id in submitted_server_ids if self.get_mcp_server_by_id(server_id) is not None] @@ -2647,7 +2647,7 @@ class MCPServerManager: ) return tool_permissions except Exception as e: - verbose_logger.warning(f"Failed to resolve toolset permissions: {e!s}") + verbose_logger.warning(f"Failed to resolve toolset permissions: {e}") return {} def invalidate_toolset_cache(self, toolset_id: str | None = None) -> None: @@ -2764,7 +2764,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}: {e!s}") + verbose_logger.warning(f"Failed to get tools from server {server_id}: {e}") return [] async def list_tools( @@ -2822,7 +2822,7 @@ class MCPServerManager: return tools except Exception as e: verbose_logger.warning( - f"Failed to list tools from server {server.name}: {e!s}. Continuing with other servers." + f"Failed to list tools from server {server.name}: {e}. Continuing with other servers." ) return [] @@ -3476,12 +3476,12 @@ class MCPServerManager: www_authenticate=None if server.is_dcr_bridge else challenge_header, server_name=server.name, ) from e - verbose_logger.warning(f"Failed to get tools from server {server.name}: {e!s}") + verbose_logger.warning(f"Failed to get tools from server {server.name}: {e}") raise MCPServerListError(ServerListFault(tag="internal", status_code=e.status_code), server.name) from e except MCPServerListError: raise except Exception as e: - verbose_logger.warning(f"Failed to get tools from server {server.name}: {e!s}") + verbose_logger.warning(f"Failed to get tools from server {server.name}: {e}") raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge) async def get_prompts_from_server( @@ -3532,7 +3532,7 @@ class MCPServerManager: return prefixed_or_original_prompts except Exception as e: - verbose_logger.warning(f"Failed to get prompts from server {server.name}: {e!s}") + verbose_logger.warning(f"Failed to get prompts from server {server.name}: {e}") return [] async def get_resources_from_server( @@ -3574,7 +3574,7 @@ class MCPServerManager: return prefixed_resources except Exception as e: - verbose_logger.warning(f"Failed to get resources from server {server.name}: {e!s}") + verbose_logger.warning(f"Failed to get resources from server {server.name}: {e}") return [] async def get_resource_templates_from_server( @@ -3618,7 +3618,7 @@ class MCPServerManager: return prefixed_templates except Exception as e: - verbose_logger.warning(f"Failed to get resource templates from server {server.name}: {e!s}") + verbose_logger.warning(f"Failed to get resource templates from server {server.name}: {e}") return [] async def read_resource_from_server( @@ -4215,10 +4215,10 @@ class MCPServerManager: verbose_logger.warning(f"Task cancelled while listing tools from {server_name}") raise MCPServerListError(ServerListFault(tag="internal"), server_name) from e except ConnectionError as e: - verbose_logger.warning(f"Connection error while listing tools from {server_name}: {e!s}") + verbose_logger.warning(f"Connection error while listing tools from {server_name}: {e}") raise MCPServerListError(ServerListFault(tag="unreachable"), server_name) from e except Exception as e: - verbose_logger.warning(f"Error listing tools from {server_name}: {e!s}") + verbose_logger.warning(f"Error listing tools from {server_name}: {e}") raise_classified_list_failure(e, server_name) _SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024 @@ -4533,7 +4533,7 @@ class MCPServerManager: return result except Exception as e: - error_msg = f"Error calling OpenAPI tool {tool_name}: {e!s}" + error_msg = f"Error calling OpenAPI tool {tool_name}: {e}" verbose_logger.error(error_msg) return CallToolResult( content=[TextContent(type="text", text=error_msg)], @@ -4639,7 +4639,7 @@ class MCPServerManager: HTTPException, ) as e: # Re-raise guardrail exceptions to properly fail the MCP call - verbose_logger.error(f"Guardrail blocked MCP tool call pre call: {e!s}") + verbose_logger.error(f"Guardrail blocked MCP tool call pre call: {e}") raise e return hook_result @@ -4995,7 +4995,7 @@ class MCPServerManager: GuardrailRaisedException, HTTPException, ) as e: - verbose_logger.error(f"Guardrail blocked MCP tool call during result check: {e!s}") + verbose_logger.error(f"Guardrail blocked MCP tool call during result check: {e}") raise e # If proxy_logging_obj is None, the tool call result is at index 0 @@ -5194,7 +5194,7 @@ class MCPServerManager: GuardrailRaisedException, HTTPException, ) as e: - verbose_logger.error(f"Guardrail blocked MCP tool call during result check: {e!s}") + verbose_logger.error(f"Guardrail blocked MCP tool call during result check: {e}") raise e async def call_tool( @@ -5345,7 +5345,7 @@ class MCPServerManager: 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: {e!s}" + f"No running event loop - skipping tool name to MCP server name mapping initialization: {e}" ) async def _initialize_tool_name_to_mcp_server_name_mapping(self): @@ -5364,12 +5364,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} due to upstream auth error: {e!s}" + f"Skipping tool name mapping for server {server.name} due to upstream auth error: {e}" ) continue except Exception as e: verbose_logger.warning( - f"Failed to get tools from server {server.name} during tool name mapping initialization: {e!s}" + f"Failed to get tools from server {server.name} during tool name mapping initialization: {e}" ) continue for tool in tools: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 6f4fde6fbfe..d0458db51c6 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -654,7 +654,7 @@ if MCP_AVAILABLE: return { "tools": [], "error": "server_error", - "message": f"Failed to get tools from server {server.name}: {e!s}", + "message": f"Failed to get tools from server {server.name}: {e}", } return { "tools": list_tools_result, @@ -866,7 +866,7 @@ if MCP_AVAILABLE: errors.append( f"{get_server_prefix(server)}: {classify_list_exception(e).tag}" if isinstance(e, (MCPServerListError, MCPUpstreamAuthError)) - else f"{get_server_prefix(server)}: {e!s}" + else f"{get_server_prefix(server)}: {e}" ) continue @@ -905,7 +905,7 @@ if MCP_AVAILABLE: return { "tools": [], "error": "unexpected_error", - "message": f"An unexpected error occurred: {e!s}", + "message": f"An unexpected error occurred: {e}", } @router.post("/tools/call", dependencies=[Depends(user_api_key_auth)]) @@ -1052,7 +1052,7 @@ if MCP_AVAILABLE: }, ) except BlockedPiiEntityError as e: - verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {e!s}") + verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {e}") raise HTTPException( status_code=400, detail={ @@ -1063,7 +1063,7 @@ if MCP_AVAILABLE: }, ) except GuardrailRaisedException as e: - verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {e!s}") + verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {e}") raise HTTPException( status_code=400, detail={ @@ -1082,15 +1082,15 @@ if MCP_AVAILABLE: # Locally generated denials (tool/server permission, IP filtering, BYOK) stay at error level # so restriction probing keeps full monitoring visibility; the relayed upstream 401 above is # the only status demoted to info. - verbose_logger.error(f"HTTPException in MCP tool call: {e!s}") + verbose_logger.error(f"HTTPException in MCP tool call: {e}") raise e except Exception as e: - verbose_logger.exception(f"Unexpected error in MCP tool call: {e!s}") + verbose_logger.exception(f"Unexpected error in MCP tool call: {e}") raise HTTPException( status_code=500, detail={ "error": "internal_server_error", - "message": f"An unexpected error occurred: {e!s}", + "message": f"An unexpected error occurred: {e}", }, ) diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index 779cc5861d4..e694c2da7e3 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -1292,5 +1292,5 @@ async def handle_sampling_create_message( verbose_logger.exception("MCP sampling handler failed: %s", e) return ErrorData( code=-1, - message=f"Sampling failed: {e!s}", + message=f"Sampling failed: {e}", ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index fe47c264dfa..a894413019e 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -805,7 +805,7 @@ if MCP_AVAILABLE: } return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta}) except Exception as e: - verbose_logger.exception(f"Error in list_tools endpoint: {e!s}") + verbose_logger.exception(f"Error in list_tools endpoint: {e}") # Return empty list instead of failing completely # This prevents the HTTP stream from failing and allows the client to get a response return [] @@ -1080,26 +1080,26 @@ if MCP_AVAILABLE: isError=True, ) except BlockedPiiEntityError as e: - verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {e!s}") + verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {e}") return CallToolResult( content=[ TextContent( - text=f"Error: Blocked PII entity detected - {e!s}", + text=f"Error: Blocked PII entity detected - {e}", type="text", ) ], isError=True, ) except GuardrailRaisedException as e: - verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {e!s}") + verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {e}") return CallToolResult( - content=[TextContent(text=f"Error: Guardrail violation - {e!s}", type="text")], + content=[TextContent(text=f"Error: Guardrail violation - {e}", type="text")], isError=True, ) except HTTPException as e: - verbose_logger.error(f"HTTPException in MCP tool call: {e!s}") + verbose_logger.error(f"HTTPException in MCP tool call: {e}") return CallToolResult( - content=[TextContent(text=f"Error: {e.detail!s}", type="text")], + content=[TextContent(text=f"Error: {e.detail}", type="text")], isError=True, ) except MCPUpstreamAuthError as e: @@ -1121,7 +1121,7 @@ if MCP_AVAILABLE: except Exception as e: verbose_logger.exception(f"MCP mcp_server_tool_call - error: {e}") return CallToolResult( - content=[TextContent(text=f"Error: {e!s}", type="text")], + content=[TextContent(text=f"Error: {e}", type="text")], isError=True, ) @@ -1173,7 +1173,7 @@ if MCP_AVAILABLE: 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: {e!s}") + verbose_logger.exception(f"Error in list_prompts endpoint: {e}") # Return empty list instead of failing completely # This prevents the HTTP stream from failing and allows the client to get a response return [] @@ -1265,7 +1265,7 @@ if MCP_AVAILABLE: 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: {e!s}") + verbose_logger.exception(f"Error in list_resources endpoint: {e}") return [] finally: if _session_reset_token is not None: @@ -1310,7 +1310,7 @@ if MCP_AVAILABLE: ) return resource_templates except Exception as e: - verbose_logger.exception(f"Error in list_resource_templates endpoint: {e!s}") + verbose_logger.exception(f"Error in list_resource_templates endpoint: {e}") return [] finally: if _session_reset_token is not None: @@ -2036,7 +2036,7 @@ if MCP_AVAILABLE: verbose_logger.debug(f"MCP list_tools: omitting {server.name}; it needs upstream auth") return [], classify_list_exception(e) except Exception as e: - verbose_logger.exception(f"Error getting tools from server {server.name}: {e!s}") + verbose_logger.exception(f"Error getting tools from server {server.name}: {e}") return [], classify_list_exception(e) # Fetch tools from all servers in parallel @@ -2169,7 +2169,7 @@ if MCP_AVAILABLE: 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}: {e!s}") + verbose_logger.exception(f"Error getting prompts from server {server.name}: {e}") # Continue with other servers instead of failing completely verbose_logger.info(f"Successfully fetched {len(all_prompts)} prompts total from all MCP servers") @@ -2221,7 +2221,7 @@ if MCP_AVAILABLE: 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}: {e!s}") + verbose_logger.exception(f"Error getting resources from server {server.name}: {e}") verbose_logger.info(f"Successfully fetched {len(all_resources)} resources total from all MCP servers") @@ -2359,7 +2359,7 @@ if MCP_AVAILABLE: verbose_logger.debug(f"Successfully fetched {len(listing.tools)} tools from managed MCP servers") return listing except Exception as e: - verbose_logger.exception(f"Error getting tools from managed MCP servers: {e!s}") + verbose_logger.exception(f"Error getting tools from managed MCP servers: {e}") # Continue with an empty listing instead of failing completely return AggregateToolListing(tools=[], outcomes={}) @@ -2398,7 +2398,7 @@ if MCP_AVAILABLE: ) 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: {e!s}") + verbose_logger.exception(f"Error getting tools from managed MCP servers: {e}") # Continue with empty managed tools list instead of failing completely return managed_prompts @@ -2428,7 +2428,7 @@ if MCP_AVAILABLE: ) 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: {e!s}") + verbose_logger.exception(f"Error getting resources from managed MCP servers: {e}") return managed_resources @@ -3335,8 +3335,8 @@ if MCP_AVAILABLE: result = tool.handler(**arguments) return [TextContent(text=str(result), type="text")] except Exception as e: - verbose_logger.exception(f"Error executing local tool {name}: {e!s}") - return [TextContent(text=f"Error: {e!s}", type="text")] + verbose_logger.exception(f"Error executing local tool {name}: {e}") + return [TextContent(text=f"Error: {e}", type="text")] def _get_mcp_servers_in_path(path: str) -> list[str] | None: """ diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index 62733edf378..a321c40b9e2 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -55,7 +55,7 @@ async def list_mcp_toolsets( rows = await MCPToolsetRepository(prisma_client).table.find_many(where=where) return [_toolset_from_row(r) for r in rows] except Exception as e: - verbose_proxy_logger.warning(f"litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - {e!s}") + verbose_proxy_logger.warning(f"litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - {e}") return [] diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3ccf3ea9952..268b08d0f2e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -543,6 +543,7 @@ class LiteLLMRoutes(enum.Enum): "/v2/team/list", "/organization/list", "/team/available", + "/team/metadata_schema", "/user/info", "/v2/user/info", "/model/info", @@ -609,6 +610,7 @@ class LiteLLMRoutes(enum.Enum): "/team/block", "/team/unblock", "/team/available", + "/team/metadata_schema", "/team/permissions_list", "/team/permissions_update", "/team/permissions_bulk_update", diff --git a/litellm/proxy/a2a/discovery.py b/litellm/proxy/a2a/discovery.py index 66c661972a8..ffd85331bfd 100644 --- a/litellm/proxy/a2a/discovery.py +++ b/litellm/proxy/a2a/discovery.py @@ -122,11 +122,11 @@ async def fetch_well_known_card( # dict so production (``user_url_validation=True``) doesn't 500. response = await async_safe_get(client, url, headers=headers or {}) except SSRFError as exc: - last_error = f"{url}: {exc!s}" + last_error = f"{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}" + last_error = f"{url}: {exc}" verbose_proxy_logger.debug("A2A discovery failed for %s: %s", url, exc) continue @@ -138,7 +138,7 @@ async def fetch_well_known_card( try: card = response.json() except Exception as exc: - last_error = f"{url}: invalid JSON ({exc!s})" + last_error = f"{url}: invalid JSON ({exc})" continue if not isinstance(card, dict): diff --git a/litellm/proxy/a2a/endpoints.py b/litellm/proxy/a2a/endpoints.py index bcc07629ab1..cd5024a8456 100644 --- a/litellm/proxy/a2a/endpoints.py +++ b/litellm/proxy/a2a/endpoints.py @@ -104,7 +104,7 @@ async def discover_agent_card( raise HTTPException(status_code=400, detail=str(exc)) except Exception as exc: verbose_proxy_logger.exception("Unexpected error during A2A discovery: %s", exc) - raise HTTPException(status_code=500, detail=f"Discovery failed: {exc!s}") + raise HTTPException(status_code=500, detail=f"Discovery failed: {exc}") return JSONResponse( content={"url": request.url, "agent_card": card}, diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 79808c06daa..6d48b31658b 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -441,7 +441,7 @@ async def _handle_stream_message( "message": getattr( proxy_exc, "message", - f"Streaming error: {proxy_exc!s}", + f"Streaming error: {proxy_exc}", ), }, } @@ -491,7 +491,7 @@ async def _handle_stream_message( "id": request_id, "error": { "code": -32603, - "message": f"Streaming error: {e!s}", + "message": f"Streaming error: {e}", }, } ) @@ -974,4 +974,4 @@ async def invoke_agent_a2a( ) except Exception: pass - return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e!s}", 500) + return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 60367178c7f..5ae992648d7 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -271,7 +271,7 @@ class AgentRegistry: 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: {e!s}") + raise Exception(f"Error adding agent to DB: {e}") async def delete_agent_from_db(self, agent_id: str, prisma_client: PrismaClient) -> Mapping[str, object]: """ @@ -281,7 +281,7 @@ class AgentRegistry: deleted_agent = await agents_table(prisma_client).delete(where={"agent_id": agent_id}) return dict(deleted_agent) except Exception as e: - raise Exception(f"Error deleting agent from DB: {e!s}") + raise Exception(f"Error deleting agent from DB: {e}") async def patch_agent_in_db( self, @@ -363,7 +363,7 @@ class AgentRegistry: 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: {e!s}") + raise Exception(f"Error patching agent in DB: {e}") async def update_agent_in_db( self, @@ -450,7 +450,7 @@ class AgentRegistry: 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: {e!s}") + raise Exception(f"Error updating agent in DB: {e}") @staticmethod async def get_all_agents_from_db( @@ -478,7 +478,7 @@ class AgentRegistry: return agents except Exception as e: - raise Exception(f"Error getting agents from DB: {e!s}") + raise Exception(f"Error getting agents from DB: {e}") def get_agent_by_id( self, @@ -494,7 +494,7 @@ class AgentRegistry: return None except Exception as e: - raise Exception(f"Error getting agent from DB: {e!s}") + raise Exception(f"Error getting agent from DB: {e}") def get_agent_by_name(self, agent_name: str) -> AgentResponse | None: """ @@ -507,7 +507,7 @@ class AgentRegistry: return None except Exception as e: - raise Exception(f"Error getting agent from DB: {e!s}") + raise Exception(f"Error getting agent from DB: {e}") global_agent_registry = AgentRegistry() diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 8acff11b009..6999228c83d 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -59,7 +59,7 @@ class AgentRequestHandler: return list(set(allowed_agents)) except Exception as e: - verbose_logger.warning(f"Failed to get allowed agents: {e!s}") + verbose_logger.warning(f"Failed to get allowed agents: {e}") return [] @staticmethod @@ -179,7 +179,7 @@ class AgentRequestHandler: return list(set(all_agents)) except Exception as e: - verbose_logger.warning(f"Failed to get allowed agents for key: {e!s}") + verbose_logger.warning(f"Failed to get allowed agents for key: {e}") return [] @staticmethod @@ -255,7 +255,7 @@ 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: {e!s}") + verbose_logger.warning(f"Failed to get allowed agents for team: {e}") return [] @staticmethod @@ -310,7 +310,7 @@ class AgentRequestHandler: return list(agent_ids) except Exception as e: - verbose_logger.warning(f"Failed to get agents from access groups: {e!s}") + verbose_logger.warning(f"Failed to get agents from access groups: {e}") return [] @staticmethod @@ -369,7 +369,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: {e!s}") + verbose_logger.warning(f"Failed to get agent access groups for key: {e}") return [] @staticmethod @@ -412,5 +412,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: {e!s}") + verbose_logger.warning(f"Failed to get agent access groups for team: {e}") return [] diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 1efbdeb0132..db5341dbe5a 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -316,8 +316,8 @@ async def get_agents( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.agent_endpoints.get_agents(): Exception occurred - {e!s}") - raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e!s}"}) + verbose_proxy_logger.exception(f"litellm.proxy.agent_endpoints.get_agents(): Exception occurred - {e}") + raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) #### CRUD ENDPOINTS FOR AGENTS #### 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 03fdede0cf4..bf797b92850 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 @@ -121,7 +121,7 @@ async def get_marketplace(): verbose_proxy_logger.exception(f"Error generating marketplace: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to generate marketplace: {e!s}"}, + detail={"error": f"Failed to generate marketplace: {e}"}, ) @@ -304,7 +304,7 @@ async def register_plugin( verbose_proxy_logger.exception(f"Error registering plugin: {e}") raise HTTPException( status_code=500, - detail={"error": f"Registration failed: {e!s}"}, + detail={"error": f"Registration failed: {e}"}, ) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 4b566caf2b1..5535928dfac 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -142,7 +142,7 @@ async def anthropic_response( _usage = _blocked_response_usage(e.original_response) _anthropic_response = AnthropicMessagesResponse( - id=f"msg_{uuid.uuid4()!s}", + id=f"msg_{uuid.uuid4()}", type="message", role="assistant", content=[{"type": "text", "text": e.message}], @@ -189,7 +189,7 @@ async def anthropic_response( 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(f"litellm.proxy.proxy_server.anthropic_response(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.anthropic_response(): Exception occured - {e}") # Extract model_id from request metadata (same as success path) litellm_metadata = data.get("litellm_metadata", {}) or {} @@ -209,7 +209,7 @@ async def anthropic_response( litellm_logging_obj=None, ) - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -301,8 +301,8 @@ async def count_tokens( detail=detail, ) except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - {e!s}") - raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e!s}"}) + verbose_proxy_logger.exception(f"litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - {e}") + raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) @router.post( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f876b303510..6b1d845ec95 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -246,7 +246,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None 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}: {e!s}, assuming it has cost") + verbose_proxy_logger.debug(f"Error checking cost for model {model_name}: {e}, assuming it has cost") return False # All models checked have zero cost @@ -973,7 +973,7 @@ async def get_default_end_user_budget( return _budget_obj except Exception as e: - verbose_proxy_logger.error(f"Error fetching default end user budget: {e!s}") + verbose_proxy_logger.error(f"Error fetching default end user budget: {e}") return None @@ -1620,6 +1620,34 @@ async def _get_fuzzy_user_object( return response +async def _backfill_null_user_email( + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + user_row: LiteLLM_UserTable, + user_email: str | None, +) -> LiteLLM_UserTable: + if user_email is None or user_row.user_email is not None or prisma_client is None: + return user_row + + user_repo = UserRepository(prisma_client) + await user_repo.backfill_null_user_email( + user_id=user_row.user_id, + user_email=user_email, + ) + db_row = await user_repo.find_by_id(user_row.user_id) + if db_row is None: + return user_row + email_update = {"user_email": db_row.user_email} # mutable-ok: model_copy update payload is dict-shaped + updated_row = user_row.model_copy(update=email_update) + await user_api_key_cache.async_set_cache( + key=user_row.user_id, + value=updated_row, + model_type=LiteLLM_UserTable, + ttl=get_management_object_ttl(user_api_key_cache), + ) + return updated_row + + @log_db_metrics async def get_user_object( user_id: str | None, @@ -1648,7 +1676,12 @@ async def get_user_object( model_type=LiteLLM_UserTable, ) if cached_user_obj is not None: - return cached_user_obj + return await _backfill_null_user_email( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_row=cached_user_obj, + user_email=user_email, + ) # else, check db if prisma_client is None: raise Exception("No db connected") @@ -1732,6 +1765,12 @@ async def get_user_object( response.organization_memberships = _dumped_memberships _response = LiteLLM_UserTable.model_validate(dict(response)) + _response = await _backfill_null_user_email( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_row=_response, + user_email=user_email, + ) response_dict = _response.model_dump() # save the user object to cache @@ -2238,7 +2277,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}': {e!s}"}, + detail={"error": f"Error looking up team by alias '{team_alias}': {e}"}, ) @@ -2324,7 +2363,7 @@ async def get_org_object_by_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}': {e!s}"}, + detail={"error": f"Error looking up organization by alias '{org_alias}': {e}"}, ) diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index e96d3db65ff..681647814e7 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -95,7 +95,7 @@ class UserAPIKeyAuthExceptionHandler: use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False), ) verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {e!s}\nRequester IP Address:{requester_ip}", + f"litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {e}\nRequester IP Address:{requester_ip}", extra={"requester_ip": requester_ip}, ) @@ -150,7 +150,7 @@ class UserAPIKeyAuthExceptionHandler: ) if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_401_UNAUTHORIZED), diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index dfa5b22d285..a03ed13180c 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -582,7 +582,7 @@ def route_in_additonal_public_routes(current_route: str): return False except Exception as e: - verbose_proxy_logger.error(f"route_in_additonal_public_routes: {e!s}") + verbose_proxy_logger.error(f"route_in_additonal_public_routes: {e}") return False @@ -619,7 +619,7 @@ def get_request_route(request: Request) -> str: return raw_path except Exception as e: verbose_proxy_logger.debug( - f"error on get_request_route: {e!s}, defaulting to request.url.path={request.url.path}" + f"error on get_request_route: {e}, defaulting to request.url.path={request.url.path}" ) return str(request.url.path) @@ -639,7 +639,7 @@ def get_request_route_template(request: Request) -> str | None: template = getattr(route, "path", None) return template if isinstance(template, str) and template else None except Exception as e: - verbose_proxy_logger.debug(f"error on get_request_route_template: {e!s}") + verbose_proxy_logger.debug(f"error on get_request_route_template: {e}") return None diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index b1ccdc87830..cf4b47e3180 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -777,8 +777,8 @@ class JWTHandler: return userinfo except Exception as e: - verbose_proxy_logger.error(f"Error fetching OIDC UserInfo: {e!s}") - raise Exception(f"Failed to fetch OIDC UserInfo: {e!s}") + verbose_proxy_logger.error(f"Error fetching OIDC UserInfo: {e}") + raise Exception(f"Failed to fetch OIDC UserInfo: {e}") _unscoped_jwt_warning_emitted = False @@ -987,7 +987,7 @@ class JWTHandler: code=status.HTTP_401_UNAUTHORIZED, ) except Exception as e: - raise Exception(f"Validation fails: {e!s}") + raise Exception(f"Validation fails: {e}") return self._apply_issuer_claim_mappings( token=payload, @@ -1032,7 +1032,7 @@ class JWTHandler: code=status.HTTP_401_UNAUTHORIZED, ) except Exception as e: - raise Exception(f"Validation fails: {e!s}") + raise Exception(f"Validation fails: {e}") raise Exception("Invalid JWT Submitted") diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index a25f3e58d2c..1f61ef7ea28 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -48,7 +48,7 @@ class LicenseCheck: else: self.public_key = None except Exception as e: - verbose_proxy_logger.error(f"Error reading public key: {e!s}") + verbose_proxy_logger.error(f"Error reading public key: {e}") def _verify(self, license_str: str) -> bool: verbose_proxy_logger.debug( @@ -84,7 +84,7 @@ class LicenseCheck: return premium except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.auth.litellm_license.py::_verify - Unable to verify License={license_str} via api. - {e!s}" + f"litellm.proxy.auth.litellm_license.py::_verify - Unable to verify License={license_str} via api. - {e}" ) return False @@ -187,6 +187,6 @@ class LicenseCheck: except Exception as e: verbose_proxy_logger.debug( - f"litellm.proxy.auth.litellm_license.py::verify_license_without_api_request - Unable to verify License locally. - {e!s}" + f"litellm.proxy.auth.litellm_license.py::verify_license_without_api_request - Unable to verify License locally. - {e}" ) return False diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 72450453174..286837c8909 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1232,6 +1232,26 @@ async def _user_api_key_auth_builder( valid_token.jwt_claims = jwt_claims do_standard_jwt_auth = False # Fall through to virtual key checks + if valid_token.user_id is not None and valid_token.user_email is None: + mapped_claims = jwt_claims or {} # mutable-ok: empty-dict fallback for the None-claims case + mapped_user_email = jwt_handler.get_user_email(token=mapped_claims, default_value=None) + mapped_jwt_user_id = jwt_handler.get_user_id(token=mapped_claims, default_value=None) + if mapped_user_email is not None and mapped_jwt_user_id == valid_token.user_id: + try: + mapped_user_obj = await get_user_object( + user_id=valid_token.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + user_email=mapped_user_email, + ) + except Exception as e: + verbose_proxy_logger.debug(f"JWT mapped-key user_email backfill skipped: {e}") + else: + if mapped_user_obj is not None: + valid_token.user_email = mapped_user_obj.user_email elif isinstance(resolve_result, _PendingAutoRegister): # Run full JWT policy (RBAC, scope, custom_validate, # email-domain) via auth_builder, then create the key @@ -1481,7 +1501,7 @@ async def _user_api_key_auth_builder( except Exception as e: if isinstance(e, litellm.BudgetExceededError): raise e - verbose_proxy_logger.debug(f"Unable to find user in db. Error - {e!s}") + verbose_proxy_logger.debug(f"Unable to find user in db. Error - {e}") ### CHECK IF ADMIN ### # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead @@ -1729,7 +1749,7 @@ async def _user_api_key_auth_builder( ) except Exception as e: verbose_logger.debug( - f"litellm.proxy.auth.user_api_key_auth.py::user_api_key_auth() - Unable to get user from db/cache. Setting user_obj to None. Exception received - {e!s}" + f"litellm.proxy.auth.user_api_key_auth.py::user_api_key_auth() - Unable to get user from db/cache. Setting user_obj to None. Exception received - {e}" ) user_obj = None @@ -2754,7 +2774,7 @@ async def _lookup_end_user_and_apply_budget( except Exception as e: if isinstance(e, litellm.BudgetExceededError): raise e - verbose_proxy_logger.debug(f"Unable to find user in db. Error - {e!s}") + verbose_proxy_logger.debug(f"Unable to find user in db. Error - {e}") return valid_token, end_user_object diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 8b2a437009d..0c2764db33f 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -340,7 +340,7 @@ async def create_batch( 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(f"litellm.proxy.proxy_server.create_batch(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.create_batch(): Exception occured - {e}") raise handle_exception_on_proxy(e) @@ -592,7 +592,7 @@ async def retrieve_batch( 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(f"litellm.proxy.proxy_server.retrieve_batch(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.retrieve_batch(): Exception occured - {e}") raise handle_exception_on_proxy(e) @@ -773,7 +773,7 @@ async def list_batches( original_exception=e, request_data={"after": after, "limit": limit}, ) - verbose_proxy_logger.error(f"litellm.proxy.proxy_server.retrieve_batch(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.retrieve_batch(): Exception occured - {e}") raise handle_exception_on_proxy(e) @@ -982,7 +982,7 @@ async def cancel_batch( 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(f"litellm.proxy.proxy_server.create_batch(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.create_batch(): Exception occured - {e}") raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py index e64f3e9e7e3..50b2f63e18a 100644 --- a/litellm/proxy/caching_routes.py +++ b/litellm/proxy/caching_routes.py @@ -43,7 +43,7 @@ def _extract_cache_params() -> dict[str, Any]: 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: {e!s}") + verbose_proxy_logger.debug(f"Error extracting cache params: {e}") return {} @@ -158,7 +158,7 @@ async def cache_delete(request: Request): except Exception as e: raise HTTPException( status_code=500, - detail=f"Cache Delete Failed({e!s})", + detail=f"Cache Delete Failed({e})", ) @@ -173,7 +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): {e!s}") + verbose_proxy_logger.warning(f"CLIENT LIST command failed (likely restricted on managed Redis): {e}") return ["CLIENT LIST command not available on this Redis instance"], -1 @@ -209,7 +209,7 @@ async def cache_redis_info(): except Exception as e: raise HTTPException( status_code=503, - detail=f"Service Unhealthy ({e!s})", + detail=f"Service Unhealthy ({e})", ) @@ -245,5 +245,5 @@ async def cache_flushall(): except Exception as e: raise HTTPException( status_code=503, - detail=f"Service Unhealthy ({e!s})", + detail=f"Service Unhealthy ({e})", ) diff --git a/litellm/proxy/client/cli/commands/chat.py b/litellm/proxy/client/cli/commands/chat.py index f0c91be686d..3e86c79a90b 100644 --- a/litellm/proxy/client/cli/commands/chat.py +++ b/litellm/proxy/client/cli/commands/chat.py @@ -386,5 +386,5 @@ def _stream_response( console.print(f"[red]{e.response.text}[/red]") return None except Exception as e: - console.print(f"\n[red]Error: {e!s}[/red]") + console.print(f"\n[red]Error: {e}[/red]") return None diff --git a/litellm/proxy/client/cli/commands/credentials.py b/litellm/proxy/client/cli/commands/credentials.py index 8187f811778..cdfa4c5cd69 100644 --- a/litellm/proxy/client/cli/commands/credentials.py +++ b/litellm/proxy/client/cli/commands/credentials.py @@ -71,7 +71,7 @@ def create(ctx: click.Context, credential_name: str, info: str, values: str): credential_info = json.loads(info) credential_values = json.loads(values) except json.JSONDecodeError as e: - raise click.BadParameter(f"Invalid JSON: {e!s}") + raise click.BadParameter(f"Invalid JSON: {e}") try: response = client.create(credential_name, credential_info, credential_values) diff --git a/litellm/proxy/client/cli/commands/keys.py b/litellm/proxy/client/cli/commands/keys.py index ec5dca25518..8ebed1749f4 100644 --- a/litellm/proxy/client/cli/commands/keys.py +++ b/litellm/proxy/client/cli/commands/keys.py @@ -122,7 +122,7 @@ def generate( aliases_dict = json.loads(aliases) if aliases else None config_dict = json.loads(config) if config else None except json.JSONDecodeError as e: - raise click.BadParameter(f"Invalid JSON: {e!s}") + raise click.BadParameter(f"Invalid JSON: {e}") try: response = client.generate( models=models_list, @@ -316,7 +316,7 @@ def _import_keys_to_destination( except Exception as e: failed_count += 1 key_alias = key.get("key_alias", "N/A") - click.echo(f"Failed to import key {key_alias}: {e!s}", err=True) + click.echo(f"Failed to import key {key_alias}: {e}", err=True) return imported_count, failed_count @@ -389,5 +389,5 @@ def import_keys( click.echo(e.response.text, err=True) raise click.Abort() except Exception as e: - click.echo(f"Error: {e!s}", err=True) + click.echo(f"Error: {e}", err=True) raise click.Abort() diff --git a/litellm/proxy/client/cli/commands/teams.py b/litellm/proxy/client/cli/commands/teams.py index 442ac40a775..2d88e4bbce2 100644 --- a/litellm/proxy/client/cli/commands/teams.py +++ b/litellm/proxy/client/cli/commands/teams.py @@ -76,7 +76,7 @@ def list(ctx: click.Context): click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) raise click.Abort() except Exception as e: - click.echo(f"Error: {e!s}", err=True) + click.echo(f"Error: {e}", err=True) raise click.Abort() @@ -99,7 +99,7 @@ def available(ctx: click.Context): error_body = e.response.json() click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) except Exception as e: - click.echo(f"Error: {e!s}", err=True) + click.echo(f"Error: {e}", err=True) raise click.Abort() @@ -158,5 +158,5 @@ def assign_key(ctx: click.Context, team_id: str | None): click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) raise click.Abort() except Exception as e: - click.echo(f"Error: {e!s}", err=True) + click.echo(f"Error: {e}", err=True) raise click.Abort() diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4d4f459a080..bc3bcd233f0 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -43,6 +43,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import can_key_call_resolved_model from litellm.proxy.auth.auth_utils import check_response_size_is_safe from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, @@ -53,6 +54,7 @@ from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.types.guardrails import GuardrailEventHooks from litellm.types.router import RouterRateLimitError from litellm.types.utils import ServerToolUse @@ -384,6 +386,39 @@ async def _authorize_response_file_search_vector_stores( ) +async def _resolve_per_request_model_group_alias( + requested_model: object, + router_settings: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router, +) -> str | None: + """ + Resolve ``router_settings.model_group_alias`` coming from a key or team. + + The Router only ever resolves aliases from its own instance attribute, which + holds the global config map and is shared across requests, so a per-request + map has to be applied here instead of being forwarded to the Router. + + Model access was authorized against the requested group, so the target is + authorized in its own right before the rewrite; a key that may not call the + target gets the usual 403 rather than being quietly served it. + + Returns the target model group, or None when no alias applies. + """ + if not isinstance(requested_model, str): + return None + target = resolve_model_group_alias(router_settings.get("model_group_alias"), requested_model) + if target is None or target == requested_model: + return None + await can_key_call_resolved_model( + model=target, + llm_model_list=llm_router.model_list, + valid_token=user_api_key_dict, + llm_router=llm_router, + ) + return target + + async def _parse_event_data_for_error(event_line: str | bytes) -> int | None: """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 @@ -908,7 +943,7 @@ def _log_llm_api_exception(e: Exception) -> None: "litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, upstream LLM request cancelled" ) return - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {e}") async def _cancel_llm_call_on_client_disconnect( @@ -1285,6 +1320,35 @@ class ProxyBaseLLMRequestProcessing: ): self.data["model"] = user_api_key_dict.aliases[self.data["model"]] + # Apply hierarchical router_settings (Key > Team) + # Global router_settings are already on the Router object itself. + # This sits with the other alias rewrites, and ahead of the guardrail + # merge and the pre-call hooks, so everything that keys off the model + # group -- model-level guardrails, per-model budgets and rate limits, + # the logging object -- sees the group that will actually serve. + if llm_router is not None and proxy_config is not None: + from litellm.proxy.proxy_server import prisma_client + + router_settings = await proxy_config._get_hierarchical_router_settings( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + + # If router_settings found (from key or team), apply them + # Pass settings as per-request overrides instead of creating a new Router + # This avoids expensive Router instantiation on each request + if router_settings is not None: + self.data["router_settings_override"] = router_settings + alias_target = await _resolve_per_request_model_group_alias( + requested_model=self.data.get("model"), + router_settings=router_settings, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + ) + if alias_target is not None: + self.data["model"] = alias_target + 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( @@ -1339,23 +1403,6 @@ class ProxyBaseLLMRequestProcessing: call_type=route_type, # type: ignore ) - # Apply hierarchical router_settings (Key > Team) - # Global router_settings are already on the Router object itself. - if llm_router is not None and proxy_config is not None: - from litellm.proxy.proxy_server import prisma_client - - router_settings = await proxy_config._get_hierarchical_router_settings( - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - ) - - # If router_settings found (from key or team), apply them - # Pass settings as per-request overrides instead of creating a new Router - # This avoids expensive Router instantiation on each request - if router_settings is not None: - self.data["router_settings_override"] = router_settings - if "messages" in self.data and self.data["messages"]: logging_obj.update_messages(self.data["messages"]) @@ -2696,7 +2743,7 @@ class ProxyBaseLLMRequestProcessing: status_code=http_status_error.response.status_code, detail={"error": error_text}, ) - error_msg = f"{e!s}" + error_msg = f"{e}" # Check for AttributeError in the exception chain. # The AttributeError may be wrapped in multiple layers # (e.g. AttributeError -> OpenAIException -> APIConnectionError), @@ -2898,7 +2945,7 @@ class ProxyBaseLLMRequestProcessing: raise except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.async_data_generator(): Exception occured - {e!s}" + f"litellm.proxy.proxy_server.async_data_generator(): Exception occured - {e}" ) transformed_exception = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -2914,7 +2961,7 @@ class ProxyBaseLLMRequestProcessing: if isinstance(e, HTTPException): raise e error_traceback = _redact_string(traceback.format_exc()) - error_msg = f"{e!s}\n\n{error_traceback}" + error_msg = f"{e}\n\n{error_traceback}" proxy_exception = ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index 55d6f083fdd..a884eab462a 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -273,7 +273,7 @@ class CustomOpenAPISpec: except Exception as e: # If schema addition fails, continue without it - verbose_proxy_logger.debug(f"Failed to add {operation_name} request schema: {e!s}") + verbose_proxy_logger.debug(f"Failed to add {operation_name} request schema: {e}") return openapi_schema @@ -302,7 +302,7 @@ class CustomOpenAPISpec: operation_name="chat completion", ) except ImportError as e: - verbose_proxy_logger.debug(f"Failed to import ProxyChatCompletionRequest: {e!s}") + verbose_proxy_logger.debug(f"Failed to import ProxyChatCompletionRequest: {e}") return openapi_schema @staticmethod @@ -328,7 +328,7 @@ class CustomOpenAPISpec: operation_name="embedding", ) except ImportError as e: - verbose_proxy_logger.debug(f"Failed to import EmbeddingRequest: {e!s}") + verbose_proxy_logger.debug(f"Failed to import EmbeddingRequest: {e}") return openapi_schema @staticmethod @@ -356,7 +356,7 @@ class CustomOpenAPISpec: operation_name="responses API", ) except ImportError as e: - verbose_proxy_logger.debug(f"Failed to import ResponsesAPIRequestParams: {e!s}") + verbose_proxy_logger.debug(f"Failed to import ResponsesAPIRequestParams: {e}") return openapi_schema @staticmethod diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 7d3150a3e72..7d2b303a7ca 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -653,7 +653,7 @@ async def configure_gc_thresholds_endpoint( ) 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: {e!s}") + raise HTTPException(status_code=500, detail=f"Failed to set GC thresholds: {e}") # Get current object count to show immediate impact current_count = gc.get_count()[0] @@ -783,4 +783,4 @@ def init_verbose_loggers(): except Exception as e: import logging - logging.warning(f"Failed to init verbose loggers: {e!s}") + logging.warning(f"Failed to init verbose loggers: {e}") diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index b7b8bfd1eea..651e59ef959 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -145,7 +145,7 @@ def decrypt_value_helper( # if it's not str - do not decrypt it, return the value return value except Exception as e: - error_message = f"Error decrypting value for key: {key}, Did your master_key/salt key change recently? \nError: {e!s}\nSet permanent salt key - https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" + error_message = f"Error decrypting value for key: {key}, Did your master_key/salt key change recently? \nError: {e}\nSet permanent salt key - https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" if exception_type == "debug": verbose_proxy_logger.debug(error_message) return value if return_original_value else None diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 0dd910e1901..67212539cc4 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -98,9 +98,9 @@ async def _read_request_body(request: Request | None) -> dict: # Above the configured size, skip the repair and raise the 400 now. 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: {e!s}") + verbose_proxy_logger.error(f"Invalid JSON payload received: {e}") raise ProxyException( - message=f"Invalid JSON payload: {e!s}", + message=f"Invalid JSON payload: {e}", type="invalid_request_error", param="request_body", code=status.HTTP_400_BAD_REQUEST, @@ -120,9 +120,9 @@ async def _read_request_body(request: Request | None) -> dict: parsed_body = json.loads(body_str) except json.JSONDecodeError: # If both orjson and json.loads fail, throw a proper error - verbose_proxy_logger.error(f"Invalid JSON payload received: {e!s}") + verbose_proxy_logger.error(f"Invalid JSON payload received: {e}") raise ProxyException( - message=f"Invalid JSON payload: {e!s}", + message=f"Invalid JSON payload: {e}", type="invalid_request_error", param="request_body", code=status.HTTP_400_BAD_REQUEST, @@ -134,7 +134,7 @@ async def _read_request_body(request: Request | None) -> dict: except (json.JSONDecodeError, orjson.JSONDecodeError, ProxyException) as e: # Re-raise ProxyException as-is - verbose_proxy_logger.error(f"Invalid JSON payload received: {e!s}") + verbose_proxy_logger.error(f"Invalid JSON payload received: {e}") raise except Exception as e: # Catch unexpected errors to avoid crashes @@ -426,7 +426,7 @@ def extract_nested_form_metadata(form_data: dict[str, Any], prefix: str = "litel 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}': {e!s}") + verbose_proxy_logger.error(f"Error parsing metadata key '{key}': {e}") continue return metadata diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 225f7cfdf6c..56aedb76590 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -34,9 +34,9 @@ def get_file_contents_from_s3(bucket_name, object_key): except ImportError as e: # this is most likely if a user is not using the litellm docker container - verbose_proxy_logger.error(f"ImportError: {e!s}") + verbose_proxy_logger.error(f"ImportError: {e}") except Exception as e: - verbose_proxy_logger.error(f"Error retrieving file contents: {e!s}") + verbose_proxy_logger.error(f"Error retrieving file contents: {e}") return None @@ -57,7 +57,7 @@ async def get_config_file_contents_from_gcs(bucket_name, object_key): return config except Exception as e: - verbose_proxy_logger.error(f"Error retrieving file contents: {e!s}") + verbose_proxy_logger.error(f"Error retrieving file contents: {e}") return None @@ -111,10 +111,10 @@ def download_python_file_from_s3( return True except ImportError as e: - verbose_proxy_logger.error(f"ImportError: {e!s}") + verbose_proxy_logger.error(f"ImportError: {e}") return False except Exception as e: - verbose_proxy_logger.exception(f"Error downloading Python file: {e!s}") + verbose_proxy_logger.exception(f"Error downloading Python file: {e}") return False @@ -158,7 +158,7 @@ async def download_python_file_from_gcs( return True except Exception as e: - verbose_proxy_logger.exception(f"Error downloading Python file from GCS: {e!s}") + verbose_proxy_logger.exception(f"Error downloading Python file from GCS: {e}") return False 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 b5ff2ffa9cc..d7c70bdb20c 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -68,7 +68,7 @@ class SpendLogCleanup: return True except ValueError as e: verbose_proxy_logger.warning( - f"Invalid maximum_spend_logs_retention_period value: {retention_setting}, error: {e!s}" + f"Invalid maximum_spend_logs_retention_period value: {retention_setting}, error: {e}" ) return False diff --git a/litellm/proxy/example_config_yaml/custom_team_metadata_validate.py b/litellm/proxy/example_config_yaml/custom_team_metadata_validate.py new file mode 100644 index 00000000000..1ca21bc3104 --- /dev/null +++ b/litellm/proxy/example_config_yaml/custom_team_metadata_validate.py @@ -0,0 +1,40 @@ +"""Example validator for `general_settings.custom_team_metadata_validate`. + +Wire it up in the proxy config: + +```yaml +general_settings: + custom_team_metadata_validate: custom_team_metadata_validate.validate_team_metadata + team_metadata_validation_timeout: 5 + team_metadata_validation_error_message: "Validation service unavailable, contact your admin." +``` + +Return `valid=False` with an `error_message` to reject the write with that +message (HTTP 400). Raise any exception (for example, when the upstream +validation service is unreachable) to fail closed with the generic +`team_metadata_validation_error_message` (HTTP 503). +""" + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataValidationPayload, + TeamMetadataValidationResult, +) + +VALID_COST_CENTERS = frozenset({"CC-1001", "CC-1002", "CC-2001"}) + + +async def validate_team_metadata( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + cost_center = payload.metadata.get("cost_center") + if cost_center is None: + return TeamMetadataValidationResult( + valid=False, + error_message="Team metadata must include a cost_center. Contact the FinOps team.", + ) + if cost_center not in VALID_COST_CENTERS: + return TeamMetadataValidationResult( + valid=False, + error_message=f"Cost center {cost_center} is not recognized. Contact the FinOps team.", + ) + return TeamMetadataValidationResult(valid=True) diff --git a/litellm/proxy/example_config_yaml/store_model_db_config.yaml b/litellm/proxy/example_config_yaml/store_model_db_config.yaml index 5b77a53b4b7..3f19ab39a02 100644 --- a/litellm/proxy/example_config_yaml/store_model_db_config.yaml +++ b/litellm/proxy/example_config_yaml/store_model_db_config.yaml @@ -7,4 +7,7 @@ model_list: general_settings: store_model_in_db: true + custom_team_metadata_validate: team_metadata_validator_e2e.validate_team_metadata + team_metadata_validation_timeout: 5 + team_metadata_validation_error_message: "Cost center validation is unavailable right now; the team was not saved. Contact FinOps." diff --git a/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py b/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py new file mode 100644 index 00000000000..23390ffec1c --- /dev/null +++ b/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py @@ -0,0 +1,107 @@ +"""Dispatching team metadata validator for the store_model_in_db e2e suite. + +The suite runs one proxy with one config, so a single registered validator +dispatches to one of three independent implementations chosen per request via +the `_e2e_validator_impl` metadata key: + +- `allowlist`: requires `cost_center` and checks it against a static set +- `http`: POSTs the metadata to the cost center service at + `TEAM_METADATA_VALIDATION_SERVICE_URL`; transport errors raise (fail closed) +- `http_down`: like `http` but targets a closed port, proving the 503 path +- `immutable`: requires `cost_center` and forbids changing it once set + +A request whose metadata carries no `_e2e_validator_impl` key is accepted +untouched, so the rest of the suite's team operations are unaffected. An +unknown impl value raises, which the proxy converts to the fail-closed 503. +""" + +import os + +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataValidationPayload, + TeamMetadataValidationResult, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + +ALLOWED_COST_CENTERS = frozenset({"CC-1001", "CC-1002"}) +CLOSED_PORT_URL = "http://127.0.0.1:9/validate" + + +async def _validate_allowlist(payload: TeamMetadataValidationPayload) -> TeamMetadataValidationResult: + cost_center = payload.metadata.get("cost_center") + if cost_center is None: + return TeamMetadataValidationResult( + valid=False, + error_message="cost_center is required in team metadata. Contact the FinOps team.", + ) + if cost_center not in ALLOWED_COST_CENTERS: + return TeamMetadataValidationResult( + valid=False, + error_message=f"Cost center {cost_center} is not recognized. Contact the FinOps team.", + ) + return TeamMetadataValidationResult(valid=True) + + +async def _validate_via_http(payload: TeamMetadataValidationPayload, service_url: str) -> TeamMetadataValidationResult: + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + response = await client.post( + service_url, + json={ # mutable-ok: httpx serializes the request body from a plain dict + "operation": payload.operation, + "metadata": payload.metadata, + }, + timeout=2.0, + ) + body = response.json() + if body.get("ok") is True: + return TeamMetadataValidationResult(valid=True) + return TeamMetadataValidationResult( + valid=False, + error_message=body.get("reason", "Rejected by the cost center service."), + ) + + +class _ImmutableCostCenterValidator: + def __init__(self, immutable_key: str = "cost_center") -> None: + self.immutable_key = immutable_key + + async def __call__(self, payload: TeamMetadataValidationPayload) -> TeamMetadataValidationResult: + current = payload.metadata.get(self.immutable_key) + if current is None: + return TeamMetadataValidationResult( + valid=False, + error_message=f"{self.immutable_key} is required in team metadata. Contact the FinOps team.", + ) + if payload.operation == "update" and payload.existing_metadata is not None: + prior = payload.existing_metadata.get(self.immutable_key) + if prior is not None and prior != current: + return TeamMetadataValidationResult( + valid=False, + error_message=( + f"{self.immutable_key} is immutable once set " + f"(stored: {prior}, requested: {current}). Contact the FinOps team." + ), + ) + return TeamMetadataValidationResult(valid=True) + + +_IMMUTABLE_VALIDATOR = _ImmutableCostCenterValidator() + + +async def validate_team_metadata( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + impl = payload.metadata.get("_e2e_validator_impl") + if impl is None: + return TeamMetadataValidationResult(valid=True) + if impl == "allowlist": + return await _validate_allowlist(payload) + if impl == "http": + service_url = os.environ.get("TEAM_METADATA_VALIDATION_SERVICE_URL", "http://localhost:9414/validate") + return await _validate_via_http(payload, service_url) + if impl == "http_down": + return await _validate_via_http(payload, CLOSED_PORT_URL) + if impl == "immutable": + return await _IMMUTABLE_VALIDATOR(payload) + raise ValueError(f"unknown _e2e_validator_impl: {impl}") diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index 778ea729e32..4daab1caf96 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -199,9 +199,7 @@ async def create_fine_tuning_job( 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( - f"litellm.proxy.proxy_server.create_fine_tuning_job(): Exception occurred - {e!s}" - ) + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.create_fine_tuning_job(): Exception occurred - {e}") raise handle_exception_on_proxy(e) @@ -340,7 +338,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( - f"litellm.proxy.proxy_server.retrieve_fine_tuning_job(): Exception occurred - {e!s}" + f"litellm.proxy.proxy_server.retrieve_fine_tuning_job(): Exception occurred - {e}" ) raise handle_exception_on_proxy(e) @@ -468,9 +466,7 @@ async def list_fine_tuning_jobs( 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( - f"litellm.proxy.proxy_server.list_fine_tuning_jobs(): Exception occurred - {e!s}" - ) + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.list_fine_tuning_jobs(): Exception occurred - {e}") raise handle_exception_on_proxy(e) @@ -608,7 +604,5 @@ async def cancel_fine_tuning_job( 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( - f"litellm.proxy.proxy_server.cancel_fine_tuning_job(): Exception occurred - {e!s}" - ) + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.cancel_fine_tuning_job(): Exception occurred - {e}") raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 3d8ed8dc1e2..12373b7fb97 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1420,7 +1420,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: {e!s}") + raise HTTPException(status_code=500, detail=f"Error reading category file: {e}") @router.get( @@ -1452,7 +1452,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: {e!s}") from e + raise HTTPException(status_code=500, detail=f"Error reading major_airlines.json: {e}") from e @router.post( @@ -1540,10 +1540,10 @@ async def validate_blocked_words_file(request: dict[str, str]): "message": f"Valid YAML file with {len(blocked_words_list)} blocked word(s)", } except yaml.YAMLError as e: - return {"valid": False, "error": f"Invalid YAML syntax: {e!s}"} + return {"valid": False, "error": f"Invalid YAML syntax: {e}"} except Exception as e: verbose_proxy_logger.exception("Error validating blocked words file") - return {"valid": False, "error": f"Validation error: {e!s}"} + return {"valid": False, "error": f"Validation error: {e}"} def _get_field_type_from_annotation(field_annotation: Any) -> str: diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index f8fedb22872..5aae14b83e6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -2266,4 +2266,4 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): raise except Exception as e: verbose_proxy_logger.error("Bedrock Guardrail: Failed to apply guardrail: %s", str(e)) - raise Exception(f"Bedrock guardrail failed: {e!s}") + raise Exception(f"Bedrock guardrail failed: {e}") diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index c019d445fd4..43a3671ad97 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -480,10 +480,10 @@ async def http_request( return _http_success_response(e.response) except httpx.RequestError as e: verbose_proxy_logger.warning(f"Custom code http_request error: {e}") - return _http_error_response(f"Request failed: {e!s}") + return _http_error_response(f"Request failed: {e}") except Exception as e: verbose_proxy_logger.warning(f"Custom code http_request unexpected error: {e}") - return _http_error_response(f"Unexpected error: {e!s}") + return _http_error_response(f"Unexpected error: {e}") async def _execute_http_request( diff --git a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py index c7ed4028218..96e7e605349 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py +++ b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py @@ -228,7 +228,7 @@ class DeepKeepGuardrail(CustomGuardrail): **({"http_status_code": http_status_code} if http_status_code else {}), ) verbose_proxy_logger.error("DeepKeep guardrail API error: %s", str(error)) - raise DeepKeepGuardrailAPIError(f"DeepKeep guardrail API failed: {error!s}") + raise DeepKeepGuardrailAPIError(f"DeepKeep guardrail API failed: {error}") @staticmethod def _build_return_inputs( 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 a694ef897ab..f1f37263eae 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 @@ -357,7 +357,7 @@ class GenericGuardrailAPI(CustomGuardrail): **({"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)) - raise Exception(f"Generic Guardrail API failed: {error!s}") + raise Exception(f"Generic Guardrail API failed: {error}") @log_guardrail_information async def apply_guardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index dba82eeb32c..90d131893c6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -299,8 +299,8 @@ class LassoGuardrail(CustomGuardrail): except Exception as e: if isinstance(e, HTTPException): raise e - verbose_proxy_logger.error(f"Error in post-call Lasso masking: {e!s}") - raise LassoGuardrailAPIError(f"Failed to apply post-call masking: {e!s}") + verbose_proxy_logger.error(f"Error in post-call Lasso masking: {e}") + raise LassoGuardrailAPIError(f"Failed to apply post-call masking: {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") @@ -599,7 +599,7 @@ class LassoGuardrail(CustomGuardrail): # Log error with context verbose_proxy_logger.error( - f"Error calling Lasso API: {error!s}", + f"Error calling Lasso API: {error}", extra={ "guardrail_name": getattr(self, "guardrail_name", "unknown"), "message_type": message_type, @@ -620,7 +620,7 @@ class LassoGuardrail(CustomGuardrail): raise LassoGuardrailAPIError(f"API error: {error.response.status_code}") # Generic error handling - raise LassoGuardrailAPIError(f"Failed to verify request safety with Lasso API: {error!s}") + raise LassoGuardrailAPIError(f"Failed to verify request safety with Lasso API: {error}") def _log_masking_applied( self, 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 5650a6b07ac..c6900c38cbf 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 @@ -750,7 +750,7 @@ class ContentFilterGuardrail(CustomGuardrail): 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}: {e!s}") + raise Exception(f"Error loading blocked words file {file_path}: {e}") 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.""" 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 d30de723443..8292f575c74 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py @@ -168,7 +168,7 @@ def get_available_content_categories() -> list[dict[str, str]]: # 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}: {e!s}") + verbose_proxy_logger.warning(f"Failed to load category file {filename}: {e}") continue elif filename.endswith(".json"): # JSON category files (e.g. harm_toxic_abuse.json) - no YAML header, use filename diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index 642c1dcbca8..b2f91083cc0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -163,7 +163,7 @@ class NomaGuardrail(CustomGuardrail): try: asyncio.create_task(coro) except Exception as e: - verbose_proxy_logger.error(f"Failed to create background Noma task: {e!s}") + verbose_proxy_logger.error(f"Failed to create background Noma task: {e}") async def _process_user_message_check( self, @@ -348,7 +348,7 @@ class NomaGuardrail(CustomGuardrail): return "guardrail_failed_to_respond" except Exception as e: - verbose_proxy_logger.error(f"Error determining NOMA guardrail status: {e!s}") + verbose_proxy_logger.error(f"Error determining NOMA guardrail status: {e}") return "guardrail_failed_to_respond" def _should_only_sensitive_data_failed(self, classification_obj: dict) -> bool: @@ -513,7 +513,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: {e!s}") + verbose_proxy_logger.error(f"Noma background user message check failed: {e}") async def _check_llm_response_background( self, @@ -525,7 +525,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: {e!s}") + verbose_proxy_logger.error(f"Noma background response check failed: {e}") async def _handle_verdict_background( self, @@ -547,7 +547,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: {e!s}") + verbose_proxy_logger.error(f"Noma background verdict handling failed: {e}") async def async_pre_call_hook( self, @@ -570,7 +570,7 @@ class NomaGuardrail(CustomGuardrail): try: 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: {e!s}") + verbose_proxy_logger.error(f"Failed to start background Noma pre-call check: {e}") return data try: @@ -594,7 +594,7 @@ class NomaGuardrail(CustomGuardrail): event_type=GuardrailEventHooks.pre_call, ) - verbose_proxy_logger.error(f"Noma pre-call hook failed: {e!s}") + verbose_proxy_logger.error(f"Noma pre-call hook failed: {e}") if self.block_failures: raise @@ -618,7 +618,7 @@ class NomaGuardrail(CustomGuardrail): try: 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: {e!s}") + verbose_proxy_logger.error(f"Failed to start background Noma moderation check: {e}") return data try: @@ -642,7 +642,7 @@ class NomaGuardrail(CustomGuardrail): event_type=GuardrailEventHooks.during_call, ) - verbose_proxy_logger.error(f"Noma moderation hook failed: {e!s}") + verbose_proxy_logger.error(f"Noma moderation hook failed: {e}") if self.block_failures: raise @@ -665,7 +665,7 @@ class NomaGuardrail(CustomGuardrail): 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: {e!s}") + verbose_proxy_logger.error(f"Failed to start background Noma post-call check: {e}") return response try: @@ -689,7 +689,7 @@ class NomaGuardrail(CustomGuardrail): event_type=GuardrailEventHooks.post_call, ) - verbose_proxy_logger.error(f"Noma post-call hook failed: {e!s}") + verbose_proxy_logger.error(f"Noma post-call hook failed: {e}") if self.block_failures: raise return response @@ -828,7 +828,7 @@ class NomaGuardrail(CustomGuardrail): except Exception as e: if self.block_failures: raise - verbose_proxy_logger.error(f"Noma streaming post-call hook failed: {e!s}") + verbose_proxy_logger.error(f"Noma streaming post-call hook failed: {e}") for chunk in all_chunks: yield chunk return diff --git a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py index b86273f754a..37ce84b8e6e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py +++ b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py @@ -118,7 +118,7 @@ class OnyxGuardrail(CustomGuardrail): payload = parsed.get("response", {}) except Exception as e: verbose_proxy_logger.error( - f"Error in converting request_data to ModelResponse: {e!s}", + f"Error in converting request_data to ModelResponse: {e}", extra={ "conversation_id": conversation_id, "input_type": input_type, @@ -133,7 +133,7 @@ class OnyxGuardrail(CustomGuardrail): raise e except Exception as e: verbose_proxy_logger.error( - f"Error in apply_guardrail guard: {e!s}", + f"Error in apply_guardrail guard: {e}", extra={"conversation_id": conversation_id, "input_type": input_type}, ) return inputs diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 947a81c1b79..fe2e87ff661 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -250,7 +250,7 @@ class OvalixGuardrail(CustomGuardrail): 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}", + message=f"Ovalix guardrail error: {e}", should_wrap_with_default_message=False, ) from 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 5d134fd01c2..782ffef61cf 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 @@ -231,7 +231,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): 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: {e!s}") + verbose_proxy_logger.error(f"PANW Prisma AIRS: Error extracting response text: {e}") return "" async def _call_panw_api( @@ -433,7 +433,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): } except httpx.TimeoutException as e: - verbose_proxy_logger.error(f"PANW Prisma AIRS: Timeout error: {e!s}") + verbose_proxy_logger.error(f"PANW Prisma AIRS: Timeout error: {e}") return { "action": "block", "category": "timeout_error", @@ -441,7 +441,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): } except httpx.RequestError as e: - verbose_proxy_logger.error(f"PANW Prisma AIRS: Network/request error: {e!s}") + verbose_proxy_logger.error(f"PANW Prisma AIRS: Network/request error: {e}") return { "action": "block", "category": "network_error", @@ -449,7 +449,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): } except Exception as e: - verbose_proxy_logger.error(f"PANW Prisma AIRS: Unexpected error: {e!s}") + verbose_proxy_logger.error(f"PANW Prisma AIRS: Unexpected error: {e}") return {"action": "block", "category": "api_error", "_is_transient": True} @staticmethod @@ -1056,7 +1056,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"PANW Prisma AIRS scan failed: {e!s}") + verbose_proxy_logger.error(f"PANW Prisma AIRS scan failed: {e}") raise HTTPException( status_code=500, detail={ @@ -1170,7 +1170,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"PANW Prisma AIRS scan failed: {e!s}") + verbose_proxy_logger.error(f"PANW Prisma AIRS scan failed: {e}") raise HTTPException( status_code=500, detail={ @@ -1366,7 +1366,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): error_obj["code"] = e.status_code yield f"data: {json.dumps({'error': error_obj})}\n\n" except Exception as e: - verbose_proxy_logger.error(f"PANW Prisma AIRS streaming error: {e!s}") + verbose_proxy_logger.error(f"PANW Prisma AIRS streaming error: {e}") yield f"data: {json.dumps({'error': {'message': 'Security scan failed - streaming response blocked for safety', 'type': 'guardrail_scan_error', 'code': 500, 'guardrail': self.guardrail_name}})}\n\n" async def _scan_tool_calls_for_guardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index f96fca11abc..77767c8c61b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -457,7 +457,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 - {e!s}") + verbose_proxy_logger.error(f"Pillar Guardrail: API communication failed - {e}") return self._handle_api_error(e, data) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 9d0b8d2777c..7a38c4087c6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -139,9 +139,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: {e!s}, file_path={ad_hoc_recognizers}") + raise Exception(f"Error decoding JSON file: {e}, file_path={ad_hoc_recognizers}") except Exception as e: - raise Exception(f"An error occurred: {e!s}, file_path={ad_hoc_recognizers}") + raise Exception(f"An error occurred: {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, @@ -1124,7 +1124,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield mock_response_stream except Exception as e: - verbose_proxy_logger.error(f"Error masking streaming PII output: {e!s}") + verbose_proxy_logger.error(f"Error masking streaming PII output: {e}") for chunk in all_chunks: yield chunk @@ -1253,7 +1253,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield mock_response_stream except Exception as e: - verbose_proxy_logger.error(f"Error in PII streaming processing: {e!s}") + verbose_proxy_logger.error(f"Error in PII streaming processing: {e}") for chunk in remaining_chunks: yield chunk 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 9b55fcd8062..0f3a817b12c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -326,7 +326,7 @@ class PromptSecurityGuardrail(CustomGuardrail): except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error processing image: {e!s}") + verbose_proxy_logger.error(f"Error processing image: {e}") @staticmethod def _resolve_key_alias_from_request_data(request_data: dict) -> str | None: @@ -481,8 +481,8 @@ class PromptSecurityGuardrail(CustomGuardrail): except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error sanitizing image file: {e!s}") - raise HTTPException(status_code=500, detail=f"File sanitization failed: {e!s}") + verbose_proxy_logger.error(f"Error sanitizing image file: {e}") + raise HTTPException(status_code=500, detail=f"File sanitization failed: {e}") async def _process_document_item(self, item: dict, user_api_key_alias: str | None) -> dict: """Process and sanitize document/file items.""" @@ -554,8 +554,8 @@ class PromptSecurityGuardrail(CustomGuardrail): except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error sanitizing document: {e!s}") - raise HTTPException(status_code=500, detail=f"Document sanitization failed: {e!s}") + verbose_proxy_logger.error(f"Error sanitizing document: {e}") + raise HTTPException(status_code=500, detail=f"Document sanitization failed: {e}") async def process_message_files(self, messages: list, user_api_key_alias: str | None = None) -> list: """Process messages and sanitize any file content (images, documents, PDFs, etc.).""" diff --git a/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py index 4ad29bbeae8..2f2228ae312 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py @@ -10,6 +10,18 @@ if TYPE_CHECKING: def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> RubrikLogger: + """Create and register a RubrikLogger instance. + + The ``mode`` field in the guardrail config controls which surfaces are + moderated: + - ``pre_call`` (or a mode that includes it): prompt moderation via the + ``/v1/before_prompt/openai/v1`` webhook. + - ``post_call`` (the default when ``mode`` is omitted): response and tool + call moderation via the ``/v1/after_completion/openai/v1`` webhook. + + Both hooks are active when ``mode`` covers both ``pre_call`` and + ``post_call``. + """ import litellm rubrik_callback = RubrikLogger( 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 09e5ffff193..9e16e9d5786 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 @@ -351,7 +351,7 @@ class ZscalerAIGuard(CustomGuardrail): return self._handle_response(response, direction) except Exception as e: verbose_proxy_logger.error(f"{e}. Blocking request.") - user_facing_error = self._create_user_facing_error(f"{e!s}") + user_facing_error = self._create_user_facing_error(f"{e}") raise HTTPException(status_code=500, detail=user_facing_error) @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index aaaef95f4a4..b0e16c0ed2e 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -285,7 +285,7 @@ class GuardrailRegistry: return guardrail_dict except Exception as e: - raise Exception(f"Error adding guardrail to DB: {e!s}") + raise Exception(f"Error adding guardrail to DB: {e}") async def delete_guardrail_from_db(self, guardrail_id: str, prisma_client: PrismaClient): """ @@ -297,7 +297,7 @@ class GuardrailRegistry: return {"message": f"Guardrail {guardrail_id} deleted successfully"} except Exception as e: - raise Exception(f"Error deleting guardrail from DB: {e!s}") + raise Exception(f"Error deleting guardrail from DB: {e}") async def update_guardrail_in_db(self, guardrail_id: str, guardrail: Guardrail, prisma_client: PrismaClient): """ @@ -328,7 +328,7 @@ class GuardrailRegistry: # Convert to dict and return return dict(updated_guardrail) except Exception as e: - raise Exception(f"Error updating guardrail in DB: {e!s}") + raise Exception(f"Error updating guardrail in DB: {e}") @staticmethod async def get_all_guardrails_from_db( @@ -350,7 +350,7 @@ class GuardrailRegistry: return guardrails except Exception as e: - raise Exception(f"Error getting guardrails from DB: {e!s}") + raise Exception(f"Error getting guardrails from DB: {e}") async def get_guardrail_by_id_from_db(self, guardrail_id: str, prisma_client: PrismaClient) -> Guardrail | None: """ @@ -366,7 +366,7 @@ class GuardrailRegistry: return Guardrail(**(dict(guardrail))) # type: ignore except Exception as e: - raise Exception(f"Error getting guardrail from DB: {e!s}") + raise Exception(f"Error getting guardrail from DB: {e}") async def get_guardrail_by_name_from_db(self, guardrail_name: str, prisma_client: PrismaClient) -> Guardrail | None: """ @@ -382,7 +382,7 @@ class GuardrailRegistry: return Guardrail(**(dict(guardrail))) # type: ignore except Exception as e: - raise Exception(f"Error getting guardrail from DB: {e!s}") + raise Exception(f"Error getting guardrail from DB: {e}") class InMemoryGuardrailHandler: diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index 71ffc9d36ef..036ee5dca78 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -141,5 +141,5 @@ def initialize_guardrails( return litellm.guardrail_name_config_map except Exception as e: - verbose_proxy_logger.exception(f"error initializing guardrails {e!s}") + verbose_proxy_logger.exception(f"error initializing guardrails {e}") raise e diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 03645c0b2fa..c3bce5e8370 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -425,11 +425,11 @@ async def health_services_endpoint( } except Exception as e: - verbose_proxy_logger.error(f"litellm.proxy.proxy_server.health_services_endpoint(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.health_services_endpoint(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), @@ -1069,7 +1069,7 @@ async def health_endpoint( ) return _post_process(router_result) except Exception as e: - verbose_proxy_logger.error(f"litellm.proxy.proxy_server.py::health_endpoint(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.py::health_endpoint(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) raise e @@ -1110,7 +1110,7 @@ async def health_check_history_endpoint( verbose_proxy_logger.error(f"Error getting health check history: {e}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Failed to retrieve health check history: {e!s}"}, + detail={"error": f"Failed to retrieve health check history: {e}"}, ) @@ -1142,7 +1142,7 @@ async def latest_health_checks_endpoint( verbose_proxy_logger.error(f"Error getting latest health checks: {e}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Failed to retrieve latest health checks: {e!s}"}, + detail={"error": f"Failed to retrieve latest health checks: {e}"}, ) @@ -1185,7 +1185,7 @@ 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: {e!s}"}, + detail={"error": f"Failed to retrieve shared health check status: {e}"}, ) @@ -1473,7 +1473,7 @@ async def _get_health_readiness_details( "is_detailed_debug": is_detailed_debug, } except Exception as e: - raise HTTPException(status_code=503, detail=f"Service Unhealthy ({e!s})") + raise HTTPException(status_code=503, detail=f"Service Unhealthy ({e})") def _allow_public_health_readiness_details() -> bool: @@ -1897,10 +1897,8 @@ async def test_model_connection( except HTTPException as e: raise e except Exception as e: - verbose_proxy_logger.debug( - f"litellm.proxy.health_endpoints.test_model_connection(): Exception occurred - {e!s}" - ) + verbose_proxy_logger.debug(f"litellm.proxy.health_endpoints.test_model_connection(): Exception occurred - {e}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Failed to test connection: {e!s}"}, + detail={"error": f"Failed to test connection: {e}"}, ) diff --git a/litellm/proxy/hooks/azure_content_safety.py b/litellm/proxy/hooks/azure_content_safety.py index 75ce0dff59c..3c7713b2819 100644 --- a/litellm/proxy/hooks/azure_content_safety.py +++ b/litellm/proxy/hooks/azure_content_safety.py @@ -123,7 +123,7 @@ class _PROXY_AzureContentSafety( raise e except Exception as e: verbose_proxy_logger.error( - f"litellm.proxy.hooks.azure_content_safety.py::async_pre_call_hook(): Exception occured - {e!s}" + f"litellm.proxy.hooks.azure_content_safety.py::async_pre_call_hook(): Exception occured - {e}" ) verbose_proxy_logger.debug(traceback.format_exc()) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 651ede6f5bc..ce4ff2cb370 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -600,7 +600,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) raise except Exception as e: - verbose_proxy_logger.error(f"Error counting input file usage for {file_id}: {e!s}") + verbose_proxy_logger.error(f"Error counting input file usage for {file_id}: {e}") raise async def _enforce_batch_file_model_access( @@ -704,7 +704,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): detail={ "error": ( "Batch input file references a model the caller is " - f"not authorized to use: model={model_to_check}, reason={e!s}" + f"not authorized to use: model={model_to_check}, reason={e}" ) }, ) @@ -734,7 +734,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): from litellm.proxy.proxy_server import llm_router, proxy_logging_obj except ImportError as e: raise ValueError( - f"Cannot import proxy_server dependencies: {e!s}. Managed files require proxy_server to be initialized." + f"Cannot import proxy_server dependencies: {e}. Managed files require proxy_server to be initialized." ) # Get the managed files hook @@ -846,6 +846,6 @@ class _PROXY_BatchRateLimiter(CustomLogger): # Re-raise HTTP exceptions (rate limit exceeded) raise except Exception as e: - verbose_proxy_logger.error(f"Error in batch rate limiting: {e!s}", exc_info=True) + verbose_proxy_logger.error(f"Error in batch rate limiting: {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 effafdbcf35..377cd8d3d45 100644 --- a/litellm/proxy/hooks/batch_redis_get.py +++ b/litellm/proxy/hooks/batch_redis_get.py @@ -84,7 +84,7 @@ class _PROXY_BatchRedisRequests(CustomLogger): raise e except Exception as e: verbose_proxy_logger.error( - f"litellm.proxy.hooks.batch_redis_get.py::async_pre_call_hook(): Exception occured - {e!s}" + f"litellm.proxy.hooks.batch_redis_get.py::async_pre_call_hook(): Exception occured - {e}" ) verbose_proxy_logger.debug(traceback.format_exc()) diff --git a/litellm/proxy/hooks/cache_control_check.py b/litellm/proxy/hooks/cache_control_check.py index a5c26e0dad8..f2a0f06b95b 100644 --- a/litellm/proxy/hooks/cache_control_check.py +++ b/litellm/proxy/hooks/cache_control_check.py @@ -52,5 +52,5 @@ class _PROXY_CacheControlCheck(CustomLogger): raise e except Exception as e: verbose_logger.exception( - f"litellm.proxy.hooks.cache_control_check.py::async_pre_call_hook(): Exception occured - {e!s}" + f"litellm.proxy.hooks.cache_control_check.py::async_pre_call_hook(): Exception occured - {e}" ) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter.py b/litellm/proxy/hooks/dynamic_rate_limiter.py index d08c3488348..5d890b6787c 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter.py @@ -68,7 +68,7 @@ class DynamicRateLimiterCache: await self.cache.async_set_cache_sadd(key=key_name, value=value, ttl=self.ttl) except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.hooks.dynamic_rate_limiter.py::async_set_cache_sadd(): Exception occured - {e!s}" + f"litellm.proxy.hooks.dynamic_rate_limiter.py::async_set_cache_sadd(): Exception occured - {e}" ) raise e @@ -172,7 +172,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): ) except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.hooks.dynamic_rate_limiter.py::check_available_usage: Exception occurred - {e!s}" + f"litellm.proxy.hooks.dynamic_rate_limiter.py::check_available_usage: Exception occurred - {e}" ) return None, None, None, None, None @@ -263,6 +263,6 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): ) except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.hooks.dynamic_rate_limiter.py::async_post_call_success_hook(): Exception occured - {e!s}" + f"litellm.proxy.hooks.dynamic_rate_limiter.py::async_post_call_success_hook(): Exception occured - {e}" ) return response diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index cee11ff22ae..773abed1785 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -282,7 +282,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return max_saturation except Exception as e: - verbose_proxy_logger.error(f"Error checking saturation for {model}: {e!s}") + verbose_proxy_logger.error(f"Error checking saturation for {model}: {e}") # Fail open: assume not saturated on error return 0.0 @@ -640,7 +640,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error in dynamic rate limiter: {e!s}, allowing request") + verbose_proxy_logger.error(f"Error in dynamic rate limiter: {e}, allowing request") # Fail open on unexpected errors return None @@ -676,7 +676,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: {e!s}") + verbose_proxy_logger.exception(f"Error in dynamic rate limiter v3 post-call hook: {e}") return response async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -791,4 +791,4 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ) except Exception as e: - verbose_proxy_logger.exception(f"Error in dynamic rate limiter success event: {e!s}") + verbose_proxy_logger.exception(f"Error in dynamic rate limiter success event: {e}") diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index 2fc9779e5fd..983a59657ce 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -587,7 +587,7 @@ class SkillsInjectionHook(CustomLogger): return result or "Code executed successfully" except Exception as e: - return f"Code execution failed: {e!s}" + return f"Code execution failed: {e}" async def _execute_skill_tool( self, @@ -821,7 +821,7 @@ print('No executable skill module found') except Exception as e: verbose_proxy_logger.error(f"SkillsInjectionHook: Code execution failed: {e}") - return f"Code execution failed: {e!s}" + return f"Code execution failed: {e}" def _attach_files_to_response( self, diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 0a1a09d0792..4a768b4e7de 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -75,5 +75,5 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): raise e except Exception as e: verbose_logger.exception( - f"litellm.proxy.hooks.max_budget_limiter.py::async_pre_call_hook(): Exception occured - {e!s}" + f"litellm.proxy.hooks.max_budget_limiter.py::async_pre_call_hook(): Exception occured - {e}" ) diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index b41fd960aac..04f34d0e9cf 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -776,7 +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(f"Inside Parallel Request Limiter: An exception occurred - {e!s}") + verbose_proxy_logger.exception(f"Inside Parallel Request Limiter: An exception occurred - {e}") async def get_internal_user_object( self, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 7253a684b3c..98f1e650845 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -490,7 +490,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: {e!s}") + verbose_proxy_logger.debug(f"Could not load batch rate limiter: {e}") return self._batch_rate_limiter def _get_current_time(self) -> datetime: @@ -808,7 +808,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}: {e!s}") + verbose_proxy_logger.warning(f"Redis Lua script failed for hash tag {hash_tag}: {e}") # Fallback to in-memory cache for this group group_cache_values = await self.in_memory_cache_sliding_window( keys=group_keys, @@ -1055,7 +1055,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) counts = [max(0, int(value)) for value in raw_counts] except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror, never a 500 - verbose_proxy_logger.warning(f"parallel_count_script failed, using local mirror: {e!s}") + verbose_proxy_logger.warning(f"parallel_count_script failed, using local mirror: {e}") counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span) else: counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span) @@ -1085,7 +1085,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ], ) except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to in-memory enforcement, never a 500 - verbose_proxy_logger.warning(f"parallel_acquire_script failed, falling back to in-memory gauge: {e!s}") + verbose_proxy_logger.warning(f"parallel_acquire_script failed, falling back to in-memory gauge: {e}") async with self._check_and_increment_lock: return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span) if int(raw[0]) == 1: @@ -1212,9 +1212,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) return except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the in-memory release, never a 500 - verbose_proxy_logger.warning( - f"parallel_release_script failed, falling back to in-memory release: {e!s}" - ) + verbose_proxy_logger.warning(f"parallel_release_script failed, falling back to in-memory release: {e}") async with self._check_and_increment_lock: for counter_key in counter_keys: @@ -2240,7 +2238,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return False except Exception as e: - verbose_proxy_logger.debug(f"Error checking model failure status: {e!s}, defaulting to enforce limits") + verbose_proxy_logger.debug(f"Error checking model failure status: {e}, defaulting to enforce limits") # Fail safe: enforce limits if we can't check return True @@ -2746,7 +2744,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) except Exception as e: - verbose_proxy_logger.warning(f"TTL preservation failed, falling back to regular pipeline: {e!s}") + verbose_proxy_logger.warning(f"TTL preservation failed, falling back to regular pipeline: {e}") # Fallback to regular pipeline on error await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( increment_list=pipeline_operations, @@ -3003,7 +3001,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) except Exception as e: - verbose_proxy_logger.exception(f"Error in rate limit success event: {e!s}") + verbose_proxy_logger.exception(f"Error in rate limit success event: {e}") async def async_logging_hook( self, @@ -3120,7 +3118,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if stash is not None and reserved_tokens > 0: stash.reservation_released = True except Exception as e: - verbose_proxy_logger.exception(f"Error in rate limit failure event: {e!s}") + verbose_proxy_logger.exception(f"Error in rate limit failure event: {e}") async def async_release_max_parallel_requests_on_disconnect( self, @@ -3185,7 +3183,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) except Exception as e: - verbose_proxy_logger.exception(f"Error in rate limit post-call hook: {e!s}") + verbose_proxy_logger.exception(f"Error in rate limit post-call hook: {e}") async def async_post_call_failure_hook( self, diff --git a/litellm/proxy/hooks/prompt_injection_detection.py b/litellm/proxy/hooks/prompt_injection_detection.py index 3e8518d55dc..e7192b9b063 100644 --- a/litellm/proxy/hooks/prompt_injection_detection.py +++ b/litellm/proxy/hooks/prompt_injection_detection.py @@ -197,7 +197,7 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger): raise e except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.hooks.prompt_injection_detection.py::async_pre_call_hook(): Exception occured - {e!s}" + f"litellm.proxy.hooks.prompt_injection_detection.py::async_pre_call_hook(): Exception occured - {e}" ) async def async_moderation_hook( # type: ignore diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 857429fa89f..0319a680714 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -308,7 +308,7 @@ class _ProxyDBLogger(CustomLogger): 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" ) except Exception as e: - error_msg = f"Error in tracking cost callback - {e!s}\n Traceback:{traceback.format_exc()}" + error_msg = f"Error in tracking cost callback - {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", {}) diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py index 444c39340a0..b242f763fcb 100644 --- a/litellm/proxy/hooks/user_management_event_hooks.py +++ b/litellm/proxy/hooks/user_management_event_hooks.py @@ -71,7 +71,7 @@ class UserManagementEventHooks: ) ) except Exception as e: - verbose_proxy_logger.warning(f"Unable to create audit log for user on `/user/new` - {e!s}") + verbose_proxy_logger.warning(f"Unable to create audit log for user on `/user/new` - {e}") @staticmethod async def async_send_user_invitation_email( diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 7666ad0f065..36f702e1b4b 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -185,7 +185,7 @@ async def image_generation( 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(f"litellm.proxy.proxy_server.image_generation(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.image_generation(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -195,7 +195,7 @@ async def image_generation( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d58f953d2ef..6864caccea2 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4,6 +4,7 @@ import json import re import time from collections import OrderedDict +from collections.abc import Mapping from typing import TYPE_CHECKING, Any from fastapi import HTTPException, Request @@ -20,6 +21,8 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + TRUSTED_CALLBACK_VARS_FIELD, + _request_blocked_callback_params, iter_client_callback_metadata_dicts, ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -46,6 +49,12 @@ from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_head # 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()) +_REDACTED_HEADER_VALUE = "***REDACTED***" +_CREDENTIAL_HEADER_NAMES = SpecialHeaders.litellm_credential_header_names() | frozenset( + {"cookie", "proxy-authorization"} +) +_TRANSPORT_ONLY_CREDENTIAL_KEYS = frozenset({"provider_specific_header", "headers", "api_key"}) + # Matches any header of the form x--session-id (case-insensitive). # Excludes the two explicit litellm headers which are handled with higher priority. _GENERIC_SESSION_ID_HEADER_RE = re.compile(r"^x-.+-session-id$", re.IGNORECASE) @@ -356,6 +365,39 @@ def _strip_client_message_redaction_opt_out(data: dict[str, Any]) -> None: ) +def _strip_client_callback_credentials( + data: dict[str, Any], # mutable-ok: strips in place on the request body the pre-call pipeline threads through +) -> None: + """Drop callback credentials and destinations supplied by the caller. + + ``_request_blocked_callback_params`` (Datadog + GCS credentials, sites and agent + hosts) are already ignored when building ``standard_callback_dynamic_params``. + Strip them from the body and every client metadata slot as well, so a caller + cannot pair its own ``dd_site``/``dd_agent_host`` with the team's admin-configured + ``dd_api_key`` and have the resulting logs shipped to a host it controls. + + ``TRUSTED_CALLBACK_VARS_FIELD`` is proxy-owned; it is cleared here and repopulated + from team/key callback settings in ``add_litellm_data_to_request``. + """ + containers = (("body", data), *iter_client_callback_metadata_dicts(data)) + stripped = tuple( + f"{label}.{field}" + for label, container in containers + for field in _request_blocked_callback_params + if field in container + ) + for _, container in containers: + for field in _request_blocked_callback_params: + container.pop(field, None) + data.pop(TRUSTED_CALLBACK_VARS_FIELD, None) + if stripped: + verbose_proxy_logger.debug( + "Stripped client-supplied callback credentials from request: %s. " + "Configure these on the team or key callback settings instead.", + ", ".join(sorted(stripped)), + ) + + def _strip_client_pricing_overrides(data: dict[str, Any]) -> None: """Drop pricing overrides from the request body and any metadata variant. @@ -524,7 +566,8 @@ def safe_add_api_version_from_query_params(data: dict, request: Request): def convert_key_logging_metadata_to_callback( - data: AddTeamCallback, team_callback_settings_obj: TeamCallbackMetadata | None + data: AddTeamCallback, + team_callback_settings_obj: TeamCallbackMetadata | None, ) -> TeamCallbackMetadata: if team_callback_settings_obj is None: team_callback_settings_obj = TeamCallbackMetadata() @@ -711,6 +754,30 @@ def clean_headers( return clean_headers +def _is_credential_header(header: str) -> bool: + """Whether `header` carries a caller credential rather than request context.""" + return header.lower() in _CREDENTIAL_HEADER_NAMES + + +def redact_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + """Return a copy of `headers` with credential-bearing values masked. + + `clean_headers` deliberately preserves some credential headers so they can be + forwarded to the upstream provider; an Anthropic subscription OAuth token in + `Authorization`, or a client-supplied provider key in `x-api-key`. Those values + must never reach a logging callback or a spend log, so every observability-facing + copy of the header dict is built through this helper while the copy that is + forwarded upstream keeps the real values. + + The returned object is a plain dict; guardrail hooks stamp their own headers onto + the stored copy and the logging callbacks JSON-serialize it. + """ + return { + header: (_REDACTED_HEADER_VALUE if _is_credential_header(header) else value) + for header, value in headers.items() + } + + class LiteLLMProxyRequestSetup: @staticmethod def _get_timeout_from_request(headers: dict) -> float | None: @@ -1407,7 +1474,8 @@ async def add_litellm_data_to_request( _headers, allow_client_message_redaction_opt_out=_allow_client_message_redaction_opt_out, ) - verbose_proxy_logger.debug(f"Request Headers: {_headers}") + _logging_safe_headers = redact_credential_headers(_headers) + verbose_proxy_logger.debug(f"Request Headers: {_logging_safe_headers}") verbose_proxy_logger.debug(f"Raw Headers: {_raw_headers}") if forward_llm_auth and "x-api-key" in _headers: @@ -1428,7 +1496,7 @@ async def add_litellm_data_to_request( data["proxy_server_request"] = { "url": str(request.url), "method": request.method, - "headers": _headers, + "headers": _logging_safe_headers, "body": None, # filled in post-strip; see below "arrival_time": arrival_time, # Track when request arrived at proxy } @@ -1454,7 +1522,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): - data[_metadata_variable_name]["headers"] = _headers + data[_metadata_variable_name]["headers"] = _logging_safe_headers # check for forwardable headers data = LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group( @@ -1563,6 +1631,10 @@ async def add_litellm_data_to_request( if not _key_or_team_allows_client_pricing_override(user_api_key_dict): _strip_client_pricing_overrides(data) + # Same reason as the strips above: runs after the metadata string-to-dict parse + # so JSON-string metadata cannot smuggle callback credentials past the dict guard. + _strip_client_callback_credentials(data) + if not _allow_client_message_redaction_opt_out and litellm.turn_off_message_logging is True: _strip_client_message_redaction_opt_out(data) @@ -1579,7 +1651,7 @@ async def add_litellm_data_to_request( # self-reference — body.proxy_server_request.body would be the same # dict as body, producing an infinite traversal loop for any consumer # that walks the structure. - _body_snapshot_exclude = {"secret_fields", "proxy_server_request"} + _body_snapshot_exclude = frozenset({"secret_fields", "proxy_server_request"}) | _TRANSPORT_ONLY_CREDENTIAL_KEYS _body_snapshot = {k: v for k, v in data.items() if k not in _body_snapshot_exclude} data["proxy_server_request"]["body"] = _body_snapshot @@ -1686,7 +1758,7 @@ async def add_litellm_data_to_request( data[_metadata_variable_name]["user_api_key_team_object_permission_id"] = getattr( user_api_key_dict, "team_object_permission_id", None ) - data[_metadata_variable_name]["headers"] = _headers + data[_metadata_variable_name]["headers"] = _logging_safe_headers data[_metadata_variable_name]["endpoint"] = str(request.url) # Carry the proxy-receive instant via metadata (like `endpoint`) so the # OTel layer can compute pre-request latency, including on the failure @@ -1771,6 +1843,9 @@ async def add_litellm_data_to_request( # unpack callback_vars in data for k, v in callback_settings_obj.callback_vars.items(): data[k] = v + # Callbacks that must not honour request-supplied credentials read this + # proxy-owned field instead of the raw request kwargs. + data[TRUSTED_CALLBACK_VARS_FIELD] = callback_settings_obj.callback_vars # Add disabled callbacks from key metadata if user_api_key_dict.metadata and "litellm_disabled_callbacks" in user_api_key_dict.metadata: diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index 439449403c4..22438e6f336 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -386,7 +386,7 @@ class CacheSettingsManager: verbose_proxy_logger.info("Cache settings initialized from database") except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.management_endpoints.cache_settings_endpoints.py::CacheSettingsManager::init_cache_settings_in_db - {e!s}" + f"litellm.proxy.management_endpoints.cache_settings_endpoints.py::CacheSettingsManager::init_cache_settings_in_db - {e}" ) @staticmethod @@ -480,8 +480,8 @@ async def get_cache_settings( redis_type_descriptions=REDIS_TYPE_DESCRIPTIONS, ) except Exception as e: - verbose_proxy_logger.error(f"Error fetching cache settings: {e!s}") - raise HTTPException(status_code=500, detail=f"Error fetching cache settings: {e!s}") + verbose_proxy_logger.error(f"Error fetching cache settings: {e}") + raise HTTPException(status_code=500, detail=f"Error fetching cache settings: {e}") @router.post( @@ -539,10 +539,10 @@ async def test_cache_connection( return CacheTestResponse(**result) except Exception as e: - verbose_proxy_logger.error(f"Error testing cache connection: {e!s}") + verbose_proxy_logger.error(f"Error testing cache connection: {e}") return CacheTestResponse( status="failed", - message=f"Cache connection test failed: {e!s}", + message=f"Cache connection test failed: {e}", error=str(e), ) @@ -652,5 +652,5 @@ async def update_cache_settings( "settings": _redact_credentials(cache_settings), } except Exception as e: - verbose_proxy_logger.error(f"Error updating cache settings: {e!s}") - raise HTTPException(status_code=500, detail=f"Error updating cache settings: {e!s}") + verbose_proxy_logger.error(f"Error updating cache settings: {e}") + raise HTTPException(status_code=500, detail=f"Error updating cache settings: {e}") diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 9bd8db16769..0dc85f98786 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -993,10 +993,10 @@ async def get_daily_activity( ) except Exception as e: - verbose_proxy_logger.exception(f"Error fetching daily activity: {e!s}") + verbose_proxy_logger.exception(f"Error fetching daily activity: {e}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Failed to fetch analytics: {e!s}"}, + detail={"error": f"Failed to fetch analytics: {e}"}, ) @@ -1082,8 +1082,8 @@ async def get_daily_activity_aggregated( ) except Exception as e: - verbose_proxy_logger.exception(f"Error fetching aggregated daily activity: {e!s}") + verbose_proxy_logger.exception(f"Error fetching aggregated daily activity: {e}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Failed to fetch analytics: {e!s}"}, + detail={"error": f"Failed to fetch analytics: {e}"}, ) diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 9d985e48a60..f2295452f4d 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -129,7 +129,7 @@ async def get_cost_discount_config( return {"values": cost_discount_config} except Exception as e: - verbose_proxy_logger.error(f"Error fetching cost discount config: {e!s}") + verbose_proxy_logger.error(f"Error fetching cost discount config: {e}") return {"values": {}} @@ -224,10 +224,10 @@ async def update_cost_discount_config( "values": cost_discount_config, } except Exception as e: - verbose_proxy_logger.error(f"Error updating cost discount config: {e!s}") + verbose_proxy_logger.error(f"Error updating cost discount config: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to update cost discount config: {e!s}"}, + detail={"error": f"Failed to update cost discount config: {e}"}, ) @@ -262,7 +262,7 @@ async def get_cost_margin_config( return {"values": cost_margin_config} except Exception as e: - verbose_proxy_logger.error(f"Error fetching cost margin config: {e!s}") + verbose_proxy_logger.error(f"Error fetching cost margin config: {e}") return {"values": {}} @@ -398,10 +398,10 @@ async def update_cost_margin_config( "values": cost_margin_config, } except Exception as e: - verbose_proxy_logger.error(f"Error updating cost margin config: {e!s}") + verbose_proxy_logger.error(f"Error updating cost margin config: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to update cost margin config: {e!s}"}, + detail={"error": f"Failed to update cost margin config: {e}"}, ) @@ -484,7 +484,7 @@ async def estimate_cost( raise HTTPException( status_code=404, detail={ - "error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e!s}" + "error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e}" }, ) diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 09977fdce40..ff384190d31 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -103,7 +103,7 @@ async def block_user(data: BlockUsers): return {"blocked_users": records} except Exception as e: - verbose_proxy_logger.error(f"An error occurred - {e!s}") + verbose_proxy_logger.error(f"An error occurred - {e}") raise HTTPException(status_code=500, detail={"error": str(e)}) @@ -390,7 +390,7 @@ async def new_end_user( return _to_customer_response(end_user_record) except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.management_endpoints.customer_endpoints.new_end_user(): Exception occured - {e!s}" + f"litellm.proxy.management_endpoints.customer_endpoints.new_end_user(): Exception occured - {e}" ) if "Unique constraint failed on the fields: (`user_id`)" in str(e): raise ProxyException( @@ -455,7 +455,7 @@ async def end_user_info( except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.management_endpoints.customer_endpoints.end_user_info(): Exception occured - {e!s}" + f"litellm.proxy.management_endpoints.customer_endpoints.end_user_info(): Exception occured - {e}" ) raise handle_exception_on_proxy(e) @@ -636,7 +636,7 @@ async def update_end_user( # update based on remaining passed in values except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.update_end_user(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.update_end_user(): Exception occured - {e}") raise handle_exception_on_proxy(e) @@ -711,7 +711,7 @@ async def delete_end_user( # update based on remaining passed in values except Exception as e: - verbose_proxy_logger.error(f"litellm.proxy.proxy_server.delete_end_user(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.delete_end_user(): Exception occured - {e}") raise handle_exception_on_proxy(e) @@ -767,7 +767,7 @@ async def list_end_user( except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.management_endpoints.customer_endpoints.list_end_user(): Exception occured - {e!s}" + f"litellm.proxy.management_endpoints.customer_endpoints.list_end_user(): Exception occured - {e}" ) raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/management_endpoints/fallback_management_endpoints.py b/litellm/proxy/management_endpoints/fallback_management_endpoints.py index f765cf379e4..3df5384b551 100644 --- a/litellm/proxy/management_endpoints/fallback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/fallback_management_endpoints.py @@ -182,10 +182,10 @@ async def create_fallback( except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error creating fallback: {e!s}", exc_info=True) + verbose_proxy_logger.error(f"Error creating fallback: {e}", exc_info=True) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Failed to create fallback: {e!s}"}, + detail={"error": f"Failed to create fallback: {e}"}, ) @@ -239,10 +239,10 @@ async def get_fallback( except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error getting fallback: {e!s}", exc_info=True) + verbose_proxy_logger.error(f"Error getting fallback: {e}", exc_info=True) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Failed to get fallback: {e!s}"}, + detail={"error": f"Failed to get fallback: {e}"}, ) @@ -350,8 +350,8 @@ async def delete_fallback( except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error deleting fallback: {e!s}", exc_info=True) + verbose_proxy_logger.error(f"Error deleting fallback: {e}", exc_info=True) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Failed to delete fallback: {e!s}"}, + detail={"error": f"Failed to delete fallback: {e}"}, ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d87a0b3d096..8e31b1f6e62 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -330,7 +330,7 @@ async def _add_user_to_team( except HTTPException as e: if e.status_code == 400 and ("already exists" in str(e) or "doesn't exist" in str(e)): verbose_proxy_logger.debug( - f"litellm.proxy.management_endpoints.internal_user_endpoints.new_user(): User already exists in team - {e!s}" + f"litellm.proxy.management_endpoints.internal_user_endpoints.new_user(): User already exists in team - {e}" ) else: verbose_proxy_logger.error( @@ -348,7 +348,7 @@ async def _add_user_to_team( and ProxyErrorTypes.team_member_already_in_team in e.type ): verbose_proxy_logger.debug( - f"litellm.proxy.management_endpoints.internal_user_endpoints.new_user(): User already exists in team - {e!s}" + f"litellm.proxy.management_endpoints.internal_user_endpoints.new_user(): User already exists in team - {e}" ) else: verbose_proxy_logger.error( @@ -605,7 +605,7 @@ async def new_user( return new_user_response except Exception as e: - verbose_proxy_logger.exception(f"/user/new: Exception occured - {e!s}") + verbose_proxy_logger.exception(f"/user/new: Exception occured - {e}") raise handle_exception_on_proxy(e) @@ -900,7 +900,7 @@ async def user_info( return response_data except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.user_info(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.user_info(): Exception occured - {e}") raise handle_exception_on_proxy(e) @@ -1050,7 +1050,7 @@ async def user_info_v2( object_permission=user_data.get("object_permission"), ) except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.user_info_v2(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.user_info_v2(): Exception occured - {e}") raise handle_exception_on_proxy(e) @@ -1320,7 +1320,7 @@ async def _invalidate_cached_user_entitlement(user_id: str | None, object_permis try: await user_api_key_cache.async_delete_cache(key=key) except Exception as e: # noqa: BLE001 # a cache we cannot clear still expires; never fail the write - verbose_proxy_logger.warning(f"Failed to invalidate cached entitlement key {key!r}: {e!s}") + verbose_proxy_logger.warning(f"Failed to invalidate cached entitlement key {key!r}: {e}") async def _update_single_user_helper( @@ -1569,11 +1569,11 @@ async def user_update( ) return response except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.user_update(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.user_update(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), @@ -2395,7 +2395,7 @@ async def add_internal_user_to_organization( return new_membership except Exception as e: - raise Exception(f"Failed to add user to organization: {e!s}") + raise Exception(f"Failed to add user to organization: {e}") async def _resolve_org_filter_for_user_search( @@ -2593,8 +2593,8 @@ async def ui_view_users( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception(f"Error searching users: {e!s}") - raise HTTPException(status_code=500, detail=f"Error searching users: {e!s}") + verbose_proxy_logger.exception(f"Error searching users: {e}") + raise HTTPException(status_code=500, detail=f"Error searching users: {e}") # Using shared metric helper implementations from common_daily_activity @@ -2716,10 +2716,10 @@ async def get_user_daily_activity( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception(f"/spend/daily/analytics: Exception occured - {e!s}") + verbose_proxy_logger.exception(f"/spend/daily/analytics: Exception occured - {e}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Failed to fetch analytics: {e!s}"}, + detail={"error": f"Failed to fetch analytics: {e}"}, ) @@ -2808,8 +2808,8 @@ async def get_user_daily_activity_aggregated( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception(f"/user/daily/activity/aggregated: Exception occured - {e!s}") + verbose_proxy_logger.exception(f"/user/daily/activity/aggregated: Exception occured - {e}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Failed to fetch analytics: {e!s}"}, + detail={"error": f"Failed to fetch analytics: {e}"}, ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index fb9e6fd739e..d4403b3f5db 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -940,7 +940,7 @@ async def _common_key_generation_helper( data = apply_enterprise_key_management_params(data, team_table) except Exception as e: verbose_proxy_logger.debug( - f"litellm.proxy.proxy_server.generate_key_fn(): Enterprise key management params not applied - {e!s}" + f"litellm.proxy.proxy_server.generate_key_fn(): Enterprise key management params not applied - {e}" ) # TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable @@ -1732,7 +1732,7 @@ async def generate_key_fn( ) except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.generate_key_fn(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.generate_key_fn(): Exception occured - {e}") raise handle_exception_on_proxy(e) @@ -1934,9 +1934,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ casted_metadata[k] = v except Exception as e: - verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.prepare_metadata_fields(): Exception occured - {e!s}" - ) + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.prepare_metadata_fields(): Exception occured - {e}") non_default_values["metadata"] = encrypt_callback_vars(casted_metadata) return non_default_values @@ -2799,10 +2797,10 @@ async def update_key_fn( return {"key": key, **response["data"]} # update based on remaining passed in values except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.update_key_fn(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.update_key_fn(): Exception occured - {e}") if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), @@ -3368,7 +3366,7 @@ async def delete_key_fn( return {"deleted_keys": deleted_keys} except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.delete_key_fn(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.delete_key_fn(): Exception occured - {e}") raise handle_exception_on_proxy(e) @@ -3908,7 +3906,7 @@ async def generate_key_helper_fn( # If it's not valid JSON/YAML, keep as is or set to empty dict key_data["router_settings"] = {} except Exception as e: - verbose_proxy_logger.error(f"litellm.proxy.proxy_server.generate_key_helper_fn(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.generate_key_helper_fn(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise e @@ -4115,7 +4113,7 @@ async def delete_verification_tokens( raise Exception("DB not connected. prisma_client is None") except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.delete_verification_tokens(): Exception occured - {e!s}" + f"litellm.proxy.proxy_server.delete_verification_tokens(): Exception occured - {e}" ) verbose_proxy_logger.debug(traceback.format_exc()) raise e @@ -4387,7 +4385,7 @@ async def _rotate_master_key( }, ) except Exception as e: - verbose_proxy_logger.error(f"Failed to re-encrypt credential {cred.credential_name}: {e!s}") + verbose_proxy_logger.error(f"Failed to re-encrypt credential {cred.credential_name}: {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") @@ -5451,7 +5449,7 @@ async def list_keys( verbose_proxy_logger.exception(f"Error in list_keys: {e}") if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"error({e!s})"), + message=getattr(e, "detail", f"error({e})"), type=ProxyErrorTypes.internal_server_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR), @@ -5603,7 +5601,7 @@ async def key_aliases( verbose_proxy_logger.exception(f"Error in key_aliases: {e}") if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"error({e!s})"), + message=getattr(e, "detail", f"error({e})"), type=ProxyErrorTypes.internal_server_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), @@ -6340,7 +6338,7 @@ async def key_health( except Exception as e: raise ProxyException( - message=f"Key health check failed: {e!s}", + message=f"Key health check failed: {e}", type=ProxyErrorTypes.internal_server_error, param=getattr(e, "param", "None"), code=status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -6425,7 +6423,7 @@ async def test_key_logging( return LoggingCallbackStatus( callbacks=logging_callbacks, status="unhealthy", - details=f"Logging test failed: {e!s}", + details=f"Logging test failed: {e}", ) await asyncio.sleep(2) # wait for callbacks to run, callbacks use batching so wait for the flush event @@ -6556,5 +6554,5 @@ def validate_model_max_budget(model_max_budget: dict | None) -> None: BudgetConfig(**_info) except Exception as e: raise ValueError( - f"Invalid model_max_budget: {e!s}. Example of valid model_max_budget: https://docs.litellm.ai/docs/proxy/users" + f"Invalid model_max_budget: {e}. Example of valid model_max_budget: https://docs.litellm.ai/docs/proxy/users" ) diff --git a/litellm/proxy/management_endpoints/management_v1/budgets.py b/litellm/proxy/management_endpoints/management_v1/budgets.py index 8ecd7b1fa30..1bd47a940be 100644 --- a/litellm/proxy/management_endpoints/management_v1/budgets.py +++ b/litellm/proxy/management_endpoints/management_v1/budgets.py @@ -191,7 +191,7 @@ async def list_budgets( raise except Exception as e: # noqa: BLE001 # a driver error answers as a problem document, not the OpenAI error shape verbose_proxy_logger.exception( - f"litellm.proxy.management_endpoints.management_v1.budgets.list_budgets(): Exception occured - {e!s}" + f"litellm.proxy.management_endpoints.management_v1.budgets.list_budgets(): Exception occured - {e}" ) raise ManagementProblem( ProblemDetail( diff --git a/litellm/proxy/management_endpoints/management_v1/spend_logs.py b/litellm/proxy/management_endpoints/management_v1/spend_logs.py index 1927e94d01b..403e2760fb9 100644 --- a/litellm/proxy/management_endpoints/management_v1/spend_logs.py +++ b/litellm/proxy/management_endpoints/management_v1/spend_logs.py @@ -188,7 +188,7 @@ async def list_spend_log_end_users( except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.management_endpoints.management_v1.spend_logs.list_spend_log_end_users(): " - f"Exception occured - {e!s}" + f"Exception occured - {e}" ) raise ManagementProblem( ProblemDetail( diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 2ae9da576b2..da86a7f06f2 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -399,7 +399,7 @@ 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: {e!s}") + verbose_proxy_logger.debug(f"Failed to encrypt temporary MCP server payload for Redis cache: {e}") return if not isinstance(encrypted_payload, str): @@ -413,7 +413,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: {e!s}") + verbose_proxy_logger.debug(f"Failed to write temporary MCP server to Redis cache: {e}") async def _get_temporary_mcp_server_from_redis( server_id: str, @@ -435,7 +435,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: {e!s}") + verbose_proxy_logger.debug(f"Failed reading temporary MCP server from Redis cache: {e}") return None if not isinstance(cached_server, str): @@ -454,7 +454,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: {e!s}") + verbose_proxy_logger.debug(f"Invalid decrypted temporary MCP payload in Redis cache: {e}") return None if not isinstance(loaded, dict): return None @@ -463,7 +463,7 @@ if MCP_AVAILABLE: try: return MCPServer.model_validate(payload_dict) except Exception as e: - verbose_proxy_logger.debug(f"Invalid temporary MCP server payload in Redis cache: {e!s}") + verbose_proxy_logger.debug(f"Invalid temporary MCP server payload in Redis cache: {e}") return None async def get_cached_temporary_mcp_server( @@ -1183,10 +1183,10 @@ if MCP_AVAILABLE: touched_by=user_api_key_dict.user_id or user_api_key_dict.team_id, ) except Exception as e: - verbose_proxy_logger.exception(f"Error registering mcp server: {e!s}") + verbose_proxy_logger.exception(f"Error registering mcp server: {e}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Error registering mcp server: {e!s}"}, + detail={"error": f"Error registering mcp server: {e}"}, ) # Do NOT add to runtime registry — pending servers are not active return _redact_mcp_credentials(new_mcp_server) @@ -1483,10 +1483,10 @@ if MCP_AVAILABLE: touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, ) except Exception as e: - verbose_proxy_logger.exception(f"Error creating mcp server: {e!s}") + verbose_proxy_logger.exception(f"Error creating mcp server: {e}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Error creating mcp server: {e!s}"}, + detail={"error": f"Error creating mcp server: {e}"}, ) # Registry refresh is best-effort: the row is already committed, so a @@ -1498,7 +1498,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 registry refresh failed: {e!s}" + f"MCP server {new_mcp_server.server_id} created but in-memory registry refresh failed: {e}" ) return _redact_mcp_credentials(new_mcp_server) @@ -1559,10 +1559,10 @@ if MCP_AVAILABLE: ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS, ) except Exception as e: - verbose_proxy_logger.exception(f"Error caching temporary mcp server: {e!s}") + verbose_proxy_logger.exception(f"Error caching temporary mcp server: {e}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Error caching temporary mcp server: {e!s}"}, + detail={"error": f"Error caching temporary mcp server: {e}"}, ) return _redact_mcp_credentials(temp_record) 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 2ac0b32ec13..b294b2674e4 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -439,10 +439,10 @@ async def create_model_group( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception(f"Error creating access group '{data.access_group}': {e!s}") + verbose_proxy_logger.exception(f"Error creating access group '{data.access_group}': {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to create access group: {e!s}"}, + detail={"error": f"Failed to create access group: {e}"}, ) @@ -489,10 +489,10 @@ async def list_access_groups( return ListAccessGroupsResponse(access_groups=access_groups_list) except Exception as e: - verbose_proxy_logger.exception(f"Error listing access groups: {e!s}") + verbose_proxy_logger.exception(f"Error listing access groups: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to list access groups: {e!s}"}, + detail={"error": f"Failed to list access groups: {e}"}, ) @@ -546,10 +546,10 @@ 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}': {e!s}") + verbose_proxy_logger.exception(f"Error getting access group info for '{access_group}': {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to get access group info: {e!s}"}, + detail={"error": f"Failed to get access group info: {e}"}, ) @@ -627,7 +627,7 @@ async def update_access_group( except Exception as e: raise HTTPException( status_code=500, - detail={"error": f"Failed to check access group existence: {e!s}"}, + detail={"error": f"Failed to check access group existence: {e}"}, ) # Validation: Check if all new models exist (only if using model_names path) @@ -699,10 +699,10 @@ async def update_access_group( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception(f"Error updating access group '{access_group}': {e!s}") + verbose_proxy_logger.exception(f"Error updating access group '{access_group}': {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to update access group: {e!s}"}, + detail={"error": f"Failed to update access group: {e}"}, ) @@ -759,7 +759,7 @@ async def delete_access_group( except Exception as e: raise HTTPException( status_code=500, - detail={"error": f"Failed to check access group existence: {e!s}"}, + detail={"error": f"Failed to check access group existence: {e}"}, ) try: @@ -800,8 +800,8 @@ async def delete_access_group( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception(f"Error deleting access group '{access_group}': {e!s}") + verbose_proxy_logger.exception(f"Error deleting access group '{access_group}': {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to delete access group: {e!s}"}, + detail={"error": f"Failed to delete access group: {e}"}, ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 35b64963d4d..475c5813cdd 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -61,6 +61,7 @@ from litellm.repositories.team_repository import TeamRepository from litellm.router import Router from litellm.router_utils.auto_router_model_naming import ( STRATEGY_ROUTER_PARAM_FIELDS, + validate_complexity_router_config_write, validate_strategy_router_model_write, ) from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -117,11 +118,20 @@ def _strategy_router_write_violation( An auto-router deployment's ``litellm_params.model`` (``auto_router/...``) is the discriminator the router loads it by; a write that mangles it makes the router drop the deployment silently under ``ignore_invalid_deployments``. - Only writes that supply ``litellm_params.model`` are judged, against the - merged (stored + incoming) params, so partial patches and restores of an - already-corrupted row stay legal. Returns the violation, or None. + Only writes that supply ``litellm_params.model`` are judged on the naming + contract, against the merged (stored + incoming) params, so partial patches + and restores of an already-corrupted row stay legal. A config is judged only + when the write carries one, for the same reason: a rename must not be held + hostage by a stored config it does not touch. Returns the violation, or None. """ - if incoming_params is None or incoming_params.model is None: + if incoming_params is None: + return None + config_violation = validate_complexity_router_config_write( + complexity_router_config=incoming_params.complexity_router_config + ) + if config_violation is not None: + return config_violation + if incoming_params.model is None: return None present_fields = frozenset( field @@ -355,13 +365,13 @@ async def patch_model( return updated_model except Exception as e: - verbose_proxy_logger.exception(f"Error in patch_model: {e!s}") + verbose_proxy_logger.exception(f"Error in patch_model: {e}") if isinstance(e, (HTTPException, ProxyException)): raise e raise ProxyException( - message=f"Error updating model: {e!s}", + message=f"Error updating model: {e}", type=ProxyErrorTypes.internal_server_error, code=status.HTTP_500_INTERNAL_SERVER_ERROR, param=None, @@ -462,13 +472,13 @@ async def _set_model_blocked_status( return updated_model except Exception as e: - verbose_proxy_logger.exception(f"Error in model {action}: {e!s}") + verbose_proxy_logger.exception(f"Error in model {action}: {e}") if isinstance(e, (HTTPException, ProxyException)): raise e raise ProxyException( - message=f"Error updating model blocked status: {e!s}", + message=f"Error updating model blocked status: {e}", type=ProxyErrorTypes.internal_server_error, code=status.HTTP_500_INTERNAL_SERVER_ERROR, param=None, @@ -1223,10 +1233,10 @@ async def delete_model( ) except Exception as e: - verbose_proxy_logger.exception(f"Failed to delete model. Due to error - {e!s}") + verbose_proxy_logger.exception(f"Failed to delete model. Due to error - {e}") if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), @@ -1429,10 +1439,10 @@ async def add_new_model( return model_response except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.add_new_model(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.add_new_model(): Exception occured - {e}") if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), @@ -1582,10 +1592,10 @@ async def update_model( return model_response except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.update_model(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.update_model(): Exception occured - {e}") if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), @@ -1675,13 +1685,13 @@ async def update_public_model_groups( } except Exception as e: - verbose_proxy_logger.exception(f"Error updating public model groups: {e!s}") + verbose_proxy_logger.exception(f"Error updating public model groups: {e}") if isinstance(e, HTTPException): raise e raise ProxyException( - message=f"Error updating public model groups: {e!s}", + message=f"Error updating public model groups: {e}", type=ProxyErrorTypes.internal_server_error, code=status.HTTP_500_INTERNAL_SERVER_ERROR, param=None, @@ -1743,13 +1753,13 @@ async def update_useful_links( } except Exception as e: - verbose_proxy_logger.exception(f"Error updating public model groups: {e!s}") + verbose_proxy_logger.exception(f"Error updating public model groups: {e}") if isinstance(e, HTTPException): raise e raise ProxyException( - message=f"Error updating public model groups: {e!s}", + message=f"Error updating public model groups: {e}", type=ProxyErrorTypes.internal_server_error, code=status.HTTP_500_INTERNAL_SERVER_ERROR, param=None, @@ -1970,5 +1980,5 @@ async def clear_cache() -> frozenset[str] | None: ) return still_desired_ids except Exception as e: - verbose_proxy_logger.exception(f"Failed to clear cache and reload models. Due to error - {e!s}") + verbose_proxy_logger.exception(f"Failed to clear cache and reload models. Due to error - {e}") return None diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 3ec728c7f79..949c35e4182 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -1261,7 +1261,7 @@ async def organization_member_add( verbose_proxy_logger.exception(f"Error adding member to organization: {e}") if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 061b820093c..0adc0610c60 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -120,7 +120,7 @@ async def get_router_settings( routing_strategy_descriptions=ROUTING_STRATEGY_DESCRIPTIONS, ) except Exception as e: - verbose_proxy_logger.error(f"Error fetching router settings: {e!s}") + verbose_proxy_logger.error(f"Error fetching router settings: {e}") raise @@ -168,5 +168,5 @@ async def get_router_fields( routing_strategy_descriptions=ROUTING_STRATEGY_DESCRIPTIONS, ) except Exception as e: - verbose_proxy_logger.error(f"Error fetching router fields: {e!s}") + verbose_proxy_logger.error(f"Error fetching router fields: {e}") raise diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index fecb14b08d3..8e701fa9e20 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -201,7 +201,7 @@ async def _get_model_names(prisma_client: "PrismaClient", model_ids: Sequence[st models = await _table(ModelRepository(prisma_client)).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: {e!s}") + verbose_proxy_logger.error(f"Error getting model names: {e}") return {} @@ -331,7 +331,7 @@ async def new_tag( "tag": tag_config, } except Exception as e: - verbose_proxy_logger.exception(f"Error creating tag: {e!s}") + verbose_proxy_logger.exception(f"Error creating tag: {e}") raise HTTPException(status_code=500, detail=str(e)) @@ -372,7 +372,7 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): data={"litellm_params": json.dumps(existing_params)}, ) except Exception as e: - verbose_proxy_logger.exception(f"Error adding tag to deployment: {e!s}") + verbose_proxy_logger.exception(f"Error adding tag to deployment: {e}") raise HTTPException(status_code=500, detail=str(e)) @@ -461,7 +461,7 @@ async def update_tag( "tag": tag_config, } except Exception as e: - verbose_proxy_logger.exception(f"Error updating tag: {e!s}") + verbose_proxy_logger.exception(f"Error updating tag: {e}") raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 8fab485cc7a..32b22dd6ade 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -37,7 +37,10 @@ from litellm.proxy.litellm_pre_call_utils import ( _get_validated_callback_metadata, convert_key_logging_metadata_to_callback, ) -from litellm.proxy.management_endpoints.team_endpoints import _verify_team_access +from litellm.proxy.management_endpoints.team_endpoints import ( + _refresh_cached_team, + _verify_team_access, +) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.repositories.team_repository import TeamRepository @@ -262,7 +265,11 @@ async def add_team_callbacks( """ try: from litellm.proxy._types import CommonProxyErrors - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException( @@ -316,6 +323,17 @@ async def add_team_callbacks( new_team_row = await TeamRepository(prisma_client).table.update( where={"team_id": team_id}, data={"metadata": team_metadata_json}, # type: ignore + # `object_permission` is included so `_refresh_cached_team` doesn't + # write a cached team with the relation nulled out — see + # team_model_add for the full rationale. + include={"object_permission": True}, # mutable-ok: prisma include takes a dict literal + ) + + # Without this a newly registered callback stays dormant for existing keys. + await _refresh_cached_team( + team_row=new_team_row, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) await _emit_team_callback_audit_log( @@ -336,7 +354,7 @@ async def add_team_callbacks( except ProxyException as e: raise e except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.add_team_callbacks(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.add_team_callbacks(): Exception occured - {e}") raise ProxyException( message="Internal Server Error, " + str(e), type=ProxyErrorTypes.internal_server_error.value, @@ -363,6 +381,9 @@ async def disable_team_logging( """ Disable all logging callbacks for a team + Callbacks registered through POST /team/{team_id}/callback and the Admin UI are cleared, so + re-enabling logging means registering them again with their callback_vars + Parameters: - team_id (str, required): The unique identifier for the team @@ -375,7 +396,11 @@ async def disable_team_logging( """ try: - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -408,6 +433,9 @@ async def disable_team_logging( # Update metadata team_metadata["callback_settings"] = team_callback_settings_obj.model_dump() + # _get_dynamic_logging_metadata stops at metadata["logging"], where the API + # and Admin UI register callbacks, without ever reading callback_settings. + team_metadata["logging"] = [] # mutable-ok: the disabled state is persisted as an empty JSON array team_metadata = encrypt_callback_vars(team_metadata) team_metadata_json = json.dumps(team_metadata) @@ -415,6 +443,10 @@ async def disable_team_logging( updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": team_id}, data={"metadata": team_metadata_json}, # type: ignore + # `object_permission` is included so `_refresh_cached_team` doesn't + # write a cached team with the relation nulled out — see + # team_model_add for the full rationale. + include={"object_permission": True}, # mutable-ok: prisma include takes a dict literal ) if updated_team is None: @@ -423,6 +455,14 @@ async def disable_team_logging( detail={"error": f"Team id = {team_id} does not exist. Error updating team logging"}, ) + # Request-time callback resolution reads the cached team, so without this + # the DB says logging is off while live keys keep sending until it expires. + await _refresh_cached_team( + team_row=updated_team, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + # Disabling a team's logging callbacks is itself a logging-control # action — emit an audit-log row so the action remains traceable # even though the team's own observability is now off. @@ -452,7 +492,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 - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.disable_team_logging(): Exception occurred - {e}") verbose_proxy_logger.debug(traceback.format_exc()) raise ProxyException( message="Internal Server Error, " + str(e), @@ -545,11 +585,11 @@ async def get_team_callbacks( except ProxyException: raise except Exception as e: - verbose_proxy_logger.error(f"litellm.proxy.proxy_server.get_team_callbacks(): Exception occurred - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.get_team_callbacks(): Exception occurred - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Internal Server Error({e!s})"), + message=getattr(e, "detail", f"Internal Server Error({e})"), type=ProxyErrorTypes.internal_server_error.value, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index def58794040..bdd7e18a850 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -76,6 +76,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ( + OrganizationNotFoundError, _cache_team_object, allowed_route_check_inside_route, can_org_access_model, @@ -112,6 +113,10 @@ from litellm.proxy.management_helpers.object_permission_utils import ( from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, ) +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + validate_team_metadata_if_configured, +) from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, @@ -146,6 +151,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( TeamListResponse, TeamMemberAddResult, TeamMemberInfoResponse, + TeamMetadataSchemaResponse, UpdateTeamMemberPermissionsRequest, ) @@ -1210,24 +1216,10 @@ async def new_team( 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 - if data.organization_id is not None and prisma_client is not None: - org_table = await get_org_object( - org_id=data.organization_id, - user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, - ) - if org_table is None: - raise HTTPException( - status_code=400, - detail=f"Organization not found for organization_id={data.organization_id}", - ) - - await _check_org_team_limits( - org_table=org_table, - data=data, - prisma_client=prisma_client, - ) + if data.organization_id is None: + default_organization_id = _get_default_team_param("organization_id") + if isinstance(default_organization_id, str): + data.organization_id = default_organization_id # Apply defaults from litellm.default_team_params for any fields # not explicitly provided in the request. @@ -1255,6 +1247,29 @@ async def new_team( if default_budget is not None: data.max_budget = default_budget + # check org key limits - done here to handle inheriting org id from team + if data.organization_id is not None and prisma_client is not None: + try: + org_table = await get_org_object( + org_id=data.organization_id, + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + include_budget_table=True, + ) + except OrganizationNotFoundError: + org_table = None + if org_table is None: + raise HTTPException( + status_code=400, + detail=f"Organization not found for organization_id={data.organization_id}", + ) + + await _check_org_team_limits( + org_table=org_table, + data=data, + prisma_client=prisma_client, + ) + if ( user_api_key_dict.user_role is None or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN ): # don't restrict proxy admin @@ -1277,6 +1292,18 @@ async def new_team( _check_passthrough_routes_caller_permission(data, user_api_key_dict, entity="team") + if isinstance(data.metadata, dict): + TeamMemberBudgetHandler.strip_system_managed_metadata_keys(data.metadata) + + await validate_team_metadata_if_configured( + operation="create", + metadata=data.metadata, + existing_metadata=None, + team_id=data.team_id, + team_alias=data.team_alias, + user_api_key_dict=user_api_key_dict, + ) + ## ADD TO MODEL TABLE _model_id = None if data.model_aliases is not None and isinstance(data.model_aliases, dict): @@ -1291,9 +1318,6 @@ async def new_team( _model_id = model_dict.id - ## Create Team Member Budget Table - if isinstance(data.metadata, dict): - TeamMemberBudgetHandler.strip_system_managed_metadata_keys(data.metadata) data_json = data.json() ## Handle Object Permission - MCP, Vector Stores etc. @@ -1955,6 +1979,25 @@ async def update_team( if isinstance(updated_kv.get("metadata"), dict): TeamMemberBudgetHandler.strip_system_managed_metadata_keys(updated_kv["metadata"]) + if "metadata" in updated_kv: + stored_metadata = ( + { # mutable-ok: the validator payload's isinstance guard requires a plain dict + key: value + for key, value in existing_team_row.metadata.items() + if key not in TeamMemberBudgetHandler.SYSTEM_MANAGED_METADATA_KEYS + } + if isinstance(existing_team_row.metadata, dict) + else None + ) + await validate_team_metadata_if_configured( + operation="update", + metadata=updated_kv.get("metadata"), + existing_metadata=stored_metadata, + team_id=data.team_id, + team_alias=data.team_alias if data.team_alias is not None else existing_team_row.team_alias, + user_api_key_dict=user_api_key_dict, + ) + # Check budget_duration and budget_reset_at _set_budget_reset_at(data, updated_kv) @@ -2411,7 +2454,7 @@ async def _process_team_members( except Exception as e: raise HTTPException( status_code=500, - detail={"error": f"Unable to add user - {data.member}, to team - {data.team_id}, for reason - {e!s}"}, + detail={"error": f"Unable to add user - {data.member}, to team - {data.team_id}, for reason - {e}"}, ) updated_users.append(updated_user) if updated_tm is not None: @@ -2433,7 +2476,7 @@ async def _process_team_members( except Exception as e: raise HTTPException( status_code=500, - detail={"error": f"Unable to add user - {m}, to team - {data.team_id}, for reason - {e!s}"}, + detail={"error": f"Unable to add user - {m}, to team - {data.team_id}, for reason - {e}"}, ) updated_users.append(updated_user) if updated_tm is not None: @@ -3936,7 +3979,7 @@ async def team_info( ) if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), @@ -4169,6 +4212,24 @@ async def unblock_team( return record +@router.get( + "/team/metadata_schema", + tags=["team management"], # mutable-ok: fastapi's decorator signature types tags as a list + dependencies=(Depends(user_api_key_auth),), + response_model=TeamMetadataSchemaResponse, +) +async def get_team_metadata_schema(): + """ + Get the team metadata fields declared in ``general_settings.team_metadata_schema``. + + The UI uses this to prepopulate the team metadata form with the declared + keys. Returns an empty ``fields`` list when no schema is configured. This + schema is advisory; server-side enforcement stays with + ``custom_team_metadata_validate``. + """ + return TeamMetadataSchemaResponse(fields=TEAM_METADATA_SCHEMA_REGISTRY.get()) + + @router.get("/team/available") async def list_available_teams( http_request: Request, @@ -4808,7 +4869,7 @@ async def list_team( ) except Exception as e: team_exception = f"""Invalid team object for team_id: {team.team_id}. team_object={team.model_dump()}. - Error: {e!s} + Error: {e} """ verbose_proxy_logger.exception(team_exception) continue @@ -4925,7 +4986,7 @@ async def ui_view_teams( return teams except Exception as e: - raise HTTPException(status_code=500, detail=f"Error searching teams: {e!s}") + raise HTTPException(status_code=500, detail=f"Error searching teams: {e}") def add_new_models_to_team(team_obj: LiteLLM_TeamTable, new_models: list[str]) -> list[str]: diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d274879f82a..bc05f72ae14 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -2197,7 +2197,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: {e!s}") + raise HTTPException(status_code=500, detail=f"Failed to process CLI SSO: {e}") @router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False) @@ -2320,7 +2320,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: {e!s}") + raise HTTPException(status_code=500, detail=f"Error checking session status: {e}") async def insert_sso_user( @@ -4479,7 +4479,7 @@ async def debug_sso_callback(request: Request): # Try to convert to string or another JSON serializable format filtered_result[key] = str(value) except Exception as e: - filtered_result[key] = f"Complex value (not displayable): {e!s}" + filtered_result[key] = f"Complex value (not displayable): {e}" # Defense-in-depth: ensure no bearer tokens leak into the rendered HTML even if # a non-conforming IdP places them in its userinfo response. diff --git a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py index 20f4e91b030..939fe7300f7 100644 --- a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py +++ b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py @@ -150,7 +150,7 @@ async def get_distinct_user_agent_tags( except Exception as e: raise HTTPException( status_code=500, - detail=f"Failed to fetch distinct user agent tags: {e!s}", + detail=f"Failed to fetch distinct user agent tags: {e}", ) @@ -243,7 +243,7 @@ async def get_daily_active_users( except Exception as e: raise HTTPException( status_code=500, - detail=f"Failed to fetch DAU analytics: {e!s}", + detail=f"Failed to fetch DAU analytics: {e}", ) @@ -364,7 +364,7 @@ async def get_weekly_active_users( except Exception as e: raise HTTPException( status_code=500, - detail=f"Failed to fetch WAU analytics: {e!s}", + detail=f"Failed to fetch WAU analytics: {e}", ) @@ -485,7 +485,7 @@ async def get_monthly_active_users( except Exception as e: raise HTTPException( status_code=500, - detail=f"Failed to fetch MAU analytics: {e!s}", + detail=f"Failed to fetch MAU analytics: {e}", ) @@ -585,12 +585,12 @@ async def get_tag_summary( except ValueError as e: raise HTTPException( status_code=400, - detail=f"Invalid date format. Use YYYY-MM-DD: {e!s}", + detail=f"Invalid date format. Use YYYY-MM-DD: {e}", ) except Exception as e: raise HTTPException( status_code=500, - detail=f"Failed to fetch tag summary analytics: {e!s}", + detail=f"Failed to fetch tag summary analytics: {e}", ) @@ -740,5 +740,5 @@ async def get_per_user_analytics( except Exception as e: raise HTTPException( status_code=500, - detail=f"Failed to fetch per-user analytics: {e!s}", + detail=f"Failed to fetch per-user analytics: {e}", ) diff --git a/litellm/proxy/management_helpers/team_metadata_validation.py b/litellm/proxy/management_helpers/team_metadata_validation.py new file mode 100644 index 00000000000..2b8ba05c602 --- /dev/null +++ b/litellm/proxy/management_helpers/team_metadata_validation.py @@ -0,0 +1,193 @@ +"""Custom validation of team metadata on team create/update. + +Operators point `general_settings.custom_team_metadata_validate` at an async +Python function (loaded via `get_instance_fn`, like `custom_key_generate`). +The function receives a `TeamMetadataValidationPayload` and returns a +`TeamMetadataValidationResult`. The proxy awaits it before committing a team +write and fails closed: a rejected value surfaces the function's own message +(HTTP 400), while any raised exception or timeout blocks the write with a +generic message (HTTP 503). +""" + +import asyncio +import inspect +from collections.abc import Awaitable, Mapping +from types import MappingProxyType +from typing import Literal, Protocol + +from fastapi import HTTPException, status +from pydantic import BaseModel, JsonValue, TypeAdapter + +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.types.proxy.management_endpoints.team_endpoints import ( + TeamMetadataFieldSchema, +) + +DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS = 5.0 +DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE = ( + "Team metadata validation is currently unavailable, so the team was not saved. Contact your proxy admin." +) +DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE = "Team metadata failed validation." + + +class TeamMetadataRequester(BaseModel): + user_id: str | None = None + user_email: str | None = None + user_role: str | None = None + + +class TeamMetadataValidationPayload(BaseModel): + operation: Literal["create", "update"] + metadata: Mapping[str, JsonValue] + existing_metadata: Mapping[str, JsonValue] | None = None + team_id: str | None = None + team_alias: str | None = None + requester: TeamMetadataRequester + + +class TeamMetadataValidationResult(BaseModel): + valid: bool + error_message: str | None = None + + +_EMPTY_METADATA: Mapping[str, JsonValue] = MappingProxyType({}) + + +class TeamMetadataValidator(Protocol): + def __call__(self, payload: TeamMetadataValidationPayload, /) -> Awaitable[TeamMetadataValidationResult]: ... + + +class TeamMetadataValidatorRegistry: + def __init__(self) -> None: + self._validator: TeamMetadataValidator | None = None + + def set(self, validator: TeamMetadataValidator | None) -> None: + self._validator = validator + + def get(self) -> TeamMetadataValidator | None: + return self._validator + + +TEAM_METADATA_VALIDATOR_REGISTRY = TeamMetadataValidatorRegistry() + +_TEAM_METADATA_SCHEMA_ADAPTER: TypeAdapter[tuple[TeamMetadataFieldSchema, ...]] = TypeAdapter( + tuple[TeamMetadataFieldSchema, ...] +) + + +def parse_team_metadata_schema(raw_schema: object) -> tuple[TeamMetadataFieldSchema, ...]: + """Parse ``general_settings.team_metadata_schema``; raises on a malformed schema so config load fails fast.""" + if raw_schema is None: + return () + fields = _TEAM_METADATA_SCHEMA_ADAPTER.validate_python(raw_schema) + keys = tuple(field.key for field in fields) + duplicate_keys = sorted(frozenset(key for key in keys if keys.count(key) > 1)) + if duplicate_keys: + raise ValueError(f"team_metadata_schema contains duplicate keys: {', '.join(duplicate_keys)}") + return fields + + +class TeamMetadataSchemaRegistry: + def __init__(self) -> None: + self._fields: tuple[TeamMetadataFieldSchema, ...] = () + + def set(self, fields: tuple[TeamMetadataFieldSchema, ...]) -> None: + self._fields = fields + + def get(self) -> tuple[TeamMetadataFieldSchema, ...]: + return self._fields + + +TEAM_METADATA_SCHEMA_REGISTRY = TeamMetadataSchemaRegistry() + + +async def run_team_metadata_validation( + validator: TeamMetadataValidator, + payload: TeamMetadataValidationPayload, + premium_user: bool, + timeout_seconds: float, + unavailable_message: str, +) -> None: + if premium_user is not True: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: HTTPException.detail has no immutable form + "error": f"custom_team_metadata_validate is an Enterprise feature. {CommonProxyErrors.not_premium_user.value}" + }, + ) + if not ( + inspect.iscoroutinefunction(validator) or inspect.iscoroutinefunction(getattr(validator, "__call__", None)) + ): + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ # mutable-ok: HTTPException.detail has no immutable form + "error": "custom_team_metadata_validate must be an async function" + }, + ) + + try: + raw_result = await asyncio.wait_for(validator(payload), timeout=timeout_seconds) + result = TeamMetadataValidationResult.model_validate(raw_result) + except Exception: # noqa: BLE001 # fail closed: any validator failure must block the team write + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={"error": unavailable_message}, # mutable-ok: HTTPException.detail has no immutable form + ) + + if not result.valid: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: HTTPException.detail has no immutable form + "error": result.error_message or DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE + }, + ) + + +def _read_timeout_seconds(general_settings: Mapping[str, object]) -> float: + raw_timeout = general_settings.get("team_metadata_validation_timeout") + if isinstance(raw_timeout, (int, float)) and not isinstance(raw_timeout, bool) and raw_timeout > 0: + return float(raw_timeout) + return DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS + + +def _read_unavailable_message(general_settings: Mapping[str, object]) -> str: + raw_message = general_settings.get("team_metadata_validation_error_message") + if isinstance(raw_message, str) and raw_message.strip(): + return raw_message + return DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE + + +async def validate_team_metadata_if_configured( + operation: Literal["create", "update"], + metadata: Mapping[str, JsonValue] | None, + existing_metadata: Mapping[str, JsonValue] | None, + team_id: str | None, + team_alias: str | None, + user_api_key_dict: UserAPIKeyAuth, + registry: TeamMetadataValidatorRegistry = TEAM_METADATA_VALIDATOR_REGISTRY, +) -> None: + from litellm.proxy.proxy_server import general_settings, premium_user + + validator = registry.get() + if validator is None: + return + + payload = TeamMetadataValidationPayload( + operation=operation, + metadata=metadata if isinstance(metadata, dict) else _EMPTY_METADATA, + existing_metadata=existing_metadata if isinstance(existing_metadata, dict) else None, + team_id=team_id, + team_alias=team_alias, + requester=TeamMetadataRequester( + user_id=user_api_key_dict.user_id, + user_email=user_api_key_dict.user_email, + user_role=user_api_key_dict.user_role.value if user_api_key_dict.user_role is not None else None, + ), + ) + await run_team_metadata_validation( + validator=validator, + payload=payload, + premium_user=premium_user, + timeout_seconds=_read_timeout_seconds(general_settings), + unavailable_message=_read_unavailable_message(general_settings), + ) diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index df8b0725257..ca25be9d92c 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -55,7 +55,7 @@ async def _parse_multipart_form(request: Request) -> dict[str, Any]: form = await request.form() except Exception as e: raise ValueError( - f"Failed to parse multipart form data: {e!s}. " + f"Failed to parse multipart form data: {e}. " "When using curl with --form/-F, do NOT set the Content-Type header " "manually — curl will set it automatically with the required boundary." ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 4e4718272bd..5d4c3c04818 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -549,7 +549,7 @@ async def create_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.exception(f"litellm.proxy.proxy_server.create_file(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.create_file(): Exception occured - {e}") if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), @@ -558,7 +558,7 @@ async def create_file( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -690,7 +690,7 @@ async def get_file_content( ) except ValueError as e: raise ProxyException( - message=f"Storage backend error: {e!s}", + message=f"Storage backend error: {e}", type="invalid_request_error", param="file_id", code=400, @@ -845,7 +845,7 @@ async def get_file_content( 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(f"litellm.proxy.proxy_server.retrieve_file_content(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.retrieve_file_content(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -855,7 +855,7 @@ async def get_file_content( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -1032,7 +1032,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(f"litellm.proxy.proxy_server.retrieve_file(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.retrieve_file(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -1042,7 +1042,7 @@ async def get_file( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -1238,7 +1238,7 @@ async def delete_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.exception(f"litellm.proxy.proxy_server.delete_file(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.delete_file(): Exception occured - {e}") if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), @@ -1247,7 +1247,7 @@ async def delete_file( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -1427,7 +1427,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(f"litellm.proxy.proxy_server.list_files(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.list_files(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -1437,7 +1437,7 @@ async def list_files( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 1395fc9d32f..0d9b0ab9c49 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -857,14 +857,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: {e!s}") + verbose_proxy_logger.error(f"BedrockError in handle_bedrock_count_tokens: {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: {e!s}") - raise HTTPException(status_code=500, detail={"error": f"CountTokens processing error: {e!s}"}) + verbose_proxy_logger.error(f"Error in handle_bedrock_count_tokens: {e}") + raise HTTPException(status_code=500, detail={"error": f"CountTokens processing error: {e}"}) async def bedrock_llm_proxy_route( 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 ddf86e9cd80..5d045ff2852 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 @@ -935,7 +935,7 @@ class AnthropicPassthroughLoggingHandler: index=0, message={ "role": "assistant", - "content": f"Error creating batch job: {e!s}", + "content": f"Error creating batch job: {e}", "tool_calls": None, "function_call": None, "provider_specific_fields": { 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 93bcac704e5..397f1d94a34 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 @@ -203,7 +203,7 @@ class AssemblyAIPassthroughLoggingHandler: return response.json() except Exception as e: - verbose_proxy_logger.exception(f"[Non blocking logging error] Error getting AssemblyAI transcript: {e!s}") + verbose_proxy_logger.exception(f"[Non blocking logging error] Error getting AssemblyAI transcript: {e}") return None def _poll_assembly_for_transcript_response( @@ -275,7 +275,7 @@ class AssemblyAIPassthroughLoggingHandler: return None except Exception as e: - verbose_proxy_logger.exception(f"[Non blocking logging error] Error getting AssemblyAI model info: {e!s}") + verbose_proxy_logger.exception(f"[Non blocking logging error] Error getting AssemblyAI model info: {e}") return None @staticmethod 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 e878f2a544d..63414a1c19e 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 @@ -183,7 +183,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): return cost except Exception as e: - verbose_proxy_logger.warning(f"Error calculating image generation cost: {e!s}") + verbose_proxy_logger.warning(f"Error calculating image generation cost: {e}") return 0.0 @staticmethod @@ -217,7 +217,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): return cost except Exception as e: - verbose_proxy_logger.warning(f"Error calculating image editing cost: {e!s}") + verbose_proxy_logger.warning(f"Error calculating image editing cost: {e}") return 0.0 @staticmethod @@ -445,7 +445,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): } except Exception as e: - verbose_proxy_logger.error(f"Error in OpenAI passthrough cost tracking: {e!s}") + verbose_proxy_logger.error(f"Error in OpenAI passthrough cost tracking: {e}") # Fall back to base handler without cost tracking base_handler = OpenAIPassthroughLoggingHandler() return base_handler.passthrough_chat_handler( @@ -514,7 +514,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): return complete_streaming_response except Exception as e: - verbose_proxy_logger.error(f"Error building complete streaming response: {e!s}") + verbose_proxy_logger.error(f"Error building complete streaming response: {e}") return None @staticmethod @@ -608,7 +608,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): } except Exception as e: - verbose_proxy_logger.error(f"Error in OpenAI streaming passthrough cost tracking: {e!s}") + verbose_proxy_logger.error(f"Error in OpenAI streaming passthrough cost tracking: {e}") return { "result": None, "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 a2f17eb8911..233127c3fef 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 @@ -759,7 +759,7 @@ class VertexPassthroughLoggingHandler: index=0, message={ "role": "assistant", - "content": f"Error creating batch prediction job: {e!s}", + "content": f"Error creating batch prediction job: {e}", "tool_calls": None, "function_call": None, "provider_specific_fields": { diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index b957618d776..b8aba215d10 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -294,8 +294,8 @@ 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(f"litellm.proxy.proxy_server.completion(): Exception occured - {e!s}") - error_msg = f"{e!s}" + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.completion(): Exception occured - {e}") + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -1502,7 +1502,7 @@ async def pass_through_request( ) else: verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.pass_through_endpoint(): Exception occured - {e!s}" + f"litellm.proxy.proxy_server.pass_through_endpoint(): Exception occured - {e}" ) ######################################################### @@ -1544,7 +1544,7 @@ async def pass_through_request( headers=custom_headers, ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 010cf8a7561..9a4a28c7678 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -89,7 +89,7 @@ class PassThroughStreamingHandler: yield chunk except Exception as e: - verbose_proxy_logger.error(f"Error in chunk_processor: {e!s}") + verbose_proxy_logger.error(f"Error in chunk_processor: {e}") raise finally: # GeneratorExit (raised on client disconnect) is not caught by @@ -115,7 +115,7 @@ class PassThroughStreamingHandler: ) ) except Exception as e: - verbose_proxy_logger.error(f"Error scheduling chunk_processor logging: {e!s}") + verbose_proxy_logger.error(f"Error scheduling chunk_processor logging: {e}") @staticmethod async def _route_streaming_logging_to_handler( @@ -165,7 +165,7 @@ class PassThroughStreamingHandler: **kwargs, ) except Exception as e: - verbose_proxy_logger.error(f"Error in _route_streaming_logging_to_handler: {e!s}") + verbose_proxy_logger.error(f"Error in _route_streaming_logging_to_handler: {e}") @staticmethod def _build_passthrough_logging_result( diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index ed0d98c6e6a..797f72f7667 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -60,8 +60,8 @@ class AttachmentRegistry: self._attachments.append(attachment) verbose_proxy_logger.debug(f"Loaded attachment for policy: {attachment.policy}") except Exception as e: - verbose_proxy_logger.error(f"Error loading attachment: {e!s}") - raise ValueError(f"Invalid attachment: {e!s}") from e + verbose_proxy_logger.error(f"Error loading attachment: {e}") + raise ValueError(f"Invalid attachment: {e}") from e self._config_attachments = tuple(self._attachments) self._initialized = True @@ -318,7 +318,7 @@ class AttachmentRegistry: ) except Exception as e: verbose_proxy_logger.exception(f"Error adding attachment to DB: {e}") - raise Exception(f"Error adding attachment to DB: {e!s}") + raise Exception(f"Error adding attachment to DB: {e}") async def delete_attachment_from_db( self, @@ -354,7 +354,7 @@ class AttachmentRegistry: return {"message": f"Attachment {attachment_id} deleted successfully"} except Exception as e: verbose_proxy_logger.exception(f"Error deleting attachment from DB: {e}") - raise Exception(f"Error deleting attachment from DB: {e!s}") + raise Exception(f"Error deleting attachment from DB: {e}") async def get_attachment_by_id_from_db( self, @@ -394,7 +394,7 @@ class AttachmentRegistry: ) except Exception as e: verbose_proxy_logger.exception(f"Error getting attachment from DB: {e}") - raise Exception(f"Error getting attachment from DB: {e!s}") + raise Exception(f"Error getting attachment from DB: {e}") async def get_all_attachments_from_db( self, @@ -432,7 +432,7 @@ class AttachmentRegistry: ] except Exception as e: verbose_proxy_logger.exception(f"Error getting attachments from DB: {e}") - raise Exception(f"Error getting attachments from DB: {e!s}") + raise Exception(f"Error getting attachments from DB: {e}") async def sync_attachments_from_db( self, @@ -468,7 +468,7 @@ class AttachmentRegistry: ) except Exception as e: verbose_proxy_logger.exception(f"Error syncing attachments from DB: {e}") - raise Exception(f"Error syncing attachments from DB: {e!s}") + raise Exception(f"Error syncing attachments from DB: {e}") # Global singleton instance diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py index 9fb700770d2..1facec0898f 100644 --- a/litellm/proxy/policy_engine/init_policies.py +++ b/litellm/proxy/policy_engine/init_policies.py @@ -167,7 +167,7 @@ async def init_policies( policy_registry.load_policies(policies_config) verbose_proxy_logger.info(f"Successfully loaded {len(policies_config)} policies") except Exception as e: - verbose_proxy_logger.error(f"Failed to load policies: {e!s}") + verbose_proxy_logger.error(f"Failed to load policies: {e}") raise # Load attachments if provided @@ -176,7 +176,7 @@ async def init_policies( attachment_registry.load_attachments(policy_attachments_config) 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: {e!s}") + verbose_proxy_logger.error(f"Failed to load policy attachments: {e}") raise return validation_result diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index 32dfc44b8ba..07a4c2abac6 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -187,8 +187,8 @@ 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}': {e!s}") - raise ValueError(f"Invalid policy '{policy_name}': {e!s}") from e + verbose_proxy_logger.error(f"Error loading policy '{policy_name}': {e}") + raise ValueError(f"Invalid policy '{policy_name}': {e}") from e self._config_policies = dict(self._policies) self._sources = {policy_name: "config" for policy_name in self._policies} @@ -433,7 +433,7 @@ class PolicyRegistry: return _row_to_policy_db_response(created_policy) except Exception as e: verbose_proxy_logger.exception(f"Error adding policy to DB: {e}") - raise Exception(f"Error adding policy to DB: {e!s}") + raise Exception(f"Error adding policy to DB: {e}") async def update_policy_in_db( self, @@ -497,7 +497,7 @@ class PolicyRegistry: return _row_to_policy_db_response(updated_policy) except Exception as e: verbose_proxy_logger.exception(f"Error updating policy in DB: {e}") - raise Exception(f"Error updating policy in DB: {e!s}") + raise Exception(f"Error updating policy in DB: {e}") async def delete_policy_from_db( self, @@ -547,7 +547,7 @@ class PolicyRegistry: return result except Exception as e: verbose_proxy_logger.exception(f"Error deleting policy from DB: {e}") - raise Exception(f"Error deleting policy from DB: {e!s}") + raise Exception(f"Error deleting policy from DB: {e}") async def get_policy_by_id_from_db( self, @@ -573,7 +573,7 @@ class PolicyRegistry: return _row_to_policy_db_response(policy) except Exception as e: verbose_proxy_logger.exception(f"Error getting policy from DB: {e}") - raise Exception(f"Error getting policy from DB: {e!s}") + raise Exception(f"Error getting policy from DB: {e}") def get_policy_by_id_for_request(self, policy_id: str) -> tuple[str, Policy] | None: """ @@ -620,7 +620,7 @@ class PolicyRegistry: return [_row_to_policy_db_response(p) for p in policies] except Exception as e: verbose_proxy_logger.exception(f"Error getting policies from DB: {e}") - raise Exception(f"Error getting policies from DB: {e!s}") + raise Exception(f"Error getting policies from DB: {e}") async def sync_policies_from_db( self, @@ -689,7 +689,7 @@ class PolicyRegistry: ) except Exception as e: verbose_proxy_logger.exception(f"Error syncing policies from DB: {e}") - raise Exception(f"Error syncing policies from DB: {e!s}") + raise Exception(f"Error syncing policies from DB: {e}") async def resolve_guardrails_from_db( self, @@ -742,7 +742,7 @@ class PolicyRegistry: return sorted(resolved_policy.guardrails) except Exception as e: verbose_proxy_logger.exception(f"Error resolving guardrails from DB: {e}") - raise Exception(f"Error resolving guardrails from DB: {e!s}") + raise Exception(f"Error resolving guardrails from DB: {e}") async def get_versions_by_policy_name( self, @@ -772,7 +772,7 @@ class PolicyRegistry: ) except Exception as e: verbose_proxy_logger.exception(f"Error getting versions: {e}") - raise Exception(f"Error getting versions: {e!s}") + raise Exception(f"Error getting versions: {e}") async def create_new_version( self, @@ -858,7 +858,7 @@ class PolicyRegistry: return _row_to_policy_db_response(created) except Exception as e: verbose_proxy_logger.exception(f"Error creating new version: {e}") - raise Exception(f"Error creating new version: {e!s}") + raise Exception(f"Error creating new version: {e}") async def update_version_status( self, @@ -963,7 +963,7 @@ class PolicyRegistry: return _row_to_policy_db_response(updated) except Exception as e: verbose_proxy_logger.exception(f"Error updating version status: {e}") - raise Exception(f"Error updating version status: {e!s}") + raise Exception(f"Error updating version status: {e}") async def compare_versions( self, @@ -1016,7 +1016,7 @@ class PolicyRegistry: ) except Exception as e: verbose_proxy_logger.exception(f"Error comparing versions: {e}") - raise Exception(f"Error comparing versions: {e!s}") + raise Exception(f"Error comparing versions: {e}") async def delete_all_versions( self, @@ -1047,7 +1047,7 @@ class PolicyRegistry: return {"message": message} except Exception as e: verbose_proxy_logger.exception(f"Error deleting all versions: {e}") - raise Exception(f"Error deleting all versions: {e!s}") + raise Exception(f"Error deleting all versions: {e}") # Global singleton instance diff --git a/litellm/proxy/policy_engine/policy_validator.py b/litellm/proxy/policy_engine/policy_validator.py index 824f009c474..67f7b37472c 100644 --- a/litellm/proxy/policy_engine/policy_validator.py +++ b/litellm/proxy/policy_engine/policy_validator.py @@ -78,7 +78,7 @@ 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")} except Exception as e: - verbose_proxy_logger.warning(f"Could not get guardrails from registry: {e!s}") + verbose_proxy_logger.warning(f"Could not get guardrails from registry: {e}") return set() async def check_team_alias_exists(self, team_alias: str) -> bool: @@ -100,7 +100,7 @@ class PolicyValidator: ) return team is not None except Exception as e: - verbose_proxy_logger.warning(f"Could not check team alias '{team_alias}': {e!s}") + verbose_proxy_logger.warning(f"Could not check team alias '{team_alias}': {e}") return True # Assume valid on error async def check_key_alias_exists(self, key_alias: str) -> bool: @@ -122,7 +122,7 @@ class PolicyValidator: ) return key is not None except Exception as e: - verbose_proxy_logger.warning(f"Could not check key alias '{key_alias}': {e!s}") + verbose_proxy_logger.warning(f"Could not check key alias '{key_alias}': {e}") return True # Assume valid on error def check_model_exists(self, model: str) -> bool: @@ -151,7 +151,7 @@ class PolicyValidator: return False except Exception as e: - verbose_proxy_logger.warning(f"Could not check model '{model}': {e!s}") + verbose_proxy_logger.warning(f"Could not check model '{model}': {e}") return True # Assume valid on error @staticmethod @@ -436,7 +436,7 @@ class PolicyValidator: PolicyValidationError( policy_name=policy_name, error_type=PolicyValidationErrorType.INVALID_SYNTAX, - message=f"Failed to parse policy: {e!s}", + message=f"Failed to parse policy: {e}", ) ) diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index c0c8ef2de54..89087a3fdd5 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -1304,7 +1304,7 @@ async def convert_prompt_file_to_json( } except Exception as e: - raise HTTPException(status_code=500, detail=f"Error converting prompt file: {e!s}") + raise HTTPException(status_code=500, detail=f"Error converting prompt file: {e}") finally: # Clean up temp file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 901ca39326b..2ff39164d80 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -461,6 +461,11 @@ from litellm.proxy.management_helpers.audit_logs import ( create_audit_log_for_update, create_object_audit_log, ) +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + TEAM_METADATA_VALIDATOR_REGISTRY, + parse_team_metadata_schema, +) from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.middleware.billable_request_metrics_middleware import ( BillableRequestMetricsMiddleware, @@ -783,6 +788,8 @@ def cleanup_router_config_variables(): user_custom_auth_path = None user_custom_key_generate = None user_custom_key_update = None + TEAM_METADATA_VALIDATOR_REGISTRY.set(None) + TEAM_METADATA_SCHEMA_REGISTRY.set(()) user_custom_sso = None user_custom_ui_sso_sign_in_handler = None use_background_health_checks = None @@ -3499,6 +3506,7 @@ _DB_OVERLAY_REMOTE_MODULE_STR_FIELDS: dict[str, tuple[str, ...]] = { "custom_auth", "custom_key_generate", "custom_key_update", + "custom_team_metadata_validate", "custom_sso", "custom_ui_sso_sign_in_handler", ), @@ -3848,7 +3856,7 @@ class ProxyConfig: with open(file_path, "r") as file: return yaml.safe_load(file) or {} except Exception as e: - raise Exception(f"Error loading yaml file {file_path}: {e!s}") + raise Exception(f"Error loading yaml file {file_path}: {e}") async def _get_config_from_file(self, config_file_path: str | None = None) -> dict: """ @@ -4286,7 +4294,7 @@ class ProxyConfig: 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}: {e!s}") + verbose_proxy_logger.error(f"Error parsing search tool {search_tool_name}: {e}") continue return search_tools_parsed if search_tools_parsed else None @@ -4829,6 +4837,14 @@ class ProxyConfig: if custom_key_update is not None: user_custom_key_update = get_instance_fn(value=custom_key_update, config_file_path=config_file_path) + custom_team_metadata_validate = general_settings.get("custom_team_metadata_validate", None) + TEAM_METADATA_VALIDATOR_REGISTRY.set( + get_instance_fn(value=custom_team_metadata_validate, config_file_path=config_file_path) + if custom_team_metadata_validate is not None + else None + ) + TEAM_METADATA_SCHEMA_REGISTRY.set(parse_team_metadata_schema(general_settings.get("team_metadata_schema"))) + 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) @@ -5499,7 +5515,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: {e!s}") + verbose_proxy_logger.exception(f"Error adding/deleting model to llm_router: {e}") if llm_router is not None: llm_model_list = llm_router.get_model_list() @@ -6143,7 +6159,7 @@ class ProxyConfig: return new_models except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy_server.py::add_deployment() - Error getting new models from DB - {e!s}" + f"litellm.proxy_server.py::add_deployment() - Error getting new models from DB - {e}" ) return None @@ -6200,7 +6216,7 @@ class ProxyConfig: await self._init_non_llm_objects_in_db(prisma_client=prisma_client) except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:add_deployment - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:add_deployment - {e}") return still_desired_ids @@ -6375,9 +6391,7 @@ class ProxyConfig: 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( - f"litellm.proxy.proxy_server.py::ProxyConfig:_init_sso_settings_in_db - {e!s}" - ) + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_sso_settings_in_db - {e}") async def _init_hashicorp_vault_config_override(self, prisma_client: PrismaClient): """ @@ -6534,7 +6548,7 @@ class ProxyConfig: ) except Exception as e: - verbose_proxy_logger.exception(f"Error in _check_and_reload_model_cost_map: {e!s}") + verbose_proxy_logger.exception(f"Error in _check_and_reload_model_cost_map: {e}") async def _check_and_reload_anthropic_beta_headers(self, prisma_client: PrismaClient): """ @@ -6631,7 +6645,7 @@ class ProxyConfig: ) except Exception as e: - verbose_proxy_logger.exception(f"Error in _check_and_reload_anthropic_beta_headers: {e!s}") + verbose_proxy_logger.exception(f"Error in _check_and_reload_anthropic_beta_headers: {e}") def _get_prompt_spec_for_db_prompt(self, db_prompt): """ @@ -6660,7 +6674,7 @@ class ProxyConfig: prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=prompt_spec) except Exception as e: - verbose_proxy_logger.debug(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - {e!s}") + verbose_proxy_logger.debug(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - {e}") async def _init_guardrails_in_db(self, prisma_client: PrismaClient): from litellm.proxy.guardrails.guardrail_registry import ( @@ -6687,7 +6701,7 @@ class ProxyConfig: # pod. Config-loaded entries are never touched. IN_MEMORY_GUARDRAIL_HANDLER.reconcile_db_guardrails(db_guardrail_ids=db_guardrail_ids) except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - {e}") async def _init_policies_in_db(self, prisma_client: PrismaClient): """ @@ -6711,7 +6725,7 @@ class ProxyConfig: verbose_proxy_logger.debug("Successfully synced policies and attachments from DB") except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_policies_in_db - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_policies_in_db - {e}") async def _init_tool_policy_in_db(self, prisma_client: PrismaClient): """ @@ -6725,9 +6739,7 @@ class ProxyConfig: await registry.sync_tool_policy_from_db(prisma_client=prisma_client) verbose_proxy_logger.debug("Successfully synced tool policy from DB") except Exception as e: - verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.py::ProxyConfig:_init_tool_policy_in_db - {e!s}" - ) + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_tool_policy_in_db - {e}") async def _init_vector_stores_in_db(self, prisma_client: PrismaClient): from litellm.vector_stores.vector_store_registry import VectorStoreRegistry @@ -6745,7 +6757,7 @@ class ProxyConfig: litellm.vector_store_registry.add_vector_store_to_registry(vector_store=vector_store) except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.py::ProxyConfig:_init_vector_stores_in_db - {e!s}" + f"litellm.proxy.proxy_server.py::ProxyConfig:_init_vector_stores_in_db - {e}" ) async def _init_vector_store_indexes_in_db(self, prisma_client: PrismaClient): @@ -6769,7 +6781,7 @@ class ProxyConfig: litellm.vector_store_index_registry.upsert_vector_store_index(vector_store_index=vector_store_index) except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.py::ProxyConfig:_init_vector_stores_in_db - {e!s}" + f"litellm.proxy.proxy_server.py::ProxyConfig:_init_vector_stores_in_db - {e}" ) async def _init_mcp_servers_in_db(self): @@ -6794,7 +6806,7 @@ class ProxyConfig: await backfill_null_oauth2_flows(prisma_client) except Exception as e: # noqa: BLE001 verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db backfill - {e!s}" + f"litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db backfill - {e}" ) try: @@ -6802,15 +6814,13 @@ class ProxyConfig: await backfill_discovery_stamped_issuers(prisma_client) except Exception as e: # noqa: BLE001 verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db issuer stamp backfill - {e!s}" + f"litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db issuer stamp backfill - {e}" ) try: await global_mcp_server_manager.reload_servers_from_database() except Exception as e: - verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db - {e!s}" - ) + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db - {e}") async def init_mcp_servers_from_db(self) -> None: if self._should_load_db_object(object_type="mcp"): @@ -6838,7 +6848,7 @@ class ProxyConfig: await global_mcp_server_manager.reload_servers_from_database() except Exception as e: # noqa: BLE001 # scheduled job: a reload failure must not kill the recurring retry verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.py::ProxyConfig:reload_mcp_servers_from_db - {e!s}" + f"litellm.proxy.proxy_server.py::ProxyConfig:reload_mcp_servers_from_db - {e}" ) async def _init_agents_in_db(self, prisma_client: PrismaClient): @@ -6850,7 +6860,7 @@ class ProxyConfig: 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) except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_agents_in_db - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_agents_in_db - {e}") async def _init_search_tools_in_db(self, prisma_client: PrismaClient): """ @@ -6890,9 +6900,7 @@ class ProxyConfig: ) except Exception as e: - verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.py::ProxyConfig:_init_search_tools_in_db - {e!s}" - ) + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.py::ProxyConfig:_init_search_tools_in_db - {e}") @staticmethod def _merge_config_and_db_search_tools( @@ -6958,7 +6966,7 @@ class ProxyConfig: CredentialAccessor.upsert_credentials(credentials) # upsert credentials that are in the all-up list except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy_server.py::get_credentials() - Error getting credentials from DB - {e!s}" + f"litellm.proxy_server.py::get_credentials() - Error getting credentials from DB - {e}" ) return [] @@ -7138,14 +7146,14 @@ async def async_assistants_data_generator(response, user_api_key_dict: UserAPIKe try: yield f"data: {c}\n\n" except Exception as e: - yield f"data: {e!s}\n\n" + yield f"data: {e}\n\n" # Streaming is done, yield the [DONE] chunk done_message = "[DONE]" yield f"data: {done_message}\n\n" except Exception as e: verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server.async_assistants_data_generator(): Exception occured - {e!s}" + f"litellm.proxy.proxy_server.async_assistants_data_generator(): Exception occured - {e}" ) await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -7586,7 +7594,7 @@ async def async_data_generator( try: yield _format_streaming_sse_chunk(chunk=chunk) except Exception as e: - yield f"data: {e!s}\n\n" + yield f"data: {e}\n\n" if pending_fallback_event: yield _format_fallback_metadata_sse_event( @@ -7624,7 +7632,7 @@ async def async_data_generator( client_disconnected = True raise except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.async_data_generator(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.async_data_generator(): Exception occured - {e}") await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, @@ -9375,8 +9383,8 @@ 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(f"litellm.proxy.proxy_server.completion(): Exception occured - {e!s}") - error_msg = f"{e!s}" + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.completion(): Exception occured - {e}") + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -9614,7 +9622,7 @@ async def moderations( 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(f"litellm.proxy.proxy_server.moderations(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.moderations(): Exception occured - {e}") if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), @@ -9623,7 +9631,7 @@ async def moderations( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -9760,7 +9768,7 @@ async def audio_speech( original_exception=e, request_data=data, ) - verbose_proxy_logger.error(f"litellm.proxy.proxy_server.audio_speech(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.audio_speech(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) raise e @@ -9902,7 +9910,7 @@ async def audio_transcriptions( 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(f"litellm.proxy.proxy_server.audio_transcription(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.audio_transcription(): Exception occured - {e}") if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), @@ -9911,7 +9919,7 @@ async def audio_transcriptions( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -10188,7 +10196,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(f"litellm.proxy.proxy_server.get_assistants(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.get_assistants(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10198,7 +10206,7 @@ async def get_assistants( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -10279,7 +10287,7 @@ async def create_assistant( 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(f"litellm.proxy.proxy_server.create_assistant(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.create_assistant(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10289,7 +10297,7 @@ async def create_assistant( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -10368,7 +10376,7 @@ async def delete_assistant( 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(f"litellm.proxy.proxy_server.delete_assistant(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.delete_assistant(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10378,7 +10386,7 @@ async def delete_assistant( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -10457,7 +10465,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(f"litellm.proxy.proxy_server.create_threads(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.create_threads(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10467,7 +10475,7 @@ async def create_threads( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -10544,7 +10552,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(f"litellm.proxy.proxy_server.get_thread(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.get_thread(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10554,7 +10562,7 @@ async def get_thread( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -10635,7 +10643,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(f"litellm.proxy.proxy_server.add_messages(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.add_messages(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10645,7 +10653,7 @@ async def add_messages( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -10722,7 +10730,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(f"litellm.proxy.proxy_server.get_messages(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.get_messages(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10732,7 +10740,7 @@ async def get_messages( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -10823,7 +10831,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(f"litellm.proxy.proxy_server.run_thread(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.run_thread(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10833,7 +10841,7 @@ async def run_thread( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -11760,7 +11768,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: {e!s}") + verbose_proxy_logger.exception(f"Error querying database models with search: {e}") search_total_count = router_models_count else: search_total_count = router_models_count @@ -11895,7 +11903,7 @@ def _sort_models( sorted_models = sorted(all_models, key=get_sort_key, reverse=reverse) return sorted_models except Exception as e: - verbose_proxy_logger.exception(f"Error sorting models by {sort_by}: {e!s}") + verbose_proxy_logger.exception(f"Error sorting models by {sort_by}: {e}") return all_models @@ -11975,7 +11983,7 @@ async def _load_team_object_for_model_filter(team_id: str, prisma_client: Prisma return None return LiteLLM_TeamTable.model_validate(team_db_object.model_dump()) except Exception as e: - verbose_proxy_logger.exception(f"Error fetching team {team_id}: {e!s}") + verbose_proxy_logger.exception(f"Error fetching team {team_id}: {e}") return None @@ -12025,7 +12033,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}: {e!s}") + verbose_proxy_logger.debug(f"Error querying database models for team {team_id}: {e}") return team_accessible_model_ids @@ -12163,7 +12171,7 @@ async def _find_model_by_id( if decrypted_models: found_model = decrypted_models[0] except Exception as e: - verbose_proxy_logger.exception(f"Error querying database for modelId {model_id}: {e!s}") + verbose_proxy_logger.exception(f"Error querying database for modelId {model_id}: {e}") # If model found, verify search filter if provided if found_model is not None: @@ -13613,7 +13621,7 @@ async def async_queue_request( ) if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), @@ -13779,7 +13787,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(f"litellm.proxy.proxy_server.login_v2(): Exception occurred - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.login_v2(): Exception occurred - {e}") if isinstance(e, ProxyException): raise e elif isinstance(e, HTTPException): @@ -13790,7 +13798,7 @@ async def login_v2(request: Request): code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=error_msg, type=ProxyErrorTypes.auth_error, @@ -13856,7 +13864,7 @@ async def login_v3(request: Request): status_code=status.HTTP_200_OK, ) except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.login_v3(): Exception occurred - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.login_v3(): Exception occurred - {e}") if isinstance(e, ProxyException): raise e elif isinstance(e, HTTPException): @@ -13867,7 +13875,7 @@ async def login_v3(request: Request): code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=error_msg, type=ProxyErrorTypes.auth_error, @@ -13929,7 +13937,7 @@ async def login_v3_exchange(request: Request): except ProxyException: raise except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.login_v3_exchange(): Exception occurred - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.login_v3_exchange(): Exception occurred - {e}") raise ProxyException( message=str(e), type=ProxyErrorTypes.auth_error, @@ -14756,11 +14764,11 @@ async def update_config( return {"message": "Config updated successfully"} except Exception as e: - verbose_proxy_logger.error(f"litellm.proxy.proxy_server.update_config(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.update_config(): Exception occured - {e}") verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), @@ -15584,7 +15592,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 - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.delete_callback(): Exception occurred - {e}") verbose_proxy_logger.debug(traceback.format_exc()) raise ProxyException( message="Error deleting callback: " + str(e), @@ -15708,10 +15716,10 @@ async def get_config( "available_callbacks": all_available_callbacks, } except Exception as e: - verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.get_config(): Exception occured - {e!s}") + verbose_proxy_logger.exception(f"litellm.proxy.proxy_server.get_config(): Exception occured - {e}") if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({e!s})"), + message=getattr(e, "detail", f"Authentication Error({e})"), type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), @@ -15826,8 +15834,8 @@ async def reload_model_cost_map( "timestamp": current_time.isoformat(), } except Exception as e: - verbose_proxy_logger.exception(f"Failed to reload model cost map: {e!s}") - raise HTTPException(status_code=500, detail=f"Failed to reload model cost map: {e!s}") + verbose_proxy_logger.exception(f"Failed to reload model cost map: {e}") + raise HTTPException(status_code=500, detail=f"Failed to reload model cost map: {e}") @router.post( @@ -15883,10 +15891,10 @@ 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: {e!s}") + verbose_proxy_logger.exception(f"Failed to schedule model cost map reload: {e}") raise HTTPException( status_code=500, - detail=f"Failed to schedule model cost map reload: {e!s}", + detail=f"Failed to schedule model cost map reload: {e}", ) @@ -15928,8 +15936,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: {e!s}") - raise HTTPException(status_code=500, detail=f"Failed to cancel model cost map reload: {e!s}") + verbose_proxy_logger.exception(f"Failed to cancel model cost map reload: {e}") + raise HTTPException(status_code=500, detail=f"Failed to cancel model cost map reload: {e}") @router.get( @@ -16015,10 +16023,10 @@ 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: {e!s}") + verbose_proxy_logger.exception(f"Failed to get model cost map reload status: {e}") raise HTTPException( status_code=500, - detail=f"Failed to get model cost map reload status: {e!s}", + detail=f"Failed to get model cost map reload status: {e}", ) @@ -16063,10 +16071,10 @@ 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: {e!s}") + verbose_proxy_logger.exception(f"Failed to get model cost map source info: {e}") raise HTTPException( status_code=500, - detail=f"Failed to get model cost map source info: {e!s}", + detail=f"Failed to get model cost map source info: {e}", ) @@ -16142,8 +16150,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: {e!s}") - raise HTTPException(status_code=500, detail=f"Failed to reload anthropic beta headers: {e!s}") + verbose_proxy_logger.exception(f"Failed to reload anthropic beta headers: {e}") + raise HTTPException(status_code=500, detail=f"Failed to reload anthropic beta headers: {e}") @router.post( @@ -16199,10 +16207,10 @@ 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: {e!s}") + verbose_proxy_logger.exception(f"Failed to schedule anthropic beta headers reload: {e}") raise HTTPException( status_code=500, - detail=f"Failed to schedule anthropic beta headers reload: {e!s}", + detail=f"Failed to schedule anthropic beta headers reload: {e}", ) @@ -16244,10 +16252,10 @@ 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: {e!s}") + verbose_proxy_logger.exception(f"Failed to cancel anthropic beta headers reload: {e}") raise HTTPException( status_code=500, - detail=f"Failed to cancel anthropic beta headers reload: {e!s}", + detail=f"Failed to cancel anthropic beta headers reload: {e}", ) @@ -16336,10 +16344,10 @@ 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: {e!s}") + verbose_proxy_logger.exception(f"Failed to get anthropic beta headers reload status: {e}") raise HTTPException( status_code=500, - detail=f"Failed to get anthropic beta headers reload status: {e!s}", + detail=f"Failed to get anthropic beta headers reload status: {e}", ) diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 6b8227ee94f..5aef914d178 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -364,7 +364,7 @@ async def get_litellm_model_cost_map(): except Exception as e: raise HTTPException( status_code=500, - detail=f"Internal Server Error ({e!s})", + detail=f"Internal Server Error ({e})", ) diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index 69a5a9861d2..f1c138fa1bf 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -103,7 +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(f"litellm.proxy.proxy_server.rerank(): Exception occured - {e!s}") + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.rerank(): Exception occured - {e}") if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), @@ -112,7 +112,7 @@ async def rerank( code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index f17d546b88b..9fa634dc12e 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -250,9 +250,7 @@ async def responses_api( f"Stored background response {response.id} in managed objects table with unified_id={response.id}" ) except Exception as e: - verbose_proxy_logger.error( - f"Failed to store background response in managed objects table: {e!s}" - ) + verbose_proxy_logger.error(f"Failed to store background response in managed objects table: {e}") return response except ModifyResponseException as e: diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 84dcc5718e7..b744396e850 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -328,7 +328,7 @@ async def background_streaming_task( ) except Exception as e: - verbose_proxy_logger.error(f"Error in background streaming task for {polling_id}: {e!s}") + verbose_proxy_logger.error(f"Error in background streaming task for {polling_id}: {e}") import traceback verbose_proxy_logger.error(traceback.format_exc()) diff --git a/litellm/proxy/search_endpoints/endpoints.py b/litellm/proxy/search_endpoints/endpoints.py index 7c3a924b3b5..0032083b09c 100644 --- a/litellm/proxy/search_endpoints/endpoints.py +++ b/litellm/proxy/search_endpoints/endpoints.py @@ -170,7 +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}: {e!s}") + verbose_proxy_logger.error(f"Search tool authorization failed for {search_tool_name_value}: {e}") raise if llm_router is not None and hasattr(llm_router, "search_tools"): diff --git a/litellm/proxy/search_endpoints/search_tool_registry.py b/litellm/proxy/search_endpoints/search_tool_registry.py index be4a588660c..d7e5efa6d1e 100644 --- a/litellm/proxy/search_endpoints/search_tool_registry.py +++ b/litellm/proxy/search_endpoints/search_tool_registry.py @@ -78,8 +78,8 @@ class SearchToolRegistry: return search_tool_dict except Exception as e: - verbose_proxy_logger.exception(f"Error adding search tool to DB: {e!s}") - raise Exception(f"Error adding search tool to DB: {e!s}") + verbose_proxy_logger.exception(f"Error adding search tool to DB: {e}") + raise Exception(f"Error adding search tool to DB: {e}") async def delete_search_tool_from_db(self, search_tool_id: str, prisma_client: PrismaClient): """ @@ -109,8 +109,8 @@ class SearchToolRegistry: "search_tool_name": existing_tool.search_tool_name, } except Exception as e: - verbose_proxy_logger.exception(f"Error deleting search tool from DB: {e!s}") - raise Exception(f"Error deleting search tool from DB: {e!s}") + verbose_proxy_logger.exception(f"Error deleting search tool from DB: {e}") + raise Exception(f"Error deleting search tool from DB: {e}") async def update_search_tool_in_db(self, search_tool_id: str, search_tool: SearchTool, prisma_client: PrismaClient): """ @@ -143,8 +143,8 @@ 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: {e!s}") - raise Exception(f"Error updating search tool in DB: {e!s}") + verbose_proxy_logger.exception(f"Error updating search tool in DB: {e}") + raise Exception(f"Error updating search tool in DB: {e}") @staticmethod async def get_all_search_tools_from_db( @@ -176,8 +176,8 @@ class SearchToolRegistry: return search_tools except Exception as e: - verbose_proxy_logger.exception(f"Error getting search tools from DB: {e!s}") - raise Exception(f"Error getting search tools from DB: {e!s}") + verbose_proxy_logger.exception(f"Error getting search tools from DB: {e}") + raise Exception(f"Error getting search tools from DB: {e}") async def get_search_tool_by_id_from_db( self, search_tool_id: str, prisma_client: PrismaClient @@ -204,8 +204,8 @@ 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: {e!s}") - raise Exception(f"Error getting search tool from DB: {e!s}") + verbose_proxy_logger.exception(f"Error getting search tool from DB: {e}") + raise Exception(f"Error getting search tool from DB: {e}") async def get_search_tool_by_name_from_db( self, search_tool_name: str, prisma_client: PrismaClient @@ -232,5 +232,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: {e!s}") - raise Exception(f"Error getting search tool from DB: {e!s}") + verbose_proxy_logger.exception(f"Error getting search tool from DB: {e}") + raise Exception(f"Error getting search tool from DB: {e}") diff --git a/litellm/proxy/spend_tracking/cloudzero_endpoints.py b/litellm/proxy/spend_tracking/cloudzero_endpoints.py index 37a53d06b0b..7b573b2fad7 100644 --- a/litellm/proxy/spend_tracking/cloudzero_endpoints.py +++ b/litellm/proxy/spend_tracking/cloudzero_endpoints.py @@ -161,10 +161,10 @@ async def get_cloudzero_settings( # Re-raise HTTPExceptions as-is raise e except Exception as e: - verbose_proxy_logger.error(f"Error retrieving CloudZero settings: {e!s}") + verbose_proxy_logger.error(f"Error retrieving CloudZero settings: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to retrieve CloudZero settings: {e!s}"}, + detail={"error": f"Failed to retrieve CloudZero settings: {e}"}, ) @@ -238,10 +238,10 @@ async def update_cloudzero_settings( ) raise e except Exception as e: - verbose_proxy_logger.error(f"Error updating CloudZero settings: {e!s}") + verbose_proxy_logger.error(f"Error updating CloudZero settings: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to update CloudZero settings: {e!s}"}, + detail={"error": f"Failed to update CloudZero settings: {e}"}, ) @@ -275,7 +275,7 @@ async def is_cloudzero_setup_in_db() -> bool: return cloudzero_config is not None and cloudzero_config.param_value is not None except Exception as e: - verbose_proxy_logger.error(f"Error checking CloudZero status: {e!s}") + verbose_proxy_logger.error(f"Error checking CloudZero status: {e}") return False @@ -317,7 +317,7 @@ async def is_cloudzero_setup() -> bool: return False except Exception as e: - verbose_proxy_logger.error(f"Error checking CloudZero setup: {e!s}") + verbose_proxy_logger.error(f"Error checking CloudZero setup: {e}") return False @@ -364,10 +364,10 @@ async def init_cloudzero_settings( return CloudZeroInitResponse(message="CloudZero settings initialized successfully", status="success") except Exception as e: - verbose_proxy_logger.error(f"Error initializing CloudZero settings: {e!s}") + verbose_proxy_logger.error(f"Error initializing CloudZero settings: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to initialize CloudZero settings: {e!s}"}, + detail={"error": f"Failed to initialize CloudZero settings: {e}"}, ) @@ -422,10 +422,10 @@ async def cloudzero_dry_run_export( ) except Exception as e: - verbose_proxy_logger.error(f"Error performing CloudZero dry run export: {e!s}") + verbose_proxy_logger.error(f"Error performing CloudZero dry run export: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to perform CloudZero dry run export: {e!s}"}, + detail={"error": f"Failed to perform CloudZero dry run export: {e}"}, ) @@ -487,10 +487,10 @@ async def cloudzero_export( ) except Exception as e: - verbose_proxy_logger.error(f"Error performing CloudZero export: {e!s}") + verbose_proxy_logger.error(f"Error performing CloudZero export: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to perform CloudZero export: {e!s}"}, + detail={"error": f"Failed to perform CloudZero export: {e}"}, ) @@ -550,8 +550,8 @@ async def delete_cloudzero_settings( except HTTPException as e: raise e except Exception as e: - verbose_proxy_logger.error(f"Error deleting CloudZero settings: {e!s}") + verbose_proxy_logger.error(f"Error deleting CloudZero settings: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to delete CloudZero settings: {e!s}"}, + detail={"error": f"Failed to delete CloudZero settings: {e}"}, ) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index d2c3b0d9391..0bcc2b9994b 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -440,7 +440,7 @@ async def view_spend_tags( except Exception as e: if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"/spend/tags Error({e!s})"), + message=getattr(e, "detail", f"/spend/tags Error({e})"), type="internal_error", param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), @@ -1492,7 +1492,7 @@ async def global_get_all_tag_names(): except Exception as e: if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"/spend/all_tag_names Error({e!s})"), + message=getattr(e, "detail", f"/spend/all_tag_names Error({e})"), type="internal_error", param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), @@ -1648,7 +1648,7 @@ async def _get_spend_report_for_time_range( return response, spend_per_tag except Exception as e: - verbose_proxy_logger.error(f"Exception in _get_daily_spend_reports {e!s}") + verbose_proxy_logger.error(f"Exception in _get_daily_spend_reports {e}") @router.post( @@ -1798,7 +1798,7 @@ async def calculate_spend(request: SpendCalculateRequest): param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) - error_msg = f"{e!s}" + error_msg = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -2667,7 +2667,7 @@ async def view_spend_logs( except Exception as e: if isinstance(e, HTTPException): raise ProxyException( - message=getattr(e, "detail", f"/spend/logs Error({e!s})"), + message=getattr(e, "detail", f"/spend/logs Error({e})"), type="internal_error", param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), @@ -2789,7 +2789,7 @@ async def global_spend_refresh(): } except Exception as e: - verbose_proxy_logger.exception(f"Failed to refresh materialized view - {e!s}") + verbose_proxy_logger.exception(f"Failed to refresh materialized view - {e}") return { "message": "Failed to refresh materialized view", "status": "failure", @@ -2830,7 +2830,7 @@ async def global_spend_for_internal_user( return response except Exception as e: - verbose_proxy_logger.error(f"/global/spend/logs Error: {e!s}") + verbose_proxy_logger.error(f"/global/spend/logs Error: {e}") raise e @@ -3387,7 +3387,7 @@ 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(f"/provider/budgets: Exception occured - {e!s}") + verbose_proxy_logger.exception(f"/provider/budgets: Exception occured - {e}") raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index 195731c3ed1..ac45594de22 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -166,10 +166,10 @@ async def get_vantage_settings( except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error retrieving Vantage settings: {e!s}") + verbose_proxy_logger.error(f"Error retrieving Vantage settings: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to retrieve Vantage settings: {e!s}"}, + detail={"error": f"Failed to retrieve Vantage settings: {e}"}, ) @@ -235,10 +235,10 @@ async def update_vantage_settings( except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error updating Vantage settings: {e!s}") + verbose_proxy_logger.error(f"Error updating Vantage settings: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to update Vantage settings: {e!s}"}, + detail={"error": f"Failed to update Vantage settings: {e}"}, ) @@ -257,7 +257,7 @@ async def is_vantage_setup_in_db() -> bool: return vantage_config is not None and vantage_config.param_value is not None except Exception as e: - verbose_proxy_logger.error(f"Error checking Vantage status: {e!s}") + verbose_proxy_logger.error(f"Error checking Vantage status: {e}") return False @@ -280,7 +280,7 @@ async def is_vantage_setup() -> bool: return True return False except Exception as e: - verbose_proxy_logger.error(f"Error checking Vantage setup: {e!s}") + verbose_proxy_logger.error(f"Error checking Vantage setup: {e}") return False @@ -324,10 +324,10 @@ async def init_vantage_settings( except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error initializing Vantage settings: {e!s}") + verbose_proxy_logger.error(f"Error initializing Vantage settings: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to initialize Vantage settings: {e!s}"}, + detail={"error": f"Failed to initialize Vantage settings: {e}"}, ) @@ -415,10 +415,10 @@ async def vantage_dry_run_export( except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error performing Vantage dry run export: {e!s}") + verbose_proxy_logger.error(f"Error performing Vantage dry run export: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to perform Vantage dry run export: {e!s}"}, + detail={"error": f"Failed to perform Vantage dry run export: {e}"}, ) @@ -488,10 +488,10 @@ async def vantage_export( except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error performing Vantage export: {e!s}") + verbose_proxy_logger.error(f"Error performing Vantage export: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to perform Vantage export: {e!s}"}, + detail={"error": f"Failed to perform Vantage export: {e}"}, ) @@ -548,8 +548,8 @@ async def delete_vantage_settings( except HTTPException: raise except Exception as e: - verbose_proxy_logger.error(f"Error deleting Vantage settings: {e!s}") + verbose_proxy_logger.error(f"Error deleting Vantage settings: {e}") raise HTTPException( status_code=500, - detail={"error": f"Failed to delete Vantage settings: {e!s}"}, + detail={"error": f"Failed to delete Vantage settings: {e}"}, ) diff --git a/litellm/proxy/types_utils/utils.py b/litellm/proxy/types_utils/utils.py index e9fb18b258e..e61fcdd859b 100644 --- a/litellm/proxy/types_utils/utils.py +++ b/litellm/proxy/types_utils/utils.py @@ -176,7 +176,7 @@ def _load_instance_from_remote_storage(remote_url: str, config_file_path: str | return instance except Exception as e: - raise ImportError(f"Failed to load custom logger from {remote_url}: {e!s}") from e + raise ImportError(f"Failed to load custom logger from {remote_url}: {e}") from e async def _download_gcs_file_wrapper(bucket_name: str, object_key: str, local_file_path: str) -> bool: @@ -190,7 +190,7 @@ async def _download_gcs_file_wrapper(bucket_name: str, object_key: str, local_fi except Exception as e: from litellm._logging import verbose_proxy_logger - verbose_proxy_logger.error(f"Error downloading from GCS: {e!s}") + verbose_proxy_logger.error(f"Error downloading from GCS: {e}") return False diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 69178cea55e..8ed848ac1bf 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy.config_resolvers.sso import ( ) from litellm.proxy.utils import invalidate_config_param from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( SSOConfigRepository, UISettingsRepository, @@ -636,6 +637,36 @@ async def _validate_default_teams_exist(teams: list[str] | list[NewUserRequestTe ) +async def _validate_default_organization_exists(organization_id: str) -> None: + """Reject a default organization that cannot be assigned. + + Teams are created from these settings long after they are saved, and an unknown + organization id would fail every future team creation instead of here, where the + admin who typed it can still fix it. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={ # mutable-ok: HTTPException detail must be a plain dict for FastAPI JSON serialization + "error": "Database not connected. Please connect a database." + }, + ) + + organization_exists = await OrganizationRepository(prisma_client).exists( + organization_id, id_field="organization_id" + ) + if not organization_exists: + raise HTTPException( + status_code=400, + detail={ # mutable-ok: HTTPException detail must be a plain dict for FastAPI JSON serialization + "error": f"Organization not found: {organization_id}. " + "An organization must exist before it can be set as the default organization for new teams." + }, + ) + + async def update_default_team_member_budget(teams: list[NewUserRequestTeam], user_api_key_dict: UserAPIKeyAuth): """ 1. Update the max member budget for the team @@ -774,6 +805,9 @@ async def update_default_team_settings( Update the default team parameters for SSO users. These settings will be applied to new teams created from SSO. """ + if settings.organization_id is not None: + await _validate_default_organization_exists(settings.organization_id) + return await _update_litellm_setting( settings=settings, settings_key="default_team_params", @@ -966,7 +1000,7 @@ async def update_sso_settings( except Exception as e: raise HTTPException( status_code=500, - detail={"error": f"Error updating environment_variables: {e!s}"}, + detail={"error": f"Error updating environment_variables: {e}"}, ) return { diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 68d75384452..5f18189b6b3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3341,7 +3341,7 @@ class PrismaClient: reason=f"prisma_get_generic_data_{table_name}_lookup_failure", ) except Exception as e: - error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {e!s}" + error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {e}" verbose_proxy_logger.error(error_msg) error_msg = error_msg + f"\nException Type: {type(e)}" error_traceback = error_msg + "\n" + traceback.format_exc() @@ -3956,7 +3956,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception in insert_data: {e!s}" + error_msg = f"LiteLLM Prisma Client Exception in insert_data: {e}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() end_time = time.time() @@ -4205,7 +4205,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception - update_data: {e!s}" + error_msg = f"LiteLLM Prisma Client Exception - update_data: {e}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() end_time = time.time() @@ -4271,7 +4271,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception - delete_data: {e!s}" + error_msg = f"LiteLLM Prisma Client Exception - delete_data: {e}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() end_time = time.time() @@ -4304,7 +4304,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception connect(): {e!s}" + error_msg = f"LiteLLM Prisma Client Exception connect(): {e}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() end_time = time.time() @@ -4334,7 +4334,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception disconnect(): {e!s}" + error_msg = f"LiteLLM Prisma Client Exception disconnect(): {e}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() end_time = time.time() @@ -5023,7 +5023,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception disconnect(): {e!s}" + error_msg = f"LiteLLM Prisma Client Exception disconnect(): {e}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() end_time = time.time() @@ -5833,7 +5833,7 @@ def _raise_failed_update_spend_exception(e: Exception, start_time: float, proxy_ """ import traceback - error_msg = f"[Non-Blocking]LiteLLM Prisma Client Exception - update spend logs: {e!s}" + error_msg = f"[Non-Blocking]LiteLLM Prisma Client Exception - update spend logs: {e}" error_traceback = error_msg + "\n" + traceback.format_exc() end_time = time.time() _duration = end_time - start_time @@ -6125,7 +6125,7 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException: if isinstance(e, HTTPException): return ProxyException( - message=getattr(e, "detail", f"error({e!s})"), + message=getattr(e, "detail", f"error({e})"), type=ProxyErrorTypes.internal_server_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 811597f3821..6176ae03d3d 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -211,7 +211,7 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d ) return embedding_config except Exception as e: - verbose_proxy_logger.debug(f"Error resolving embedding config from router for model {model_name}: {e!s}") + verbose_proxy_logger.debug(f"Error resolving embedding config from router for model {model_name}: {e}") continue return None @@ -299,7 +299,7 @@ async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) ) return embedding_config except Exception as e: - verbose_proxy_logger.debug(f"Error resolving embedding config for model {model_name}: {e!s}") + verbose_proxy_logger.debug(f"Error resolving embedding config for model {model_name}: {e}") continue return None @@ -542,7 +542,7 @@ async def new_vector_store( "vector_store": response_vs, } except Exception as e: - verbose_proxy_logger.exception(f"Error creating vector store: {e!s}") + verbose_proxy_logger.exception(f"Error creating vector store: {e}") raise HTTPException(status_code=500, detail=str(e)) @@ -647,7 +647,7 @@ async def list_vector_stores( return response except Exception as e: - verbose_proxy_logger.exception(f"Error listing vector stores: {e!s}") + verbose_proxy_logger.exception(f"Error listing vector stores: {e}") raise HTTPException(status_code=500, detail=str(e)) @@ -727,7 +727,7 @@ async def delete_vector_store( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception(f"Error deleting vector store: {e!s}") + verbose_proxy_logger.exception(f"Error deleting vector store: {e}") raise HTTPException(status_code=500, detail=str(e)) @@ -799,7 +799,7 @@ async def get_vector_store_info( # the catch-all below would otherwise rewrite them as 500. raise except Exception as e: - verbose_proxy_logger.exception(f"Error getting vector store info: {e!s}") + verbose_proxy_logger.exception(f"Error getting vector store info: {e}") raise HTTPException(status_code=500, detail=str(e)) @@ -888,5 +888,5 @@ async def update_vector_store( # as 500 with the original status code embedded in the detail. raise except Exception as e: - verbose_proxy_logger.exception(f"Error updating vector store: {e!s}") + verbose_proxy_logger.exception(f"Error updating vector store: {e}") raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py index 767e526804c..d46c93a2038 100644 --- a/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py +++ b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py @@ -60,7 +60,7 @@ def _normalize_langfuse_base_url(base_target_url: str) -> str: except Exception as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={"error": f"Invalid Langfuse host: {e!s}"}, + detail={"error": f"Invalid Langfuse host: {e}"}, ) if base_url.scheme not in ("http", "https") or not base_url.host: @@ -137,7 +137,7 @@ def _build_langfuse_proxy_target( except SSRFError as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={"error": f"Invalid Langfuse host: {e!s}"}, + detail={"error": f"Invalid Langfuse host: {e}"}, ) custom_headers["Host"] = host_header return target_url, custom_headers diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 5eb326bda18..2b567e8b52a 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -195,6 +195,17 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): return await self.update(user_id, data, id_field="user_id") + async def backfill_null_user_email(self, user_id: str, user_email: str) -> int: + """Set user_email only when the stored value is null, atomically at the database. + + Returns the number of rows updated: 0 means another writer already set an email. + """ + updated_count: int = await self.table.update_many( + where={"user_id": user_id, "user_email": None}, # mutable-ok: Prisma query filters are dict-shaped + data={"user_email": user_email}, # mutable-ok: Prisma update payloads are dict-shaped + ) + return updated_count + async def delete_user(self, user_id: str) -> LiteLLM_UserTable | None: """Delete a user.""" return await self.delete(user_id, id_field="user_id") diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 03c13e504ac..2733fed744a 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -534,5 +534,5 @@ def rerank( # Placeholder return return response except Exception as e: - verbose_logger.error(f"Error in rerank: {e!s}") + verbose_logger.error(f"Error in rerank: {e}") raise exception_type(model=model, custom_llm_provider=custom_llm_provider, original_exception=e) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index c20b35b6bbc..0241453c15f 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -313,7 +313,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _default_response_created_event_data(self) -> dict: # Use cached response ID if available, otherwise generate a new one if self._cached_response_id is None: - self._cached_response_id = f"resp_{uuid.uuid4()!s}" + self._cached_response_id = f"resp_{uuid.uuid4()}" response_created_event_data = { "id": self._cached_response_id, @@ -386,7 +386,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def create_output_item_added_event(self) -> OutputItemAddedEvent: if self._cached_item_id is None: - self._cached_item_id = f"msg_{uuid.uuid4()!s}" + self._cached_item_id = f"msg_{uuid.uuid4()}" self._sequence_number += 1 event = OutputItemAddedEvent( @@ -407,7 +407,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def create_content_part_added_event(self) -> ContentPartAddedEvent: if self._cached_item_id is None: - self._cached_item_id = f"msg_{uuid.uuid4()!s}" + self._cached_item_id = f"msg_{uuid.uuid4()}" self._sequence_number += 1 event = ContentPartAddedEvent( @@ -528,7 +528,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): 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_{uuid.uuid4()!s}" + self._cached_item_id = f"msg_{uuid.uuid4()}" return OutputTextDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, @@ -541,7 +541,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): 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_{uuid.uuid4()!s}" + self._cached_item_id = f"msg_{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 @@ -577,7 +577,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): 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_{uuid.uuid4()!s}" + self._cached_item_id = f"msg_{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 diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 50744a7b93f..39881277a10 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -844,8 +844,8 @@ class LiteLLM_Proxy_MCP_Handler: request_data=logging_request_data, error=e, ) - verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {e!s}") - error_message = f"Tool call blocked: PII entity '{getattr(e, 'entity_type', 'unknown')}' detected by guardrail '{getattr(e, 'guardrail_name', 'unknown')}'. {e!s}" + verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {e}") + error_message = f"Tool call blocked: PII entity '{getattr(e, 'entity_type', 'unknown')}' detected by guardrail '{getattr(e, 'guardrail_name', 'unknown')}'. {e}" tool_results.append( { "tool_call_id": tool_call_id, @@ -860,9 +860,9 @@ class LiteLLM_Proxy_MCP_Handler: request_data=logging_request_data, error=e, ) - verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {e!s}") + verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {e}") error_message = ( - f"Tool call blocked: Guardrail '{getattr(e, 'guardrail_name', 'unknown')}' violation. {e!s}" + f"Tool call blocked: Guardrail '{getattr(e, 'guardrail_name', 'unknown')}' violation. {e}" ) tool_results.append( { @@ -878,7 +878,7 @@ class LiteLLM_Proxy_MCP_Handler: request_data=logging_request_data, error=e, ) - verbose_logger.error(f"HTTPException in MCP tool call: {e!s}") + verbose_logger.error(f"HTTPException in MCP tool call: {e}") error_message = f"Tool call failed: {str(e.detail) if hasattr(e, 'detail') else str(e)}" tool_results.append( { @@ -898,7 +898,7 @@ class LiteLLM_Proxy_MCP_Handler: tool_results.append( { "tool_call_id": tool_call_id, - "result": f"Error executing tool: {e!s}", + "result": f"Error executing tool: {e}", "name": tool_name, } ) diff --git a/litellm/router.py b/litellm/router.py index 37190bdbf38..6bf1bdfc670 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -115,6 +115,7 @@ from litellm.router_utils.common_utils import ( _is_proxy_admin_request, filter_team_based_models, filter_web_search_deployments, + resolve_model_group_alias, ) from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.cooldown_handlers import ( @@ -1720,7 +1721,7 @@ class Router: return _deployment_copy except Exception as e: - verbose_router_logger.debug(f"Error occurred while printing deployment - {e!s}") + verbose_router_logger.debug(f"Error occurred while printing deployment - {e}") raise e ### COMPLETION, EMBEDDING, IMG GENERATION FUNCTIONS @@ -1828,7 +1829,7 @@ class Router: return response except Exception as e: - verbose_router_logger.info(f"litellm.completion(model={model_name})\033[31m Exception {e!s}\033[0m") + verbose_router_logger.info(f"litellm.completion(model={model_name})\033[31m Exception {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) @@ -1923,7 +1924,7 @@ class Router: finally: loop.close() except Exception as e: - verbose_router_logger.error(f"Silent experiment failed for model {silent_model}: {e!s}") + verbose_router_logger.error(f"Silent experiment failed for model {silent_model}: {e}") # fmt: off @@ -2754,7 +2755,7 @@ class Router: **silent_kwargs, ) except Exception as e: - verbose_router_logger.error(f"Silent experiment failed for model {silent_model}: {e!s}") + verbose_router_logger.error(f"Silent experiment failed for model {silent_model}: {e}") async def _acompletion( self, model: str, messages: list[dict[str, str]], **kwargs @@ -2907,7 +2908,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 {e!s}\033[0m") + verbose_router_logger.info(f"litellm.acompletion(model={model_name})\033[31m Exception {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 @@ -3696,7 +3697,7 @@ class Router: 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(f"litellm.image_generation(model={model_name})\033[31m Exception {e!s}\033[0m") + verbose_router_logger.info(f"litellm.image_generation(model={model_name})\033[31m Exception {e}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -3780,7 +3781,7 @@ class Router: 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(f"litellm.aimage_generation(model={model_name})\033[31m Exception {e!s}\033[0m") + verbose_router_logger.info(f"litellm.aimage_generation(model={model_name})\033[31m Exception {e}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -3884,7 +3885,7 @@ class Router: 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 {e!s}\033[0m") + verbose_router_logger.info(f"litellm.atranscription(model={model_name})\033[31m Exception {e}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -3998,7 +3999,7 @@ class Router: 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 {e!s}\033[0m") + verbose_router_logger.info(f"litellm.aspeech(model={model_name})\033[31m Exception {e}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -4056,7 +4057,7 @@ class Router: 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 {e!s}\033[0m") + verbose_router_logger.info(f"litellm.arerank(model={model_name})\033[31m Exception {e}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -4190,7 +4191,7 @@ class Router: 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 {e!s}\033[0m") + verbose_router_logger.info(f"litellm.atext_completion(model={model})\033[31m Exception {e}\033[0m") if model is not None: self.fail_calls[model] += 1 raise e @@ -4280,7 +4281,7 @@ class Router: 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 {e!s}\033[0m") + verbose_router_logger.info(f"litellm.aadapter_completion(model={model})\033[31m Exception {e}\033[0m") if model is not None: self.fail_calls[model] += 1 raise e @@ -4539,9 +4540,7 @@ class Router: return response except Exception as e: - verbose_router_logger.info( - f"ageneric_api_call_with_fallbacks(model={model})\033[31m Exception {e!s}\033[0m" - ) + verbose_router_logger.info(f"ageneric_api_call_with_fallbacks(model={model})\033[31m Exception {e}\033[0m") if model is not None: self.fail_calls[model] += 1 raise e @@ -4661,7 +4660,7 @@ class Router: 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 {e!s}\033[0m") + verbose_router_logger.info(f"{handler_name}(model={model})\033[31m Exception {e}\033[0m") if model is not None: self.fail_calls[model] += 1 raise e @@ -4726,7 +4725,7 @@ class Router: 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 {e!s}\033[0m") + verbose_router_logger.info(f"litellm.embedding(model={model_name})\033[31m Exception {e}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -4813,7 +4812,7 @@ class Router: 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 {e!s}\033[0m") + verbose_router_logger.info(f"litellm.aembedding(model={model_name})\033[31m Exception {e}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -4966,7 +4965,7 @@ class Router: return returned_response except Exception as e: verbose_router_logger.exception( - f"litellm.acreate_file(model={model}, {kwargs})\033[31m Exception {e!s}\033[0m" + f"litellm.acreate_file(model={model}, {kwargs})\033[31m Exception {e}\033[0m" ) if model is not None: self.fail_calls[model] += 1 @@ -5061,9 +5060,7 @@ class Router: return response except Exception as e: - verbose_router_logger.exception( - f"litellm.avector_store_create(model={model})\033[31m Exception {e!s}\033[0m" - ) + verbose_router_logger.exception(f"litellm.avector_store_create(model={model})\033[31m Exception {e}\033[0m") if model is not None: self.fail_calls[model] += 1 raise e @@ -5178,7 +5175,7 @@ class Router: return response # type: ignore except Exception as e: verbose_router_logger.exception( - f"litellm._acreate_batch(model={model}, {kwargs})\033[31m Exception {e!s}\033[0m" + f"litellm._acreate_batch(model={model}, {kwargs})\033[31m Exception {e}\033[0m" ) if model is not None: self.fail_calls[model] += 1 @@ -5400,7 +5397,7 @@ class Router: return response # type: ignore except Exception as e: verbose_router_logger.exception( - f"litellm._acancel_batch(model={model}, {kwargs})\033[31m Exception {e!s}\033[0m" + f"litellm._acancel_batch(model={model}, {kwargs})\033[31m Exception {e}\033[0m" ) if model is not None: self.fail_calls[model] += 1 @@ -6948,7 +6945,7 @@ class Router: except Exception as e: verbose_router_logger.debug( - f"litellm.router.Router::deployment_callback_on_success(): Exception occured - {e!s}" + f"litellm.router.Router::deployment_callback_on_success(): Exception occured - {e}" ) def sync_deployment_callback_on_success( @@ -9014,7 +9011,7 @@ class Router: custom_llm_provider=litellm_params.custom_llm_provider, ) except litellm.exceptions.BadRequestError as e: - verbose_router_logger.error(f"litellm.router.py::get_model_group_info() - {e!s}") + verbose_router_logger.error(f"litellm.router.py::get_model_group_info() - {e}") if model_info is None: supported_openai_params = litellm.get_supported_openai_params( @@ -10228,7 +10225,7 @@ class Router: ) except Exception as e: verbose_router_logger.error( - f"litellm.router.py::_pre_call_checks: failed to count tokens. Returning initial list of deployments. Got - {e!s}" + f"litellm.router.py::_pre_call_checks: failed to count tokens. Returning initial list of deployments. Got - {e}" ) return _returned_deployments if input_tokens > max_input_tokens: @@ -10239,7 +10236,7 @@ class Router: ) continue except Exception as e: - verbose_router_logger.exception(f"An error occurs - {e!s}") + verbose_router_logger.exception(f"An error occurs - {e}") model_id = _model_info.get("id", "") ## RPM CHECK ## @@ -10335,16 +10332,7 @@ class Router: - str, the litellm model name - None, if model is not in model group alias """ - if model not in self.model_group_alias: - return None - - _item = self.model_group_alias[model] - if isinstance(_item, str): - model = _item - else: - model = _item["model"] - - return model + return resolve_model_group_alias(self.model_group_alias, model) def _get_deployment_by_litellm_model(self, model: str) -> list: """ @@ -11623,7 +11611,7 @@ class Router: if model_id is not None: 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: {e!s}") + verbose_router_logger.error(f"Error in _track_deployment_metrics: {e}") def get_num_retries_from_retry_policy(self, exception: Exception, model_group: str | None = None): return _get_num_retries_from_retry_policy( diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index ff395828b2a..70e1c12665d 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -97,7 +97,7 @@ class BaseRoutingStrategy(ABC): default_sync_interval ) # Wait for DEFAULT_REDIS_SYNC_INTERVAL seconds before next sync except Exception as e: - verbose_router_logger.error(f"Error in periodic sync task: {e!s}") + verbose_router_logger.error(f"Error in periodic sync task: {e}") await asyncio.sleep( default_sync_interval ) # Still wait DEFAULT_REDIS_SYNC_INTERVAL seconds on error before retrying @@ -146,7 +146,7 @@ class BaseRoutingStrategy(ABC): return return_result except Exception as e: - verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {e!s}") + verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {e}") self.redis_increment_operation_queue = [] def add_to_in_memory_keys_to_update(self, key: str): @@ -226,4 +226,4 @@ class BaseRoutingStrategy(ABC): 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: {e!s}") + verbose_router_logger.exception(f"Error syncing in-memory cache with Redis: {e}") diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 619f1fc4629..3b8a75f4e49 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -514,7 +514,7 @@ class RouterBudgetLimiting(CustomLogger): DEFAULT_REDIS_SYNC_INTERVAL ) # Wait for DEFAULT_REDIS_SYNC_INTERVAL seconds before next sync except Exception as e: - verbose_router_logger.error(f"Error in periodic sync task: {e!s}") + verbose_router_logger.error(f"Error in periodic sync task: {e}") await asyncio.sleep( DEFAULT_REDIS_SYNC_INTERVAL ) # Still wait DEFAULT_REDIS_SYNC_INTERVAL seconds on error before retrying @@ -545,7 +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: {e!s}") + verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {e}") async def _sync_in_memory_spend_with_redis(self): """ @@ -600,7 +600,7 @@ class RouterBudgetLimiting(CustomLogger): 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: {e!s}") + verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {e}") def _get_budget_config_for_deployment( self, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 970d8de8575..1fa98f13c25 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -426,13 +426,13 @@ class ComplexityRouterConfig(BaseModel): # Session affinity: pin the first turn's routed model for the rest of the session session_affinity: bool = Field( - default=True, + default=False, description=( "When True and a session_id is resolvable on the request, pin the model chosen on the " "session's first turn and reuse it for every later turn, skipping re-classification. " - "On by default so multi-turn sessions stay on one model, preserving provider prompt " - "caches and avoiding cross-model conversation-history errors. Set False to reclassify " - "every turn." + "Off by default so every turn is classified on its own merits and routed to the cheapest " + "adequate tier. Set True to keep a multi-turn session on one model, which preserves " + "provider prompt caches and avoids cross-model conversation-history errors." ), ) session_affinity_ttl_seconds: int = Field( diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index 12820ae1237..ba7d32c42ad 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -91,7 +91,7 @@ class LowestCostLoggingHandler(CustomLogger): self.logged_success += 1 except Exception as e: verbose_logger.exception( - f"litellm.router_strategy.lowest_cost.py::log_success_event(): Exception occured - {e!s}" + f"litellm.router_strategy.lowest_cost.py::log_success_event(): Exception occured - {e}" ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -170,7 +170,7 @@ class LowestCostLoggingHandler(CustomLogger): self.logged_success += 1 except Exception as e: verbose_logger.exception( - f"litellm.proxy.hooks.prompt_injection_detection.py::async_pre_call_hook(): Exception occured - {e!s}" + f"litellm.proxy.hooks.prompt_injection_detection.py::async_pre_call_hook(): Exception occured - {e}" ) async def async_get_available_deployments( diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 2f73450b8d2..0adcdebcbf2 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -160,7 +160,7 @@ class LowestLatencyLoggingHandler(CustomLogger): self.logged_success += 1 except Exception as e: verbose_logger.exception( - f"litellm.proxy.hooks.prompt_injection_detection.py::async_pre_call_hook(): Exception occured - {e!s}" + f"litellm.proxy.hooks.prompt_injection_detection.py::async_pre_call_hook(): Exception occured - {e}" ) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): @@ -217,7 +217,7 @@ class LowestLatencyLoggingHandler(CustomLogger): return except Exception as e: verbose_logger.exception( - f"litellm.proxy.hooks.prompt_injection_detection.py::async_pre_call_hook(): Exception occured - {e!s}" + f"litellm.proxy.hooks.prompt_injection_detection.py::async_pre_call_hook(): Exception occured - {e}" ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -350,7 +350,7 @@ class LowestLatencyLoggingHandler(CustomLogger): self.logged_success += 1 except Exception as e: verbose_logger.exception( - f"litellm.router_strategy.lowest_latency.py::async_log_success_event(): Exception occured - {e!s}" + f"litellm.router_strategy.lowest_latency.py::async_log_success_event(): Exception occured - {e}" ) def _get_available_deployments( diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 4a4352fe19d..f8e7e93eb54 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -73,7 +73,7 @@ class LowestTPMLoggingHandler(CustomLogger): self.logged_success += 1 except Exception as e: verbose_router_logger.error( - f"litellm.router_strategy.lowest_tpm_rpm.py::async_log_success_event(): Exception occured - {e!s}" + f"litellm.router_strategy.lowest_tpm_rpm.py::async_log_success_event(): Exception occured - {e}" ) verbose_router_logger.debug(traceback.format_exc()) @@ -135,7 +135,7 @@ class LowestTPMLoggingHandler(CustomLogger): self.logged_success += 1 except Exception as e: verbose_router_logger.exception( - f"litellm.router_strategy.lowest_tpm_rpm.py::async_log_success_event(): Exception occured - {e!s}" + f"litellm.router_strategy.lowest_tpm_rpm.py::async_log_success_event(): Exception occured - {e}" ) verbose_router_logger.debug(traceback.format_exc()) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 03793c5577c..a81428fd5fa 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -245,7 +245,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): self.logged_success += 1 except Exception as e: verbose_logger.exception( - f"litellm.proxy.hooks.lowest_tpm_rpm_v2.py::log_success_event(): Exception occured - {e!s}" + f"litellm.proxy.hooks.lowest_tpm_rpm_v2.py::log_success_event(): Exception occured - {e}" ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -289,7 +289,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): self.logged_success += 1 except Exception as e: verbose_logger.exception( - f"litellm.proxy.hooks.lowest_tpm_rpm_v2.py::async_log_success_event(): Exception occured - {e!s}" + f"litellm.proxy.hooks.lowest_tpm_rpm_v2.py::async_log_success_event(): Exception occured - {e}" ) def _return_potential_deployments( diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py index b89e1c154f0..51a3f127b29 100644 --- a/litellm/router_utils/auto_router_model_naming.py +++ b/litellm/router_utils/auto_router_model_naming.py @@ -62,12 +62,42 @@ def classify_strategy_router_model(model: str) -> StrategyRouterKind | None: return "semantic" +def validate_complexity_router_config_write(complexity_router_config: Mapping[str, object] | None) -> str | None: + """Reject a complexity config the router would refuse to build a deployment from. + + Parsed with the router's own ``ComplexityRouterConfig`` rather than a copy of + its rules, so the boundary rejects exactly what the load would. Plugins are + resolved from dotted paths only on the config.yaml path, so a written config + reaches this function in the same shape the load hands to the same model. + Judged on the config alone: a patch may write one without naming a model, and + the stored model is encrypted at rest, so it cannot be classified here. + """ + from pydantic import ValidationError + + from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig + + if complexity_router_config is None: + return None + try: + _ = ComplexityRouterConfig.model_validate(complexity_router_config) + except ValidationError as exc: + first = exc.errors()[0] + location = ".".join(str(part) for part in first.get("loc", ())) or "complexity_router_config" + return ( + f"complexity_router_config is invalid at {location}: {first.get('msg', 'invalid value')}. " + "The router would drop this deployment at load time, so the write is rejected instead." + ) + return None + + def validate_strategy_router_model_write(model: str, present_fields: frozenset[str]) -> str | None: """Check that writing ``model`` leaves a deployment the router can load. ``present_fields`` is the set of strategy-router param fields that are non-None on the deployment after the write (stored fields merged with the incoming ones). Returns a human-readable violation, or None when coherent. + A config's contents are ``validate_complexity_router_config_write``'s to + judge, since a write may carry one without naming a model at all. """ kind = classify_strategy_router_model(model) if kind is None: diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 189296a8955..bf1c814c049 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -22,6 +22,27 @@ def _is_proxy_admin_request(request_kwargs: Mapping[str, object] | None) -> bool return getattr(user_api_key_auth, "user_role", None) == "proxy_admin" +def resolve_model_group_alias(model_group_alias: object, model: str) -> str | None: + """ + Resolve ``model`` through a ``model_group_alias`` map. + + Handles both supported entry shapes, the plain string form + ``{"alias": "target"}`` and the item form + ``{"alias": {"model": "target", "hidden": true}}``, and tolerates malformed + entries: the map can come from a key or team row rather than from validated + config, so a bad value must not raise mid-request. + + Returns the target model group, or None when the map does not rewrite ``model``. + """ + if not isinstance(model_group_alias, Mapping): + return None + entry = model_group_alias.get(model) + target = entry.get("model") if isinstance(entry, Mapping) else entry + if not isinstance(target, str) or not target: + return None + return target + + 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 diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index ef62a5d8c6c..4e9a11a4bfd 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -58,7 +58,7 @@ class CooldownCache: return cooldown_key, cooldown_data except Exception as e: - verbose_logger.error(f"CooldownCache::_common_add_cooldown_logic - Exception occurred - {e!s}") + verbose_logger.error(f"CooldownCache::_common_add_cooldown_logic - Exception occurred - {e}") raise e def add_deployment_to_cooldown( @@ -92,7 +92,7 @@ class CooldownCache: ttl=_cooldown_time, ) except Exception as e: - verbose_logger.error(f"CooldownCache::add_deployment_to_cooldown - Exception occurred - {e!s}") + verbose_logger.error(f"CooldownCache::add_deployment_to_cooldown - Exception occurred - {e}") raise e @staticmethod diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 0c92a6fa2ab..3fad860fa7d 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -190,7 +190,7 @@ async def log_success_fallback_event(original_model_group: str, kwargs: dict, or original_exception=original_exception, ) except Exception as e: - verbose_router_logger.error(f"Error in log_success_fallback_event: {e!s}") + verbose_router_logger.error(f"Error in log_success_fallback_event: {e}") async def log_failure_fallback_event(original_model_group: str, kwargs: dict, original_exception: Exception): @@ -218,7 +218,7 @@ async def log_failure_fallback_event(original_model_group: str, kwargs: dict, or original_exception=original_exception, ) except Exception as e: - verbose_router_logger.error(f"Error in log_failure_fallback_event: {e!s}") + verbose_router_logger.error(f"Error in log_failure_fallback_event: {e}") def _check_non_standard_fallback_format(fallbacks: list[Any] | None) -> bool: diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 004f7b53869..42704cea826 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -150,7 +150,7 @@ class PatternMatchRouter: matched_pattern=pattern_match, deployments=llm_deployments ) except Exception as e: - verbose_router_logger.debug(f"Error in PatternMatchRouter.route: {e!s}") + verbose_router_logger.debug(f"Error in PatternMatchRouter.route: {e}") return None # No matching pattern found 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 d67f2a2bf47..da8b452fa8a 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 @@ -212,7 +212,7 @@ class ModelRateLimitingCheck(CustomLogger): self._refund_io_token_reservation_if_any() raise except Exception as e: - verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.pre_call_check: {e!s}") + verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.pre_call_check: {e}") # Don't fail the request if rate limit check fails return deployment @@ -300,7 +300,7 @@ class ModelRateLimitingCheck(CustomLogger): await self._async_refund_io_token_reservation_if_any(parent_otel_span=parent_otel_span) raise except Exception as e: - verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.async_pre_call_check: {e!s}") + verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.async_pre_call_check: {e}") # Don't fail the request if rate limit check fails return deployment @@ -360,7 +360,7 @@ class ModelRateLimitingCheck(CustomLogger): ) except Exception as e: - verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.async_log_success_event: {e!s}") + verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.async_log_success_event: {e}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): from litellm.litellm_core_utils.core_helpers import ( @@ -418,7 +418,7 @@ class ModelRateLimitingCheck(CustomLogger): ) except Exception as e: - verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.log_success_event: {e!s}") + verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.log_success_event: {e}") def log_failure_event(self, kwargs, response_obj, start_time, end_time): with contextlib.suppress(Exception): diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index 531b2b577b1..0ce0d4229c1 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -77,7 +77,7 @@ class SearchAPIRouter: 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: {e!s}") + verbose_router_logger.exception(f"Error updating router with search tools: {e}") raise e @staticmethod @@ -226,6 +226,6 @@ class SearchAPIRouter: except Exception as e: verbose_router_logger.error( - f"Error in SearchAPIRouter.async_search_with_fallbacks_helper for {search_tool_name}: {e!s}" + f"Error in SearchAPIRouter.async_search_with_fallbacks_helper for {search_tool_name}: {e}" ) raise e diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index a05ea367b19..2982d30274b 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -282,7 +282,7 @@ def get_secret( raise ValueError("Azure OIDC provider returned None token") return oidc_token except Exception as e: - error_msg = f"Azure OIDC provider failed: {e!s}" + error_msg = f"Azure OIDC provider failed: {e}" verbose_logger.error(error_msg) raise ValueError(error_msg) with open(azure_federated_token_file, "r") as f: @@ -335,7 +335,7 @@ def get_secret( ) except Exception as e: # check if it's in os.environ verbose_logger.error( - f"Defaulting to os.environ value for key={secret_name}. An exception occurred - {e!s}.\n\n{traceback.format_exc()}" + f"Defaulting to os.environ value for key={secret_name}. An exception occurred - {e}.\n\n{traceback.format_exc()}" ) secret = os.getenv(secret_name) try: diff --git a/litellm/secret_managers/secret_manager_handler.py b/litellm/secret_managers/secret_manager_handler.py index 2acb154dd59..64a00f0df58 100644 --- a/litellm/secret_managers/secret_manager_handler.py +++ b/litellm/secret_managers/secret_manager_handler.py @@ -119,7 +119,7 @@ def get_secret_from_manager( if secret is None: raise ValueError(f"No secret found in Google Secret Manager for {secret_name}") except Exception as e: - print_verbose(f"An error occurred - {e!s}") + print_verbose(f"An error occurred - {e}") raise e elif key_manager == KeyManagementSystem.HASHICORP_VAULT.value: @@ -128,7 +128,7 @@ def get_secret_from_manager( if secret is None: raise ValueError(f"No secret found in Hashicorp Secret Manager for {secret_name}") except Exception as e: - print_verbose(f"An error occurred - {e!s}") + print_verbose(f"An error occurred - {e}") raise e elif key_manager == KeyManagementSystem.CYBERARK.value: @@ -137,7 +137,7 @@ def get_secret_from_manager( if secret is None: raise ValueError(f"No secret found in CyberArk Secret Manager for {secret_name}") except Exception as e: - print_verbose(f"An error occurred - {e!s}") + print_verbose(f"An error occurred - {e}") raise e elif key_manager == KeyManagementSystem.CUSTOM.value: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index d9f8229dbed..21de28a9b29 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -301,7 +301,7 @@ class BedrockToolSpec(dict): "name": name, "description": description, } - if supports_strict_tools and strict is not None: + if supports_strict_tools and strict: tool_spec["strict"] = strict super().__init__(toolSpec=tool_spec) diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 2d9387da956..f3d0529a7ea 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,13 +1,12 @@ from typing import Any, Dict, List, Literal, Optional, Union -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, Field from litellm.proxy._types import ( KeyManagementRoutes, LiteLLM_DeletedTeamTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, - LiteLLM_UserTable, Member, ) @@ -125,3 +124,22 @@ class TeamMemberInfoResponse(LiteLLM_TeamMembership): role: Optional[str] = None user_email: Optional[str] = None team_alias: Optional[str] = None + + +class TeamMetadataFieldSchema(BaseModel): + """One declared team metadata field from ``general_settings.team_metadata_schema``. + + Advisory only: the UI uses it to prepopulate the team metadata form. + Enforcement stays with ``custom_team_metadata_validate``. + """ + + model_config = ConfigDict(extra="forbid") + + key: str = Field(min_length=1) + label: Optional[str] = None + + +class TeamMetadataSchemaResponse(BaseModel): + """Response for GET /team/metadata_schema; ``fields`` is empty when no schema is configured.""" + + fields: tuple[TeamMetadataFieldSchema, ...] diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index d4b1d98f957..f68d818d991 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -229,3 +229,7 @@ class DefaultTeamSSOParams(LiteLLMPydanticObjectBase): default=None, description="Default permissions granted to members of newly created teams (e.g. /key/generate, /key/update, /key/delete). /key/info and /key/health are always included.", ) + organization_id: str | None = Field( + default=None, + description="Default organization for new teams created without an explicit organization", + ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 18991f53e6f..3539ac0f27a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3286,8 +3286,15 @@ agentic_loop_internal_litellm_params = [ "_code_interpreter_interception_converted_stream", ] +# Proxy-owned callback credentials, stamped from admin-configured team/key callback +# settings. Listed in all_litellm_params for the same reason as the agentic-loop +# fields above: an unrecognized top-level key is swept into extra_body and sent to +# the provider. +TRUSTED_CALLBACK_VARS_FIELD = "litellm_trusted_callback_vars" + all_litellm_params = ( agentic_loop_internal_litellm_params + + [TRUSTED_CALLBACK_VARS_FIELD] + [ "metadata", "litellm_metadata", diff --git a/litellm/utils.py b/litellm/utils.py index 6ef3871a3c1..eb3e578b7e8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -947,7 +947,7 @@ 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: {e!s}") + verbose_logger.warning(f"Error removing thought signatures from tool call IDs: {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: @@ -1004,7 +1004,7 @@ def function_setup( else: messages = "default-message-value" except Exception as e: - verbose_logger.debug(f"Error extracting messages from Google contents: {e!s}") + verbose_logger.debug(f"Error extracting messages from Google contents: {e}") messages = "default-message-value" else: messages = "default-message-value" @@ -1410,7 +1410,7 @@ def client(original_function): ) kwargs["max_tokens"] = modified_max_tokens except Exception as e: - print_verbose(f"Error while checking max token limit: {e!s}") + print_verbose(f"Error while checking max token limit: {e}") # MODEL CALL result = original_function(*args, **kwargs) end_time = datetime.datetime.now() @@ -1675,7 +1675,7 @@ def client(original_function): ) kwargs["max_tokens"] = modified_max_tokens except Exception as e: - print_verbose(f"Error while checking max token limit: {e!s}") + print_verbose(f"Error while checking max token limit: {e}") # MODEL CALL result = await original_function(*args, **kwargs) @@ -2224,7 +2224,7 @@ def supports_native_streaming(model: str, custom_llm_provider: str | None) -> bo return supports_native_streaming except Exception as e: verbose_logger.debug( - f"Model not found or error in checking supports_native_streaming support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {e!s}" + f"Model not found or error in checking supports_native_streaming support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {e}" ) return False @@ -2248,7 +2248,7 @@ def supports_response_schema(model: str, custom_llm_provider: str | None = None) 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: {e!s}" + f"Model not found or error in checking response schema support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {e}" ) return False @@ -2362,7 +2362,7 @@ def _supports_factory(model: str, custom_llm_provider: str | None, key: str) -> return False except Exception as e: verbose_logger.debug( - f"Model not found or error in checking {key} support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {e!s}" + f"Model not found or error in checking {key} support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {e}" ) supported_by_provider = _supports_provider_info_factory(model, custom_llm_provider, key) @@ -2404,7 +2404,7 @@ def _is_explicitly_disabled_factory(model: str, custom_llm_provider: str | None, verbose_logger.debug( f"Model not found or error in checking {key} disabled state. " f"You passed model={model}, custom_llm_provider={custom_llm_provider}. " - f"Error: {e!s}" + f"Error: {e}" ) return False @@ -2537,7 +2537,7 @@ def get_supported_regions(model: str, custom_llm_provider: str | None = None) -> return None except Exception as e: verbose_logger.debug( - f"Model not found or error in checking supported_regions support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {e!s}" + f"Model not found or error in checking supported_regions support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {e}" ) return None @@ -6542,7 +6542,7 @@ class TextCompletionStreamWrapper: return response except Exception as e: - raise Exception(f"Error occurred converting to text completion object - chunk: {chunk}; Error: {e!s}") + raise Exception(f"Error occurred converting to text completion object - chunk: {chunk}; Error: {e}") def __next__(self): # model_response = ModelResponse(stream=True, model=self.model) @@ -6868,7 +6868,7 @@ def trim_messages( return final_messages, response_tokens return final_messages except Exception as e: # [NON-Blocking, if error occurs just return final_messages - verbose_logger.exception(f"Got exception while token trimming - {e!s}") + verbose_logger.exception(f"Got exception while token trimming - {e}") return original_messages diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 4abd587bce5..1350e2b187e 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -235,7 +235,7 @@ class VectorStoreRegistry: 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: {e!s}") + verbose_logger.debug(f"Error fetching vector store from database: {e}") return None @@ -346,7 +346,7 @@ class VectorStoreRegistry: 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: {e!s}") + verbose_logger.debug(f"Error verifying vector store {vector_store_id} in database: {e}") # Fall back to database if not found in memory (or was deleted) if vector_store is None and prisma_client is not None: @@ -355,7 +355,7 @@ class VectorStoreRegistry: 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: {e!s}") + verbose_logger.debug(f"Error fetching vector store {vector_store_id} from database: {e}") if vector_store is not None: # Create a copy to avoid modifying the registry diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 346f613ea3e..56e0391f419 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1975,6 +1975,7 @@ "prompt_cache_min_tokens": 2048 }, "anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -2011,6 +2012,7 @@ "prompt_cache_min_tokens": 1024 }, "global.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -2047,6 +2049,7 @@ "prompt_cache_min_tokens": 1024 }, "us.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2083,6 +2086,7 @@ "prompt_cache_min_tokens": 1024 }, "eu.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2119,6 +2123,7 @@ "prompt_cache_min_tokens": 1024 }, "au.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2155,6 +2160,7 @@ "prompt_cache_min_tokens": 1024 }, "jp.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 386108169bc..d27b168d6ca 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -24,7 +24,7 @@ "limit": 130 }, "ANN401": { - "limit": 1849 + "limit": 1848 }, "ASYNC230": { "limit": 14 diff --git a/tests/e2e/ui/tests/mcp/mcpServers.spec.ts b/tests/e2e/ui/tests/mcp/mcpServers.spec.ts index 37aabf9c057..43f21e77fbd 100644 --- a/tests/e2e/ui/tests/mcp/mcpServers.spec.ts +++ b/tests/e2e/ui/tests/mcp/mcpServers.spec.ts @@ -36,7 +36,7 @@ test.describe("MCP Servers", () => { await formModal.locator('input[id="url"]').fill("https://e2e-fake-mcp.test.local/mcp"); // Authentication: None - // The auth_type Form.Item has no label prop (create_mcp_server.tsx:795), so + // The auth_type Form.Item has no label prop (CreateMCPServer.tsx), so // it can't be anchored by label text. Scope via the enclosing Collapse // panel ("Authentication") instead — that anchor is stable even if the // placeholder copy changes. diff --git a/tests/proxy_behavior/management/test_team_budget_limits.py b/tests/proxy_behavior/management/test_team_budget_limits.py index dad775370ad..96a6fe7234a 100644 --- a/tests/proxy_behavior/management/test_team_budget_limits.py +++ b/tests/proxy_behavior/management/test_team_budget_limits.py @@ -10,14 +10,14 @@ Pins the five helpers Driven through /team/new + /team/update. -Structural finding pinned here, identical in shape to F1's org aggregate: -both call sites (lines 985 + 1751) load the org via `get_org_object` -WITHOUT `include_budget_table=True`, so `org_table.litellm_budget_table` -is `None` and the org max_budget / org tpm / org rpm guards inside -`_check_org_team_limits` (lines 641–694, 670–694) silently no-op. The -`models` subset guard (lines 654–667) IS reachable because it reads -`org_table.models` directly. The `_check_user_team_limits` guards reach -all branches through `user_api_key_dict`, no relation include needed. +Structural finding, updated: /team/new loads the org via `get_org_object` +WITH `include_budget_table=True`, so the org max_budget / org tpm / org rpm +guards inside `_check_org_team_limits` are live there and are pinned as +enforced below. /team/update still loads the org without the budget +relation, so its budget guards remain no-ops. The `models` subset guard IS +reachable on both because it reads `org_table.models` directly. The +`_check_user_team_limits` guards reach all branches through +`user_api_key_dict`, no relation include needed. """ import uuid @@ -132,48 +132,67 @@ async def test_check_org_team_limits_models_subset( headers={"Authorization": f"Bearer {seeder}"}, json=body, ) - assert ( - resp.status_code == expected_status - ), f"{body!r} → {resp.status_code}: {resp.text}" + assert resp.status_code == expected_status, f"{body!r} → {resp.status_code}: {resp.text}" rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id}) assert len(rows) == (1 if expected_status == 200 else 0) # --------------------------------------------------------------------------- -# _check_org_team_limits — budget / tpm / rpm structurally unreachable -# (org_table.litellm_budget_table is None at guard time). Pin the -# no-op behavior so a future change that flips include_budget_table=True -# turns these into reds. +# _check_org_team_limits — budget / tpm / rpm live on /team/new since its +# get_org_object call passes include_budget_table=True. (/team/update still +# loads the org without the budget relation, so its guards remain no-ops.) # --------------------------------------------------------------------------- -_ORG_BUDGET_DEAD_SCENARIOS = [ +_ORG_BUDGET_ENFORCED_SCENARIOS = [ ( - "org_budget/over_max_budget_unenforced", + "org_budget/over_max_budget_rejected", {"max_budget": 100, "tpm_limit": None, "rpm_limit": None}, {"max_budget": 999_999}, + 400, ), ( - "org_tpm/over_unenforced", + "org_budget/within_max_budget_accepted", + {"max_budget": 100, "tpm_limit": None, "rpm_limit": None}, + {"max_budget": 50}, + 200, + ), + ( + "org_tpm/over_rejected", {"max_budget": None, "tpm_limit": 100, "rpm_limit": None}, {"tpm_limit": 999_999}, + 400, ), ( - "org_rpm/over_unenforced", + "org_tpm/within_accepted", + {"max_budget": None, "tpm_limit": 100, "rpm_limit": None}, + {"tpm_limit": 50}, + 200, + ), + ( + "org_rpm/over_rejected", {"max_budget": None, "tpm_limit": None, "rpm_limit": 100}, {"rpm_limit": 999_999}, + 400, + ), + ( + "org_rpm/within_accepted", + {"max_budget": None, "tpm_limit": None, "rpm_limit": 100}, + {"rpm_limit": 50}, + 200, ), ] @pytest.mark.parametrize( - "org_budget,body_extras", - [(b, c) for (_id, b, c) in _ORG_BUDGET_DEAD_SCENARIOS], - ids=[s[0] for s in _ORG_BUDGET_DEAD_SCENARIOS], + "org_budget,body_extras,expected_status", + [(b, c, d) for (_id, b, c, d) in _ORG_BUDGET_ENFORCED_SCENARIOS], + ids=[s[0] for s in _ORG_BUDGET_ENFORCED_SCENARIOS], ) -async def test_check_org_team_limits_budget_dead_code_pin( +async def test_check_org_team_limits_budget_enforced( org_budget, body_extras: Dict[str, Any], + expected_status: int, proxy_client, prisma, scratch, @@ -192,9 +211,9 @@ async def test_check_org_team_limits_budget_dead_code_pin( **body_extras, }, ) - assert resp.status_code == 200, resp.text + assert resp.status_code == expected_status, f"{body_extras!r} → {resp.status_code}: {resp.text}" rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id}) - assert len(rows) == 1 + assert len(rows) == (1 if expected_status == 200 else 0) # --------------------------------------------------------------------------- @@ -279,9 +298,9 @@ async def test_check_user_team_limits( **body_extras, }, ) - assert ( - resp.status_code == expected_status - ), f"caps={actor_caps} body={body_extras} → {resp.status_code}: {resp.text}" + assert resp.status_code == expected_status, ( + f"caps={actor_caps} body={body_extras} → {resp.status_code}: {resp.text}" + ) rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id}) assert len(rows) == (1 if expected_status == 200 else 0) @@ -376,9 +395,7 @@ async def test_proxy_admin_raise_budget_allowed(proxy_client, prisma, scratch): async def test_team_admin_remove_budget_cap_blocked(proxy_client, prisma, scratch): """A team admin cannot strip the team's cap (max_budget=null); removing the ceiling is the strongest possible raise -> proxy-admin only.""" - caller_cleartext = await _seed_scratch_actor_with_caps( - prisma, scratch.prefix, max_budget=100000.0 - ) + caller_cleartext = await _seed_scratch_actor_with_caps(prisma, scratch.prefix, max_budget=100000.0) team_id = await create_scratch_team( prisma, team_id=scratch.tag("team"), diff --git a/tests/proxy_behavior/management/test_team_new.py b/tests/proxy_behavior/management/test_team_new.py index 7b07f259641..9846566d0b9 100644 --- a/tests/proxy_behavior/management/test_team_new.py +++ b/tests/proxy_behavior/management/test_team_new.py @@ -72,13 +72,9 @@ async def test_team_new_authz_matrix( headers={"Authorization": f"Bearer {caller.cleartext}"}, json=body, ) - assert ( - resp.status_code == expected_status - ), f"{actor.value} org={org_target}: {resp.status_code} {resp.text}" + assert resp.status_code == expected_status, f"{actor.value} org={org_target}: {resp.status_code} {resp.text}" - row = await prisma.db.litellm_teamtable.find_unique( - where={"team_id": scratch.prefix} - ) + row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix}) if expected_status == 200: assert row is not None assert row.organization_id == org_id @@ -94,9 +90,7 @@ async def test_team_new_rejects_negative_budget(proxy_client, prisma, scratch, w json={"team_id": scratch.prefix, "max_budget": -1}, ) assert resp.status_code == 400, resp.text - row = await prisma.db.litellm_teamtable.find_unique( - where={"team_id": scratch.prefix} - ) + row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix}) assert row is None @@ -118,12 +112,10 @@ async def test_team_new_rejects_duplicate_team_id(proxy_client, prisma, scratch, assert second.status_code == 400, second.text -async def test_team_new_unknown_organization_is_500( - proxy_client, prisma, scratch, world -): - """SURFACED, NOT ENDORSED: a /team/new with an organization_id that does - not exist currently fails 500 (the role-resolution layer raises before - the handler's own 400 'Organization not found' check is reached).""" +async def test_team_new_unknown_organization_is_400(proxy_client, prisma, scratch, world): + """A /team/new with an organization_id that does not exist fails 400: + OrganizationNotFoundError is routed into the handler's own + 'Organization not found' guard instead of escaping as a 500.""" resp = await proxy_client.post( "/team/new", headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, @@ -132,8 +124,7 @@ async def test_team_new_unknown_organization_is_500( "organization_id": scratch.tag("no-such-org"), }, ) - assert resp.status_code == 500, resp.text - row = await prisma.db.litellm_teamtable.find_unique( - where={"team_id": scratch.prefix} - ) + assert resp.status_code == 400, resp.text + assert "Organization not found" in resp.text + row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix}) assert row is None diff --git a/tests/store_model_in_db_tests/cost_center_service.py b/tests/store_model_in_db_tests/cost_center_service.py new file mode 100644 index 00000000000..c2820dfbe7e --- /dev/null +++ b/tests/store_model_in_db_tests/cost_center_service.py @@ -0,0 +1,54 @@ +"""Stand-in cost center validation service for the team metadata e2e tests. + +Accepts POST /validate with {"operation": ..., "metadata": {...}} and answers +{"ok": true} or {"ok": false, "reason": ...} based on a static allowlist. +GET /health answers 200 for the CI wait loop. +""" + +import argparse +import json +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +ALLOWED_COST_CENTERS = {"CC-1001", "CC-1002"} + + +class CostCenterHandler(BaseHTTPRequestHandler): + def _respond(self, status: int, body: dict) -> None: + payload = json.dumps(body).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def do_GET(self): + if self.path == "/health": + self._respond(200, {"status": "healthy"}) + return + self._respond(404, {"error": "not found"}) + + def do_POST(self): + if self.path != "/validate": + self._respond(404, {"error": "not found"}) + return + length = int(self.headers.get("Content-Length", 0)) + body = json.loads(self.rfile.read(length) or b"{}") + cost_center = (body.get("metadata") or {}).get("cost_center") + if cost_center is None: + self._respond(200, {"ok": False, "reason": "cost_center missing per cost center service"}) + elif cost_center not in ALLOWED_COST_CENTERS: + self._respond(200, {"ok": False, "reason": f"cost center {cost_center} rejected by cost center service"}) + else: + self._respond(200, {"ok": True}) + + def log_message(self, format, *args): + pass + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=9414) + args = parser.parse_args() + print(f"cost center service listening on {args.host}:{args.port}") + ThreadingHTTPServer((args.host, args.port), CostCenterHandler).serve_forever() diff --git a/tests/store_model_in_db_tests/test_team_metadata_validation_e2e.py b/tests/store_model_in_db_tests/test_team_metadata_validation_e2e.py new file mode 100644 index 00000000000..050b59ef426 --- /dev/null +++ b/tests/store_model_in_db_tests/test_team_metadata_validation_e2e.py @@ -0,0 +1,171 @@ +"""E2E matrix for custom team metadata validation against the DB-backed proxy. + +The proxy (see store_model_db_config.yaml) registers +team_metadata_validator_e2e.validate_team_metadata, which dispatches per +request to one of three independent implementations via the +`_e2e_validator_impl` metadata key: a static allowlist function, an +HTTP-backed function calling the cost center service started by CI, and an +immutability-enforcing class instance. Metadata without the dispatch key is +accepted untouched, so the rest of this suite is unaffected. +""" + +import os +import uuid + +import httpx +import pytest + +PROXY_BASE_URL = os.getenv("PROXY_BASE_URL", "http://localhost:4000") +MASTER_KEY = os.getenv("LITELLM_MASTER_KEY", "sk-1234") +HEADERS = {"Authorization": f"Bearer {MASTER_KEY}", "Content-Type": "application/json"} + +UNAVAILABLE_MESSAGE = "Cost center validation is unavailable right now; the team was not saved. Contact FinOps." + +IMPLS = ["allowlist", "http", "immutable"] + +REQUIRED_MESSAGES = { + "allowlist": "cost_center is required in team metadata", + "http": "cost_center missing per cost center service", + "immutable": "cost_center is required in team metadata", +} +UNKNOWN_MESSAGES = { + "allowlist": "is not recognized", + "http": "rejected by cost center service", +} + + +def _meta(impl, **fields): + return {"_e2e_validator_impl": impl, **fields} + + +def _create_team(metadata, team_id=None): + body = {"team_alias": f"meta-validate-{uuid.uuid4().hex[:8]}"} + if team_id is not None: + body["team_id"] = team_id + if metadata is not None: + body["metadata"] = metadata + return httpx.post(f"{PROXY_BASE_URL}/team/new", headers=HEADERS, json=body, timeout=30) + + +def _patch_team(team_id, body): + return httpx.patch(f"{PROXY_BASE_URL}/team/{team_id}", headers=HEADERS, json=body, timeout=30) + + +def _post_update(team_id, body): + return httpx.post(f"{PROXY_BASE_URL}/team/update", headers=HEADERS, json={"team_id": team_id, **body}, timeout=30) + + +def _team_info(team_id): + return httpx.get(f"{PROXY_BASE_URL}/team/info", headers=HEADERS, params={"team_id": team_id}, timeout=30) + + +def _delete_team(team_id): + httpx.post(f"{PROXY_BASE_URL}/team/delete", headers=HEADERS, json={"team_ids": [team_id]}, timeout=30) + + +@pytest.fixture +def team_with_cost_center(request): + impl = request.param + team_id = f"meta-validate-{impl}-{uuid.uuid4().hex[:8]}" + response = _create_team(metadata=_meta(impl, cost_center="CC-1001"), team_id=team_id) + assert response.status_code == 200, response.text + yield impl, team_id + _delete_team(team_id) + + +@pytest.mark.parametrize("impl", IMPLS) +def test_create_with_valid_cost_center_succeeds(impl): + response = _create_team(metadata=_meta(impl, cost_center="CC-1001")) + assert response.status_code == 200, response.text + team_id = response.json()["team_id"] + try: + assert response.json()["metadata"]["cost_center"] == "CC-1001" + finally: + _delete_team(team_id) + + +@pytest.mark.parametrize("impl", IMPLS) +def test_create_without_cost_center_is_rejected(impl): + team_id = f"meta-validate-reject-{impl}-{uuid.uuid4().hex[:8]}" + response = _create_team(metadata=_meta(impl), team_id=team_id) + assert response.status_code == 400, response.text + assert REQUIRED_MESSAGES[impl] in response.text + info = _team_info(team_id) + assert info.status_code == 404, "rejected create must not leave a team row behind" + + +@pytest.mark.parametrize("impl", IMPLS) +def test_create_with_unknown_cost_center(impl): + response = _create_team(metadata=_meta(impl, cost_center="CC-9999")) + if impl == "immutable": + assert response.status_code == 200, response.text + _delete_team(response.json()["team_id"]) + return + assert response.status_code == 400, response.text + assert UNKNOWN_MESSAGES[impl] in response.text + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_patch_changing_cost_center(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _patch_team(team_id, {"metadata": {"cost_center": "CC-1002"}}) + if impl == "immutable": + assert response.status_code == 400, response.text + assert "immutable once set" in response.text + info = _team_info(team_id).json()["team_info"]["metadata"] + assert info["cost_center"] == "CC-1001", "blocked update must leave stored metadata intact" + return + assert response.status_code == 200, response.text + assert response.json()["metadata"]["cost_center"] == "CC-1002" + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_patch_unrelated_key_validates_merged_result(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _patch_team(team_id, {"metadata": {"team_notes": "hello"}}) + assert response.status_code == 200, response.text + merged = response.json()["metadata"] + assert merged["cost_center"] == "CC-1001" + assert merged["team_notes"] == "hello" + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_patch_null_deleting_cost_center_is_rejected(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _patch_team(team_id, {"metadata": {"cost_center": None}}) + assert response.status_code == 400, response.text + assert REQUIRED_MESSAGES[impl] in response.text + info = _team_info(team_id).json()["team_info"]["metadata"] + assert info["cost_center"] == "CC-1001" + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_post_update_dropping_cost_center_is_rejected(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _post_update(team_id, {"metadata": _meta(impl, team_notes="only-notes")}) + assert response.status_code == 400, response.text + assert REQUIRED_MESSAGES[impl] in response.text + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_update_without_metadata_skips_validation(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _post_update(team_id, {"tpm_limit": 55}) + assert response.status_code == 200, response.text + + +def test_http_service_outage_fails_closed_with_configured_message(): + response = _create_team(metadata=_meta("http_down", cost_center="CC-1001")) + assert response.status_code == 503, response.text + assert UNAVAILABLE_MESSAGE in response.text + + +def test_metadata_without_dispatch_key_is_untouched(): + response = _create_team(metadata={"any_key": "any_value"}) + assert response.status_code == 200, response.text + team_id = response.json()["team_id"] + try: + update = _post_update(team_id, {"metadata": {"any_key": "changed"}}) + assert update.status_code == 200, update.text + finally: + _delete_team(team_id) diff --git a/tests/test_litellm/caching/test_evicted_client_closer.py b/tests/test_litellm/caching/test_evicted_client_closer.py new file mode 100644 index 00000000000..a08fd58079d --- /dev/null +++ b/tests/test_litellm/caching/test_evicted_client_closer.py @@ -0,0 +1,409 @@ +""" +Tests for EvictedClientCloser. + +An evicted client must stay open long enough for a request that already holds it +to finish, and must then actually be closed, otherwise its connection pool is +retained until a generational collection runs. A client the caller supplied is +never closed, because litellm does not own its lifecycle. +""" + +import asyncio +import gc +import weakref + +import httpx +import pytest + +from litellm.caching.evicted_client_closer import EvictedClientCloser +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + +class FakeClock: + """Hand-advanced monotonic clock, so grace windows need no real waiting.""" + + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +class AsyncClient: + def __init__(self) -> None: + self.closed = False + + async def close(self) -> None: + self.closed = True + + +class SyncClient: + def __init__(self) -> None: + self.closed = False + + def close(self) -> None: + self.closed = True + + +class CountingDeadline(float): + """A clock reading that tallies every deadline comparison made against it. + + Deadline comparisons are the work a reap does, so counting them says whether + that work tracks the entries that are due or the size of the whole queue. + """ + + comparisons = 0 + + def __add__(self, other: float) -> "CountingDeadline": + return CountingDeadline(float(self) + other) + + def __le__(self, other: float) -> bool: + CountingDeadline.comparisons += 1 + return float(self) <= float(other) + + def __gt__(self, other: float) -> bool: + CountingDeadline.comparisons += 1 + return float(self) > float(other) + + +def make_closer(clock: FakeClock, grace_seconds: float = 60.0) -> EvictedClientCloser: + return EvictedClientCloser(grace_seconds=grace_seconds, clock=clock) + + +async def _trickling_upstream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + """Serves a chunked body slowly, so a request stays on the wire long enough to observe.""" + await reader.read(4096) + writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n") + await writer.drain() + for _ in range(6): + writer.write(b"5\r\nhello\r\n") + await writer.drain() + await asyncio.sleep(0.1) + writer.write(b"0\r\n\r\n") + await writer.drain() + + +@pytest.mark.asyncio +async def test_owned_client_is_closed_once_the_grace_window_elapses(): + clock = FakeClock() + closer = make_closer(clock) + client = AsyncClient() + + closer.mark_owned(client) + closer.schedule(client) + clock.advance(61.0) + closer.reap() + await asyncio.sleep(0.05) + + assert client.closed is True + assert closer.pending_count == 0 + + +@pytest.mark.asyncio +async def test_owned_client_stays_open_inside_the_grace_window(): + """A request handed the client just before eviction is still using it.""" + clock = FakeClock() + closer = make_closer(clock) + client = AsyncClient() + + closer.mark_owned(client) + closer.schedule(client) + clock.advance(59.0) + closer.reap() + await asyncio.sleep(0.05) + + assert client.closed is False + assert closer.pending_count == 1 + + +@pytest.mark.asyncio +async def test_caller_supplied_client_is_never_closed(): + clock = FakeClock() + closer = make_closer(clock) + client = AsyncClient() + + closer.schedule(client) + clock.advance(3600.0) + closer.reap() + await asyncio.sleep(0.05) + + assert client.closed is False + assert closer.pending_count == 0 + + +@pytest.mark.asyncio +async def test_sync_client_is_closed_once_the_grace_window_elapses(): + clock = FakeClock() + closer = make_closer(clock) + client = SyncClient() + + closer.mark_owned(client) + closer.schedule(client) + clock.advance(61.0) + closer.reap() + + assert client.closed is True + + +@pytest.mark.asyncio +async def test_a_failing_close_does_not_propagate_or_block_the_others(): + class ExplodingClient: + async def close(self) -> None: + raise RuntimeError("connection already gone") + + clock = FakeClock() + closer = make_closer(clock) + exploding, healthy = ExplodingClient(), AsyncClient() + + for client in (exploding, healthy): + closer.mark_owned(client) + closer.schedule(client) + clock.advance(61.0) + closer.reap() + await asyncio.sleep(0.05) + + assert healthy.closed is True + + +@pytest.mark.asyncio +async def test_an_unhashable_cached_value_does_not_break_eviction(): + """The cache holds arbitrary values; an ownership test must never raise on one.""" + + class Unhashable: + __hash__ = None # pyright: ignore[reportAssignmentType] # unhashable by construction + + clock = FakeClock() + closer = make_closer(clock) + + closer.mark_owned(Unhashable()) + closer.schedule(Unhashable()) + + assert closer.pending_count == 0 + + +@pytest.mark.asyncio +async def test_values_with_nothing_to_close_are_never_queued(): + """The cache holds plain values too; those have nothing to reclaim.""" + + class NotAClient: + pass + + clock = FakeClock() + closer = make_closer(clock) + value = NotAClient() + + closer.mark_owned(value) + closer.schedule(value) + + assert closer.pending_count == 0 + + +@pytest.mark.asyncio +async def test_a_queued_client_is_not_kept_alive_by_the_queue(): + """Waiting out a grace window must not retain what the collector would free first.""" + clock = FakeClock() + closer = make_closer(clock) + client = AsyncClient() + gone = weakref.ref(client) + + closer.mark_owned(client) + closer.schedule(client) + del client + gc.collect() + + assert gone() is None, "the pending queue is holding the client alive" + + clock.advance(61.0) + closer.reap() + assert closer.pending_count == 0 + + +def test_sync_client_evicted_outside_an_event_loop_is_still_closed(): + """The sync httpx handler is cached and evicted from call sites with no loop.""" + clock = FakeClock() + closer = make_closer(clock) + client = SyncClient() + + closer.mark_owned(client) + closer.schedule(client) + assert closer.pending_count == 1 + + clock.advance(61.0) + closer.reap() + + assert client.closed is True + assert closer.pending_count == 0 + + +@pytest.mark.asyncio +async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped(): + clock = FakeClock() + closer = make_closer(clock) + client = AsyncClient() + closer.mark_owned(client) + + def schedule_outside_a_loop() -> None: + closer.schedule(client) + clock.advance(61.0) + closer.reap() + + await asyncio.to_thread(schedule_outside_a_loop) + assert client.closed is False, "no loop was running, so it could not have been closed" + assert closer.pending_count == 1 + + closer.reap() + await asyncio.sleep(0.05) + + assert client.closed is True + + +@pytest.mark.asyncio +async def test_a_client_evicted_on_another_event_loop_is_left_alone(): + """Closing a client bound to a different loop would schedule work on that loop.""" + clock = FakeClock() + closer = make_closer(clock) + client = AsyncClient() + closer.mark_owned(client) + + def schedule_on_its_own_loop() -> None: + asyncio.run(_schedule()) + + async def _schedule() -> None: + closer.schedule(client) + + await asyncio.to_thread(schedule_on_its_own_loop) + assert closer.pending_count == 1 + + clock.advance(61.0) + closer.reap() + await asyncio.sleep(0.05) + + assert client.closed is False + assert closer.pending_count == 1 + + +@pytest.mark.asyncio +async def test_a_client_serving_a_request_is_not_closed_when_its_grace_window_ends(): + """The grace window on its own cannot promise that a request has finished. + + ``litellm.request_timeout`` defaults to 6000 seconds and a streaming response + is bounded only by how long the upstream keeps sending, so a client past its + deadline is closed only once its own pool reports nothing in flight. + """ + server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + clock = FakeClock() + closer = make_closer(clock) + client = httpx.AsyncClient() + + closer.mark_owned(client) + closer.schedule(client) + + async def read_the_stream() -> int: + received = 0 + async with client.stream("GET", f"http://127.0.0.1:{port}/") as response: + async for chunk in response.aiter_bytes(): + received += len(chunk) + return received + + streaming = asyncio.create_task(read_the_stream()) + await asyncio.sleep(0.25) # the request is on the wire + clock.advance(3600.0) # and its grace window is long gone + closer.reap() + await asyncio.sleep(0.05) + + assert client.is_closed is False, "closed a client that was serving a request" + assert await streaming > 0, "the in-flight request did not survive the reap" + + clock.advance(3600.0) + closer.reap() + await asyncio.sleep(0.05) + + assert client.is_closed is True, "an idle client past its grace window must be closed" + assert closer.pending_count == 0 + server.close() + + +@pytest.mark.asyncio +async def test_the_aiohttp_backed_handler_is_not_closed_mid_request(): + """The default async path is aiohttp-backed, whose pool accounts for its own leases.""" + server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + clock = FakeClock() + closer = make_closer(clock) + handler = AsyncHTTPHandler() + + closer.mark_owned(handler) + closer.schedule(handler) + + request = asyncio.create_task(handler.get(f"http://127.0.0.1:{port}/")) + await asyncio.sleep(0.25) + clock.advance(3600.0) + closer.reap() + await asyncio.sleep(0.05) + + assert handler.client.is_closed is False, "closed a handler that was serving a request" + assert (await request).status_code == 200 + + clock.advance(3600.0) + closer.reap() + await asyncio.sleep(0.05) + + assert handler.client.is_closed is True + server.close() + + +def test_the_pending_queue_cannot_grow_past_its_bound(): + """A caller that churns the client cache must not be able to grow this queue.""" + clock = FakeClock() + closer = EvictedClientCloser(grace_seconds=60.0, max_pending=8, clock=clock) + clients = tuple(SyncClient() for _ in range(50)) + + for client in clients: + closer.mark_owned(client) + closer.schedule(client) + + assert closer.pending_count == 8, "the queue grew past max_pending" + + clock.advance(61.0) + closer.reap() + + assert closer.pending_count == 0 + assert sum(client.closed for client in clients) == 8, "everything queued should have been closed" + + +def test_a_reap_looks_at_what_is_due_rather_than_at_the_whole_queue(): + """Sustained churn evicts a client per request, and every read of the cache reaps. + + So the cost of a reap has to track the entries that are due, not the length of + the queue; a reap that filters the whole queue makes the pair quadratic. Each + bucket is ordered by deadline, so an up-to-date reap compares one entry per + bucket and stops. Counting the comparisons measures that directly, where a + wall-clock budget would only measure the machine. + """ + evictions = 1_000 + clock = FakeClock() + closer = EvictedClientCloser( + grace_seconds=60.0, + max_pending=evictions, + clock=lambda: CountingDeadline(clock.now), + ) + clients = tuple(SyncClient() for _ in range(evictions)) + for client in clients: + closer.mark_owned(client) + + CountingDeadline.comparisons = 0 + for client in clients: + closer.schedule(client) + closer.reap() # nothing is due yet, which is the hot path + clock.advance(61.0) + closer.reap() + + assert closer.pending_count == 0 + assert all(client.closed for client in clients) + assert CountingDeadline.comparisons < 10 * evictions, ( + f"{CountingDeadline.comparisons} deadline comparisons for {evictions} evictions; " + "a reap is walking the whole queue" + ) diff --git a/tests/test_litellm/caching/test_llm_caching_handler.py b/tests/test_litellm/caching/test_llm_caching_handler.py index 8e6a94945b0..5f0e82dbb80 100644 --- a/tests/test_litellm/caching/test_llm_caching_handler.py +++ b/tests/test_litellm/caching/test_llm_caching_handler.py @@ -19,6 +19,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path +from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.caching.llm_caching_handler import LLMClientCache @@ -156,6 +157,71 @@ def test_remove_key_no_event_loop(): assert "test-key" not in cache.cache_dict +class _FakeClock: + """Hand-advanced monotonic clock, so grace windows need no real waiting.""" + + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +@pytest.mark.asyncio +async def test_evicted_litellm_owned_client_is_closed_once_the_grace_window_elapses(): + """ + Eviction only drops the cache's reference. The SDK clients are reference + cycles, so without an explicit close the client keeps its connection pool + open until a generational collection runs. + """ + clock = _FakeClock() + cache = LLMClientCache( + max_size_in_memory=2, + evicted_client_closer=EvictedClientCloser(grace_seconds=60.0, clock=clock), + ) + + client = MockAsyncClient() + cache.set_cache("client-key", client, litellm_owned_client=True, ttl=600) + + cache.ttl_dict = {key: 0 for key in cache.ttl_dict} + cache.expiration_heap = [(0, key) for _, key in cache.expiration_heap] + cache.evict_cache() + await asyncio.sleep(0.1) + assert client.closed is False, "an in-flight request may still hold the client" + + clock.advance(61.0) + cache.get_cache("any-key") + await asyncio.sleep(0.1) + + assert client.closed is True + + +@pytest.mark.asyncio +async def test_evicted_caller_supplied_client_is_never_closed(): + """litellm does not own a client the caller passed in, so it must stay open.""" + clock = _FakeClock() + cache = LLMClientCache( + max_size_in_memory=2, + evicted_client_closer=EvictedClientCloser(grace_seconds=60.0, clock=clock), + ) + + client = MockAsyncClient() + cache.set_cache("client-key", client, ttl=600) + + cache.ttl_dict = {key: 0 for key in cache.ttl_dict} + cache.expiration_heap = [(0, key) for _, key in cache.expiration_heap] + cache.evict_cache() + + clock.advance(3600.0) + cache.get_cache("any-key") + await asyncio.sleep(0.1) + + assert client.closed is False + + def test_remove_key_removes_plain_values(): """ _remove_key correctly removes non-client values (strings, dicts, etc.). diff --git a/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py b/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py index 772e993c132..09d6f51e0a8 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py @@ -6,6 +6,7 @@ Verifies that DataDogLogger can be instantiated with per-team credentials and that the DataDogHandler correctly resolves and caches per-team loggers. """ +import copy from unittest.mock import patch import pytest @@ -13,7 +14,9 @@ import pytest from litellm.integrations.datadog.datadog import DataDogLogger from litellm.integrations.datadog.datadog_team_handler import ( DataDogHandler, - DatadogLoggingConfig, +) +from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + TRUSTED_CALLBACK_VARS_FIELD, ) from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import ( DynamicLoggingCache, @@ -94,9 +97,7 @@ class TestDataDogLoggerCredentialKwargs: assert logger.DD_API_KEY is None assert "attacker.example.com" in logger.intake_url - def test_direct_api_mode_does_not_leak_env_api_key_when_disallowed( - self, datadog_env - ): + def test_direct_api_mode_does_not_leak_env_api_key_when_disallowed(self, datadog_env): """With allow_env_credentials=False and no explicit key, init must fail rather than reuse env key.""" with pytest.raises(Exception, match="DD_API_KEY"): with patch("asyncio.create_task"): @@ -261,3 +262,96 @@ class TestStandardCallbackDynamicParamsIncludesDatadog: assert "dd_site" in annotations assert "dd_agent_host" in annotations assert "dd_agent_port" in annotations + + +def _build_logging_obj(kwargs: dict, *, with_datadog_callback: bool = True): + from litellm.litellm_core_utils.litellm_logging import Logging + + with patch("asyncio.create_task"): + return Logging( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time="2026-01-01", + litellm_call_id="test-call-id", + function_id="test-func", + dynamic_success_callbacks=["datadog"] if with_datadog_callback else None, + kwargs=kwargs, + ) + + +def _dd_loggers(logging_obj) -> list[DataDogLogger]: + return [cb for cb in (logging_obj.dynamic_success_callbacks or []) if isinstance(cb, DataDogLogger)] + + +class TestTeamCallbackFlowPassesDDCredentials: + """ + dd_* credentials reach DataDogHandler only from the proxy-stamped trusted field. + + Team callback_vars are admin-configured, so they must survive + _request_blocked_callback_params; anything the caller put in the request body + must not, or a caller could pair its own dd_site with the team's dd_api_key. + """ + + def test_trusted_callback_vars_reach_datadog_handler(self, datadog_env): + trusted_vars = {"dd_api_key": "team-dd-key-123", "dd_site": "us5.datadoghq.com"} + logging_obj = _build_logging_obj( + { + TRUSTED_CALLBACK_VARS_FIELD: trusted_vars, + "model": "gpt-4", + "litellm_params": {"metadata": {}}, + } + ) + + dd_loggers = _dd_loggers(logging_obj) + assert len(dd_loggers) == 1, "DataDogLogger should be initialized from team callback_vars" + assert dd_loggers[0].DD_API_KEY == "team-dd-key-123" + assert "us5.datadoghq.com" in dd_loggers[0].intake_url + + def test_request_kwargs_dd_params_are_ignored(self, datadog_env): + """Top-level dd_* in the call kwargs are caller-controlled and must never be honoured.""" + logging_obj = _build_logging_obj( + { + "dd_api_key": "caller-dd-key", + "dd_site": "attacker.example.com", + "dd_agent_host": "attacker.example.com", + "model": "gpt-4", + "litellm_params": {"metadata": {}}, + } + ) + + dd_loggers = _dd_loggers(logging_obj) + assert len(dd_loggers) == 1 + assert dd_loggers[0].DD_API_KEY == "global_api_key" + assert "attacker.example.com" not in dd_loggers[0].intake_url + assert "us1.datadoghq.com" in dd_loggers[0].intake_url + + def test_logging_object_stays_deepcopyable(self): + """The proxy deep-copies request data, and the Logging object rides along in it.""" + logging_obj = _build_logging_obj( + { + TRUSTED_CALLBACK_VARS_FIELD: {"dd_api_key": "team-dd-key-123", "dd_site": "us5.datadoghq.com"}, + "model": "gpt-4", + "litellm_params": {"metadata": {}}, + }, + with_datadog_callback=False, + ) + + assert copy.deepcopy(logging_obj)._trusted_callback_vars == logging_obj._trusted_callback_vars + + def test_caller_cannot_redirect_team_credentials(self, datadog_env): + """The exfil shape: caller's dd_site paired with the team's dd_api_key.""" + logging_obj = _build_logging_obj( + { + TRUSTED_CALLBACK_VARS_FIELD: {"dd_api_key": "team-dd-key-123"}, + "dd_site": "attacker.example.com", + "model": "gpt-4", + "litellm_params": {"metadata": {}}, + } + ) + + dd_loggers = _dd_loggers(logging_obj) + assert len(dd_loggers) == 1 + assert dd_loggers[0].DD_API_KEY == "team-dd-key-123" + assert "attacker.example.com" not in dd_loggers[0].intake_url diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py index 922d2fe8a15..4a2ee487c65 100644 --- a/tests/test_litellm/integrations/test_rubrik.py +++ b/tests/test_litellm/integrations/test_rubrik.py @@ -1,8 +1,8 @@ """ Tests for the Rubrik LiteLLM plugin. -Covers initialization, apply_guardrail tool blocking (all allowed, all blocked, -partial blocking, fail-open), batch logging, and Anthropic format handling. +Covers initialization, apply_guardrail (prompt moderation + response/tool +blocking), batch logging, and Anthropic format handling. """ import os @@ -13,8 +13,11 @@ import httpx import pytest from litellm.integrations.custom_guardrail import ModifyResponseException -from litellm.integrations.rubrik import RubrikLogger -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.integrations.rubrik import ( + RubrikLogger, + _MalformedToolBlockingResponseError, +) +from litellm.proxy._types import UserAPIKeyAuth from tests.test_litellm.integrations.rubrik_test_helpers import ( make_inputs_with_tools, @@ -42,6 +45,18 @@ def handler(mock_env): return RubrikLogger() +@pytest.fixture +def user_api_key_dict(): + """The authenticated caller the proxy hands to async_post_call_failure_hook.""" + return UserAPIKeyAuth( + api_key="sk-block-attribution-test", + key_alias="rubrik-probe-key", + user_id="probe-user-1", + team_id="probe-team-1", + org_id="probe-org-1", + ) + + # -- Initialization ----------------------------------------------------------- @@ -50,19 +65,19 @@ class TestInitialization: with patch("asyncio.create_task", Mock()): handler = RubrikLogger() assert ( - handler.tool_blocking_endpoint + handler.response_moderation_endpoint == "http://localhost:8080/v1/after_completion/openai/v1" ) assert handler.logging_endpoint == "http://localhost:8080/v1/litellm/batch" assert handler.key == "test-api-key" - assert isinstance(handler.tool_blocking_client, AsyncHTTPHandler) + assert handler.moderation_client is not None def test_init_with_constructor_params(self): with patch("asyncio.create_task", Mock()): handler = RubrikLogger(api_key="ctor-key", api_base="http://ctor-host:9090") assert handler.key == "ctor-key" assert ( - handler.tool_blocking_endpoint + handler.response_moderation_endpoint == "http://ctor-host:9090/v1/after_completion/openai/v1" ) @@ -82,7 +97,7 @@ class TestInitialization: with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://localhost:8080/"}): with patch("asyncio.create_task", Mock()): assert ( - RubrikLogger().tool_blocking_endpoint + RubrikLogger().response_moderation_endpoint == "http://localhost:8080/v1/after_completion/openai/v1" ) @@ -90,13 +105,13 @@ class TestInitialization: with patch("asyncio.create_task", Mock()): with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v1"}): assert ( - RubrikLogger().tool_blocking_endpoint + RubrikLogger().response_moderation_endpoint == "http://host/v1/after_completion/openai/v1" ) with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v11"}): assert ( - RubrikLogger().tool_blocking_endpoint + RubrikLogger().response_moderation_endpoint == "http://host/v11/v1/after_completion/openai/v1" ) @@ -155,10 +170,10 @@ class TestInitialization: # Do NOT patch asyncio.create_task — the real call should be # guarded and fall back gracefully when there is no event loop. handler = RubrikLogger() - assert handler.tool_blocking_endpoint.startswith("http://localhost:8080") + assert handler.response_moderation_endpoint.startswith("http://localhost:8080") # Without a running loop at init, the periodic flush task should be # deferred so batches still get drained once a log event arrives. - assert handler._flush_task is None + assert handler._periodic_flush_task is None @pytest.mark.asyncio async def test_periodic_flush_task_started_lazily_on_first_log(self, mock_env): @@ -170,7 +185,7 @@ class TestInitialization: side_effect=RuntimeError("no running loop"), ): handler = RubrikLogger() - assert handler._flush_task is None + assert handler._periodic_flush_task is None kwargs = { "standard_logging_object": { @@ -183,8 +198,8 @@ class TestInitialization: with patch.object(handler, "_log_batch_to_rubrik", AsyncMock()): await handler.async_log_success_event(kwargs, None, None, None) - assert handler._flush_task is not None - handler._flush_task.cancel() + assert handler._periodic_flush_task is not None + handler._periodic_flush_task.cancel() def test_event_hook_defaults_to_post_call_when_none_passed(self, mock_env): """`initialize_guardrail` always passes ``event_hook=litellm_params.mode`` @@ -204,14 +219,13 @@ class TestInitialization: handler = RubrikLogger(event_hook=GuardrailEventHooks.pre_call) assert handler.event_hook == GuardrailEventHooks.pre_call - def test_default_on_defaults_to_true_when_none_passed(self, mock_env): - """`initialize_guardrail` always passes ``default_on=litellm_params.default_on`` - (which is ``None`` when the user omits ``default_on``). The logger must - coerce a None ``default_on`` to True, otherwise ``should_run_guardrail`` - (which checks ``self.default_on is True``) silently skips the guardrail.""" + def test_default_on_defaults_to_false_when_none_passed(self, mock_env): + """Follows the standard litellm pattern: omitted ``default_on`` resolves + to ``False`` (off by default). Users must explicitly set + ``default_on: true`` to enable the guardrail for all requests.""" with patch("asyncio.create_task", Mock()): handler = RubrikLogger(default_on=None) - assert handler.default_on is True + assert handler.default_on is False def test_explicit_default_on_false_preserved(self, mock_env): """A user explicitly setting ``default_on: false`` in their guardrail @@ -421,7 +435,7 @@ class TestBatchLogging: ) assert len(handler.log_queue) == 1 msgs = handler.log_queue[0]["messages"] - assert isinstance(msgs, list) + assert isinstance(msgs, tuple) assert msgs[0]["role"] == "system" assert msgs[1] == {"role": "user", "content": "hi"} @@ -444,7 +458,10 @@ class TestBatchLogging: ) assert handler.log_queue[0]["id"] == "litellm-call-123" - async def test_non_anthropic_id_unchanged(self, handler): + async def test_litellm_call_id_always_used_as_correlation_key(self, handler): + """The merged plugin always uses litellm_call_id as the log ID for all + providers (not just Anthropic) so that logs correlate with the + moderation (_blocking) and failure logs for the same request.""" kwargs = { "standard_logging_object": { "id": "chatcmpl-original", @@ -461,7 +478,7 @@ class TestBatchLogging: await handler.async_log_success_event( kwargs=kwargs, response_obj=None, start_time=None, end_time=None ) - assert handler.log_queue[0]["id"] == "chatcmpl-original" + assert handler.log_queue[0]["id"] == "litellm-call-123" async def test_payload_deep_copied_not_mutated(self, handler): """Verify the shared standard_logging_object is not mutated.""" @@ -536,7 +553,7 @@ class TestApplyGuardrail: tc2 = make_tool_call_dict("call_2", "get_time") inputs = make_inputs_with_tools([tc1, tc2]) - handler.tool_blocking_client = _echo_service() + handler.moderation_client = _echo_service() result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -548,7 +565,7 @@ class TestApplyGuardrail: tc2 = make_tool_call_dict("call_2", "drop_database") inputs = make_inputs_with_tools([tc1, tc2]) - handler.tool_blocking_client = _mock_service_response( + handler.moderation_client = _mock_service_response( { "choices": [ { @@ -594,7 +611,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client with pytest.raises(ModifyResponseException): await handler.apply_guardrail( @@ -607,7 +624,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -618,7 +635,7 @@ class TestApplyGuardrail: tc1 = make_tool_call_dict("call_1", "test_tool") inputs = make_inputs_with_tools([tc1]) - handler.tool_blocking_client = _mock_service_response({"choices": []}) + handler.moderation_client = _mock_service_response({"choices": []}) result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -641,7 +658,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -675,7 +692,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client logging_obj = Mock() logging_obj.model_call_details = { @@ -697,7 +714,10 @@ class TestApplyGuardrail: assert req["model"] == "gpt-4" assert req["messages"] == [{"role": "user", "content": "hi"}] - async def test_proxy_server_request_headers_stripped(self, handler): + async def test_proxy_server_request_not_forwarded(self, handler): + """proxy_server_request is intentionally NOT included in the request + envelope: in litellm >=1.83 its ``body`` carries a UserAPIKeyAuth + instance that breaks json.dumps, silently fail-opening the guardrail.""" tc = make_tool_call_dict("call_1", "test_tool") inputs = make_inputs_with_tools([tc]) @@ -712,7 +732,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client logging_obj = Mock() logging_obj.model_call_details = { @@ -739,8 +759,8 @@ class TestApplyGuardrail: logging_obj=logging_obj, ) - forwarded = captured_payload["request"]["proxy_server_request"] - assert forwarded == {"url": "/chat/completions", "method": "POST"} + # proxy_server_request is deliberately excluded from the forwarded envelope + assert "proxy_server_request" not in captured_payload["request"] # -- Anthropic format ---------------------------------------------------------- @@ -760,7 +780,7 @@ class TestApplyGuardrailAnthropicFormat: ) inputs = make_inputs_with_tools([tc], texts=["I'll check the weather."]) - handler.tool_blocking_client = _echo_service() + handler.moderation_client = _echo_service() result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -771,7 +791,7 @@ class TestApplyGuardrailAnthropicFormat: tc = make_tool_call_dict("toolu_123", "dangerous_tool", '{"arg": "value"}') inputs = make_inputs_with_tools([tc]) - handler.tool_blocking_client = _mock_service_response( + handler.moderation_client = _mock_service_response( { "choices": [ { @@ -790,21 +810,21 @@ class TestApplyGuardrailAnthropicFormat: inputs=inputs, request_data={}, input_type="response" ) - async def test_text_only_response_no_blocking(self, handler): + async def test_text_only_response_sent_to_moderation(self, handler): + """Text-only responses (no tool calls) are sent to the response + moderation service to check the assistant's text content.""" from litellm.types.utils import GenericGuardrailAPIInputs inputs = GenericGuardrailAPIInputs(texts=["Hello! I'm Claude."]) - mock_client = AsyncMock() - mock_client.post = AsyncMock() - handler.tool_blocking_client = mock_client + # Service allows the response (returns the content unchanged) + handler.moderation_client = _echo_service() result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" ) assert result is inputs - mock_client.post.assert_not_called() async def test_service_failure_preserves_tools(self, handler): tc = make_tool_call_dict("toolu_123", "get_weather", '{"location": "SF"}') @@ -812,7 +832,7 @@ class TestApplyGuardrailAnthropicFormat: mock_client = AsyncMock() mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -850,10 +870,13 @@ class TestNormalizeToolCalls: RubrikLogger._normalize_tool_calls(["not_a_tool_call"]) -# -- Extract blocked tools ----------------------------------------------------- +# -- Extract response block ---------------------------------------------------- -class TestExtractBlockedTools: +class TestExtractResponseBlock: + """Tests for _extract_response_block, which replaces the upstream + _extract_blocked_tools and handles both text blocks and tool blocks.""" + def test_all_allowed_returns_none(self): from litellm.types.utils import ChatCompletionMessageToolCall, Function @@ -870,7 +893,7 @@ class TestExtractBlockedTools: } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc]) + result = RubrikLogger._extract_response_block(service_resp, [tc], "") assert result is None def test_some_blocked_returns_explanation(self): @@ -896,13 +919,13 @@ class TestExtractBlockedTools: } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2]) + result = RubrikLogger._extract_response_block(service_resp, [tc1, tc2], "") assert result is not None - assert "blocked fn2" in result + assert "blocked fn2" in result.explanation def test_empty_choices_raises(self): - with pytest.raises(Exception, match="empty response"): - RubrikLogger._extract_blocked_tools({"choices": []}, []) + with pytest.raises(_MalformedToolBlockingResponseError): + RubrikLogger._extract_response_block({"choices": []}, [], "") def test_null_tool_calls_treated_as_all_blocked(self): from litellm.types.utils import ChatCompletionMessageToolCall, Function @@ -920,36 +943,55 @@ class TestExtractBlockedTools: } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc]) + result = RubrikLogger._extract_response_block(service_resp, [tc], "") assert result is not None - assert "blocked everything" in result + assert "blocked everything" in result.explanation - def test_duplicate_ids_block_when_only_one_returned(self): + def test_text_block_detected(self): + """When the service replaces the response text wholesale, it's a text block.""" from litellm.types.utils import ChatCompletionMessageToolCall, Function - tc1 = ChatCompletionMessageToolCall( - id="call_dup", - type="function", - function=Function(name="fn", arguments="{}"), - ) - tc2 = ChatCompletionMessageToolCall( - id="call_dup", - type="function", - function=Function(name="fn", arguments="{}"), - ) service_resp = { "choices": [ { "message": { - "tool_calls": [{"id": "call_dup"}], - "content": "blocked duplicate", + "tool_calls": [], + "content": "This content violates policy.", } } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2]) + result = RubrikLogger._extract_response_block( + service_resp, [], "Original assistant text." + ) assert result is not None - assert "blocked duplicate" in result + assert "violates policy" in result.explanation + + def test_tool_block_with_appended_explanation(self): + """When the service appends an explanation to the original text, only the + appended part is returned as the explanation.""" + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="fn", arguments="{}") + ) + original_text = "Here is my response." + appended_explanation = "Tool call was blocked." + service_resp = { + "choices": [ + { + "message": { + "tool_calls": [], + "content": original_text + "\n\n" + appended_explanation, + } + } + ] + } + result = RubrikLogger._extract_response_block( + service_resp, [tc], original_text + ) + assert result is not None + assert appended_explanation in result.explanation # -- Sanitize proxy server request ------------------------------------------- @@ -1010,3 +1052,900 @@ class TestResolveModel: {"response": response}, {"model": "fallback"} ) assert result == "unknown" + + +# -- Additional Initialization edge cases ------------------------------------ + + +class TestInitializationEdgeCases: + def test_batch_size_zero_uses_default(self): + """RUBRIK_BATCH_SIZE=0 must warn and fall back to the default.""" + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "0"}, + ): + h = RubrikLogger() + # Should use default, not 0 + assert h.batch_size > 0 + + def test_batch_size_negative_uses_default(self): + """RUBRIK_BATCH_SIZE=-1 must warn and fall back to the default.""" + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "-5"}, + ): + h = RubrikLogger() + assert h.batch_size > 0 + + +# -- aclose() ----------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestAclose: + async def test_aclose_cancels_task_does_not_close_shared_client(self, mock_env): + """aclose() cancels the periodic flush task but does NOT close the shared + moderation_client — closing a shared cached client would break other + RubrikLogger instances that share the same connection pool.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger() + + mock_task = Mock() + mock_task.cancel = Mock() + handler._periodic_flush_task = mock_task + + handler.moderation_client = AsyncMock() + handler.moderation_client.close = AsyncMock() + + await handler.aclose() + + mock_task.cancel.assert_called_once() + handler.moderation_client.close.assert_not_awaited() + + async def test_aclose_with_none_task_does_not_close_client(self, mock_env): + """aclose() with no flush task still does not close the shared client.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger() + + handler._periodic_flush_task = None + handler.moderation_client = AsyncMock() + handler.moderation_client.close = AsyncMock() + + await handler.aclose() + + handler.moderation_client.close.assert_not_awaited() + + +# -- apply_guardrail edge cases ----------------------------------------------- + + +@pytest.mark.asyncio +class TestApplyGuardrailEdgeCases: + async def test_unknown_input_type_returns_inputs_unchanged(self, handler): + """When input_type is not 'request' or 'response', inputs are returned as-is.""" + inputs = make_inputs_with_tools([make_tool_call_dict("call_1", "tool")]) + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="unknown" + ) + assert result is inputs + + async def test_response_with_no_texts_and_no_tool_calls_returns_inputs(self, handler): + """_moderate_response early-returns when both texts and tool_calls are empty.""" + from litellm.types.utils import GenericGuardrailAPIInputs + + inputs = GenericGuardrailAPIInputs() + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_moderate_response_empty_call_details_emits_warning(self, handler): + """When logging_obj is present but model_call_details is empty, a warning is + logged and moderation proceeds (fail-open on HTTP error).""" + tc = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc]) + + logging_obj = Mock() + logging_obj.model_call_details = {} + + handler.moderation_client = _echo_service() + + result = await handler.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=logging_obj, + ) + assert result is inputs + + +# -- Prompt moderation -------------------------------------------------------- + + +@pytest.mark.asyncio +class TestPromptModeration: + async def test_prompt_moderation_passthrough(self, handler): + """Webhook returns {} (empty dict) → inputs returned unchanged.""" + inputs = {"structured_messages": [{"role": "user", "content": "Hello"}]} + + handler.moderation_client = _mock_service_response({}) + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + async def test_prompt_moderation_blocked_raises(self, handler): + """Webhook returns synthetic chat.completion → raises ModifyResponseException.""" + inputs = { + "structured_messages": [{"role": "user", "content": "Harmful prompt"}], + "model": "gpt-4", + } + + handler.moderation_client = _mock_service_response( + { + "choices": [ + { + "message": { + "role": "assistant", + "content": "This request violates our policy.", + } + } + ] + } + ) + + with pytest.raises(ModifyResponseException) as exc_info: + await handler.apply_guardrail( + inputs=inputs, request_data={"model": "gpt-4"}, input_type="request" + ) + assert "violates our policy" in exc_info.value.message + + async def test_prompt_moderation_no_messages_skips_moderation(self, handler): + """When structured_messages is absent/empty, moderation is skipped.""" + inputs = {"model": "gpt-4"} + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + async def test_prompt_moderation_stashes_logging_obj_on_block(self, handler): + """On a prompt block, _stash_block_context must set the blocked flag.""" + inputs = { + "structured_messages": [{"role": "user", "content": "bad prompt"}], + } + + handler.moderation_client = _mock_service_response( + { + "choices": [ + {"message": {"role": "assistant", "content": "Blocked."}} + ] + } + ) + + logging_obj = Mock() + logging_obj.model_call_details = {} + request_data: dict = {} + + with pytest.raises(ModifyResponseException): + await handler.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + logging_obj=logging_obj, + ) + + assert logging_obj.model_call_details.get("_rubrik_blocked") is True + assert request_data.get("_rubrik_logging_obj") is logging_obj + + +# -- _stash_block_context ----------------------------------------------------- + + +class TestStashBlockContext: + def test_with_non_none_logging_obj_sets_flag_and_stashes(self): + """Sets _rubrik_blocked flag and stores logging_obj on request_data.""" + logging_obj = Mock() + logging_obj.model_call_details = {} + request_data: dict = {} + + RubrikLogger._stash_block_context(logging_obj, request_data) + + assert logging_obj.model_call_details["_rubrik_blocked"] is True + assert request_data["_rubrik_logging_obj"] is logging_obj + + def test_with_none_logging_obj_stores_none_on_request_data(self): + """When logging_obj is None, stores None on request_data (logged as error).""" + request_data: dict = {"litellm_call_id": "test-id"} + + RubrikLogger._stash_block_context(None, request_data) + + assert request_data["_rubrik_logging_obj"] is None + + +# -- _normalize_tool_calls duck-typed ----------------------------------------- + + +class TestNormalizeToolCallsDuckTyped: + def test_duck_typed_object_with_id_and_function_attrs(self): + """Objects that have .id and .function attrs but are not + ChatCompletionMessageToolCall are handled by the third branch.""" + from litellm.types.utils import Function + + tc = Mock() + tc.id = "call_duck" + tc.type = "function" + tc.function = Function(name="duck_tool", arguments='{"x": 1}') + # Make isinstance(..., ChatCompletionMessageToolCall) return False + # by using a plain Mock (not a ChatCompletionMessageToolCall subclass) + + result = RubrikLogger._normalize_tool_calls([tc]) + assert len(result) == 1 + assert result[0].id == "call_duck" + assert result[0].function.name == "duck_tool" + + def test_duck_typed_without_type_defaults_to_function(self): + """getattr(tc, "type", None) falls back to "function" when absent.""" + from litellm.types.utils import Function + + tc = Mock(spec=["id", "function"]) # no .type attr + tc.id = "call_no_type" + tc.function = Function(name="fn", arguments="{}") + + result = RubrikLogger._normalize_tool_calls([tc]) + assert result[0].type == "function" + + +# -- _flatten_messages_for_moderation ----------------------------------------- + + +class TestFlattenMessagesForModeration: + def test_plain_string_content_preserved(self): + messages = [{"role": "user", "content": "Hello world"}] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 1 + assert result[0]["role"] == "user" + assert result[0]["content"] == "Hello world" + + def test_content_list_flattened_to_string(self): + """Content as a list of parts (e.g. Anthropic multi-part) is flattened.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello from parts"}, + ], + } + ] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 1 + assert result[0]["role"] == "user" + assert "Hello from parts" in result[0]["content"] + + def test_non_dict_messages_skipped(self): + """Non-dict entries in the messages list are silently skipped.""" + messages = [ + "raw string message", + {"role": "user", "content": "valid"}, + ] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 1 + assert result[0]["content"] == "valid" + + def test_none_messages_returns_empty(self): + result = RubrikLogger._flatten_messages_for_moderation(None) + assert result == () + + def test_multiple_messages_preserved_in_order(self): + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Question?"}, + ] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 2 + assert result[0]["role"] == "system" + assert result[1]["role"] == "user" + + +# -- _build_prompt_moderation_payload ----------------------------------------- + + +class TestBuildPromptModerationPayload: + def test_payload_includes_tools_when_present(self): + inputs = { + "model": "gpt-4", + "structured_messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "function", "function": {"name": "fn"}}], + } + payload = RubrikLogger._build_prompt_moderation_payload(inputs, {}) + assert payload["tools"] == [{"type": "function", "function": {"name": "fn"}}] + + def test_payload_includes_user_when_present(self): + inputs = { + "structured_messages": [{"role": "user", "content": "hi"}], + } + request_data = {"user": "alice"} + payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + assert payload["user"] == "alice" + + def test_payload_uses_explicit_correlation_key(self): + inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} + request_data = { + "correlation_key": "corr-123", + "litellm_call_id": "litellm-456", + } + payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + assert payload["correlation_key"] == "corr-123" + + def test_payload_falls_back_to_litellm_call_id(self): + """When correlation_key is absent, litellm_call_id is used.""" + inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} + request_data = {"litellm_call_id": "litellm-789"} + payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + assert payload["correlation_key"] == "litellm-789" + + def test_payload_omits_optional_fields_when_absent(self): + inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} + payload = RubrikLogger._build_prompt_moderation_payload(inputs, {}) + assert "tools" not in payload + assert "user" not in payload + assert "correlation_key" not in payload + + +# -- _extract_request_data tools preference ----------------------------------- + + +class TestExtractRequestDataToolsPreference: + def test_prefers_tools_from_request_data_over_optional_params(self): + """When 'tools' key exists in request_data, it wins over optional_params.""" + call_details = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-4", + "optional_params": { + "tools": [{"type": "function", "function": {"name": "from_optional"}}] + }, + } + request_data = { + "tools": [{"type": "function", "function": {"name": "from_request"}}] + } + result = RubrikLogger._extract_request_data(call_details, request_data) + assert result["tools"] == [ + {"type": "function", "function": {"name": "from_request"}} + ] + + def test_falls_back_to_optional_params_when_not_in_request_data(self): + call_details = { + "optional_params": { + "tools": [{"type": "function", "function": {"name": "from_optional"}}] + } + } + result = RubrikLogger._extract_request_data(call_details, {}) + assert result["tools"] == [ + {"type": "function", "function": {"name": "from_optional"}} + ] + + def test_explicit_empty_list_in_request_data_is_forwarded(self): + """An explicit empty tools list signals 'no tools' to the moderation service.""" + call_details = { + "optional_params": { + "tools": [{"type": "function", "function": {"name": "from_optional"}}] + } + } + request_data = {"tools": []} + result = RubrikLogger._extract_request_data(call_details, request_data) + assert result["tools"] == [] + + +# -- _extract_prompt_refusal -------------------------------------------------- + + +class TestExtractPromptRefusal: + def test_passthrough_response_returns_none(self): + """Empty dict (passthrough) → None.""" + assert RubrikLogger._extract_prompt_refusal({}) is None + + def test_no_choices_returns_none(self): + assert RubrikLogger._extract_prompt_refusal({"choices": []}) is None + + def test_block_response_returns_content(self): + service_response = { + "choices": [{"message": {"content": "Request blocked by Rubrik."}}] + } + result = RubrikLogger._extract_prompt_refusal(service_response) + assert result == "Request blocked by Rubrik." + + def test_empty_content_falls_back_to_default_message(self): + """When content is empty string or falsy, falls back to default refusal.""" + service_response = {"choices": [{"message": {"content": ""}}]} + result = RubrikLogger._extract_prompt_refusal(service_response) + assert result == "Request blocked by policy." + + def test_none_content_falls_back_to_default_message(self): + service_response = {"choices": [{"message": {"content": None}}]} + result = RubrikLogger._extract_prompt_refusal(service_response) + assert result == "Request blocked by policy." + + +# -- _prepend_system_prompt exception path ------------------------------------ + + +class TestPrependSystemPromptException: + def test_exception_during_unpack_is_caught_and_logged(self): + """When an exception is raised inside _prepend_system_prompt, it is swallowed.""" + + class ExplodingList(list): + def __iter__(self): + raise RuntimeError("iteration error!") + + payload = {"messages": ExplodingList()} + source = {"system": "You are an assistant."} + + # Must not raise + RubrikLogger._prepend_system_prompt(payload, source) + + +# -- _append_and_maybe_flush batch trigger ------------------------------------ + + +@pytest.mark.asyncio +class TestAppendAndMaybeFlush: + async def test_flush_triggered_when_queue_reaches_batch_size(self, handler): + """flush_queue is called when the queue length reaches batch_size.""" + handler.batch_size = 2 + handler.flush_queue = AsyncMock() + + await handler._append_and_maybe_flush({"msg": "a"}) + handler.flush_queue.assert_not_called() + + await handler._append_and_maybe_flush({"msg": "b"}) + handler.flush_queue.assert_called_once() + + async def test_no_flush_before_batch_size(self, handler): + handler.batch_size = 5 + handler.flush_queue = AsyncMock() + + for i in range(4): + await handler._append_and_maybe_flush({"msg": str(i)}) + + handler.flush_queue.assert_not_called() + + +# -- _enqueue_log_event exception handling ------------------------------------ + + +@pytest.mark.asyncio +class TestEnqueueLogEventExceptions: + async def test_exception_from_prepare_log_payload_is_caught(self, handler): + """Exceptions raised by _prepare_log_payload are caught and logged.""" + handler._prepare_log_payload = AsyncMock( + side_effect=RuntimeError("payload error") + ) + + # Must not raise + await handler._enqueue_log_event( + {"standard_logging_object": {"messages": [], "response": ""}}, "test" + ) + assert len(handler.log_queue) == 0 + + +# -- async_log_success_event skip when _rubrik_blocked ------------------------ + + +@pytest.mark.asyncio +class TestSuccessEventBlockedSkip: + async def test_skips_enqueue_when_rubrik_blocked_flag_set(self, handler): + """When kwargs['_rubrik_blocked'] is True, the event is not enqueued.""" + kwargs = { + "_rubrik_blocked": True, + "litellm_call_id": "blocked-call-123", + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 0 + + +# -- async_post_call_failure_hook --------------------------------------------- + + +@pytest.mark.asyncio +class TestPostCallFailureHook: + async def test_non_modify_exception_returns_immediately(self, handler, user_api_key_dict): + """Non-ModifyResponseException causes a no-op.""" + await handler.async_post_call_failure_hook( + request_data={"litellm_call_id": "test"}, + original_exception=ValueError("unrelated error"), + user_api_key_dict=user_api_key_dict, + ) + assert len(handler.log_queue) == 0 + + async def test_modify_exception_without_stashed_logging_obj_emits_warning( + self, handler, user_api_key_dict + ): + """ModifyResponseException with no _rubrik_logging_obj → warning, no enqueue.""" + request_data = {"litellm_call_id": "test-123", "model": "gpt-4"} + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data=request_data, + guardrail_name="rubrik", + ) + + await handler.async_post_call_failure_hook( + request_data=request_data, + original_exception=exc, + user_api_key_dict=user_api_key_dict, + ) + assert len(handler.log_queue) == 0 + + async def test_modify_exception_with_valid_logging_obj_enqueues_payload( + self, handler, user_api_key_dict + ): + """ModifyResponseException + stashed logging_obj → builds and enqueues.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-abc", + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "standard_logging_object": { + "id": "chatcmpl-original", + "model": "gpt-4", + "response": "original", + "messages": [{"role": "user", "content": "hi"}], + }, + "metadata": {}, + } + + request_data = {"_rubrik_logging_obj": logging_obj} + exc = ModifyResponseException( + message="blocked by policy", + model="gpt-4", + request_data=request_data, + guardrail_name="rubrik", + ) + + handler.batch_size = 10**6 # disable auto-flush + await handler.async_post_call_failure_hook( + request_data=request_data, + original_exception=exc, + user_api_key_dict=user_api_key_dict, + ) + assert len(handler.log_queue) == 1 + assert "ModifyResponseException" in handler.log_queue[0]["response"] + + async def test_logging_obj_popped_from_request_data(self, handler, user_api_key_dict): + """_rubrik_logging_obj must be popped from request_data so it is not + forwarded downstream.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-pop", + "model": "gpt-4", + "messages": [], + "standard_logging_object": { + "id": "chatcmpl-pop", + "model": "gpt-4", + "response": "text", + "messages": [], + }, + "metadata": {}, + } + + request_data = {"_rubrik_logging_obj": logging_obj} + exc = ModifyResponseException( + message="popped", + model="gpt-4", + request_data=request_data, + guardrail_name="rubrik", + ) + + handler.batch_size = 10**6 + await handler.async_post_call_failure_hook( + request_data=request_data, + original_exception=exc, + user_api_key_dict=user_api_key_dict, + ) + assert "_rubrik_logging_obj" not in request_data + + async def test_build_and_enqueue_swallows_attribute_error_from_prepare_payload( + self, handler, user_api_key_dict + ): + """When _prepare_block_failure_payload raises AttributeError/KeyError/TypeError, + the error is logged and the event is silently dropped (lines 806-812).""" + logging_obj = Mock() + # Make model_call_details.get() raise TypeError + logging_obj.model_call_details = None # .get() will raise AttributeError + + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + # Must not raise + await handler._build_and_enqueue_block_event(logging_obj, exc, None, user_api_key_dict) + assert len(handler.log_queue) == 0 + + async def test_build_and_enqueue_swallows_flush_exception(self, handler, user_api_key_dict): + """When _append_and_maybe_flush raises, the error is logged (lines 816-817).""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-flush-err", + "model": "gpt-4", + "messages": [], + "standard_logging_object": { + "id": "id-flush-err", + "model": "gpt-4", + "response": "text", + "messages": [], + }, + "metadata": {}, + } + + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + handler._append_and_maybe_flush = AsyncMock( + side_effect=RuntimeError("flush failed") + ) + + # Must not raise + await handler._build_and_enqueue_block_event(logging_obj, exc, None, user_api_key_dict) + + +# -- _prepare_block_failure_payload and _build_fallback_payload --------------- + + +class TestPrepareBlockFailurePayload: + def test_uses_standard_logging_object_when_present(self, handler, user_api_key_dict): + """When standard_logging_object is on model_call_details, it is used as base.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-slo", + "model": "gpt-4", + "standard_logging_object": { + "id": "chatcmpl-original", + "model": "gpt-4", + "response": "original response", + "messages": [{"role": "user", "content": "hi"}], + }, + "metadata": {}, + } + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + payload = handler._prepare_block_failure_payload(logging_obj, exc, user_api_key_dict) + + assert "ModifyResponseException: blocked" in payload["response"] + assert payload["id"] == "call-slo" + + def test_standard_logging_object_identity_is_not_overwritten(self, handler, user_api_key_dict): + """A streamed block can arrive with the object populated; it keeps its own identity.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-slo-identity", + "model": "gpt-4", + "standard_logging_object": { + "id": "chatcmpl-original", + "model": "gpt-4", + "response": "original response", + "messages": [], + "metadata": {"user_api_key_hash": "hash-from-standard-logging-object"}, + }, + } + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + payload = handler._prepare_block_failure_payload(logging_obj, exc, user_api_key_dict) + + assert payload["metadata"]["user_api_key_hash"] == "hash-from-standard-logging-object" + + def test_uses_fallback_when_standard_logging_object_absent(self, handler, user_api_key_dict): + """When standard_logging_object is absent, _build_fallback_payload is used.""" + from datetime import datetime + + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-fallback", + "model": "claude-3", + "messages": [{"role": "user", "content": "question"}], + "optional_params": {"temperature": 0.5}, + "metadata": {"headers": {"host": "127.0.0.1:4000", "user-agent": "curl/8.7.1"}}, + "start_time": datetime(2024, 6, 1), + } + exc = ModifyResponseException( + message="prompt blocked", + model="claude-3", + request_data={}, + guardrail_name="rubrik", + ) + + payload = handler._prepare_block_failure_payload(logging_obj, exc, user_api_key_dict) + + assert payload["id"] == "call-fallback" + assert payload["model"] == "claude-3" + assert payload["model_group"] == "claude-3" + assert "ModifyResponseException: prompt blocked" in payload["response"] + assert payload["status"] == "failure" + + def test_fallback_payload_without_start_time(self, handler, user_api_key_dict): + """_build_fallback_payload handles missing start_time gracefully.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-notime", + "model": "gpt-4", + "messages": [], + "optional_params": {}, + "metadata": {}, + } + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + payload = handler._prepare_block_failure_payload(logging_obj, exc, user_api_key_dict) + assert payload["startTime"] is None + + +class TestBlockPayloadCallerAttribution: + """A block log must identify the caller that triggered it. + + The enriched litellm metadata lives under + ``model_call_details["litellm_params"]["metadata"]``, never at the top + level, so sourcing identity from ``call_details["metadata"]`` yielded an + empty string for every block. Identity comes from the authenticated + ``user_api_key_dict`` the failure hook is handed. + """ + + def _blocked_call_details(self): + return { + "litellm_call_id": "call-attr", + "model": "claude-3", + "messages": [{"role": "user", "content": "question"}], + "optional_params": {}, + "metadata": {"headers": {"host": "127.0.0.1:4000", "user-agent": "curl/8.7.1"}}, + } + + def test_fallback_payload_identifies_the_caller(self, handler, user_api_key_dict): + payload = handler._build_fallback_payload(self._blocked_call_details(), user_api_key_dict) + + metadata = payload["metadata"] + assert metadata["user_api_key_hash"] == user_api_key_dict.api_key + assert metadata["user_api_key_alias"] == "rubrik-probe-key" + assert metadata["user_api_key_user_id"] == "probe-user-1" + assert metadata["user_api_key_team_id"] == "probe-team-1" + assert metadata["user_api_key_org_id"] == "probe-org-1" + + def test_metadata_covers_the_full_caller_key_set(self, handler, user_api_key_dict): + """A block log and a success log agree on the caller key set.""" + from litellm.types.utils import StandardLoggingUserAPIKeyMetadata + + payload = handler._build_fallback_payload(self._blocked_call_details(), user_api_key_dict) + + expected = StandardLoggingUserAPIKeyMetadata.__required_keys__ | StandardLoggingUserAPIKeyMetadata.__optional_keys__ + assert set(payload["metadata"]) == set(expected) + + def test_virtual_key_is_logged_hashed(self, handler): + """A virtual key reaches the webhook as its hash, never as the raw token.""" + raw = "sk-block-attribution-test" + payload = handler._build_fallback_payload( + self._blocked_call_details(), UserAPIKeyAuth(api_key=raw) + ) + + assert payload["metadata"]["user_api_key_hash"] not in (raw, "") + + def test_request_header_metadata_is_not_used_as_identity(self, handler, user_api_key_dict): + """The pre-fix source is present and misleading; it must not win.""" + call_details = self._blocked_call_details() + call_details["metadata"]["user_api_key_hash"] = "stale-hash-from-request-metadata" + + payload = handler._build_fallback_payload(call_details, user_api_key_dict) + + assert payload["metadata"]["user_api_key_hash"] == user_api_key_dict.api_key + + async def test_enqueued_block_event_carries_attribution(self, handler, user_api_key_dict): + """End of the real hook chain: what actually lands on the Rubrik queue.""" + logging_obj = Mock() + logging_obj.model_call_details = self._blocked_call_details() + request_data = {"litellm_call_id": "call-attr", "_rubrik_logging_obj": logging_obj} + exc = ModifyResponseException( + message="prompt blocked", + model="claude-3", + request_data=request_data, + guardrail_name="rubrik", + ) + + await handler.async_post_call_failure_hook( + request_data=request_data, + original_exception=exc, + user_api_key_dict=user_api_key_dict, + ) + + assert len(handler.log_queue) == 1 + metadata = handler.log_queue[0]["metadata"] + assert metadata["user_api_key_hash"] == user_api_key_dict.api_key + assert metadata["user_api_key_user_id"] == "probe-user-1" + + +# -- async_send_batch empty queue and flush_queue edge cases ------------------ + + +@pytest.mark.asyncio +class TestQueueEdgeCases: + async def test_async_send_batch_returns_early_on_empty_queue(self, handler): + """async_send_batch is a no-op when the queue is empty.""" + handler.async_httpx_client = AsyncMock() + await handler.async_send_batch() + handler.async_httpx_client.post.assert_not_called() + + async def test_flush_queue_returns_early_when_flush_lock_is_none(self, handler): + """flush_queue is a no-op when flush_lock is None.""" + handler.flush_lock = None + handler.log_queue = [{"msg": "a"}] + handler.async_httpx_client = AsyncMock() + + await handler.flush_queue() + handler.async_httpx_client.post.assert_not_called() + + async def test_flush_queue_returns_early_when_queue_empty_inside_lock(self, handler): + """flush_queue acquires the lock then no-ops when the queue is empty.""" + handler.log_queue = [] + handler.async_httpx_client = AsyncMock() + + await handler.flush_queue() + handler.async_httpx_client.post.assert_not_called() + + +# -- _post_json non-dict response --------------------------------------------- + + +@pytest.mark.asyncio +class TestPostJson: + async def test_raises_type_error_for_list_response(self, handler): + """When the service returns a JSON array instead of a dict, TypeError is raised.""" + mock_client = AsyncMock() + mock_resp = Mock() + mock_resp.json.return_value = ["not", "a", "dict"] + mock_resp.raise_for_status = Mock() + mock_client.post = AsyncMock(return_value=mock_resp) + handler.moderation_client = mock_client + + with pytest.raises(TypeError, match="non-dict JSON"): + await handler._post_json( + handler.prompt_moderation_endpoint, {}, "Test service" + ) + + async def test_raises_type_error_for_string_response(self, handler): + """A bare string response also raises TypeError.""" + mock_client = AsyncMock() + mock_resp = Mock() + mock_resp.json.return_value = "blocked" + mock_resp.raise_for_status = Mock() + mock_client.post = AsyncMock(return_value=mock_resp) + handler.moderation_client = mock_client + + with pytest.raises(TypeError, match="non-dict JSON"): + await handler._post_json( + handler.response_moderation_endpoint, {}, "Test service" + ) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py index 791982fc3dc..370ec4b6f60 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py @@ -4,6 +4,11 @@ Bedrock Converse routes Claude Opus 4.7/4.8 and Claude Sonnet 4 through an Anthropic-compatible validator that rejects ``toolSpec.strict`` even though Anthropic's native API accepts ``strict`` as a top-level tool field. See BerriAI/litellm#31582. + +That per-model gate only covers models whose cost-map entry carries the flag, so a +``strict: false`` that litellm itself synthesized still broke unflagged models. Since +``strict: false`` is the Chat Completions default, it is now dropped for every model +rather than forwarded as a no-op the provider can reject. See BerriAI/litellm#33193. """ import pytest @@ -31,6 +36,16 @@ _STRICT_TOOL = [ } ] +_NON_STRICT_TOOL = [ + { + "type": "function", + "function": { + **_STRICT_TOOL[0]["function"], + "strict": False, + }, + } +] + @pytest.mark.parametrize( "model_id", @@ -48,12 +63,17 @@ _STRICT_TOOL = [ "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0", "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0", "bedrock/apac.anthropic.claude-sonnet-4-20250514-v1:0", + # Sonnet 5 rejects it too, verified live against Bedrock in us-east-1 + "anthropic.claude-sonnet-5", + "bedrock/us.anthropic.claude-sonnet-5", + "bedrock/eu.anthropic.claude-sonnet-5", + "bedrock/jp.anthropic.claude-sonnet-5", ], ) def test_bedrock_tools_pt_strict_dropped_for_strict_unsupported_models( model_id: str, ) -> None: - """Opus 4.7/4.8 and Sonnet 4 reject toolSpec.strict and additionalProperties.""" + """Opus 4.7/4.8, Sonnet 4 and Sonnet 5 reject toolSpec.strict and additionalProperties.""" result = _bedrock_tools_pt(_STRICT_TOOL, model=model_id) tool_spec = result[0]["toolSpec"] assert ( @@ -81,6 +101,55 @@ def test_bedrock_tools_pt_strict_kept_for_other_anthropic(model_id: str) -> None ), f"strict missing for {model_id}: {result[0]['toolSpec']}" +@pytest.mark.parametrize( + "model_id", + [ + "bedrock/us.anthropic.claude-sonnet-5", + "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "bedrock/us.anthropic.claude-sonnet-4-6", + "bedrock/us.anthropic.claude-opus-4-8", + ], +) +def test_bedrock_tools_pt_falsy_strict_always_dropped(model_id: str) -> None: + """``strict: false`` is the Chat Completions default, so forwarding it says nothing + the provider does not already assume. Bedrock Converse rejects the key's presence + for a growing set of Claude models, so it is dropped for every model, including the + ones whose cost-map entry still allows ``strict: true`` through.""" + result = _bedrock_tools_pt(_NON_STRICT_TOOL, model=model_id) + tool_spec = result[0]["toolSpec"] + assert ( + "strict" not in tool_spec + ), f"no-op strict: false leaked into toolSpec for {model_id}: {tool_spec}" + + +def test_responses_bridge_function_tool_does_not_reach_bedrock_with_strict() -> None: + """The Responses-to-Chat-Completions bridge stamps ``strict: false`` onto every + function tool even when the caller never sent one, which is how Codex CLI requests + acquired the key. Assert the fabricated value does not survive to toolSpec.""" + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + responses_tool = { + "type": "function", + "name": "get_weather", + "description": "Get the weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } + chat_tools, _ = ( + LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + [responses_tool] + ) + ) + result = _bedrock_tools_pt(chat_tools, model="bedrock/us.anthropic.claude-sonnet-5") + assert "strict" not in result[0]["toolSpec"] + + @pytest.mark.parametrize( "model_id", [ @@ -129,6 +198,15 @@ def test_bedrock_converse_supports_strict_tools_helper() -> None: ) is False ) + assert bedrock_converse_supports_strict_tools("anthropic.claude-sonnet-5") is False + assert ( + bedrock_converse_supports_strict_tools("bedrock/us.anthropic.claude-sonnet-5") + is False + ) + assert ( + bedrock_converse_supports_strict_tools("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") + is True + ) @pytest.mark.parametrize( @@ -143,6 +221,12 @@ def test_bedrock_converse_supports_strict_tools_helper() -> None: "us.anthropic.claude-sonnet-4-20250514-v1:0", "eu.anthropic.claude-sonnet-4-20250514-v1:0", "apac.anthropic.claude-sonnet-4-20250514-v1:0", + "anthropic.claude-sonnet-5", + "global.anthropic.claude-sonnet-5", + "us.anthropic.claude-sonnet-5", + "eu.anthropic.claude-sonnet-5", + "au.anthropic.claude-sonnet-5", + "jp.anthropic.claude-sonnet-5", ], ) def test_strict_tools_flag_set_in_model_cost_map(cost_map_key: str) -> None: diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index a3280b90fe3..c5d4bd044cc 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -2034,3 +2034,74 @@ def test_azure_traditional_api_uses_azure_openai_client(): assert isinstance( async_client, AsyncAzureOpenAI ), f"Expected AsyncAzureOpenAI client for api_version={api_version}" + + +def test_evicting_an_azure_client_built_on_the_callers_session_leaves_it_open(monkeypatch): + """`initialize_azure_sdk_client` puts `litellm.aclient_session` on the SDK client. + + That session belongs to the caller. `AsyncAzureOpenAI.close()` closes whatever + http client it was handed, so treating the wrapper as litellm's to close would + close the caller's shared session out from under them. + """ + import httpx + + from litellm.caching.evicted_client_closer import EvictedClientCloser + from litellm.caching.llm_caching_handler import LLMClientCache + + shared_session = httpx.AsyncClient() + closer = EvictedClientCloser(grace_seconds=0.0) + monkeypatch.setattr(litellm, "aclient_session", shared_session) + monkeypatch.setattr( + litellm, + "in_memory_llm_clients_cache", + LLMClientCache(evicted_client_closer=closer), + ) + + wrapper = BaseAzureLLM().get_azure_openai_client( + api_key="not-a-real-key", + api_base="https://litellm.openai.azure.com", + api_version="2024-02-01", + litellm_params={}, + _is_async=True, + ) + + assert wrapper is not None + assert wrapper._client is shared_session, "the wrapper should be built on the caller's session" + + closer.schedule(wrapper) + closer.reap() + + assert closer.pending_count == 0, "a wrapper around the caller's session must never be queued" + assert shared_session.is_closed is False, "closed the session the caller configured" + + +def test_an_azure_client_litellm_built_its_own_http_client_for_is_still_closed(monkeypatch): + """The ownership check must not turn the reclaim off for the ordinary case.""" + from litellm.caching.evicted_client_closer import EvictedClientCloser + from litellm.caching.llm_caching_handler import LLMClientCache + + closer = EvictedClientCloser(grace_seconds=0.0) + monkeypatch.setattr(litellm, "aclient_session", None) + monkeypatch.setattr(litellm, "client_session", None) + monkeypatch.setattr( + litellm, + "in_memory_llm_clients_cache", + LLMClientCache(evicted_client_closer=closer), + ) + + wrapper = BaseAzureLLM().get_azure_openai_client( + api_key="not-a-real-key", + api_base="https://litellm.openai.azure.com", + api_version="2024-02-01", + litellm_params={}, + _is_async=False, + ) + + assert wrapper is not None + closer.schedule(wrapper) + + assert closer.pending_count == 1, "litellm built this client's http client, so it owns it" + + closer.reap() + + assert wrapper.is_closed() is True diff --git a/tests/test_litellm/llms/openai/test_openai_common_utils.py b/tests/test_litellm/llms/openai/test_openai_common_utils.py index ce25f7e9af6..a099b5c659f 100644 --- a/tests/test_litellm/llms/openai/test_openai_common_utils.py +++ b/tests/test_litellm/llms/openai/test_openai_common_utils.py @@ -175,3 +175,75 @@ def test_get_openai_client_cache_key(client_type): ) assert isinstance(key, str) assert "api_key=sk-test" in key + + +def test_evicting_a_client_built_on_the_callers_session_leaves_that_session_open(monkeypatch): + """`litellm.aclient_session` belongs to the caller, who goes on using it. + + `_get_async_http_client` hands that session straight back, so the SDK client + litellm builds around it is only a wrapper. The SDK's `close()` closes + whatever http client it was given, so treating the wrapper as litellm's to + close would close the caller's shared session out from under them. + """ + import httpx + + from litellm.caching.evicted_client_closer import EvictedClientCloser + from litellm.caching.llm_caching_handler import LLMClientCache + from litellm.llms.openai.openai import OpenAIChatCompletion + + shared_session = httpx.AsyncClient() + closer = EvictedClientCloser(grace_seconds=0.0) + monkeypatch.setattr(litellm, "aclient_session", shared_session) + monkeypatch.setattr( + litellm, + "in_memory_llm_clients_cache", + LLMClientCache(evicted_client_closer=closer), + ) + + wrapper = OpenAIChatCompletion()._get_openai_client( + is_async=True, + api_key="sk-not-a-real-key", + api_base="https://api.openai.com/v1", + max_retries=2, + ) + + assert wrapper is not None + assert wrapper._client is shared_session, "the wrapper should be built on the caller's session" + + closer.schedule(wrapper) + closer.reap() + + assert closer.pending_count == 0, "a wrapper around the caller's session must never be queued" + assert shared_session.is_closed is False, "closed the session the caller configured" + + +def test_a_client_litellm_built_its_own_http_client_for_is_still_closed(monkeypatch): + """The ownership check must not turn the reclaim off for the ordinary case.""" + from litellm.caching.evicted_client_closer import EvictedClientCloser + from litellm.caching.llm_caching_handler import LLMClientCache + from litellm.llms.openai.openai import OpenAIChatCompletion + + closer = EvictedClientCloser(grace_seconds=0.0) + monkeypatch.setattr(litellm, "aclient_session", None) + monkeypatch.setattr(litellm, "client_session", None) + monkeypatch.setattr( + litellm, + "in_memory_llm_clients_cache", + LLMClientCache(evicted_client_closer=closer), + ) + + wrapper = OpenAIChatCompletion()._get_openai_client( + is_async=False, + api_key="sk-not-a-real-key", + api_base="https://api.openai.com/v1", + max_retries=2, + ) + + assert wrapper is not None + closer.schedule(wrapper) + + assert closer.pending_count == 1, "litellm built this client's http client, so it owns it" + + closer.reap() + + assert wrapper.is_closed() is True diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5f3b0f36b95..f5aa695cb78 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -883,6 +883,229 @@ async def test_get_user_object_upsert_includes_user_email(): assert creation_args["user_id"] == "new_test_user" +@pytest.mark.asyncio +async def test_get_user_object_backfills_null_email_from_cache_hit(): + """ + Regression (LIT-4710): an existing user row with a null user_email must be + backfilled from the JWT-provided email even when served from cache, so the + JWT-to-virtual-key path (which resolves straight to the cached user) stops + logging user_api_key_user_email=null forever. Before the fix the cached row + was returned unchanged and the DB was never updated. + """ + cache = UserApiKeyCache() + existing = LiteLLM_UserTable( + user_id="jwt-user-1", user_email=None, user_role="internal_user" + ) + await cache.async_set_cache( + key="jwt-user-1", value=existing, model_type=LiteLLM_UserTable + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable( + user_id="jwt-user-1", + user_email="jwt-user-1@example.com", + user_role="internal_user", + ) + ) + + result = await get_user_object( + user_id="jwt-user-1", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + user_id_upsert=False, + proxy_logging_obj=None, + user_email="jwt-user-1@example.com", + ) + + assert result is not None + assert result.user_email == "jwt-user-1@example.com" + + mock_prisma_client.db.litellm_usertable.update_many.assert_called_once() + update_kwargs = mock_prisma_client.db.litellm_usertable.update_many.call_args.kwargs + assert update_kwargs["where"] == {"user_id": "jwt-user-1", "user_email": None} + assert update_kwargs["data"]["user_email"] == "jwt-user-1@example.com" + + refreshed = await cache.async_get_cache( + key="jwt-user-1", model_type=LiteLLM_UserTable + ) + assert refreshed is not None + assert refreshed.user_email == "jwt-user-1@example.com" + + +@pytest.mark.asyncio +async def test_get_user_object_backfills_null_email_from_db_read(): + """ + Regression (LIT-4710): a user row read from the DB with a null user_email is + backfilled from the JWT-provided email before it is cached and returned. + """ + cache = UserApiKeyCache() + db_row = LiteLLM_UserTable( + user_id="jwt-user-3", user_email=None, user_role="internal_user" + ) + backfilled_row = LiteLLM_UserTable( + user_id="jwt-user-3", + user_email="jwt-user-3@example.com", + user_role="internal_user", + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + side_effect=[db_row, backfilled_row] + ) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=1) + + with patch( + "litellm.proxy.auth.auth_checks._should_check_db", return_value=True + ): + result = await get_user_object( + user_id="jwt-user-3", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + user_id_upsert=False, + proxy_logging_obj=None, + user_email="jwt-user-3@example.com", + ) + + assert result is not None + assert result.user_email == "jwt-user-3@example.com" + mock_prisma_client.db.litellm_usertable.update_many.assert_called_once() + + refreshed = await cache.async_get_cache( + key="jwt-user-3", model_type=LiteLLM_UserTable + ) + assert refreshed is not None + assert refreshed.user_email == "jwt-user-3@example.com" + + +@pytest.mark.asyncio +async def test_get_user_object_does_not_overwrite_existing_email(): + """ + LIT-4710 guardrail: backfill is scoped to null-to-value. An existing non-null + user_email (e.g. one an operator set intentionally) must never be overwritten + by the JWT-provided email. + """ + cache = UserApiKeyCache() + existing = LiteLLM_UserTable( + user_id="jwt-user-2", + user_email="operator-set@example.com", + user_role="internal_user", + ) + await cache.async_set_cache( + key="jwt-user-2", value=existing, model_type=LiteLLM_UserTable + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=0) + + result = await get_user_object( + user_id="jwt-user-2", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + user_id_upsert=False, + proxy_logging_obj=None, + user_email="different@example.com", + ) + + assert result is not None + assert result.user_email == "operator-set@example.com" + mock_prisma_client.db.litellm_usertable.update_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_user_object_backfill_race_prefers_db_email(): + """ + LIT-4710 race guard: when the null-guarded update matches 0 rows because a + concurrent writer already backfilled an email, the cache must be refreshed + with the value the DB accepted, not this request's proposed email. + """ + cache = UserApiKeyCache() + existing = LiteLLM_UserTable( + user_id="jwt-user-4", user_email=None, user_role="internal_user" + ) + await cache.async_set_cache( + key="jwt-user-4", value=existing, model_type=LiteLLM_UserTable + ) + + winner_row = LiteLLM_UserTable( + user_id="jwt-user-4", + user_email="winner@example.com", + user_role="internal_user", + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=winner_row + ) + + result = await get_user_object( + user_id="jwt-user-4", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + user_id_upsert=False, + proxy_logging_obj=None, + user_email="loser@example.com", + ) + + assert result is not None + assert result.user_email == "winner@example.com" + + refreshed = await cache.async_get_cache( + key="jwt-user-4", model_type=LiteLLM_UserTable + ) + assert refreshed is not None + assert refreshed.user_email == "winner@example.com" + + +@pytest.mark.asyncio +async def test_get_user_object_backfill_caches_persisted_email_not_proposed(): + """ + LIT-4710 cache-coherence: even when the null-guarded update succeeds, the + cache must be refreshed from the row the DB actually holds, not this + request's proposed email. A concurrent ordinary user update (not null + guarded) can change the email in the window before the cache write, so + optimistically caching the proposed email would serve a stale value. + """ + cache = UserApiKeyCache() + existing = LiteLLM_UserTable( + user_id="jwt-user-5", user_email=None, user_role="internal_user" + ) + await cache.async_set_cache( + key="jwt-user-5", value=existing, model_type=LiteLLM_UserTable + ) + + persisted_row = LiteLLM_UserTable( + user_id="jwt-user-5", + user_email="admin-edited@example.com", + user_role="internal_user", + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=persisted_row + ) + + result = await get_user_object( + user_id="jwt-user-5", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + user_id_upsert=False, + proxy_logging_obj=None, + user_email="jwt-user-5@example.com", + ) + + assert result is not None + assert result.user_email == "admin-edited@example.com" + + refreshed = await cache.async_get_cache( + key="jwt-user-5", model_type=LiteLLM_UserTable + ) + assert refreshed is not None + assert refreshed.user_email == "admin-edited@example.com" + + @pytest.mark.asyncio async def test_get_user_object_upsert_routes_default_team_to_membership(monkeypatch): """Regression for LIT-4324: a configured default team (list of NewUserRequestTeam diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index affaaa3fbf4..3177fc5ba44 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -1989,6 +1989,231 @@ class TestJWTOAuth2Coexistence: assert result.org_id == "validated-org" assert result.user_email == "validated@example.com" + @pytest.mark.asyncio + async def test_mapped_virtual_key_backfills_and_sets_user_email(self): + """ + Regression (LIT-4710): when a JWT resolves straight to an existing + virtual-key mapping (skipping auth_builder), the token's user_email must + still backfill the resolved user and be set on the returned + UserAPIKeyAuth. Before the fix the mapped path never passed the email + through, so user_api_key_user_email stayed null on every request. + """ + jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" + general_settings = {"enable_jwt_auth": True} + user_api_key_cache = DualCache() + prisma_client = MagicMock() + jwt_handler = MagicMock() + jwt_handler.is_jwt.return_value = True + jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "mapped-user"}) + jwt_handler.get_user_email = MagicMock(return_value="mapped@example.com") + jwt_handler.get_user_id = MagicMock(return_value="mapped-user") + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", + user_email_jwt_field="sub", + virtual_key_mapping_cache_ttl=300, + ) + + mapped_key = UserAPIKeyAuth( + token="hashed-mapped-key", + api_key="hashed-mapped-key", + user_id="mapped-user", + user_email=None, + ) + backfilled_user = LiteLLM_UserTable( + user_id="mapped-user", + user_email="mapped@example.com", + user_role="internal_user", + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + with ( + patch("litellm.proxy.proxy_server.general_settings", general_settings), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), + patch( + "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key", + new_callable=AsyncMock, + return_value=mapped_key, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + return_value=backfilled_user, + ) as mock_get_user_object, + ): + result = await _user_api_key_auth_builder( + request=mock_request, + api_key=jwt_token, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "gpt-4o-mini"}, + ) + + assert result.user_id == "mapped-user" + assert result.user_email == "mapped@example.com" + assert ( + mock_get_user_object.call_args_list[0].kwargs["user_email"] + == "mapped@example.com" + ) + + @pytest.mark.asyncio + async def test_mapped_virtual_key_does_not_backfill_mismatched_owner(self): + """ + LIT-4710 security guard: when an admin-created mapping points a JWT at a + virtual key owned by a different user, the JWT principal's email must not + be written onto the mapped key owner's record. Backfill only runs when the + mapped key owner is the JWT principal. + """ + jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" + general_settings = {"enable_jwt_auth": True} + user_api_key_cache = DualCache() + prisma_client = MagicMock() + jwt_handler = MagicMock() + jwt_handler.is_jwt.return_value = True + jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "jwt-principal"}) + jwt_handler.get_user_email = MagicMock(return_value="principal@example.com") + jwt_handler.get_user_id = MagicMock(return_value="jwt-principal") + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", + user_email_jwt_field="sub", + virtual_key_mapping_cache_ttl=300, + ) + + mapped_key = UserAPIKeyAuth( + token="hashed-mapped-key", + api_key="hashed-mapped-key", + user_id="other-owner", + user_email=None, + ) + other_owner = LiteLLM_UserTable( + user_id="other-owner", + user_email=None, + user_role="internal_user", + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + with ( + patch("litellm.proxy.proxy_server.general_settings", general_settings), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), + patch( + "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key", + new_callable=AsyncMock, + return_value=mapped_key, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + return_value=other_owner, + ) as mock_get_user_object, + ): + result = await _user_api_key_auth_builder( + request=mock_request, + api_key=jwt_token, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "gpt-4o-mini"}, + ) + + assert result.user_id == "other-owner" + assert result.user_email is None + assert all( + call.kwargs.get("user_email") != "principal@example.com" + for call in mock_get_user_object.call_args_list + ) + + @pytest.mark.asyncio + async def test_mapped_virtual_key_backfill_failure_does_not_break_auth(self): + """ + LIT-4710 resilience: a mapped-key request served from a valid cached key + must still authenticate when the best-effort email backfill cannot reach + the database, retaining null email rather than failing the request. + """ + jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" + general_settings = {"enable_jwt_auth": True} + user_api_key_cache = DualCache() + prisma_client = MagicMock() + jwt_handler = MagicMock() + jwt_handler.is_jwt.return_value = True + jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "mapped-user"}) + jwt_handler.get_user_email = MagicMock(return_value="mapped@example.com") + jwt_handler.get_user_id = MagicMock(return_value="mapped-user") + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", + user_email_jwt_field="sub", + virtual_key_mapping_cache_ttl=300, + ) + + mapped_key = UserAPIKeyAuth( + token="hashed-mapped-key", + api_key="hashed-mapped-key", + user_id="mapped-user", + user_email=None, + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + with ( + patch("litellm.proxy.proxy_server.general_settings", general_settings), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), + patch( + "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key", + new_callable=AsyncMock, + return_value=mapped_key, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + side_effect=Exception("can't reach database server"), + ), + ): + result = await _user_api_key_auth_builder( + request=mock_request, + api_key=jwt_token, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "gpt-4o-mini"}, + ) + + assert result.user_id == "mapped-user" + assert result.user_email is None + @pytest.mark.asyncio async def test_routing_override_routes_matching_jwt_to_oauth2(self): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 8dbf38555b3..95405a3b016 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -3477,6 +3477,140 @@ class TestStrategyRouterWriteValidation: is None ) + def test_create_with_empty_keyword_rule_rejected(self): + """LIT-5133: the router refuses to build a rule with no keyword, but only at load time. + Without this the row is written, dropped on reload, and the caller gets a 500 plus a + deployment that can never come back.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + + violation = _strategy_router_write_violation( + incoming_params=LiteLLM_Params( + model="auto_router/complexity_router", + complexity_router_default_model="gpt-4o-mini", + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": [], "tier": "COMPLEX"}], + }, + ), + existing_params=None, + ) + assert violation is not None + assert "complexity_router_config is invalid" in violation + assert "keyword_tier_rules" in violation + + def test_patch_that_only_renames_does_not_judge_the_stored_config(self): + """Only a config the write actually carries is judged. A row stored before this validation + existed is already unloadable, and holding its rename hostage would break the restore path + this function documents; the repair is a write that supplies a good config.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + stored_bad = LiteLLM_Params( + model="auto_router/complexity_router", + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": [" "], "tier": "COMPLEX"}], + }, + ) + assert ( + _strategy_router_write_violation( + incoming_params=updateLiteLLMParams(model="auto_router/complexity_router"), + existing_params=stored_bad, + ) + is None + ) + + def test_incoming_config_replaces_stored_rather_than_merging(self): + """The field is written wholesale, so a good incoming config must clear a bad stored one.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + stored_bad = LiteLLM_Params( + model="auto_router/complexity_router", + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": [], "tier": "COMPLEX"}], + }, + ) + assert ( + _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + model="auto_router/complexity_router", + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": ["invoice"], "tier": "COMPLEX"}], + }, + ), + existing_params=stored_bad, + ) + is None + ) + + def test_config_only_patch_is_judged_without_a_model_in_the_payload(self): + """A patch may carry a config and no model, which is what a caller updating only the + routing rules sends. That path skipped the naming contract, so it has to be judged on the + config alone against the stored model, or it overwrites a working router with one that + cannot load and takes it out of service.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + violation = _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": [], "tier": "COMPLEX"}], + } + ), + existing_params=self._stored_complexity_params(), + ) + assert violation is not None + assert "complexity_router_config is invalid" in violation + + def test_config_only_patch_with_a_loadable_config_is_allowed(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + assert ( + _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": ["invoice"], "tier": "COMPLEX"}], + } + ), + existing_params=self._stored_complexity_params(), + ) + is None + ) + + def test_config_only_patch_is_judged_on_the_config_alone(self): + """The stored model is encrypted at rest, so a patch that names no model cannot be + classified from the row. An unloadable config is rejected on its own merits instead, + which is also the only reading that closes the path regardless of what is stored.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + violation = _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + complexity_router_config={"keyword_tier_rules": [{"keywords": [], "tier": "COMPLEX"}]} + ), + existing_params=LiteLLM_Params(model="c2VjcmV0-encrypted-at-rest"), + ) + assert violation is not None + assert "complexity_router_config is invalid" in violation + def test_create_semantic_router_missing_embedding_rejected(self): from litellm.proxy.management_endpoints.model_management_endpoints import ( _strategy_router_write_violation, diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py index 3b2b1ccb793..21e25d30b82 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py @@ -66,6 +66,23 @@ def _admin_auth() -> UserAPIKeyAuth: ) +@pytest.fixture(autouse=True) +def stub_team_cache_refresh(): + """Keep the cached-team refresh out of the way of the mocked prisma rows. + + The endpoints under test now refresh the auth cache after their DB write. + That helper validates a real Prisma row into LiteLLM_TeamTableCachedObj, + which the MagicMock rows these tests use cannot satisfy. The refresh being + called at all is asserted explicitly in + test_disable_team_logging_refreshes_cached_team. + """ + with patch( + "litellm.proxy.management_endpoints.team_callback_endpoints._refresh_cached_team", + new_callable=AsyncMock, + ) as refresh: + yield refresh + + @pytest.fixture def unauthorized_caller(): return UserAPIKeyAuth( @@ -238,6 +255,9 @@ async def test_disable_team_logging_emits_audit_log_when_enabled(monkeypatch): assert before["metadata"]["callback_settings"]["success_callback"] == ["langfuse"] assert after["metadata"]["callback_settings"]["success_callback"] == [] assert after["metadata"]["callback_settings"]["failure_callback"] == [] + # The audit row has to show the slot the callbacks actually live in, so a + # disable of a logging-configured team does not record an empty diff. + assert after["metadata"]["logging"] == [] @pytest.mark.asyncio @@ -718,3 +738,207 @@ async def test_get_team_callbacks_reports_empty_for_team_without_callbacks(): "failure_callbacks": [], "callback_vars": {}, } + + +@pytest.mark.asyncio +async def test_disable_team_logging_stops_callbacks_registered_via_api(): + """Disabling logging must stop the callbacks that are actually running. + + Callbacks registered through the API or the Admin UI live in + metadata["logging"], and request-time resolution stops at that slot without + reading callback_settings. Clearing only callback_settings therefore reports + success while the team keeps sending to its logging destination. This drives + the endpoint and then asks the real request-time resolver what the written + row would do. + """ + from litellm.proxy.litellm_pre_call_utils import _get_dynamic_logging_metadata + + metadata = { + "logging": [ + { + "callback_name": "langsmith", + "callback_type": "success", + "callback_vars": {"langsmith_project": "tenant-project"}, + } + ] + } + mock_prisma = _patch_prisma(_team_row(team_id="team-1", metadata=metadata)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.master_key", None), + ): + response = await disable_team_logging( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + assert response["status"] == "success" + written = json.loads(mock_prisma.db.litellm_teamtable.update.await_args.kwargs["data"]["metadata"]) + assert written["logging"] == [] + + resolved = _get_dynamic_logging_metadata( + UserAPIKeyAuth(api_key="hashed", team_id="team-1", team_metadata=written), + proxy_config=MagicMock(**{"load_team_config.return_value": {}}), + ) + assert not (resolved.success_callback if resolved else None) + assert not (resolved.failure_callback if resolved else None) + + +@pytest.mark.asyncio +async def test_disable_team_logging_refreshes_cached_team(stub_team_cache_refresh): + """The DB write alone does not stop delivery. + + Auth serves a cached team object and request-time callback resolution reads + the metadata off it, so without this refresh a key that is already in flight + keeps sending to the destination until the cache entry expires. + """ + metadata = { + "logging": [ + { + "callback_name": "langsmith", + "callback_type": "success", + "callback_vars": {"langsmith_project": "tenant-project"}, + } + ] + } + mock_prisma = _patch_prisma(_team_row(team_id="team-1", metadata=metadata)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.master_key", None), + ): + await disable_team_logging( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + stub_team_cache_refresh.assert_awaited_once() + refreshed = stub_team_cache_refresh.await_args.kwargs["team_row"] + assert refreshed is mock_prisma.db.litellm_teamtable.update.return_value + # The row fed to the cache has to carry object_permission, or the refresh + # publishes a team whose tool allowlists look empty, which reads as + # unrestricted on the search-tool and MCP-tool checks. + update_kwargs = mock_prisma.db.litellm_teamtable.update.await_args.kwargs + assert update_kwargs["include"]["object_permission"] is True + + +@pytest.mark.asyncio +async def test_add_team_callbacks_refreshes_cached_team(stub_team_cache_refresh): + """Registering a callback must take effect for keys that are already live.""" + mock_prisma = _patch_prisma(_team_row(team_id="team-1", metadata={"logging": []})) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.master_key", None), + ): + await add_team_callbacks( + data=AddTeamCallback( + callback_name="langsmith", + callback_type="success", + callback_vars={"langsmith_project": "tenant-project"}, + ), + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + stub_team_cache_refresh.assert_awaited_once() + refreshed = stub_team_cache_refresh.await_args.kwargs["team_row"] + assert refreshed is mock_prisma.db.litellm_teamtable.update.return_value + update_kwargs = mock_prisma.db.litellm_teamtable.update.await_args.kwargs + assert update_kwargs["include"]["object_permission"] is True + + +@pytest.mark.asyncio +async def test_disable_team_logging_clears_both_metadata_shapes(): + """A team carrying both shapes ends up with neither active.""" + from litellm.proxy.litellm_pre_call_utils import _get_dynamic_logging_metadata + + metadata = { + "logging": [ + { + "callback_name": "langsmith", + "callback_type": "success_and_failure", + "callback_vars": {"langsmith_project": "tenant-project"}, + } + ], + "callback_settings": { + "success_callback": ["gcs_bucket"], + "failure_callback": ["langfuse"], + "callback_vars": {"gcs_bucket_name": "legacy-bucket"}, + }, + } + mock_prisma = _patch_prisma(_team_row(team_id="team-1", metadata=metadata)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.master_key", None), + ): + await disable_team_logging( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + written = json.loads(mock_prisma.db.litellm_teamtable.update.await_args.kwargs["data"]["metadata"]) + assert written["logging"] == [] + assert written["callback_settings"]["success_callback"] == [] + assert written["callback_settings"]["failure_callback"] == [] + + resolved = _get_dynamic_logging_metadata( + UserAPIKeyAuth(api_key="hashed", team_id="team-1", team_metadata=written), + proxy_config=MagicMock(**{"load_team_config.return_value": {}}), + ) + assert not (resolved.success_callback if resolved else None) + assert not (resolved.failure_callback if resolved else None) + + +@pytest.mark.asyncio +async def test_disable_team_logging_leaves_team_re_enablable(): + """The emptied slot must still accept a fresh registration afterwards.""" + metadata = { + "logging": [ + { + "callback_name": "langsmith", + "callback_type": "success", + "callback_vars": {"langsmith_project": "tenant-project"}, + } + ] + } + row = _team_row(team_id="team-1", metadata=metadata) + mock_prisma = _patch_prisma(row) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.master_key", None), + ): + await disable_team_logging( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + row.metadata = json.loads(mock_prisma.db.litellm_teamtable.update.await_args.kwargs["data"]["metadata"]) + row.model_dump.return_value["metadata"] = row.metadata + + await add_team_callbacks( + data=AddTeamCallback( + callback_name="langfuse", + callback_type="success", + callback_vars={"langfuse_public_key": "pk-lf-new"}, + ), + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + written = json.loads(mock_prisma.db.litellm_teamtable.update.await_args.kwargs["data"]["metadata"]) + assert [entry["callback_name"] for entry in written["logging"]] == ["langfuse"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_default_params.py b/tests/test_litellm/proxy/management_endpoints/test_team_default_params.py index e0b90332ca0..a485d95db06 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_default_params.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_default_params.py @@ -10,13 +10,14 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException -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 litellm from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_OrganizationTable, NewTeamRequest, + ProxyException, UserAPIKeyAuth, LitellmUserRoles, ) @@ -76,9 +77,7 @@ class TestConfigFieldsDefaultTeamParams: db_param_value=db_settings, ) - assert result["litellm_settings"]["default_team_params"] == { - "max_budget": 100.0 - } + assert result["litellm_settings"]["default_team_params"] == {"max_budget": 100.0} # Existing keys preserved assert result["litellm_settings"]["cache"] is False @@ -172,6 +171,22 @@ class TestNewTeamDefaultParamsApplied: user_role=LitellmUserRoles.PROXY_ADMIN, ) + def _make_org(self, organization_id: str, max_budget: float | None = None) -> LiteLLM_OrganizationTable: + return LiteLLM_OrganizationTable( + organization_id=organization_id, + budget_id="budget-id", + created_by="admin-user", + updated_by="admin-user", + litellm_budget_table=None if max_budget is None else LiteLLM_BudgetTable(max_budget=max_budget), + ) + + def _patch_org_lookup(self, monkeypatch, **mock_kwargs) -> AsyncMock: + from litellm.proxy.management_endpoints import team_endpoints + + lookup = AsyncMock(**mock_kwargs) + monkeypatch.setattr(team_endpoints, "get_org_object", lookup) + return lookup + @pytest.mark.asyncio async def test_all_defaults_applied_when_not_provided(self, monkeypatch): """When no budget/rate/permission fields are in the request, all defaults apply.""" @@ -312,6 +327,7 @@ class TestNewTeamDefaultParamsApplied: assert data.tpm_limit is None assert data.rpm_limit is None assert data.team_member_permissions is None + assert data.organization_id is None @pytest.mark.asyncio async def test_legacy_default_team_settings_fallback(self, monkeypatch): @@ -370,6 +386,144 @@ class TestNewTeamDefaultParamsApplied: # default_team_params wins (100.0), legacy fallback (999.0) not used assert data.max_budget == 100.0 + @pytest.mark.asyncio + async def test_default_organization_applied_and_validated(self, monkeypatch): + """The default org must land before the org-validation block, so a defaulted + org goes through the same existence + org-limit checks as an explicit one.""" + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr(litellm, "default_team_params", {"organization_id": "default-org"}) + org_lookup = self._patch_org_lookup(monkeypatch, return_value=self._make_org("default-org")) + + data = NewTeamRequest(team_alias="my-team") + + try: + await new_team( + data=data, + user_api_key_dict=self._make_admin_auth(), + http_request=MagicMock(), + ) + except Exception: + pass + + assert data.organization_id == "default-org" + org_lookup.assert_awaited_once() + assert org_lookup.await_args.kwargs["org_id"] == "default-org" + + @pytest.mark.asyncio + async def test_explicit_organization_wins_over_default(self, monkeypatch): + """An organization_id in the request must not be replaced by the default.""" + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr(litellm, "default_team_params", {"organization_id": "default-org"}) + org_lookup = self._patch_org_lookup(monkeypatch, return_value=self._make_org("explicit-org")) + + data = NewTeamRequest(team_alias="my-team", organization_id="explicit-org") + + try: + await new_team( + data=data, + user_api_key_dict=self._make_admin_auth(), + http_request=MagicMock(), + ) + except Exception: + pass + + assert data.organization_id == "explicit-org" + assert org_lookup.await_args.kwargs["org_id"] == "explicit-org" + + @pytest.mark.asyncio + async def test_nonexistent_default_organization_returns_400(self, monkeypatch): + """get_org_object raises instead of returning None, so an org that no longer + exists surfaced as a 500; team creation must report a 400 instead.""" + from litellm.proxy.auth.auth_checks import OrganizationNotFoundError + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr(litellm, "default_team_params", {"organization_id": "deleted-org"}) + self._patch_org_lookup( + monkeypatch, + side_effect=OrganizationNotFoundError("Organization doesn't exist in db. Organization=deleted-org"), + ) + + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=NewTeamRequest(team_alias="my-team"), + user_api_key_dict=self._make_admin_auth(), + http_request=MagicMock(), + ) + + assert exc_info.value.code == "400" + assert "deleted-org" in exc_info.value.message + + @pytest.mark.asyncio + async def test_defaulted_max_budget_validated_against_org_budget(self, monkeypatch): + """Defaults must be applied BEFORE _check_org_team_limits runs, or a default + max_budget above the org's cap is persisted unchecked.""" + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr( + litellm, + "default_team_params", + {"organization_id": "capped-org", "max_budget": 500.0}, + ) + self._patch_org_lookup(monkeypatch, return_value=self._make_org("capped-org", max_budget=100.0)) + + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=NewTeamRequest(team_alias="my-team"), + user_api_key_dict=self._make_admin_auth(), + http_request=MagicMock(), + ) + + assert exc_info.value.code == "400" + assert "exceeds organization's max_budget" in exc_info.value.message + + @pytest.mark.asyncio + async def test_explicit_budget_validated_against_default_org_budget(self, monkeypatch): + """The org lookup must load the budget table (include_budget_table=True); + without it litellm_budget_table is None and every budget comparison is skipped.""" + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr(litellm, "default_team_params", {"organization_id": "capped-org"}) + org_lookup = self._patch_org_lookup(monkeypatch, return_value=self._make_org("capped-org", max_budget=100.0)) + + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=NewTeamRequest(team_alias="my-team", max_budget=500.0), + user_api_key_dict=self._make_admin_auth(), + http_request=MagicMock(), + ) + + assert exc_info.value.code == "400" + assert "exceeds organization's max_budget" in exc_info.value.message + assert org_lookup.await_args.kwargs["include_budget_table"] is True + + @pytest.mark.asyncio + async def test_defaults_within_org_budget_still_created(self, monkeypatch): + """A default budget under the org cap must not be rejected by the reordered check.""" + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr( + litellm, + "default_team_params", + {"organization_id": "capped-org", "max_budget": 50.0}, + ) + self._patch_org_lookup(monkeypatch, return_value=self._make_org("capped-org", max_budget=100.0)) + + data = NewTeamRequest(team_alias="my-team") + + try: + await new_team( + data=data, + user_api_key_dict=self._make_admin_auth(), + http_request=MagicMock(), + ) + except Exception: + pass + + assert data.organization_id == "capped-org" + assert data.max_budget == 50.0 + # --------------------------------------------------------------------------- # _update_litellm_setting: setattr ordering @@ -536,18 +690,12 @@ class TestBulkUpdateTeamMemberPermissions: mock_batcher.commit = AsyncMock(return_value=None) mock_prisma = MagicMock() - mock_prisma.db.litellm_teamtable.find_many = AsyncMock( - return_value=[team_a, team_b] - ) + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_a, team_b]) mock_prisma.db.batch_ = MagicMock(return_value=mock_batcher) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) - data = BulkUpdateTeamMemberPermissionsRequest( - permissions=["/team/daily/activity"], apply_to_all_teams=True - ) - result = await bulk_update_team_member_permissions( - data=data, user_api_key_dict=self._admin_key_dict() - ) + data = BulkUpdateTeamMemberPermissionsRequest(permissions=["/team/daily/activity"], apply_to_all_teams=True) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) assert result["teams_updated"] == 2 calls = mock_batcher.litellm_teamtable.update.call_args_list @@ -555,19 +703,14 @@ class TestBulkUpdateTeamMemberPermissions: team_a_call = [c for c in calls if c.kwargs["where"]["team_id"] == "team-a"][0] assert "/key/generate" in team_a_call.kwargs["data"]["team_member_permissions"] - assert ( - "/team/daily/activity" - in team_a_call.kwargs["data"]["team_member_permissions"] - ) + assert "/team/daily/activity" in team_a_call.kwargs["data"]["team_member_permissions"] team_b_call = [c for c in calls if c.kwargs["where"]["team_id"] == "team-b"][0] assert "/key/delete" in team_b_call.kwargs["data"]["team_member_permissions"] assert "/key/update" in team_b_call.kwargs["data"]["team_member_permissions"] @pytest.mark.asyncio - async def test_all_teams_skips_teams_that_already_have_permission( - self, monkeypatch - ): + async def test_all_teams_skips_teams_that_already_have_permission(self, monkeypatch): """apply_to_all_teams: teams that already have the permission are skipped.""" from litellm.proxy.management_endpoints.team_endpoints import ( bulk_update_team_member_permissions, @@ -583,18 +726,12 @@ class TestBulkUpdateTeamMemberPermissions: mock_batcher.commit = AsyncMock(return_value=None) mock_prisma = MagicMock() - mock_prisma.db.litellm_teamtable.find_many = AsyncMock( - return_value=[team_has, team_missing] - ) + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_has, team_missing]) mock_prisma.db.batch_ = MagicMock(return_value=mock_batcher) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) - data = BulkUpdateTeamMemberPermissionsRequest( - permissions=["/team/daily/activity"], apply_to_all_teams=True - ) - result = await bulk_update_team_member_permissions( - data=data, user_api_key_dict=self._admin_key_dict() - ) + data = BulkUpdateTeamMemberPermissionsRequest(permissions=["/team/daily/activity"], apply_to_all_teams=True) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) assert result["teams_updated"] == 1 calls = mock_batcher.litellm_teamtable.update.call_args_list @@ -618,18 +755,12 @@ class TestBulkUpdateTeamMemberPermissions: mock_batcher.commit = AsyncMock(return_value=None) mock_prisma = MagicMock() - mock_prisma.db.litellm_teamtable.find_many = AsyncMock( - side_effect=[page1, page2] - ) + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(side_effect=[page1, page2]) mock_prisma.db.batch_ = MagicMock(return_value=mock_batcher) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) - data = BulkUpdateTeamMemberPermissionsRequest( - permissions=["/team/daily/activity"], apply_to_all_teams=True - ) - result = await bulk_update_team_member_permissions( - data=data, user_api_key_dict=self._admin_key_dict() - ) + data = BulkUpdateTeamMemberPermissionsRequest(permissions=["/team/daily/activity"], apply_to_all_teams=True) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) assert result["teams_updated"] == 502 find_calls = mock_prisma.db.litellm_teamtable.find_many.call_args_list @@ -656,18 +787,14 @@ class TestBulkUpdateTeamMemberPermissions: mock_batcher.commit = AsyncMock(return_value=None) mock_prisma = MagicMock() - mock_prisma.db.litellm_teamtable.find_many = AsyncMock( - return_value=[team_a, team_b] - ) + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_a, team_b]) mock_prisma.db.batch_ = MagicMock(return_value=mock_batcher) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) data = BulkUpdateTeamMemberPermissionsRequest( permissions=["/team/daily/activity"], team_ids=["team-a", "team-b"] ) - result = await bulk_update_team_member_permissions( - data=data, user_api_key_dict=self._admin_key_dict() - ) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) assert result["teams_updated"] == 2 @@ -692,18 +819,14 @@ class TestBulkUpdateTeamMemberPermissions: mock_batcher.commit = AsyncMock(return_value=None) mock_prisma = MagicMock() - mock_prisma.db.litellm_teamtable.find_many = AsyncMock( - return_value=[team_has, team_missing] - ) + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_has, team_missing]) mock_prisma.db.batch_ = MagicMock(return_value=mock_batcher) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) data = BulkUpdateTeamMemberPermissionsRequest( permissions=["/team/daily/activity"], team_ids=["team-has", "team-missing"] ) - result = await bulk_update_team_member_permissions( - data=data, user_api_key_dict=self._admin_key_dict() - ) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) assert result["teams_updated"] == 1 calls = mock_batcher.litellm_teamtable.update.call_args_list @@ -731,9 +854,7 @@ class TestBulkUpdateTeamMemberPermissions: ) with pytest.raises(HTTPException) as exc_info: - await bulk_update_team_member_permissions( - data=data, user_api_key_dict=self._admin_key_dict() - ) + await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) assert exc_info.value.status_code == 404 assert "team-b" in str(exc_info.value.detail) @@ -753,14 +874,10 @@ class TestBulkUpdateTeamMemberPermissions: mock_prisma = MagicMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) - data = BulkUpdateTeamMemberPermissionsRequest( - permissions=["/team/daily/activity"] - ) + data = BulkUpdateTeamMemberPermissionsRequest(permissions=["/team/daily/activity"]) with pytest.raises(HTTPException) as exc_info: - await bulk_update_team_member_permissions( - data=data, user_api_key_dict=self._admin_key_dict() - ) + await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) assert exc_info.value.status_code == 400 @@ -784,9 +901,7 @@ class TestBulkUpdateTeamMemberPermissions: ) with pytest.raises(HTTPException) as exc_info: - await bulk_update_team_member_permissions( - data=data, user_api_key_dict=self._admin_key_dict() - ) + await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) assert exc_info.value.status_code == 400 @@ -804,9 +919,7 @@ class TestBulkUpdateTeamMemberPermissions: monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) data = BulkUpdateTeamMemberPermissionsRequest(permissions=[]) - result = await bulk_update_team_member_permissions( - data=data, user_api_key_dict=self._admin_key_dict() - ) + result = await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._admin_key_dict()) assert result["teams_updated"] == 0 mock_prisma.db.litellm_teamtable.find_many.assert_not_called() @@ -824,14 +937,10 @@ class TestBulkUpdateTeamMemberPermissions: mock_prisma = MagicMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) - data = BulkUpdateTeamMemberPermissionsRequest( - permissions=["/team/daily/activity"], apply_to_all_teams=True - ) + data = BulkUpdateTeamMemberPermissionsRequest(permissions=["/team/daily/activity"], apply_to_all_teams=True) with pytest.raises(HTTPException) as exc_info: - await bulk_update_team_member_permissions( - data=data, user_api_key_dict=self._non_admin_key_dict() - ) + await bulk_update_team_member_permissions(data=data, user_api_key_dict=self._non_admin_key_dict()) assert exc_info.value.status_code == 403 @@ -844,6 +953,4 @@ class TestBulkUpdateTeamMemberPermissions: ) with pytest.raises(ValidationError): - BulkUpdateTeamMemberPermissionsRequest( - permissions=["/not/a/real/permission"] - ) + BulkUpdateTeamMemberPermissionsRequest(permissions=["/not/a/real/permission"]) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 0de2c1ac71d..fc018a5333e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -9863,11 +9863,14 @@ async def _drive_team_write( raw_body=None, user=None, find_returns_none=False, + mock_sink=None, ): """Drive POST ``update_team`` or PATCH ``patch_team`` against a mocked team. Returns ``(endpoint_result, update_mock)``; propagates whatever the endpoint raises. Inspect ``update_mock.call_args.kwargs["data"]`` for the DB write. + Pass a dict as ``mock_sink`` to receive the update mock even when the + endpoint raises. """ from unittest.mock import AsyncMock, MagicMock, Mock from unittest.mock import patch as _patch @@ -9913,6 +9916,8 @@ async def _drive_team_write( return_value=LiteLLM_TeamTable(team_id=_PATCH_TEAM_ID, team_alias="t") ) pc.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + if mock_sink is not None: + mock_sink["update"] = pc.db.litellm_teamtable.update req = Mock(spec=Request) if kind == "post": @@ -10175,6 +10180,249 @@ async def test_patch_returns_full_team_object_not_wrapper(): assert result.team_id == _PATCH_TEAM_ID +# --------------------------------------------------------------------------- +# custom_team_metadata_validate wiring: the configured validator must gate +# every team write path (POST /team/new, POST /team/update, PATCH /team/{id}) +# and must see the metadata that will actually be written. +# --------------------------------------------------------------------------- + +from contextlib import contextmanager + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_VALIDATOR_REGISTRY, + TeamMetadataValidationResult, +) + + +@contextmanager +def _configured_team_metadata_validator(validator): + TEAM_METADATA_VALIDATOR_REGISTRY.set(validator) + try: + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + yield + finally: + TEAM_METADATA_VALIDATOR_REGISTRY.set(None) + + +def _recording_validator(recorded, valid=True, error_message=None): + async def validator(payload): + recorded.append(payload) + return TeamMetadataValidationResult(valid=valid, error_message=error_message) + + return validator + + +@pytest.mark.asyncio +async def test_update_validator_sees_replacement_on_post_and_merged_on_patch(): + """POST hands the validator the wholesale replacement; PATCH hands it the + RFC 7386 merged result including preserved keys.""" + existing = {"cost_center": "OLD", "keep": 1} + body = {"metadata": {"cost_center": "NEW"}} + + recorded_post = [] + with _configured_team_metadata_validator(_recording_validator(recorded_post)): + await _drive_team_write("post", existing_metadata=existing, payload=body) + + recorded_patch = [] + with _configured_team_metadata_validator(_recording_validator(recorded_patch)): + await _drive_team_write("patch", existing_metadata=existing, payload=body) + + assert len(recorded_post) == 1 + assert recorded_post[0].operation == "update" + assert recorded_post[0].metadata == {"cost_center": "NEW"} + assert recorded_post[0].existing_metadata == existing + + assert len(recorded_patch) == 1 + assert recorded_patch[0].operation == "update" + assert recorded_patch[0].metadata == {"cost_center": "NEW", "keep": 1} + assert recorded_patch[0].existing_metadata == existing + + +@pytest.mark.asyncio +async def test_patch_null_delete_removes_key_from_validated_metadata(): + """Deleting a key via PATCH null must be visible to the validator as the + key's absence in the resulting metadata, so a required key cannot be + silently dropped.""" + recorded = [] + with _configured_team_metadata_validator(_recording_validator(recorded)): + await _drive_team_write( + "patch", + existing_metadata={"cost_center": "OLD", "keep": 1}, + payload={"metadata": {"cost_center": None}}, + ) + + assert len(recorded) == 1 + assert recorded[0].metadata == {"keep": 1} + assert recorded[0].existing_metadata == {"cost_center": "OLD", "keep": 1} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_update_without_metadata_skips_validator(kind): + recorded = [] + with _configured_team_metadata_validator(_recording_validator(recorded)): + await _drive_team_write(kind, existing_metadata={"k": "v"}, payload={"tpm_limit": 5}) + + assert recorded == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_update_validator_rejection_blocks_db_write(kind): + recorded = [] + validator = _recording_validator(recorded, valid=False, error_message="cost center rejected, contact FinOps") + sink = {} + + with _configured_team_metadata_validator(validator): + with pytest.raises(ProxyException) as exc_info: + await _drive_team_write( + kind, + existing_metadata={"cost_center": "OLD"}, + payload={"metadata": {"cost_center": "BAD"}}, + mock_sink=sink, + ) + + assert str(exc_info.value.code) == "400" + assert "cost center rejected, contact FinOps" in str(exc_info.value.message) + assert len(recorded) == 1 + sink["update"].assert_not_awaited() + + +@pytest.mark.asyncio +async def test_new_team_validator_runs_without_metadata_and_rejection_blocks_create(): + """Create always validates, even when the request carries no metadata, so a + required-key policy can reject a team created without one.""" + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + recorded = [] + validator = _recording_validator(recorded, valid=False, error_message="cost_center is required") + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + _configured_team_metadata_validator(validator), + ): + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_prisma.db.litellm_teamtable.create = AsyncMock() + mock_license.is_team_count_over_limit.return_value = False + + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=NewTeamRequest(team_alias="no-metadata-team"), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"), + ) + + assert str(exc_info.value.code) == "400" + assert "cost_center is required" in str(exc_info.value.message) + assert len(recorded) == 1 + assert recorded[0].operation == "create" + assert recorded[0].metadata == {} + assert recorded[0].existing_metadata is None + mock_prisma.db.litellm_teamtable.create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_new_team_validator_accept_proceeds_to_create(mock_db_client, mock_admin_auth): + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + mock_db_client.jsonify_team_object = lambda db_data: db_data + mock_db_client.get_data = AsyncMock(return_value=None) + mock_db_client.update_data = AsyncMock(return_value=MagicMock()) + mock_db_client.db = MagicMock() + + team_create_result = MagicMock(team_id="team-accept-1") + team_create_result.model_dump.return_value = {"team_id": "team-accept-1"} + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_usertable = MagicMock() + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + recorded = [] + with _configured_team_metadata_validator(_recording_validator(recorded)): + await new_team( + data=NewTeamRequest(team_alias="accepted-team", metadata={"cost_center": "CC-1001"}), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + assert len(recorded) == 1 + assert recorded[0].operation == "create" + assert recorded[0].metadata == {"cost_center": "CC-1001"} + mock_db_client.db.litellm_teamtable.create.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_new_team_rejection_precedes_model_alias_write(): + """A rejected create must not leave an orphaned LiteLLM_ModelTable row: + validation runs before the model_aliases insert.""" + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + validator = _recording_validator([], valid=False, error_message="cost_center is required") + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + _configured_team_metadata_validator(validator), + ): + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_prisma.db.litellm_teamtable.create = AsyncMock() + mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model-1")) + mock_license.is_team_count_over_limit.return_value = False + + with pytest.raises(ProxyException): + await new_team( + data=NewTeamRequest( + team_alias="alias-orphan-check", + model_aliases={"alias-model": "gpt-4o"}, + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"), + ) + + mock_prisma.db.litellm_modeltable.create.assert_not_awaited() + mock_prisma.db.litellm_teamtable.create.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_update_existing_metadata_excludes_system_managed_keys(kind): + """The validator's existing_metadata must be symmetric with metadata: + server-owned keys (team_member_budget_id) are stripped from both, so a + key-preservation validator never sees them 'disappear'.""" + recorded = [] + stored = {"cost_center": "CC-1001", "team_member_budget_id": "budget-abc"} + + with _configured_team_metadata_validator(_recording_validator(recorded)): + await _drive_team_write( + kind, + existing_metadata=dict(stored), + payload={"metadata": {"notes": "x"}}, + ) + + assert len(recorded) == 1 + assert recorded[0].existing_metadata == {"cost_center": "CC-1001"} + assert "team_member_budget_id" not in recorded[0].metadata + + # --------------------------------------------------------------------------- # PATCH body is validated through PatchTeamRequest before it is handed to # update_team. The write below must stay byte-identical to what the untyped @@ -10340,6 +10588,76 @@ async def test_list_available_teams_filters_joined_and_validates_rows(monkeypatc assert find_many_kwargs["where"] == {"team_id": {"in": ["team-open"]}} +@pytest.mark.asyncio +async def test_get_team_metadata_schema_returns_configured_fields(): + from litellm.proxy.management_endpoints.team_endpoints import get_team_metadata_schema + from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + parse_team_metadata_schema, + ) + + TEAM_METADATA_SCHEMA_REGISTRY.set( + parse_team_metadata_schema( + [ + {"key": "cost_center", "label": "Cost Center"}, + {"key": "app_name", "label": "Application Name"}, + ] + ) + ) + try: + result = await get_team_metadata_schema() + finally: + TEAM_METADATA_SCHEMA_REGISTRY.set(()) + + assert [field.key for field in result.fields] == ["cost_center", "app_name"] + assert result.fields[0].label == "Cost Center" + assert result.fields[1].label == "Application Name" + + +@pytest.mark.asyncio +async def test_get_team_metadata_schema_empty_when_unconfigured(): + from litellm.proxy.management_endpoints.team_endpoints import get_team_metadata_schema + from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + ) + + TEAM_METADATA_SCHEMA_REGISTRY.set(()) + result = await get_team_metadata_schema() + + assert result.fields == () + + +def test_get_team_metadata_schema_route_requires_auth(): + from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + parse_team_metadata_schema, + ) + + with patch("litellm.proxy.proxy_server.master_key", "sk-1234"): + response = client.get("/team/metadata_schema") + assert response.status_code == 401 + + TEAM_METADATA_SCHEMA_REGISTRY.set(parse_team_metadata_schema([{"key": "cost_center", "label": "Cost Center"}])) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1" + ) + try: + authed = client.get("/team/metadata_schema") + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + TEAM_METADATA_SCHEMA_REGISTRY.set(()) + + assert authed.status_code == 200 + assert authed.json() == {"fields": [{"key": "cost_center", "label": "Cost Center"}]} + + +def test_team_metadata_schema_route_is_readable_by_non_admins(): + from litellm.proxy._types import LiteLLMRoutes + + assert "/team/metadata_schema" in LiteLLMRoutes.info_routes.value + assert "/team/metadata_schema" in LiteLLMRoutes.management_routes.value + + def _provisioning_caller(role: LitellmUserRoles) -> UserAPIKeyAuth: return UserAPIKeyAuth(user_id="caller-1", user_role=role) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 795b7cd5a9e..979eb09d7db 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -606,6 +606,55 @@ async def test_default_team_params(team_params): assert create_call_args["models"] == ["special-gpt-5"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "team_params", + [ + DefaultTeamSSOParams(max_budget=10, budget_duration="1d", organization_id="default-org"), + {"max_budget": 10, "budget_duration": "1d", "organization_id": "default-org"}, + ], +) +async def test_default_team_params_organization_id_reaches_sso_created_team(team_params): + """The SSO auto-team path builds NewTeamRequest straight from default_team_params, + so a default organization_id must land on the created team row and be validated.""" + from litellm.proxy._types import LiteLLM_OrganizationTable + + litellm.default_team_params = team_params + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.create = AsyncMock() + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_prisma.get_data = AsyncMock(return_value=None) + mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + + mock_org = LiteLLM_OrganizationTable( + organization_id="default-org", + budget_id="budget-id", + created_by="admin", + updated_by="admin", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + AsyncMock(return_value=mock_org), + ) as mock_get_org: + team_id = str(uuid.uuid4()) + await MicrosoftSSOHandler.create_litellm_teams_from_service_principal_team_ids( + service_principal_teams=[ + MicrosoftServicePrincipalTeam( + principalId=team_id, + principalDisplayName="Test Team", + ) + ] + ) + + mock_prisma.db.litellm_teamtable.create.assert_called_once() + create_call_args = mock_prisma.db.litellm_teamtable.create.call_args.kwargs["data"] + assert create_call_args["organization_id"] == "default-org" + assert mock_get_org.call_args.kwargs["org_id"] == "default-org" + + @pytest.mark.asyncio async def test_create_team_without_default_params(): """ diff --git a/tests/test_litellm/proxy/management_helpers/team_metadata_validator_impls.py b/tests/test_litellm/proxy/management_helpers/team_metadata_validator_impls.py new file mode 100644 index 00000000000..39ebeb97ec1 --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/team_metadata_validator_impls.py @@ -0,0 +1,92 @@ +"""Three independent `custom_team_metadata_validate` implementations. + +Used by the matrix tests in `test_team_metadata_validation.py` and loadable +directly from a proxy config via `get_instance_fn` for live verification: + +- `validate_allowlist`: plain async function; requires `cost_center` and + checks it against a static allowlist. +- `validate_via_http`: async function that POSTs the metadata to an external + validation service (`TEAM_METADATA_VALIDATION_SERVICE_URL`); any transport + error or non-2xx response raises, exercising the fail-closed path. +- `IMMUTABLE_COST_CENTER_VALIDATOR`: class instance with an async + `__call__`; requires `cost_center` and forbids changing it once set, using + `existing_metadata` and `operation`. +""" + +import os + +import httpx + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataValidationPayload, + TeamMetadataValidationResult, +) + +ALLOWED_COST_CENTERS = frozenset({"CC-1001", "CC-1002"}) +DEFAULT_SERVICE_URL = "http://localhost:9414/validate" + + +async def validate_allowlist( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + cost_center = payload.metadata.get("cost_center") + if cost_center is None: + return TeamMetadataValidationResult( + valid=False, + error_message="cost_center is required in team metadata. Contact the FinOps team.", + ) + if cost_center not in ALLOWED_COST_CENTERS: + return TeamMetadataValidationResult( + valid=False, + error_message=f"Cost center {cost_center} is not recognized. Contact the FinOps team.", + ) + return TeamMetadataValidationResult(valid=True) + + +async def validate_via_http( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + service_url = os.environ.get("TEAM_METADATA_VALIDATION_SERVICE_URL", DEFAULT_SERVICE_URL) + async with httpx.AsyncClient(timeout=2.0) as client: + response = await client.post( + service_url, + json={"operation": payload.operation, "metadata": payload.metadata}, + ) + response.raise_for_status() + body = response.json() + if body.get("ok") is True: + return TeamMetadataValidationResult(valid=True) + return TeamMetadataValidationResult( + valid=False, + error_message=body.get("reason", "Rejected by the cost center service."), + ) + + +class ImmutableCostCenterValidator: + def __init__(self, immutable_key: str = "cost_center") -> None: + self.immutable_key = immutable_key + + async def __call__( + self, + payload: TeamMetadataValidationPayload, + ) -> TeamMetadataValidationResult: + current = payload.metadata.get(self.immutable_key) + if current is None: + return TeamMetadataValidationResult( + valid=False, + error_message=f"{self.immutable_key} is required in team metadata. Contact the FinOps team.", + ) + if payload.operation == "update" and payload.existing_metadata is not None: + prior = payload.existing_metadata.get(self.immutable_key) + if prior is not None and prior != current: + return TeamMetadataValidationResult( + valid=False, + error_message=( + f"{self.immutable_key} is immutable once set " + f"(stored: {prior}, requested: {current}). Contact the FinOps team." + ), + ) + return TeamMetadataValidationResult(valid=True) + + +IMMUTABLE_COST_CENTER_VALIDATOR = ImmutableCostCenterValidator() diff --git a/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py new file mode 100644 index 00000000000..d22c3db0f7e --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py @@ -0,0 +1,704 @@ +import asyncio +import os +import sys +from unittest.mock import patch + +import pytest +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.management_helpers.team_metadata_validation import ( + DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE, + DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS, + DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE, + TeamMetadataRequester, + TeamMetadataValidationPayload, + TeamMetadataValidationResult, + TeamMetadataValidatorRegistry, + _read_timeout_seconds, + _read_unavailable_message, + run_team_metadata_validation, + validate_team_metadata_if_configured, +) + + +def _registry_with(validator): + registry = TeamMetadataValidatorRegistry() + registry.set(validator) + return registry + +UNAVAILABLE_MESSAGE = "validation system down, contact ops" + + +def _payload(**overrides): + values = { + "operation": "create", + "metadata": {"cost_center": "CC-1001"}, + "existing_metadata": None, + "team_id": "team-1", + "team_alias": "alias-1", + "requester": TeamMetadataRequester(user_id="u1"), + } + values.update(overrides) + return TeamMetadataValidationPayload(**values) + + +async def _run(validator, payload=None, premium_user=True, timeout_seconds=1.0): + await run_team_metadata_validation( + validator=validator, + payload=payload or _payload(), + premium_user=premium_user, + timeout_seconds=timeout_seconds, + unavailable_message=UNAVAILABLE_MESSAGE, + ) + + +@pytest.mark.asyncio +async def test_valid_result_passes(): + async def validator(payload): + return TeamMetadataValidationResult(valid=True) + + await _run(validator) + + +@pytest.mark.asyncio +async def test_rejection_raises_400_with_validator_message(): + async def validator(payload): + return TeamMetadataValidationResult(valid=False, error_message="cost center rejected, contact FinOps") + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "cost center rejected, contact FinOps"} + + +@pytest.mark.asyncio +async def test_rejection_without_message_uses_default(): + async def validator(payload): + return TeamMetadataValidationResult(valid=False) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE} + + +@pytest.mark.asyncio +async def test_dict_shaped_return_is_accepted(): + async def validator(payload): + return {"valid": False, "error_message": "rejected via dict"} + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "rejected via dict"} + + +@pytest.mark.asyncio +async def test_validator_exception_fails_closed_with_generic_message(): + async def validator(payload): + raise RuntimeError("internal validation service is down") + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": UNAVAILABLE_MESSAGE} + + +@pytest.mark.asyncio +async def test_validator_timeout_fails_closed(): + async def validator(payload): + await asyncio.sleep(1.0) + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator, timeout_seconds=0.01) + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": UNAVAILABLE_MESSAGE} + + +@pytest.mark.asyncio +async def test_malformed_return_shape_fails_closed(): + async def validator(payload): + return "not-a-validation-result" + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": UNAVAILABLE_MESSAGE} + + +@pytest.mark.asyncio +async def test_non_premium_user_is_rejected(): + async def validator(payload): + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator, premium_user=False) + assert exc_info.value.status_code == 400 + assert CommonProxyErrors.not_premium_user.value in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_non_coroutine_validator_is_rejected(): + def validator(payload): + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == {"error": "custom_team_metadata_validate must be an async function"} + + +@pytest.mark.parametrize( + "general_settings, expected", + [ + ({}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": 2}, 2.0), + ({"team_metadata_validation_timeout": 0.5}, 0.5), + ({"team_metadata_validation_timeout": True}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": -1}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": 0}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": "3"}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ], +) +def test_read_timeout_seconds(general_settings, expected): + assert _read_timeout_seconds(general_settings) == expected + + +@pytest.mark.parametrize( + "general_settings, expected", + [ + ({}, DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE), + ( + {"team_metadata_validation_error_message": "call the help desk"}, + "call the help desk", + ), + ({"team_metadata_validation_error_message": " "}, DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE), + ({"team_metadata_validation_error_message": None}, DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE), + ], +) +def test_read_unavailable_message(general_settings, expected): + assert _read_unavailable_message(general_settings) == expected + + +@pytest.mark.asyncio +async def test_adapter_is_noop_when_unconfigured(): + calls = [] + + async def validator(payload): + calls.append(payload) + return TeamMetadataValidationResult(valid=True) + + await validate_team_metadata_if_configured( + operation="create", + metadata={"cost_center": "CC-1001"}, + existing_metadata=None, + team_id="team-1", + team_alias="alias-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="u1"), + registry=TeamMetadataValidatorRegistry(), + ) + assert calls == [] + + +@pytest.mark.asyncio +async def test_adapter_builds_payload_and_reads_settings(): + recorded = [] + + async def validator(payload): + recorded.append(payload) + return TeamMetadataValidationResult(valid=True) + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.proxy_server.general_settings", + {"team_metadata_validation_timeout": 3, "team_metadata_validation_error_message": "ops msg"}, + ), + ): + await validate_team_metadata_if_configured( + operation="update", + metadata={"cost_center": "CC-2001"}, + existing_metadata={"cost_center": "CC-1001", "keep": 1}, + team_id="team-9", + team_alias="alias-9", + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user-9", + user_email="user-9@example.com", + ), + registry=_registry_with(validator), + ) + + assert len(recorded) == 1 + payload = recorded[0] + assert payload.operation == "update" + assert payload.metadata == {"cost_center": "CC-2001"} + assert payload.existing_metadata == {"cost_center": "CC-1001", "keep": 1} + assert payload.team_id == "team-9" + assert payload.team_alias == "alias-9" + assert payload.requester.user_id == "user-9" + assert payload.requester.user_email == "user-9@example.com" + assert payload.requester.user_role == LitellmUserRoles.INTERNAL_USER.value + + +@pytest.mark.asyncio +async def test_adapter_normalizes_non_dict_metadata_to_empty_dict(): + recorded = [] + + async def validator(payload): + recorded.append(payload) + return TeamMetadataValidationResult(valid=True) + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + await validate_team_metadata_if_configured( + operation="create", + metadata=None, + existing_metadata=None, + team_id="team-1", + team_alias=None, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="u1"), + registry=_registry_with(validator), + ) + + assert len(recorded) == 1 + assert recorded[0].metadata == {} + assert recorded[0].existing_metadata is None + + +# --------------------------------------------------------------------------- +# Validator implementation matrix: three independent implementations +# (allowlist function, HTTP-service-backed function, immutability class +# instance) driven through the real team write endpoints. +# --------------------------------------------------------------------------- + +import json as _json +import socket +import threading +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from unittest.mock import AsyncMock, MagicMock, Mock + +import team_metadata_validator_impls as impls + +from litellm.proxy._types import ProxyException +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_VALIDATOR_REGISTRY, +) + + +class _CostCenterServiceHandler(BaseHTTPRequestHandler): + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + body = _json.loads(self.rfile.read(length) or b"{}") + cost_center = (body.get("metadata") or {}).get("cost_center") + if cost_center is None: + resp = {"ok": False, "reason": "cost_center missing per cost center service"} + elif cost_center not in impls.ALLOWED_COST_CENTERS: + resp = {"ok": False, "reason": f"cost center {cost_center} rejected by cost center service"} + else: + resp = {"ok": True} + payload = _json.dumps(resp).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format, *args): + pass + + +@pytest.fixture(scope="module") +def cost_center_service_url(): + server = ThreadingHTTPServer(("127.0.0.1", 0), _CostCenterServiceHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}/validate" + finally: + server.shutdown() + thread.join(timeout=5) + + +def _closed_port_url(): + probe = socket.socket() + probe.bind(("127.0.0.1", 0)) + port = probe.getsockname()[1] + probe.close() + return f"http://127.0.0.1:{port}/validate" + + +@contextmanager +def _configured(validator): + TEAM_METADATA_VALIDATOR_REGISTRY.set(validator) + try: + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + yield + finally: + TEAM_METADATA_VALIDATOR_REGISTRY.set(None) + + +async def _drive_create(metadata, mock_sink=None): + from fastapi import Request + + from litellm.proxy._types import LiteLLM_TeamTable, NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as pc, + patch("litellm.proxy.proxy_server._license_check") as lic, + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + ): + team_row = MagicMock(team_id="matrix-team-1") + team_row.model_dump.return_value = {"team_id": "matrix-team-1"} + pc.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + pc.get_data = AsyncMock(return_value=None) + pc.update_data = AsyncMock(return_value=MagicMock()) + pc.db.litellm_teamtable.create = AsyncMock(return_value=team_row) + pc.db.litellm_teamtable.count = AsyncMock(return_value=0) + pc.db.litellm_teamtable.update = AsyncMock(return_value=team_row) + pc.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + pc.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model-1")) + lic.is_team_count_over_limit.return_value = False + if mock_sink is not None: + mock_sink["team_create"] = pc.db.litellm_teamtable.create + mock_sink["model_create"] = pc.db.litellm_modeltable.create + + request_kwargs = {"team_alias": "matrix-team"} + if metadata is not None: + request_kwargs["metadata"] = metadata + return await new_team( + data=NewTeamRequest(**request_kwargs), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + +async def _drive_update(kind, existing_metadata, payload): + from fastapi import Request + + from litellm.proxy._types import LiteLLM_TeamTable, PatchTeamRequest, UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import patch_team, update_team + + team_id = "matrix-team-upd" + existing = LiteLLM_TeamTable( + team_id=team_id, + team_alias="matrix", + metadata=existing_metadata, + organization_id=None, + ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as pc, + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.team_endpoints._refresh_cached_team", + new=AsyncMock(), + ), + ): + pc.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing) + pc.db.litellm_teamtable.update = AsyncMock( + return_value=LiteLLM_TeamTable(team_id=team_id, team_alias="matrix") + ) + pc.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + + req = Mock(spec=Request) + if kind == "post": + return await update_team( + data=UpdateTeamRequest(team_id=team_id, **payload), + http_request=req, + user_api_key_dict=auth, + litellm_changed_by=None, + ) + return await patch_team( + team_id=team_id, + data=PatchTeamRequest.model_validate(dict(payload)), + http_request=req, + user_api_key_dict=auth, + litellm_changed_by=None, + ) + + +OK = ("ok", None) + +_MATRIX_IMPLS = { + "allowlist": lambda: impls.validate_allowlist, + "http": lambda: impls.validate_via_http, + "immutable_class": lambda: impls.IMMUTABLE_COST_CENTER_VALIDATOR, +} + +# (scenario, kind, existing_metadata, request payload, {impl: expected}) +# expected is ("ok", None) or ("reject", ) +_MATRIX_SCENARIOS = [ + ( + "create-valid-cost-center", + "create", + None, + {"cost_center": "CC-1001"}, + {"allowlist": OK, "http": OK, "immutable_class": OK}, + ), + ( + "create-missing-cost-center", + "create", + None, + None, + { + "allowlist": ("reject", "cost_center is required"), + "http": ("reject", "cost_center missing per cost center service"), + "immutable_class": ("reject", "cost_center is required"), + }, + ), + ( + "create-unknown-cost-center", + "create", + None, + {"cost_center": "CC-9999"}, + { + "allowlist": ("reject", "is not recognized"), + "http": ("reject", "rejected by cost center service"), + "immutable_class": OK, + }, + ), + ( + "patch-change-cost-center", + "patch", + {"cost_center": "CC-1001"}, + {"metadata": {"cost_center": "CC-1002"}}, + { + "allowlist": OK, + "http": OK, + "immutable_class": ("reject", "immutable once set"), + }, + ), + ( + "patch-unrelated-key-preserves-cost-center", + "patch", + {"cost_center": "CC-1001"}, + {"metadata": {"notes": "hello"}}, + {"allowlist": OK, "http": OK, "immutable_class": OK}, + ), + ( + "patch-null-deletes-cost-center", + "patch", + {"cost_center": "CC-1001"}, + {"metadata": {"cost_center": None}}, + { + "allowlist": ("reject", "cost_center is required"), + "http": ("reject", "cost_center missing per cost center service"), + "immutable_class": ("reject", "cost_center is required"), + }, + ), + ( + "post-replace-drops-cost-center", + "post", + {"cost_center": "CC-1001"}, + {"metadata": {"notes": "only-notes"}}, + { + "allowlist": ("reject", "cost_center is required"), + "http": ("reject", "cost_center missing per cost center service"), + "immutable_class": ("reject", "cost_center is required"), + }, + ), + ( + "update-without-metadata-skips-validation", + "post", + {"cost_center": "CC-1001"}, + {"tpm_limit": 5}, + {"allowlist": OK, "http": OK, "immutable_class": OK}, + ), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("impl_name", sorted(_MATRIX_IMPLS)) +@pytest.mark.parametrize( + "scenario, kind, existing_metadata, request_payload, expectations", + _MATRIX_SCENARIOS, + ids=[row[0] for row in _MATRIX_SCENARIOS], +) +async def test_validator_implementation_matrix( + monkeypatch, + cost_center_service_url, + impl_name, + scenario, + kind, + existing_metadata, + request_payload, + expectations, +): + monkeypatch.setenv("TEAM_METADATA_VALIDATION_SERVICE_URL", cost_center_service_url) + validator = _MATRIX_IMPLS[impl_name]() + outcome, message_part = expectations[impl_name] + + async def drive(): + if kind == "create": + return await _drive_create(metadata=request_payload) + return await _drive_update(kind, existing_metadata, request_payload) + + with _configured(validator): + if outcome == "ok": + await drive() + else: + with pytest.raises(ProxyException) as exc_info: + await drive() + assert str(exc_info.value.code) == "400", f"{scenario} x {impl_name}" + assert message_part in str(exc_info.value.message), f"{scenario} x {impl_name}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "kind, existing_metadata, request_payload", + [ + ("create", None, {"cost_center": "CC-1001"}), + ("patch", {"cost_center": "CC-1001"}, {"metadata": {"notes": "x"}}), + ], +) +async def test_http_validator_service_outage_fails_closed(monkeypatch, kind, existing_metadata, request_payload): + monkeypatch.setenv("TEAM_METADATA_VALIDATION_SERVICE_URL", _closed_port_url()) + + with _configured(impls.validate_via_http): + with pytest.raises(ProxyException) as exc_info: + if kind == "create": + await _drive_create(metadata=request_payload) + else: + await _drive_update(kind, existing_metadata, request_payload) + + assert str(exc_info.value.code) == "503" + assert DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_class_instance_with_async_call_is_accepted(): + await _run(impls.ImmutableCostCenterValidator(), payload=_payload(metadata={"cost_center": "CC-1"})) + + +@pytest.mark.asyncio +async def test_class_instance_with_sync_call_is_rejected(): + class SyncValidator: + def __call__(self, payload): + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(SyncValidator()) + assert exc_info.value.status_code == 500 + + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataSchemaRegistry, + parse_team_metadata_schema, +) + + +def test_parse_schema_none_returns_empty(): + assert parse_team_metadata_schema(None) == () + + +def test_parse_schema_round_trips_fields_in_order(): + raw = [ + {"key": "cost_center", "label": "Cost Center"}, + {"key": "app_name"}, + ] + + fields = parse_team_metadata_schema(raw) + + assert [field.key for field in fields] == ["cost_center", "app_name"] + assert fields[0].label == "Cost Center" + assert fields[1].label is None + + +@pytest.mark.parametrize( + "raw", + [ + "cost_center", + {"key": "cost_center"}, + [{"label": "missing key"}], + [{"key": ""}], + [{"key": "cost_center", "required": True}], + [{"key": "cost_center", "description": "Cost center code"}], + [{"key": "cost_center", "allowed_values": ["CC-1001"]}], + ], +) +def test_parse_schema_malformed_raises(raw): + with pytest.raises(Exception): + parse_team_metadata_schema(raw) + + +def test_parse_schema_duplicate_keys_raise(): + with pytest.raises(ValueError, match="duplicate"): + parse_team_metadata_schema([{"key": "cost_center"}, {"key": "app_name"}, {"key": "cost_center"}]) + + +def test_schema_registry_defaults_empty_and_round_trips(): + registry = TeamMetadataSchemaRegistry() + assert registry.get() == () + + fields = parse_team_metadata_schema([{"key": "cost_center"}]) + registry.set(fields) + assert registry.get() == fields + + registry.set(()) + assert registry.get() == () + + +@pytest.mark.asyncio +async def test_non_callable_validator_is_rejected_with_clean_500(): + class NotCallable: + pass + + with pytest.raises(HTTPException) as exc_info: + await _run(NotCallable()) + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == {"error": "custom_team_metadata_validate must be an async function"} + + +def test_parse_schema_duplicate_error_lists_offending_keys(): + with pytest.raises(ValueError) as exc_info: + parse_team_metadata_schema( + [{"key": "cost_center"}, {"key": "app_name"}, {"key": "cost_center"}, {"key": "app_name"}] + ) + assert str(exc_info.value) == "team_metadata_schema contains duplicate keys: app_name, cost_center" + + +@pytest.mark.asyncio +async def test_adapter_applies_configured_timeout_to_slow_validator(): + async def slow_validator(payload): + await asyncio.sleep(0.2) + return TeamMetadataValidationResult(valid=True) + + registry = TeamMetadataValidatorRegistry() + registry.set(slow_validator) + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.proxy_server.general_settings", + {"team_metadata_validation_timeout": 0.01}, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await validate_team_metadata_if_configured( + operation="create", + metadata={"cost_center": "CC-1001"}, + existing_metadata=None, + team_id="team-1", + team_alias="alias-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="u1"), + registry=registry, + ) + + assert exc_info.value.status_code == 503 diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 9a3702d2355..28d4d87e26f 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2016,6 +2016,48 @@ async def test_ProxyConfig__get_hierarchical_router_settings_missing_returns_non assert out is None +@pytest.mark.asyncio +async def test_ProxyConfig__get_hierarchical_router_settings_falls_back_to_team(monkeypatch): + """A key with no router_settings inherits the team's, so a team-level + model_group_alias reaches the request path at all.""" + pc = ProxyConfig() + fake_key = SimpleNamespace(router_settings=None, team_id="team-1") + team_settings = {"model_group_alias": {"group-a": "group-b"}} + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_team_object", + AsyncMock(return_value=SimpleNamespace(router_settings=team_settings)), + ) + + out = await pc._get_hierarchical_router_settings( + user_api_key_dict=fake_key, + prisma_client=None, + proxy_logging_obj=None, + ) + + assert out == team_settings + + +@pytest.mark.asyncio +async def test_ProxyConfig__get_hierarchical_router_settings_key_shadows_team_entirely(monkeypatch): + """Resolution returns whichever object it finds first, it does not merge + per field, so a key that sets any router setting hides every team setting + including an alias the key itself never set.""" + pc = ProxyConfig() + fake_key = SimpleNamespace(router_settings={"num_retries": 3}, team_id="team-1") + team_lookup = AsyncMock(return_value=SimpleNamespace(router_settings={"model_group_alias": {"group-a": "group-b"}})) + monkeypatch.setattr("litellm.proxy.proxy_server.get_team_object", team_lookup) + + out = await pc._get_hierarchical_router_settings( + user_api_key_dict=fake_key, + prisma_client=None, + proxy_logging_obj=None, + ) + + assert out == {"num_retries": 3} + assert "model_group_alias" not in out + team_lookup.assert_not_called() + + # --------------------------------------------------------------------------- # ProxyConfig._add_router_settings_from_db_config # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 3bb84e095a0..4d98a05da8d 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1,6 +1,7 @@ import asyncio import copy import datetime +from types import SimpleNamespace from typing import AsyncGenerator, Callable, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -28,11 +29,14 @@ from litellm.proxy.common_request_processing import ( _is_azure_model_router_request, _override_openai_response_model, _parse_event_data_for_error, + _resolve_per_request_model_group_alias, _should_return_raw_model_name, _UpstreamClosingStreamingResponse, create_response, ) from litellm.proxy.dd_span_tagger import DDSpanTagger +from litellm.proxy._types import ProxyException +from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth from litellm.proxy.utils import ProxyLogging @@ -5354,3 +5358,228 @@ class TestModelDeploymentsSupportStreamOptions: def test_non_string_model_is_not_injected(self): assert self._support(None, None) is False + + +class TestPerRequestModelGroupAlias: + """``router_settings.model_group_alias`` on a key or team has to be resolved + by the proxy: the Router resolves aliases from its own shared instance + attribute, which only ever holds the global config map.""" + + @staticmethod + def _router() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "group-a", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, + }, + { + "model_name": "group-b", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-fake"}, + }, + ] + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "alias_map, expected", + [ + ({"group-a": "group-b"}, "group-b"), + ({"group-a": {"model": "group-b", "hidden": True}}, "group-b"), + ({"group-b": "group-a"}, None), + ({"group-a": "group-a"}, None), + ({"group-a": {"hidden": True}}, None), + ({}, None), + (None, None), + ], + ) + async def test_resolves_alias_for_the_requested_model_group(self, alias_map, expected): + resolved = await _resolve_per_request_model_group_alias( + requested_model="group-a", + router_settings={"model_group_alias": alias_map}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash", models=[]), + llm_router=self._router(), + ) + + assert resolved == expected + + @pytest.mark.asyncio + async def test_alias_target_outside_the_key_allowlist_is_rejected(self): + """Access was authorized against the requested group, so a rewrite that + the key could not have requested directly must not be served.""" + with pytest.raises(ProxyException) as exc_info: + await _resolve_per_request_model_group_alias( + requested_model="group-a", + router_settings={"model_group_alias": {"group-a": "group-b"}}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash", models=["group-a"]), + llm_router=self._router(), + ) + + assert exc_info.value.code == "403" + assert "group-b" in exc_info.value.message + + @pytest.mark.asyncio + async def test_alias_target_inside_the_key_allowlist_resolves(self): + resolved = await _resolve_per_request_model_group_alias( + requested_model="group-a", + router_settings={"model_group_alias": {"group-a": "group-b"}}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash", models=["group-a", "group-b"]), + llm_router=self._router(), + ) + + assert resolved == "group-b" + + @pytest.mark.asyncio + @pytest.mark.parametrize("requested_model", [None, ["group-a", "group-b"]]) + async def test_non_string_requested_model_is_left_alone(self, requested_model): + """The routed model is not always a string (a batch request carries a + list), and an unhashable one must not blow up the alias lookup.""" + resolved = await _resolve_per_request_model_group_alias( + requested_model=requested_model, + router_settings={"model_group_alias": {"group-a": "group-b"}}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash", models=[]), + llm_router=self._router(), + ) + + assert resolved is None + + @pytest.mark.asyncio + async def test_pre_call_logic_rewrites_the_requested_model(self, monkeypatch): + """End to end through the request path: a key carrying the alias must + leave pre-call processing pointing at the alias target, not at the + group the caller asked for.""" + processing_obj = ProxyBaseLLMRequestProcessing(data={"model": "group-a"}) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def mock_add_litellm_data_to_request(*args, **kwargs): + return kwargs.get("data", {}) + + async def passthrough_pre_call_hook(user_api_key_dict, data, call_type): + return copy.deepcopy(data) + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=passthrough_pre_call_hook) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + + mock_proxy_config = MagicMock(spec=ProxyConfig) + mock_proxy_config._get_hierarchical_router_settings = AsyncMock( + return_value={"model_group_alias": {"group-a": "group-b"}} + ) + + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash", models=[]), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=mock_proxy_config, + route_type="acompletion", + llm_router=self._router(), + ) + + assert returned_data["model"] == "group-b" + assert returned_data["router_settings_override"] == {"model_group_alias": {"group-a": "group-b"}} + # The rewrite has to land before the pre-call hooks: they are where + # per-model budgets and rate limits are enforced, so resolving later + # applies the requested group's limits to a call the target serves. + assert mock_proxy_logging_obj.pre_call_hook.call_args.kwargs["data"]["model"] == "group-b" + + @pytest.mark.asyncio + async def test_team_level_alias_rewrites_the_requested_model(self, monkeypatch): + """The team path is separate resolution, not a variant of the key path: + settings are looked up on the team only when the key carries none. Runs + the real hierarchical lookup rather than mocking it, so this covers the + team half of the fix end to end.""" + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + + processing_obj = ProxyBaseLLMRequestProcessing(data={"model": "group-a"}) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def mock_add_litellm_data_to_request(*args, **kwargs): + return kwargs.get("data", {}) + + async def passthrough_pre_call_hook(user_api_key_dict, data, call_type): + return copy.deepcopy(data) + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=passthrough_pre_call_hook) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_team_object", + AsyncMock(return_value=SimpleNamespace(router_settings={"model_group_alias": {"group-a": "group-b"}})), + ) + + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash", models=[], team_id="team-1"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=RealProxyConfig(), + route_type="acompletion", + llm_router=self._router(), + ) + + assert returned_data["model"] == "group-b" + + @pytest.mark.asyncio + async def test_model_level_guardrails_resolve_against_the_alias_target(self, monkeypatch): + """Model-level guardrails are merged by model group name, so the merge + must see the target rather than the group the caller named.""" + processing_obj = ProxyBaseLLMRequestProcessing(data={"model": "group-a"}) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def mock_add_litellm_data_to_request(*args, **kwargs): + return kwargs.get("data", {}) + + async def passthrough_pre_call_hook(user_api_key_dict, data, call_type): + return copy.deepcopy(data) + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=passthrough_pre_call_hook) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + + merged_for: list = [] + + def recording_merge(data, llm_router, trust_client_model_info=True): + merged_for.append(data.get("model")) + return data + + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "_check_and_merge_model_level_guardrails", + recording_merge, + ) + + mock_proxy_config = MagicMock(spec=ProxyConfig) + mock_proxy_config._get_hierarchical_router_settings = AsyncMock( + return_value={"model_group_alias": {"group-a": "group-b"}} + ) + + await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash", models=[]), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=mock_proxy_config, + route_type="acompletion", + llm_router=self._router(), + ) + + assert merged_for == ["group-b"] diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index bceefae3a9f..0e9aac7bf85 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -28,6 +28,9 @@ from litellm.proxy.litellm_pre_call_utils import ( check_if_token_is_service_account, clean_headers, ) +from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + TRUSTED_CALLBACK_VARS_FIELD, +) from litellm.types.utils import CredentialItem sys.path.insert( @@ -5554,3 +5557,316 @@ def test_warn_stale_team_alias_once_evicts_oldest_key_beyond_cap(monkeypatch): pre_call_utils._warn_stale_team_alias_once("key-3", "stale alias") assert list(pre_call_utils._STALE_TEAM_ALIAS_WARNING_KEYS) == ["key-2", "key-3"] + + +_OAUTH_TOKEN = "Bearer sk-ant-oat01-regression-token-lit5108" + + +def _all_header_dicts(data: dict, metadata_variable_name: str) -> list[dict]: + metadata = data.get(metadata_variable_name) or {} + proxy_server_request = data["proxy_server_request"] + body = proxy_server_request["body"] + return [ + metadata.get("headers") or {}, + (metadata.get("requester_metadata") or {}).get("headers") or {}, + proxy_server_request["headers"], + (body.get(metadata_variable_name) or {}).get("headers") or {}, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path, metadata_variable_name", + [ + ("/v1/messages", "litellm_metadata"), + ("/v1/chat/completions", "metadata"), + ], +) +async def test_add_litellm_data_to_request_redacts_oauth_header_from_logging_copies(path, metadata_variable_name): + """The Anthropic subscription token is forwarded upstream but never handed to logging.""" + request_mock = _make_request_mock( + path, + { + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + "Authorization": _OAUTH_TOKEN, + "x-litellm-api-key": "Bearer sk-virtual-key", + }, + ) + + updated = await add_litellm_data_to_request( + data={"model": "anthropic-claude", "messages": [{"role": "user", "content": "hello"}]}, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"forward_client_headers_to_llm_api": True}, + version="test-version", + ) + + for header_dict in _all_header_dicts(updated, metadata_variable_name): + assert header_dict.get("Authorization") != _OAUTH_TOKEN + assert "sk-ant-oat01" not in json.dumps(header_dict) + + assert "sk-ant-oat01" not in json.dumps(updated["proxy_server_request"], default=repr) + + assert updated["proxy_server_request"]["headers"] is updated[metadata_variable_name]["headers"] + + assert updated["provider_specific_header"]["extra_headers"]["Authorization"] == _OAUTH_TOKEN + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_keeps_every_forwarded_credential_out_of_logging_copies(): + """Credentials kept for transport must not survive anywhere under proxy_server_request.""" + secrets = { + "x-api-key": "sk-byok-provider-key-lit5108", + "cookie": "litellm_jwt=session-token-lit5108", + "proxy-authorization": "Bearer proxy-token-lit5108", + } + request_mock = _make_request_mock( + "/v1/chat/completions", + { + "Content-Type": "application/json", + "x-litellm-api-key": "Bearer sk-virtual-key", + **secrets, + }, + ) + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]}, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={ + "forward_llm_provider_auth_headers": True, + "forward_client_headers_to_llm_api": True, + }, + version="test-version", + ) + + assert updated["api_key"] == secrets["x-api-key"] + assert updated["headers"]["x-api-key"] == secrets["x-api-key"] + + logged = json.dumps(updated["proxy_server_request"], default=repr) + for value in secrets.values(): + assert value not in logged + + + +@pytest.mark.parametrize( + "header, expected_redacted", + [ + ("Authorization", True), + ("X-Api-Key", True), + ("x-goog-api-key", True), + ("Ocp-Apim-Subscription-Key", True), + ("API-Key", True), + ("Cookie", True), + ("Proxy-Authorization", True), + ("anthropic-version", False), + ("user-agent", False), + ], +) +def test_redact_credential_headers_classifies_each_header(header, expected_redacted): + from litellm.proxy.litellm_pre_call_utils import redact_credential_headers + + headers = {header: "secret-value"} + + redacted = redact_credential_headers(headers) + + assert redacted[header] == ("***REDACTED***" if expected_redacted else "secret-value") + assert headers[header] == "secret-value" + + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_debug_log_does_not_print_credentials(): + """The request-header debug line carries values the stdout secret filter does not match.""" + import litellm.proxy.litellm_pre_call_utils as pre_call_utils + + request_mock = _make_request_mock( + "/v1/chat/completions", + { + "Content-Type": "application/json", + "Ocp-Apim-Subscription-Key": "apim-plaintext-token-lit5108", + "x-litellm-api-key": "Bearer sk-virtual-key", + }, + ) + + with patch.object(pre_call_utils.verbose_proxy_logger, "debug") as mock_debug: + await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]}, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"forward_llm_provider_auth_headers": True}, + version="test-version", + ) + + logged = " ".join(str(call) for call in mock_debug.call_args_list) + assert "apim-plaintext-token-lit5108" not in logged + + +def _callback_credential_request_mock() -> MagicMock: + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.path = "/v1/chat/completions" + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + return request_mock + + +_DATADOG_TEAM_KEY = UserAPIKeyAuth( + api_key="hashed-key", + team_id="team-1", + team_metadata={ + "logging": [ + { + "callback_name": "datadog", + "callback_type": "success", + "callback_vars": {"dd_api_key": "team-dd-key"}, + } + ] + }, +) + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_strips_caller_supplied_callback_credentials(): + """ + The team admin sets dd_api_key only; a caller pairing its own dd_site with that key + would ship the team's Datadog credential to a host it controls. + """ + caller_destinations = {"dd_site": "attacker.example.com", "dd_agent_host": "attacker.example.com"} + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "hello"}], + **caller_destinations, + "gcs_bucket_name": "attacker-bucket", + TRUSTED_CALLBACK_VARS_FIELD: {"dd_site": "smuggled.example.com"}, + "metadata": {**caller_destinations, "safe_user_metadata": "kept"}, + "litellm_metadata": dict(caller_destinations), + "litellm_params": {"metadata": dict(caller_destinations)}, + } + + updated = await add_litellm_data_to_request( + data=data, + request=_callback_credential_request_mock(), + user_api_key_dict=_DATADOG_TEAM_KEY, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert "dd_site" not in updated + assert "dd_agent_host" not in updated + assert "gcs_bucket_name" not in updated + assert updated["dd_api_key"] == "team-dd-key" + assert updated[TRUSTED_CALLBACK_VARS_FIELD] == {"dd_api_key": "team-dd-key"} + for metadata_key in ("metadata", "litellm_metadata"): + assert "dd_site" not in updated[metadata_key] + assert "dd_agent_host" not in updated[metadata_key] + assert "dd_site" not in updated["litellm_params"]["metadata"] + assert updated["metadata"]["safe_user_metadata"] == "kept" + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_strips_caller_supplied_callback_credentials_with_clientside_creds_allowed(): + """`allow_client_side_credentials` opens the auth-layer ban; the strip must still hold.""" + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "hello"}], + "dd_site": "attacker.example.com", + } + + updated = await add_litellm_data_to_request( + data=data, + request=_callback_credential_request_mock(), + user_api_key_dict=_DATADOG_TEAM_KEY, + proxy_config=MagicMock(), + general_settings={"allow_client_side_credentials": True}, + version="test-version", + ) + + assert "dd_site" not in updated + assert updated[TRUSTED_CALLBACK_VARS_FIELD] == {"dd_api_key": "team-dd-key"} + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_omits_trusted_callback_vars_without_team_callbacks(): + """Without team/key callback settings the trusted field must not exist for a callback to read.""" + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "hello"}], + TRUSTED_CALLBACK_VARS_FIELD: {"dd_api_key": "caller-key", "dd_site": "attacker.example.com"}, + } + + updated = await add_litellm_data_to_request( + data=data, + request=_callback_credential_request_mock(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert TRUSTED_CALLBACK_VARS_FIELD not in updated + + +def test_trusted_callback_vars_never_reach_the_provider(): + """ + The stamped field rides the request body, so it has to be a recognised litellm param; + otherwise the OpenAI param builder sweeps it into extra_body and the provider 400s. + """ + from litellm.utils import get_non_default_completion_params + + non_default = get_non_default_completion_params( + { + "model": "gpt-4", + TRUSTED_CALLBACK_VARS_FIELD: {"dd_api_key": "team-dd-key"}, + "some_provider_param": "kept", + } + ) + + assert TRUSTED_CALLBACK_VARS_FIELD not in non_default + assert non_default["some_provider_param"] == "kept" + + +@pytest.mark.asyncio +async def test_key_level_callback_vars_survive_the_strip(): + """ + Key-level callbacks configure their own destination and credentials, and they replace + team settings rather than merging with them, so only the request body is untrusted. + """ + key_with_datadog_callback = UserAPIKeyAuth( + api_key="hashed-key", + metadata={ + "logging": [ + { + "callback_name": "datadog", + "callback_type": "success", + "callback_vars": {"dd_api_key": "key-dd-key", "dd_site": "us5.datadoghq.com"}, + } + ] + }, + ) + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "hello"}], + "dd_site": "attacker.example.com", + } + + updated = await add_litellm_data_to_request( + data=data, + request=_callback_credential_request_mock(), + user_api_key_dict=key_with_datadog_callback, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated[TRUSTED_CALLBACK_VARS_FIELD] == {"dd_api_key": "key-dd-key", "dd_site": "us5.datadoghq.com"} + assert updated["dd_site"] == "us5.datadoghq.com" \ No newline at end of file diff --git a/tests/test_litellm/proxy/test_model_level_guardrails.py b/tests/test_litellm/proxy/test_model_level_guardrails.py index 48163bf5ed5..a1278e399b5 100644 --- a/tests/test_litellm/proxy/test_model_level_guardrails.py +++ b/tests/test_litellm/proxy/test_model_level_guardrails.py @@ -598,10 +598,15 @@ async def test_pre_call_merges_model_level_guardrails_before_pre_call_hook(): } ) - captured_pre_call_data: dict = {} + captured_pre_call_guardrails: list = [] async def fake_pre_call_hook(*, user_api_key_dict, data, call_type): - captured_pre_call_data.update(data) + # Snapshot the list rather than the dict: metadata is shared by + # reference, so a merge that happens after this point would otherwise + # show up here retroactively and the assertion would pass either way. + captured_pre_call_guardrails.extend( + (data.get("metadata") or {}).get("guardrails") or data.get("guardrails") or [] + ) return data proxy_logging = MagicMock() @@ -616,13 +621,9 @@ async def test_pre_call_merges_model_level_guardrails_before_pre_call_hook(): proxy_config = MagicMock() proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None) - # Stop the function before any post-pre_call_hook logic so we can keep - # the test focused. Raising _StopAfterPreCall in the next await fires - # right after the guardrail merge + pre_call_hook complete. - class _StopAfterPreCall(Exception): - pass - - proxy_config._get_hierarchical_router_settings.side_effect = _StopAfterPreCall() + # Assert on what pre_call_hook was handed rather than short-circuiting the + # function part way through: a sentinel keyed to one particular later call + # silently stops testing the ordering as soon as that call moves. with ( patch( @@ -640,30 +641,24 @@ async def test_pre_call_merges_model_level_guardrails_before_pre_call_hook(): ): from litellm.proxy._types import UserAPIKeyAuth - try: - await processing.common_processing_pre_call_logic( - request=MagicMock(headers={}, url=MagicMock(path="/v1/chat/completions")), - general_settings={}, - user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), - proxy_logging_obj=proxy_logging, - proxy_config=proxy_config, - route_type="acompletion", - version=None, - user_model=None, - user_temperature=None, - user_request_timeout=None, - user_max_tokens=None, - user_api_base=None, - model=None, - llm_router=mock_router, - ) - except _StopAfterPreCall: - pass + await processing.common_processing_pre_call_logic( + request=MagicMock(headers={}, url=MagicMock(path="/v1/chat/completions")), + general_settings={}, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + proxy_logging_obj=proxy_logging, + proxy_config=proxy_config, + route_type="acompletion", + version=None, + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model=None, + llm_router=mock_router, + ) # The pre_call_hook must have received data with the model-level # guardrail already merged in. Before the fix, this assertion fails # because pre_call_hook saw the original data without merge. - merged = (captured_pre_call_data.get("metadata") or {}).get("guardrails") or ( - captured_pre_call_data.get("guardrails") or [] - ) - assert "my-pre-call-guardrail" in merged + assert "my-pre-call-guardrail" in captured_pre_call_guardrails diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index d4fd5bc2dce..1075bffbeb2 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -2816,6 +2816,94 @@ def test_update_internal_user_settings_without_teams_skips_team_lookup(mock_prox assert mock_proxy_config["save_call_count"]() == 1 +@pytest.fixture +def mock_organization_lookup(monkeypatch): + """Back /update/default_team_settings with a fake organization table. + + Yields the set of organization ids that exist; the test mutates it before the call. + """ + from unittest.mock import AsyncMock, MagicMock + + import litellm + import litellm.proxy.proxy_server as proxy_server_module + + existing_organization_ids: set = set() + + async def _find_unique(where): + organization_id = where["organization_id"] + if organization_id not in existing_organization_ids: + return None + return {"organization_id": organization_id} + + find_unique = AsyncMock(side_effect=_find_unique) + fake_prisma = MagicMock() + fake_prisma.db.litellm_organizationtable.find_unique = find_unique + + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + monkeypatch.setattr(litellm, "default_team_params", {}) + + return { + "existing_organization_ids": existing_organization_ids, + "find_unique": find_unique, + } + + +def test_update_default_team_settings_rejects_unknown_organization( + mock_proxy_config, mock_auth, mock_organization_lookup +): + """Regression: an unknown default org saved fine here and then failed every + future team creation, far from the admin who typed it.""" + mock_organization_lookup["existing_organization_ids"].add("real-org") + + resp = client.patch( + "/update/default_team_settings", + json={"max_budget": 10.0, "organization_id": "ghost-org"}, + ) + + assert resp.status_code == 400, resp.text + assert "ghost-org" in resp.json()["detail"]["error"] + assert mock_proxy_config["save_call_count"]() == 0 + + import litellm + + assert litellm.default_team_params == {} + + +def test_update_default_team_settings_saves_when_organization_exists( + mock_proxy_config, mock_auth, mock_organization_lookup +): + """A real organization id still saves and reaches the in-memory settings.""" + mock_organization_lookup["existing_organization_ids"].add("real-org") + + resp = client.patch( + "/update/default_team_settings", + json={"max_budget": 10.0, "organization_id": "real-org"}, + ) + + assert resp.status_code == 200, resp.text + assert resp.json()["settings"]["organization_id"] == "real-org" + assert mock_proxy_config["save_call_count"]() == 1 + + import litellm + + assert litellm.default_team_params["organization_id"] == "real-org" + + +def test_update_default_team_settings_without_organization_skips_lookup( + mock_proxy_config, mock_auth, mock_organization_lookup +): + """Settings changes that don't set an organization must not pay for a DB round trip.""" + resp = client.patch( + "/update/default_team_settings", + json={"max_budget": 10.0}, + ) + + assert resp.status_code == 200, resp.text + mock_organization_lookup["find_unique"].assert_not_awaited() + assert mock_proxy_config["save_call_count"]() == 1 + + def test_update_mcp_semantic_filter_settings_requires_proxy_admin(monkeypatch): """Non-admin callers must not mutate global MCP semantic filter settings.""" from litellm.proxy._types import UserAPIKeyAuth diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index cc73273450a..3a94b1e0f85 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2903,7 +2903,7 @@ class TestRoutingDecisionCauseLogging: class TestSessionAffinity: - """Test the session_affinity sticky-routing behavior (on by default).""" + """Test the session_affinity sticky-routing behavior (off by default).""" REASONING_MESSAGE = [ { @@ -2917,18 +2917,14 @@ class TestSessionAffinity: def session_affinity_config(self, basic_config) -> Dict: return {**basic_config, "session_affinity": True} - @pytest.fixture - def session_affinity_disabled_config(self, basic_config) -> Dict: - return {**basic_config, "session_affinity": False} - @staticmethod def _request_kwargs(session_id: str) -> Dict: return {"metadata": {"session_id": session_id}} @pytest.mark.asyncio - async def test_enabled_by_default_pins_model(self, mock_router_instance, basic_config): - """Regression: session_affinity defaults to True, so a shared session_id pins the - first turn's model and later turns reuse it instead of reclassifying.""" + async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): + """Regression: session_affinity defaults to False, so a shared session_id must NOT + pin the first turn's model; every turn is classified on its own merits.""" assert "session_affinity" not in basic_config mock_router_instance.cache = DualCache() router = ComplexityRouter( @@ -2944,19 +2940,17 @@ class TestSessionAffinity: model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE ) assert first.model == "o1-preview" - assert second.model == "o1-preview" + assert second.model == "gpt-4o-mini" @pytest.mark.asyncio - async def test_can_be_disabled_reclassifies_every_turn( - self, mock_router_instance, session_affinity_disabled_config - ): - """Regression: session_affinity=False must still reclassify every turn even when a - shared session_id is present, so the opt-out keeps working.""" + async def test_can_be_enabled_to_pin_every_later_turn(self, mock_router_instance, session_affinity_config): + """Regression: session_affinity=True is the opt-in, so a shared session_id reuses the + first turn's model instead of reclassifying.""" mock_router_instance.cache = DualCache() router = ComplexityRouter( model_name="test-router", litellm_router_instance=mock_router_instance, - complexity_router_config=session_affinity_disabled_config, + complexity_router_config=session_affinity_config, ) request_kwargs = self._request_kwargs("session-1") first = await router.async_pre_routing_hook( @@ -2966,7 +2960,7 @@ class TestSessionAffinity: model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE ) assert first.model == "o1-preview" - assert second.model == "gpt-4o-mini" + assert second.model == "o1-preview" @pytest.mark.asyncio async def test_pins_model_after_first_turn(self, mock_router_instance, session_affinity_config): diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index ca290caac0b..73c9742876c 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -2,6 +2,7 @@ import pytest from litellm.router_utils.auto_router_model_naming import ( classify_strategy_router_model, + validate_complexity_router_config_write, validate_strategy_router_model_write, ) @@ -75,3 +76,74 @@ def test_validate_rejects_incoherent_writes(model, present_fields, expected_frag ) def test_validate_accepts_coherent_writes(model, present_fields): assert validate_strategy_router_model_write(model=model, present_fields=present_fields) is None + + +VALID_TIERS = { + "SIMPLE": ["gpt-4o-mini"], + "MEDIUM": ["gpt-4o-mini"], + "COMPLEX": ["gpt-4o"], + "REASONING": ["gpt-4o"], +} + + +@pytest.mark.parametrize( + "keyword_tier_rules,expected_fragment", + [ + ([{"keywords": [], "tier": "COMPLEX"}], "at least 1 item"), + ([{"keywords": [" "], "tier": "COMPLEX"}], "non-empty keyword"), + ( + [{"keywords": ["invoice"], "tier": "MEDIUM"}, {"keywords": [], "tier": "COMPLEX"}], + "at least 1 item", + ), + ], +) +def test_validate_rejects_unloadable_complexity_config(keyword_tier_rules, expected_fragment): + """A rule with no keyword makes ComplexityRouterConfig unbuildable, so the row must never be + written: without this the deployment is persisted, dropped at load, and the caller gets a 500.""" + violation = validate_complexity_router_config_write( + complexity_router_config={ + "tiers": VALID_TIERS, + "classifier_type": "heuristic", + "keyword_tier_rules": keyword_tier_rules, + } + ) + assert violation is not None + assert "complexity_router_config is invalid" in violation + assert expected_fragment in violation + + +@pytest.mark.parametrize( + "complexity_router_config", + [ + {"tiers": VALID_TIERS, "classifier_type": "heuristic"}, + { + "tiers": VALID_TIERS, + "classifier_type": "heuristic", + "keyword_tier_rules": [{"keywords": ["invoice", "refund"], "tier": "MEDIUM"}], + }, + # extra="allow" on the model, so an unrecognised key is not this gate's business + {"tiers": VALID_TIERS, "classifier_type": "heuristic", "some_future_key": "value"}, + ], +) +def test_validate_accepts_loadable_complexity_config(complexity_router_config): + assert validate_complexity_router_config_write(complexity_router_config=complexity_router_config) is None + + +def test_naming_check_ignores_the_config_entirely(): + """The naming contract and the config's contents are separate questions with separate owners; + a write may carry a config without naming a model, so neither can stand in for the other.""" + violation = validate_strategy_router_model_write( + model="auto_router/complexity_router", present_fields=frozenset() + ) + assert violation is not None + assert "requires" in violation + + +def test_config_check_ignores_the_model_entirely(): + assert validate_complexity_router_config_write(complexity_router_config=None) is None + assert ( + validate_complexity_router_config_write( + complexity_router_config={"tiers": VALID_TIERS, "keyword_tier_rules": [{"keywords": [], "tier": "COMPLEX"}]} + ) + is not None + ) diff --git a/tests/test_litellm/router_utils/test_router_utils_common_utils.py b/tests/test_litellm/router_utils/test_router_utils_common_utils.py index efa5f2382dc..7d453c72652 100644 --- a/tests/test_litellm/router_utils/test_router_utils_common_utils.py +++ b/tests/test_litellm/router_utils/test_router_utils_common_utils.py @@ -10,6 +10,7 @@ from litellm.router_utils.common_utils import ( add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, + resolve_model_group_alias, ) @@ -516,3 +517,44 @@ class TestAddModelFileIdMappings: def test_should_return_empty_mapping_when_given_empty_list(self): result = add_model_file_id_mappings([], []) assert result == {} + + +class TestResolveModelGroupAlias: + """``model_group_alias`` maps reach this helper from validated config and + from key/team rows, so both entry shapes must resolve and malformed entries + must not raise mid-request.""" + + @pytest.mark.parametrize( + "alias_map, expected", + [ + ({"group-a": "group-b"}, "group-b"), + ({"group-a": {"model": "group-b", "hidden": True}}, "group-b"), + ({"group-a": {"model": "group-b"}}, "group-b"), + ({"other": "group-b"}, None), + ({}, None), + (None, None), + ("not-a-map", None), + ({"group-a": {"hidden": True}}, None), + ({"group-a": {"model": 5}}, None), + ({"group-a": 5}, None), + ({"group-a": None}, None), + ({"group-a": ""}, None), + ], + ) + def test_resolves_both_entry_shapes_and_tolerates_malformed_entries(self, alias_map, expected): + assert resolve_model_group_alias(alias_map, "group-a") == expected + + def test_router_alias_resolution_uses_the_shared_helper(self): + router = Router( + model_list=[ + { + "model_name": "group-b", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, + } + ], + model_group_alias={"group-a": "group-b", "group-item": {"model": "group-b", "hidden": True}}, + ) + + assert router._get_model_from_alias("group-a") == "group-b" + assert router._get_model_from_alias("group-item") == "group-b" + assert router._get_model_from_alias("group-b") is None diff --git a/tests/test_litellm/test_circleci_rust_toolchain.py b/tests/test_litellm/test_circleci_rust_toolchain.py new file mode 100644 index 00000000000..c35ced51e16 --- /dev/null +++ b/tests/test_litellm/test_circleci_rust_toolchain.py @@ -0,0 +1,148 @@ +"""Static guardrails for how CircleCI provisions Rust. + +The root package builds `litellm-rust` through maturin, so any job that runs +`uv sync` or `uv build` compiles the bridge. The `cimg/python` images ship no +Rust toolchain, and when cargo is missing maturin's `puccinialin` helper +quietly provisions one itself: it fetches `rustup-init` from the unversioned +`https://static.rust-lang.org/rustup/dist//` path with no checksum and +installs a floating `stable` toolchain. uv suppresses build-backend output on a +successful sync, so that happens with nothing in the job log to show for it, +and the compiler a job builds with changes whenever upstream publishes. + +Two invariants are pinned here: + + 1. No step list (job or reusable command) reaches a `uv sync` / `uv build` + without a Rust toolchain already provisioned ahead of it. That is the + `install_rust` command on Linux and an inline pinned rustup install in the + Windows job, so the check accepts either. A new job that syncs without one + falls back to the unpinned path, which is exactly the regression a static + check catches at PR time and a green CI run does not. + 2. `install_rust` itself pins what it downloads: an explicit rustup version in + the URL, a verified SHA-256, and an exact toolchain version rather than a + channel name. + +The Windows job predates `install_rust` and provisions its toolchain inline, so +invariant 2 is scoped to `install_rust`; invariant 1 covers both. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest +import yaml + +REPO_ROOT = Path(__file__).resolve().parents[2] +CONFIG = REPO_ROOT / ".circleci" / "config.yml" + +BUILDS_WORKSPACE = re.compile(r"\buv\s+(?:sync|build)\b") +RUSTUP_ARCHIVE_URL = re.compile(r"https://static\.rust-lang\.org/rustup/archive/\d+\.\d+\.\d+/") +EXACT_TOOLCHAIN = re.compile(r"--default-toolchain\s+\"?\d+\.\d+\.\d+\"?") + + +def _config() -> dict[str, object]: + return yaml.safe_load(CONFIG.read_text()) + + +def _step_text(step: object) -> str: + """Flatten one step into the shell text it runs, or '' for a command reference.""" + if isinstance(step, dict): + run = step.get("run") + if isinstance(run, str): + return run + if isinstance(run, dict): + command = run.get("command") + return command if isinstance(command, str) else "" + return "" + + +def _without_comments(text: str) -> str: + return "\n".join(line for line in text.splitlines() if not line.lstrip().startswith("#")) + + +def _provisions_rust(step: object) -> bool: + if step == "install_rust": + return True + text = _step_text(step) + return "rustup-init" in text and ("sha256sum" in text or "SHA256" in text) + + +def _step_lists() -> dict[str, list[object]]: + config = _config() + lists: dict[str, list[object]] = {} + for kind in ("jobs", "commands"): + section = config.get(kind) + if not isinstance(section, dict): + continue + for name, body in section.items(): + steps = body.get("steps") if isinstance(body, dict) else None + if isinstance(steps, list): + lists[f"{kind[:-1]} {name}"] = steps + return lists + + +def _first_unprovisioned_build(steps: list[object]) -> str | None: + """Return the shell text of the first workspace build reached without Rust, if any.""" + rust_ready = False + for step in steps: + if _provisions_rust(step): + rust_ready = True + text = _step_text(step) + if BUILDS_WORKSPACE.search(_without_comments(text)) and not rust_ready: + return text + return None + + +def test_step_lists_exist() -> None: + lists = _step_lists() + assert "command install_rust" in lists + building = { + name + for name, steps in lists.items() + if any(BUILDS_WORKSPACE.search(_without_comments(_step_text(s))) for s in steps) + } + assert len(building) > 10, f"expected many workspace-building step lists, found {sorted(building)}" + + +def test_no_workspace_build_without_a_provisioned_rust_toolchain() -> None: + offenders = { + name: build for name, steps in _step_lists().items() if (build := _first_unprovisioned_build(steps)) is not None + } + assert not offenders, ( + "these CircleCI step lists run `uv sync`/`uv build` with no Rust toolchain provisioned first, " + "so maturin will download an unpinned rustup and a floating toolchain instead: " + f"{ {name: build.strip().splitlines()[0] for name, build in offenders.items()} }" + ) + + +@pytest.fixture(name="install_rust_command") +def _install_rust_command() -> str: + steps = _step_lists()["command install_rust"] + return "\n".join(_step_text(step) for step in steps) + + +def test_install_rust_pins_the_rustup_version_in_the_url(install_rust_command: str) -> None: + assert RUSTUP_ARCHIVE_URL.search(install_rust_command), ( + "install_rust must download rustup-init from a version-pinned /rustup/archive// URL; " + "the /rustup/dist/ path always serves whatever rustup is current" + ) + assert "/rustup/dist/" not in install_rust_command + + +def test_install_rust_verifies_the_installer_checksum(install_rust_command: str) -> None: + assert "sha256sum -c" in install_rust_command + assert re.search(r"RUSTUP_SHA256=[0-9a-f]{64}\b", install_rust_command), ( + "install_rust must compare the downloaded installer against a hardcoded SHA-256 " + "taken from rust-lang's published .sha256 sidecar" + ) + checksum_index = install_rust_command.index("sha256sum -c") + execute_index = install_rust_command.index("/tmp/rustup-init -y") + assert checksum_index < execute_index, "the checksum must be verified before the installer is executed" + + +def test_install_rust_pins_an_exact_toolchain_version(install_rust_command: str) -> None: + assert EXACT_TOOLCHAIN.search(install_rust_command), ( + "install_rust must pin an exact toolchain version (e.g. 1.97.1); a channel name like " + "stable/beta/nightly makes the compiler drift with whatever upstream published that day" + ) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index a2cf4e351f5..289c0a0afd6 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 23349 }, "LIT002": { - "limit": 27253 + "limit": 27252 }, "LIT003": { "limit": 292 diff --git a/ui/litellm-dashboard/CLAUDE.md b/ui/litellm-dashboard/CLAUDE.md index 5ec9392d2b0..701b37ec6aa 100644 --- a/ui/litellm-dashboard/CLAUDE.md +++ b/ui/litellm-dashboard/CLAUDE.md @@ -3,3 +3,9 @@ Never put LiteLLM tokens or API keys in `localStorage`. `localStorage` survives When you fix lint violations that are grandfathered in `eslint-suppressions.json`, run `eslint . --prune-suppressions` and commit the updated baseline so the gate ratchets down instead of leaving a stale suppression `src/lib/http/schema.d.ts` is generated from the proxy's OpenAPI spec; never hand-edit it. After changing a backend route or response model that the dashboard consumes, run `npm run gen:api` and commit the result (CI `Check UI API Types Sync` enforces this) + +Tests come in three tiers, named by the standard definitions. `Foo.test.tsx` is a unit test: one module, collaborators replaced by doubles, no multi-component tree, and it should run in milliseconds. `Foo.integration.test.tsx` renders a real component tree with real children and only stubs the network boundary; it costs seconds per case, so it earns its place by proving wiring that a unit test cannot reach. Browser-level tests live in `tests/e2e/ui/` as Playwright specs against a live proxy + +When a component holds logic worth asserting, extract the logic and unit-test it there rather than driving it through a render. `CreateMCPServer` is the worked example: its payload building lives in `createServerPayload.ts` with 46 unit tests that run in single-digit milliseconds, while `CreateMCPServer.integration.test.tsx` keeps only the cases that prove a form field reaches the right payload key. A test that renders a whole modal to assert the shape of one object belongs in the first category, not the second + +Most of the suite predates this split and is not yet classified, so an unsuffixed `*.test.tsx` is not evidence that a file is really a unit test. Classify what you touch diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 8c6d940cfea..50b652ff57e 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -779,6 +779,25 @@ "count": 1 } }, + "src/app/(dashboard)/mcp-servers/_components/AwsSigV4Fields.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx": { + "max-lines": { + "count": 1 + }, + "no-nested-ternary": { + "count": 1 + }, + "no-restricted-imports": { + "count": 2 + }, + "react-hooks/set-state-in-effect": { + "count": 4 + } + }, "src/app/(dashboard)/mcp-servers/_components/DcrBridgeToggle.tsx": { "no-restricted-imports": { "count": 1 @@ -856,6 +875,11 @@ "count": 1 } }, + "src/app/(dashboard)/mcp-servers/_components/OpenApiByokFields.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.test.tsx": { "no-restricted-imports": { "count": 1 @@ -900,23 +924,6 @@ "count": 1 } }, - "src/app/(dashboard)/mcp-servers/_components/create_mcp_server.tsx": { - "local/filename-pascal-case": { - "count": 1 - }, - "max-lines": { - "count": 1 - }, - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 2 - }, - "react-hooks/set-state-in-effect": { - "count": 4 - } - }, "src/app/(dashboard)/mcp-servers/_components/index.tsx": { "local/filename-pascal-case": { "count": 1 @@ -2855,6 +2862,16 @@ "count": 2 } }, + "src/components/common_components/MetadataKeyValueFields.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/common_components/MetadataKeyValueFields.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/common_components/ModelAliasManager.tsx": { "no-restricted-imports": { "count": 1 @@ -4135,11 +4152,6 @@ "count": 1 } }, - "src/components/view_logs/LogDetailsDrawer/prettyMessagesUtils.ts": { - "no-nested-ternary": { - "count": 1 - } - }, "src/components/view_logs/LogDetailsDrawer/useKeyboardNavigation.ts": { "react-hooks/immutability": { "count": 2 diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index c9953cce2ab..a1bd63151b4 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -67,7 +67,7 @@ "jsdom": "27.4.0", "knip": "5.83.1", "openapi-typescript": "7.13.0", - "postcss": "8.5.22", + "postcss": "8.5.23", "prettier": "3.2.5", "tailwindcss": "4.3.2", "tw-animate-css": "1.4.0", @@ -5529,9 +5529,9 @@ } }, "node_modules/brace-expansion": { - "version": "5.0.8", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.8.tgz", - "integrity": "sha512-JZyDyq3D4AUifKTPOB7DELf6XsB3WdPuNxCtob1vFXPsSXhdAiHBWJ/tJ8HAc9aH84BK+5JFZLNkJKx3G9kzQg==", + "version": "5.0.9", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.9.tgz", + "integrity": "sha512-ScQ4IuvIEF1TMlP7Zt+vjJ//9zlPb2SDcxWxM3bk8s6t6GGdJ7KO1dCcTidOPJKePW30LE/2cT7wCyPho9/Wxg==", "dev": true, "license": "MIT", "dependencies": { @@ -11064,9 +11064,9 @@ } }, "node_modules/postcss": { - "version": "8.5.22", - "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.22.tgz", - "integrity": "sha512-KBDEIpLrvpv16pp3K0Fw+UCoZfopFjjgeB+0tA/aaThfEE74kKDLrgg603YvOWJyg3+WYtyq3xYsQWsIyZlPqQ==", + "version": "8.5.23", + "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.23.tgz", + "integrity": "sha512-g50586zr4bZmwFiTlflMu8E0bDTb5I5gertgwAKmsdUlTQIhZtunzUlD1WSzwcVWPoAVpsrA6vlfCD7oXvRwgg==", "funding": [ { "type": "opencollective", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 32d93729dbe..4760b622f9b 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -79,7 +79,7 @@ "jsdom": "27.4.0", "knip": "5.83.1", "openapi-typescript": "7.13.0", - "postcss": "8.5.22", + "postcss": "8.5.23", "prettier": "3.2.5", "tailwindcss": "4.3.2", "tw-animate-css": "1.4.0", @@ -90,13 +90,13 @@ "overrides": { "prismjs": "1.30.0", "js-yaml": "4.3.0", - "brace-expansion": "5.0.8", + "brace-expansion": "5.0.9", "glob": "13.0.0", "minimatch": "10.2.4", "ws": "8.21.0", "braces": "3.0.3", "axios": "1.13.6", - "postcss": "8.5.22", + "postcss": "8.5.23", "esbuild": "0.28.1", "date-fns": "^4.4.0", "sharp": "^0.35.0" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.test.tsx new file mode 100644 index 00000000000..603cddb7b89 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.test.tsx @@ -0,0 +1,142 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderWithProviders } from "@/../tests/test-utils"; +import { screen, fireEvent } from "@testing-library/react"; +import { TeamGuardrailsTab } from "./TeamGuardrailsTab"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +vi.mock("@/components/networking", () => ({ + listGuardrailSubmissions: vi.fn(), + approveGuardrailSubmission: vi.fn(), + rejectGuardrailSubmission: vi.fn(), + updateGuardrailCall: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/guardrails/useRegisterGuardrail", () => ({ + useRegisterGuardrail: () => ({ + mutateAsync: vi.fn(), + isPending: false, + }), +})); + +vi.mock("@/components/common_components/team_dropdown", () => ({ + default: () => null, +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(), +})); + +import { listGuardrailSubmissions } from "@/components/networking"; + +const pendingSubmission = { + guardrail_id: "guard-1", + guardrail_name: "test-pending-guardrail", + status: "pending_review", + team_id: "team-1", + team_guardrail: true, + litellm_params: { + guardrail: "generic_guardrail_api", + mode: "pre_call", + api_base: "https://example.com/guard", + headers: { "X-API-Key": "secret" }, + extra_headers: ["x-request-id"], + }, + guardrail_info: {}, + submitted_at: "2026-05-09T00:00:00Z", +}; + +const baseAuth = { + token: "test-token", + accessToken: "test-token", + userId: "user-1", + userEmail: "user@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, +}; + +describe("TeamGuardrailsTab — approve/reject role gate", () => { + const mockUseAuthorized = vi.mocked(useAuthorized); + + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(listGuardrailSubmissions).mockResolvedValue({ + submissions: [pendingSubmission], + summary: { total: 1, pending_review: 1, active: 0, rejected: 0 }, + }); + }); + + it("hides Approve and Reject buttons for an internal user on a pending submission", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: "Internal User" }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + + expect(screen.queryByRole("button", { name: /approve/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /reject/i })).not.toBeInTheDocument(); + }); + + it("hides Approve and Reject buttons for an Admin Viewer, whom the backend rejects with 403", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: "Admin Viewer" }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + + expect(screen.queryByRole("button", { name: /approve/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /reject/i })).not.toBeInTheDocument(); + }); + + it("shows Approve and Reject buttons for an admin on a pending submission", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: "Admin" }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + + expect(screen.getByRole("button", { name: /approve/i })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /reject/i })).toBeInTheDocument(); + }); + + it("hides Approve and Reject buttons when userRole is undefined (defaults to non-admin)", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: undefined }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + + expect(screen.queryByRole("button", { name: /approve/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /reject/i })).not.toBeInTheDocument(); + }); + + it("disables all admin-only write controls for a non-admin, including the detail panel", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: "Internal User" }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + expect(screen.getByRole("switch")).toBeDisabled(); + + fireEvent.click(screen.getByRole("button", { name: "Review" })); + await screen.findByText("Forward LiteLLM API Key"); + + expect(screen.queryByRole("button", { name: /approve/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /reject/i })).not.toBeInTheDocument(); + screen.getAllByRole("switch").forEach((toggle) => expect(toggle).toBeDisabled()); + expect(screen.queryByRole("button", { name: "Add" })).not.toBeInTheDocument(); + expect(screen.queryByLabelText(/^Remove/)).not.toBeInTheDocument(); + expect(screen.queryByPlaceholderText("e.g. x-request-id")).not.toBeInTheDocument(); + }); + + it("keeps all write controls enabled for an admin in the detail panel", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: "Admin" }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + + fireEvent.click(screen.getByRole("button", { name: "Review" })); + await screen.findByText("Forward LiteLLM API Key"); + + expect(screen.getAllByRole("button", { name: /approve/i }).length).toBeGreaterThanOrEqual(2); + screen.getAllByRole("switch").forEach((toggle) => expect(toggle).toBeEnabled()); + expect(screen.getAllByRole("button", { name: "Add" })).toHaveLength(2); + expect(screen.getByLabelText("Remove X-API-Key")).toBeInTheDocument(); + expect(screen.getByLabelText("Remove x-request-id")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx index 4217a765732..496e1129371 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx @@ -27,6 +27,8 @@ import { import NotificationsManager from "@/components/molecules/notifications_manager"; import TeamDropdown from "@/components/common_components/team_dropdown"; import { useRegisterGuardrail } from "@/app/(dashboard)/hooks/guardrails/useRegisterGuardrail"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { isProxyAdminRole } from "@/utils/roles"; type GuardrailStatus = "active" | "pending" | "rejected"; @@ -188,16 +190,25 @@ function StatCard({ label, value, color }: { label: string; value: number; color ); } -function Toggle({ enabled, onToggle }: { enabled: boolean; onToggle: () => void }) { +function Toggle({ + enabled, + onToggle, + disabled = false, +}: { + enabled: boolean; + onToggle: () => void; + disabled?: boolean; +}) { return ( - {g.status === "pending" && ( + {isAdmin && g.status === "pending" && ( <> + {isAdmin && ( + + )} ))} )} -
- setNewStaticHeaderKey(e.target.value)} - placeholder="Header name (e.g. X-API-Key)" - className="flex-1 min-w-0 text-xs font-mono border border-gray-200 rounded-sm px-2 py-1.5 text-gray-700 placeholder-gray-400 focus:outline-hidden focus:ring-1 focus:ring-blue-500" - onKeyDown={(e) => { - if (e.key === "Enter") { - e.preventDefault(); + {isAdmin && ( +
+ setNewStaticHeaderKey(e.target.value)} + placeholder="Header name (e.g. X-API-Key)" + className="flex-1 min-w-0 text-xs font-mono border border-gray-200 rounded-sm px-2 py-1.5 text-gray-700 placeholder-gray-400 focus:outline-hidden focus:ring-1 focus:ring-blue-500" + onKeyDown={(e) => { + if (e.key === "Enter") { + e.preventDefault(); + const key = newStaticHeaderKey.trim(); + const value = newStaticHeaderValue.trim(); + if (key && !g.customHeaders.some((h) => h.key.toLowerCase() === key.toLowerCase())) { + onUpdateCustomHeaders([...g.customHeaders, { key, value }]); + setNewStaticHeaderKey(""); + setNewStaticHeaderValue(""); + } + } + }} + /> + setNewStaticHeaderValue(e.target.value)} + placeholder="Value" + className="flex-1 min-w-0 text-xs font-mono border border-gray-200 rounded-sm px-2 py-1.5 text-gray-700 placeholder-gray-400 focus:outline-hidden focus:ring-1 focus:ring-blue-500" + onKeyDown={(e) => { + if (e.key === "Enter") { + e.preventDefault(); + const key = newStaticHeaderKey.trim(); + const value = newStaticHeaderValue.trim(); + if (key && !g.customHeaders.some((h) => h.key.toLowerCase() === key.toLowerCase())) { + onUpdateCustomHeaders([...g.customHeaders, { key, value }]); + setNewStaticHeaderKey(""); + setNewStaticHeaderValue(""); + } + } + }} + /> + -
+ }} + className="text-xs font-medium text-blue-600 hover:text-blue-700 border border-blue-200 bg-blue-50 hover:bg-blue-100 px-2 py-1.5 rounded-sm transition-colors shrink-0" + > + Add + +
+ )}
@@ -546,50 +565,54 @@ function DetailPanel({ className="flex items-center justify-between gap-2 text-xs font-mono bg-gray-50 border border-gray-200 rounded-sm px-2 py-1.5" > {name} - + {isAdmin && ( + + )} ))} )} -
- setNewExtraHeader(e.target.value)} - placeholder="e.g. x-request-id" - className="flex-1 min-w-0 text-xs font-mono border border-gray-200 rounded-sm px-2 py-1.5 text-gray-700 placeholder-gray-400 focus:outline-hidden focus:ring-1 focus:ring-blue-500" - onKeyDown={(e) => { - if (e.key === "Enter") { - e.preventDefault(); + {isAdmin && ( +
+ setNewExtraHeader(e.target.value)} + placeholder="e.g. x-request-id" + className="flex-1 min-w-0 text-xs font-mono border border-gray-200 rounded-sm px-2 py-1.5 text-gray-700 placeholder-gray-400 focus:outline-hidden focus:ring-1 focus:ring-blue-500" + onKeyDown={(e) => { + if (e.key === "Enter") { + e.preventDefault(); + const name = newExtraHeader.trim().toLowerCase(); + if (name && !g.extraHeaders.map((h) => h.toLowerCase()).includes(name)) { + onUpdateExtraHeaders([...g.extraHeaders, name]); + setNewExtraHeader(""); + } + } + }} + /> + -
+ }} + className="text-xs font-medium text-blue-600 hover:text-blue-700 border border-blue-200 bg-blue-50 hover:bg-blue-100 px-2 py-1.5 rounded-sm transition-colors" + > + Add + +
+ )}
- {g.status === "pending" && ( + {isAdmin && g.status === "pending" && (
+
+ ), +})); + vi.mock("./ModelSelect/ModelSelect", () => { const ModelSelect = ({ value, onChange }: { value: string[]; onChange: (value: string[]) => void }) => ( + + + + + remove(fieldName)} + style={{ color: "#ef4444" }} + /> + + ))} + + + + + )} + + ); +}; + +export default MetadataKeyValueFields; diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx index db8e75302fe..456aee6330f 100644 --- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx @@ -11,6 +11,7 @@ interface OrganizationDropdownProps { disabled?: boolean; loading?: boolean; style?: React.CSSProperties; + placeholder?: string; } const OrganizationDropdown: React.FC = ({ @@ -20,11 +21,12 @@ const OrganizationDropdown: React.FC = ({ disabled, loading, style, + placeholder = "All Organizations", }) => { return (