diff --git a/.circleci/config.yml b/.circleci/config.yml index a8a33335ad7..dbeb412506f 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2690,6 +2690,122 @@ jobs: path: ui/litellm-dashboard/playwright-report destination: e2e-playwright-report + e2e_ui_testing_server_root_path: + docker: + - image: cimg/python:3.12-browsers@sha256:b432899af01c9a311bf74f4f22e9ada2e5306d4b1b4383f8d29e1228a5844ef2 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + - image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84 + environment: + POSTGRES_USER: e2euser + POSTGRES_PASSWORD: e2epassword + POSTGRES_DB: litellm_e2e + resource_class: large + working_directory: ~/project + environment: + DATABASE_URL: "postgresql://e2euser:e2epassword@localhost:5432/litellm_e2e" + CI: "true" + # The whole job exercises the proxy mounted under a prefix. SERVER_ROOT_PATH + # is read both by the proxy at boot (to rewrite the built UI bundle in place) + # and by migration.serverRootPath.config.ts, which refuses to run without it. + SERVER_ROOT_PATH: "/litellm" + steps: + - checkout + - setup_google_dns + - install_uv + - restore_cache: + keys: + - v1-uv-cache-{{ checksum "uv.lock" }} + - run: + name: Install Python dependencies + command: | + uv sync --frozen --all-groups --all-extras --python 3.12 + uv run --no-sync python -m prisma generate --schema litellm/proxy/schema.prisma + - save_cache: + key: v1-uv-cache-{{ checksum "uv.lock" }} + paths: + - ~/.cache/uv + - restore_cache: + keys: + - ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + - run: + name: Install Node dependencies and Playwright + command: | + cd ui/litellm-dashboard + npm ci + npx playwright install chromium + - save_cache: + key: ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }} + paths: + - ui/litellm-dashboard/node_modules + - ~/.cache/ms-playwright + - run: + name: Build UI from source + command: | + cd ui/litellm-dashboard + npm run build + rm -rf ../../litellm/proxy/_experimental/out + mv out ../../litellm/proxy/_experimental/out + find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do + d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html" + done + - wait_for_service: + url: tcp://localhost:5432 + timeout: "30" + - run: + name: Push Prisma schema + command: uv run --no-sync python -m prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss + - run: + name: Seed database + command: | + PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \ + -f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql + - run: + name: Start mock LLM server + command: uv run --no-sync python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py + background: true + - run: + name: Start LiteLLM proxy under a server root path + environment: + LITELLM_MASTER_KEY: "sk-1234" + MOCK_LLM_URL: "http://127.0.0.1:8090/v1" + DISABLE_SCHEMA_UPDATE: "true" + # Output flows to this step's own log, so a boot crash is visible here + # rather than swallowed by a downstream readiness probe. + command: | + LITELLM_LICENSE="$LITELLM_LICENSE" \ + uv run --no-sync python -m litellm.proxy.proxy_cli \ + --config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \ + --port 4000 + background: true + - run: + name: Wait for prefixed proxy to be ready + command: | + for i in $(seq 1 60); do + HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" --max-time 5 -H "Authorization: Bearer sk-1234" http://127.0.0.1:4000/litellm/health 2>/dev/null || true) + if [ "$HTTP_CODE" = "200" ]; then + echo "Prefixed proxy is ready" + exit 0 + fi + sleep 2 + done + echo "Prefixed proxy failed to start; see the 'Start LiteLLM proxy under a server root path' step for the boot log" + exit 1 + - run: + name: Run migration smoke under SERVER_ROOT_PATH + command: | + cd ui/litellm-dashboard + LITELLM_LICENSE="$LITELLM_LICENSE" \ + npx playwright test --config e2e_tests/migration.serverRootPath.config.ts + no_output_timeout: 10m + - store_artifacts: + path: ui/litellm-dashboard/test-results + destination: e2e-server-root-path-test-results + - store_artifacts: + path: ui/litellm-dashboard/playwright-report + destination: e2e-server-root-path-playwright-report + build_docker_database_image: machine: image: ubuntu-2204:2024.04.1 @@ -2795,6 +2911,8 @@ workflows: filters: *main_branches - e2e_ui_testing: filters: *main_branches + - e2e_ui_testing_server_root_path: + filters: *main_branches - build_and_test: requires: - build_docker_database_image diff --git a/.github/workflows/create-release.yml b/.github/workflows/create-release.yml index a726a921a2b..4834775e329 100644 --- a/.github/workflows/create-release.yml +++ b/.github/workflows/create-release.yml @@ -52,6 +52,22 @@ jobs: // are stable maintenance releases, not pre-releases. const isPrerelease = /(?:rc|nightly|alpha|beta|[-.]dev)/i.test(tag); + // A stable release should only claim the repo "latest" badge when its + // version is >= the current latest. Otherwise a backport (e.g. 1.84.6) + // would steal "latest" from a newer line (e.g. 1.88.1). + const versionKey = (rawTag) => { + const m = String(rawTag).match(/^v?(\d+)\.(\d+)\.(\d+)/); + if (!m) return null; + const maintenance = String(rawTag).match(/(?:\.post|\.patch\.)(\d+)/i); + return [Number(m[1]), Number(m[2]), Number(m[3]), maintenance ? Number(maintenance[1]) : 0]; + }; + const isAtLeast = (a, b) => { + for (let i = 0; i < a.length; i++) { + if (a[i] !== b[i]) return a[i] > b[i]; + } + return true; + }; + const cosignSection = [ `## Verify Docker Image Signature`, ``, @@ -90,6 +106,22 @@ jobs: ].join('\n'); try { + let makeLatest = "false"; + const newVersion = versionKey(tag); + if (!isPrerelease && newVersion) { + let latestVersion = null; + try { + const latest = await github.rest.repos.getLatestRelease({ + owner: context.repo.owner, + repo: context.repo.repo, + }); + latestVersion = versionKey(latest.data.tag_name); + } catch (error) { + if (error.status !== 404) throw error; + } + makeLatest = (!latestVersion || isAtLeast(newVersion, latestVersion)) ? "true" : "false"; + } + const response = await github.rest.repos.createRelease({ draft: true, generate_release_notes: true, @@ -108,6 +140,7 @@ jobs: release_id: response.data.id, body: updatedBody, draft: false, + make_latest: makeLatest, }); } catch (error) { diff --git a/.github/workflows/test-unit-misc.yml b/.github/workflows/test-unit-misc.yml index 9add77ff424..a7363ac3b43 100644 --- a/.github/workflows/test-unit-misc.yml +++ b/.github/workflows/test-unit-misc.yml @@ -28,6 +28,8 @@ jobs: tests/test_litellm/completion_extras tests/test_litellm/containers tests/test_litellm/experimental_mcp_client + tests/test_litellm/models + tests/test_litellm/repositories tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index c114a838d6d..3d2fa3e51c8 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -240,6 +240,24 @@ graph LR 7. `DBSpendUpdateWriter.update_database()` queues spend increments to Redis 8. Background job `update_spend` flushes queued spend to PostgreSQL every 60s +### Data Access Layer (Models & Repositories) + +Database entities and the operations on them live in two packages at the root of `litellm/` so both the gateway (`proxy/`) and the SDK can use them without importing proxy internals: + +- `litellm/models/` holds the canonical Pydantic definitions for every persisted entity (`LiteLLM_VerificationToken`, `LiteLLM_TeamTable`, `LiteLLM_UserTable`, etc.). `proxy/_types.py` re-exports these for backwards compatibility, so existing imports keep working. +- `litellm/repositories/` holds the data-access layer. `BaseRepository[T]` provides the generic CRUD (`find_by_id`, `find_many`, `create`, `update`, `delete`, `count`, `exists`); entity repositories such as `VerificationTokenRepository`, `TeamRepository`, and `UserRepository` add domain-specific queries and writes on top of it. + +Conventions to follow when touching this layer: + +| Concern | How it's handled | +|---------|------------------| +| JSON columns | Prisma `Json` columns are stored as JSON strings. Repositories `json.dumps()` on write and `json.loads()` on read (see `_to_model` and the `_build_*_data` helpers). | +| Archive-then-delete | `delete_team` / `delete_token` copy the row into the `LiteLLM_Deleted*` table and delete the original inside a single `prisma_client.db.tx()` transaction. Archive payloads are built explicitly so only columns that exist on the archive table are written. | +| Column vs. field names | Where a model field differs from its DB column (for example `org_id` maps to the `organization_id` column), the repository translates in both directions rather than relying on Pydantic to guess. | +| Array mutations | Adds use Prisma's atomic `push` (`add_member`, `add_admin`, `add_models`) to avoid read-modify-write races. Removals fall back to read-modify-write because Prisma has no atomic array remove. | + +To add a new entity, define the model under `litellm/models/`, re-export it from `proxy/_types.py` if existing code imports it from there, and add a repository under `litellm/repositories/` (subclass `BaseRepository` for plain CRUD, or add bespoke methods when the entity needs encryption, archiving, or atomic array updates). Mirror the tests in `tests/test_litellm/repositories/`. + --- ## 2. SDK Request Flow diff --git a/CLAUDE.md b/CLAUDE.md index 02a9630b486..758eac7e266 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -52,6 +52,19 @@ Do not put names of customers or customer company names in code, PRs, and issues CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI +Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors): + +- Composition over inheritance +- Never-nester: early returns over deep nesting +- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never) +- No mutation; instead of mutable lists and dicts, prefer tuples, NamedTuples, frozen dataclasses, etc. +- Use dependency injection +- Fully typed; no `Any` or coarse types like dict[str, Any]. Every function parameter must be strongly typed +- Use tagged unions + match +- No monster files or god objects + +Follow conventional commits for commit names and PR titles + ## Think Before Coding **Don't assume. Don't hide confusion. Surface tradeoffs** diff --git a/litellm/__init__.py b/litellm/__init__.py index f22971dfa13..e6c30e12286 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -442,6 +442,13 @@ def _dev_env_hot_reload_enabled() -> bool: custom_prometheus_tags: List[str] = [] prometheus_metrics_config: Optional[List] = None prometheus_emit_stream_label: bool = False +# Opt-in: emit `rate_limit_category` and `rate_limit_type` labels on +# `litellm_proxy_failed_requests_metric`. Off by default to preserve the +# pre-unification label set so existing dashboards / recording rules keyed on +# that metric keep matching after upgrade. Enable when downstream consumers +# are ready to split 429s by source (vendor vs. litellm) and dimension +# (RPM/TPM/concurrent/budget). +prometheus_emit_rate_limit_labels: bool = False prometheus_user_budget_label_include_email_alias: bool = False prometheus_end_user_metrics_max_series_per_metric: Optional[int] = 10000 prometheus_end_user_metrics_ttl_seconds: Optional[float] = 3600.0 @@ -1303,6 +1310,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None): NotFoundError, PermissionDeniedError, RateLimitError, + RateLimitErrorCategory, + RateLimitType, ServiceUnavailableError, BadGatewayError, OpenAIError, diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index d02afe37569..a0d63f5043c 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -129,7 +129,7 @@ "bash_20241022": null, "bash_20250124": null, "code-execution-2025-08-25": null, - "compact-2026-01-12": null, + "compact-2026-01-12": "compact-2026-01-12", "computer-use-2025-01-24": "computer-use-2025-01-24", "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 11733ce4cee..b6cfc8e7907 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -309,9 +309,13 @@ def get_cache_key(self, **kwargs) -> str: param_value = kwargs[param] cache_key += f"{str(param)}: {str(param_value)}" - verbose_logger.debug("\nCreated cache key: %s", cache_key) hashed_cache_key = Cache._get_hashed_cache_key(cache_key) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) + verbose_logger.debug( + "\nCreated cache key: %s (source material length: %d)", + hashed_cache_key, + len(cache_key), + ) # Remove preset_cache_key from kwargs to avoid "got multiple values" TypeError # when kwargs already contains preset_cache_key from upstream callers kwargs_for_preset = {k: v for k, v in kwargs.items() if k != "preset_cache_key"} @@ -497,6 +501,34 @@ def _get_cache_logic( return cached_response return cached_result + @staticmethod + def _get_safe_cache_lookup_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]: + cache_lookup_kwargs: Dict[str, Any] = {} + for prompt_kwarg in ("messages", "input"): + if prompt_kwarg in kwargs: + cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg] + + if isinstance(kwargs.get("metadata"), dict): + cache_lookup_kwargs["metadata"] = {} + + return cache_lookup_kwargs + + @staticmethod + def _update_metadata_from_cache_lookup_kwargs( + original_kwargs: Dict[str, Any], cache_lookup_kwargs: Dict[str, Any] + ) -> None: + original_metadata = original_kwargs.get("metadata") + cache_lookup_metadata = cache_lookup_kwargs.get("metadata") + if not isinstance(original_metadata, dict) or not isinstance( + cache_lookup_metadata, dict + ): + return + + if "semantic-similarity" in cache_lookup_metadata: + original_metadata["semantic-similarity"] = cache_lookup_metadata[ + "semantic-similarity" + ] + def get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): """ Retrieves the cached result for the given arguments. @@ -511,7 +543,6 @@ def get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): try: # never block execution if self.should_use_cache(**kwargs) is not True: return - messages = kwargs.get("messages", []) if "cache_key" in kwargs: cache_key = kwargs["cache_key"] else: @@ -523,12 +554,19 @@ def get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): or cache_control_args.get("s-max-age") or float("inf") ) + cache_lookup_kwargs = self._get_safe_cache_lookup_kwargs(kwargs) if dynamic_cache_object is not None: cached_result = dynamic_cache_object.get_cache( - cache_key, messages=messages + cache_key, **cache_lookup_kwargs ) else: - cached_result = self.cache.get_cache(cache_key, messages=messages) + cached_result = self.cache.get_cache( + cache_key, **cache_lookup_kwargs + ) + self._update_metadata_from_cache_lookup_kwargs( + original_kwargs=kwargs, + cache_lookup_kwargs=cache_lookup_kwargs, + ) return self._get_cache_logic( cached_result=cached_result, max_age=max_age ) @@ -549,7 +587,6 @@ async def async_get_cache( if self.should_use_cache(**kwargs) is not True: return - kwargs.get("messages", []) if "cache_key" in kwargs: cache_key = kwargs["cache_key"] else: @@ -654,6 +691,7 @@ def _convert_to_cached_embedding( self, embedding_response: Any, model: Optional[str], + prompt_tokens: Optional[int] = None, prompt_tokens_details: Optional[dict] = None, ) -> CachedEmbedding: """ @@ -666,6 +704,7 @@ def _convert_to_cached_embedding( "index": embedding_response.get("index"), "object": embedding_response.get("object"), "model": model, + "prompt_tokens": prompt_tokens, "prompt_tokens_details": prompt_tokens_details, } elif hasattr(embedding_response, "model_dump"): @@ -675,6 +714,7 @@ def _convert_to_cached_embedding( "index": data.get("index"), "object": data.get("object"), "model": model, + "prompt_tokens": prompt_tokens, "prompt_tokens_details": prompt_tokens_details, } else: @@ -684,6 +724,7 @@ def _convert_to_cached_embedding( "index": data.get("index"), "object": data.get("object"), "model": model, + "prompt_tokens": prompt_tokens, "prompt_tokens_details": prompt_tokens_details, } except KeyError as e: @@ -732,6 +773,29 @@ def _get_per_item_prompt_tokens_details( per_item[key] = value return per_item if per_item else None + def _get_per_item_prompt_tokens( + self, + result: EmbeddingResponse, + idx_in_result_data: int, + ) -> Optional[int]: + """ + Extract the per-item prompt_tokens from a response for caching. + + Single-item responses store the full usage.prompt_tokens. Multi-item + responses distribute it evenly (with remainder) so that summing all + per-item values on retrieval reconstructs the original total. + """ + if result.usage is None or result.usage.prompt_tokens is None: + return None + + total = result.usage.prompt_tokens + num_items = len(result.data) + if num_items <= 1: + return total + + quotient, remainder = divmod(total, num_items) + return quotient + (1 if idx_in_result_data < remainder else 0) + def add_embedding_response_to_cache( self, result: EmbeddingResponse, @@ -743,7 +807,11 @@ def add_embedding_response_to_cache( kwargs["cache_key"] = preset_cache_key embedding_response = result.data[idx_in_result_data] - # Extract per-item prompt_tokens_details from response usage + # Extract per-item prompt_tokens + details from response usage + prompt_tokens = self._get_per_item_prompt_tokens( + result=result, + idx_in_result_data=idx_in_result_data, + ) prompt_tokens_details = self._get_per_item_prompt_tokens_details( result=result, idx_in_result_data=idx_in_result_data, @@ -754,6 +822,7 @@ def add_embedding_response_to_cache( embedding_dict: CachedEmbedding = self._convert_to_cached_embedding( embedding_response, model_name, + prompt_tokens=prompt_tokens, prompt_tokens_details=prompt_tokens_details, ) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 3f4e54382c9..48691335b40 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -394,7 +394,7 @@ def _extract_model_from_cached_results( return cr["model"] return None - def _process_async_embedding_cached_response( + def _process_async_embedding_cached_response( # noqa: PLR0915 self, final_embedding_cached_response: Optional[EmbeddingResponse], cached_result: List[Optional[CachedEmbedding]], @@ -456,7 +456,10 @@ def _process_async_embedding_cached_response( index=idx, object="embedding", ) - if isinstance(kwargs_input_as_list[idx], str): + cached_prompt_tokens = cr.get("prompt_tokens") + if cached_prompt_tokens is not None: + prompt_tokens += cached_prompt_tokens + elif isinstance(kwargs_input_as_list[idx], str): from litellm.utils import token_counter prompt_tokens += token_counter( diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index da9e7b1e587..cce4b75795f 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -213,6 +213,78 @@ def _get_ttl(self, **kwargs) -> Optional[int]: ttl = int(ttl) return ttl + @classmethod + def _get_prompt_from_kwargs(cls, **kwargs) -> Optional[str]: + """ + Extract a semantic-cache prompt from chat or Responses API request kwargs. + """ + messages = kwargs.get("messages") + if messages: + return get_str_from_messages(messages) + + if "input" not in kwargs: + return None + + prompt_parts: List[str] = [] + cls._collect_responses_input_text(kwargs.get("input"), prompt_parts) + prompt = "\n".join(prompt_parts).strip() + return prompt or None + + @classmethod + def _collect_responses_input_text(cls, value: Any, prompt_parts: List[str]) -> None: + value = cls._coerce_response_input_value(value) + if value is None: + return + + if isinstance(value, str): + stripped_value = value.strip() + if stripped_value: + prompt_parts.append(stripped_value) + return + + if isinstance(value, (list, tuple)): + for item in value: + cls._collect_responses_input_text(item, prompt_parts) + return + + if isinstance(value, dict): + content = value.get("content") + if content is not None: + cls._collect_responses_input_text(content, prompt_parts) + return + + for text_key in ("text", "output", "input_text", "output_text"): + text_value = value.get(text_key) + if isinstance(text_value, str): + stripped_text = text_value.strip() + if stripped_text: + prompt_parts.append(stripped_text) + return + return + + content = getattr(value, "content", None) + if content is not None: + cls._collect_responses_input_text(content, prompt_parts) + return + + for text_key in ("text", "output", "input_text", "output_text"): + text_value = getattr(value, text_key, None) + if isinstance(text_value, str): + stripped_text = text_value.strip() + if stripped_text: + prompt_parts.append(stripped_text) + return + + @staticmethod + def _coerce_response_input_value(value: Any) -> Any: + model_dump = getattr(value, "model_dump", None) + if callable(model_dump): + return model_dump() + dict_method = getattr(value, "dict", None) + if callable(dict_method): + return dict_method() + return value + def _get_embedding(self, prompt: str) -> List[float]: """ Generate an embedding vector for the given prompt using the configured embedding model. @@ -278,13 +350,11 @@ def set_cache(self, key: str, value: Any, **kwargs) -> None: value_str: Optional[str] = None try: - # Extract the prompt from messages - messages = kwargs.get("messages", []) - if not messages: - print_verbose("No messages provided for semantic caching") + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + print_verbose("No prompt provided for semantic caching") return - prompt = get_str_from_messages(messages) value_str = str(value) store_kwargs: Dict[str, Any] = { @@ -315,14 +385,12 @@ def get_cache(self, key: str, **kwargs) -> Any: print_verbose(f"Redis semantic-cache get_cache, kwargs: {kwargs}") try: - # Extract the prompt from messages - messages = kwargs.get("messages", []) - if not messages: - print_verbose("No messages provided for semantic cache lookup") + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + print_verbose("No prompt provided for semantic cache lookup") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - prompt = get_str_from_messages(messages) # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. check_kwargs: Dict[str, Any] = { @@ -428,13 +496,11 @@ async def async_set_cache(self, key: str, value: Any, **kwargs) -> None: print_verbose(f"Async Redis semantic-cache set_cache, kwargs: {kwargs}") try: - # Extract the prompt from messages - messages = kwargs.get("messages", []) - if not messages: - print_verbose("No messages provided for semantic caching") + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + print_verbose("No prompt provided for semantic caching") return - prompt = get_str_from_messages(messages) value_str = str(value) # Generate embedding for the value (response) to cache @@ -471,15 +537,12 @@ async def async_get_cache(self, key: str, **kwargs) -> Any: print_verbose(f"Async Redis semantic-cache get_cache, kwargs: {kwargs}") try: - # Extract the prompt from messages - messages = kwargs.get("messages", []) - if not messages: - print_verbose("No messages provided for semantic cache lookup") + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + print_verbose("No prompt provided for semantic cache lookup") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - prompt = get_str_from_messages(messages) - # Generate embedding for the prompt prompt_embedding = await self._get_async_embedding(prompt, **kwargs) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 51abbbf729b..6d8b5cf8a57 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -402,6 +402,20 @@ def transform_request( instructions, ) = self.convert_chat_completion_messages_to_responses_api(messages) + # OpenAI's Responses API rejects an empty input. For a system-only + # request, carry the system message as a system-role input item instead + # of instructions, mirroring how non-string system content is already + # handled in convert_chat_completion_messages_to_responses_api. + if not input_items and instructions is not None: + input_items = [ + { + "type": "message", + "role": "system", + "content": [{"type": "input_text", "text": instructions}], + } + ] + instructions = None + optional_params = self._extract_extra_body_params(optional_params) # Build responses API request using the reverse transformation logic diff --git a/litellm/constants.py b/litellm/constants.py index 36e578bd323..57f55e6c177 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -831,6 +831,7 @@ "nano-gpt", # Nano-GPT - JSON-configured provider "poe", # Poe - JSON-configured provider "chutes", # Chutes - JSON-configured provider + "parasail", # Parasail - JSON-configured provider "featherless_ai", "nscale", "nebius", @@ -1157,6 +1158,7 @@ "openai.gpt-oss-120b-1:0", "anthropic.claude-haiku-4-5-20251001-v1:0", "anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-fable-5", "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6-v1:0", diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 15f6030d4a3..1cbef6b0b49 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -9,13 +9,109 @@ ## LiteLLM versions of the OpenAI Exception Types -from typing import Any, Dict, Optional +import enum +from typing import Any, Dict, Optional, Union import httpx import openai from litellm.types.utils import LiteLLMCommonStrings + +class RateLimitErrorCategory(str, enum.Enum): + """ + Category of a rate limit error, allowing callers to distinguish where the rate + limit originated. Exposed on every :class:`RateLimitError` instance via the + ``category`` attribute. + + Use these values to switch on the rate limit source, e.g.:: + + try: + ... + except litellm.RateLimitError as e: + if e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT: + ... # litellm's own limiter (key/team/user/model RPM/TPM/budget) + elif e.category == RateLimitErrorCategory.VENDOR_RATE_LIMIT: + ... # the upstream LLM provider returned 429 + """ + + VENDOR_RATE_LIMIT = "vendor_rate_limit" + """The upstream LLM provider returned a rate-limit response (e.g. OpenAI 429).""" + + VENDOR_BATCH_RATE_LIMIT = "vendor_batch_rate_limit" + """The upstream LLM provider returned a rate-limit response on a batch endpoint.""" + + LITELLM_RATE_LIMIT = "litellm_rate_limit" + """LiteLLM's own rate limiter (key/team/user/model RPM/TPM, budget, parallel-requests, etc.) blocked the request.""" + + LITELLM_BATCH_RATE_LIMIT = "litellm_batch_rate_limit" + """LiteLLM's own batch rate limiter (token/request budget across a batch input file) blocked the request.""" + + +class RateLimitType(str, enum.Enum): + """ + The dimension that was exceeded when a rate-limit error fired. + + This is orthogonal to :class:`RateLimitErrorCategory` — *category* tells + callers **who** rate-limited the request (the upstream vendor vs. one of + litellm's own limiters), while *type* tells them **which limit dimension** + was exceeded (an RPM ceiling, a TPM ceiling, a max-parallel-requests + ceiling, a budget cap, or a max-iterations cap). + + Surfaced both on every :class:`RateLimitError` instance via the + ``rate_limit_type`` attribute and on the structured + ``StandardLoggingPayload.error_information.error_rate_limit_type`` field + so custom callbacks / metrics consumers can split rate-limit failures by + cause without parsing free-text error messages. + """ + + REQUESTS = "requests" + """Requests-per-minute (RPM) or requests-per-window ceiling exceeded.""" + + TOKENS = "tokens" + """Tokens-per-minute (TPM) or tokens-per-window ceiling exceeded.""" + + CONCURRENT_REQUESTS = "concurrent_requests" + """``max_parallel_requests`` — too many in-flight requests at once.""" + + BUDGET = "budget" + """Spend budget cap reached (key, team, user, or per-session).""" + + MAX_ITERATIONS = "max_iterations" + """Per-session max-iterations cap reached (agent-style flows).""" + + +_RATE_LIMIT_CATEGORY_VALUES = frozenset(c.value for c in RateLimitErrorCategory) +_RATE_LIMIT_TYPE_VALUES = frozenset(t.value for t in RateLimitType) + + +def validate_rate_limit_category(value: Any) -> Optional[str]: + """Return ``value`` only if it matches a known :class:`RateLimitErrorCategory`. + + Used at duck-typed read sites (StandardLoggingPayload extraction, Prometheus + labels) to reject `.category` strings set by unrelated third-party exceptions + — otherwise those would leak into custom-callback payloads and Prometheus + label cardinality. + """ + if isinstance(value, RateLimitErrorCategory): + return value.value + if isinstance(value, str) and value in _RATE_LIMIT_CATEGORY_VALUES: + return value + return None + + +def validate_rate_limit_type(value: Any) -> Optional[str]: + """Return ``value`` only if it matches a known :class:`RateLimitType`. + + See :func:`validate_rate_limit_category` for the rationale. + """ + if isinstance(value, RateLimitType): + return value.value + if isinstance(value, str) and value in _RATE_LIMIT_TYPE_VALUES: + return value + return None + + _MINIMAL_ERROR_RESPONSE: Optional[httpx.Response] = None @@ -321,6 +417,18 @@ def __repr__(self): class RateLimitError(openai.RateLimitError): # type: ignore + """ + Unified rate-limit error. + + Every rate-limit condition surfaced by litellm — whether it originated from + an upstream LLM provider, a vendor batch endpoint, or one of litellm's own + proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget, + max-iterations, etc.) — is raised as an instance of this class. + + The :attr:`category` attribute lets callers distinguish the source. See + :class:`RateLimitErrorCategory` for the available values. + """ + def __init__( self, message, @@ -330,6 +438,12 @@ def __init__( litellm_debug_info: Optional[str] = None, max_retries: Optional[int] = None, num_retries: Optional[int] = None, + category: Union[str, RateLimitErrorCategory] = ( + RateLimitErrorCategory.VENDOR_RATE_LIMIT + ), + rate_limit_type: Optional[Union[str, RateLimitType]] = None, + headers: Optional[Dict[str, str]] = None, + detail: Any = None, ): self.status_code = 429 self.message = "litellm.RateLimitError: {}".format(message) @@ -338,9 +452,39 @@ def __init__( self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries + self.category = ( + category.value if isinstance(category, RateLimitErrorCategory) else category + ) + # Which dimension was exceeded — request count, token count, parallel + # requests, budget, max iterations. None when the source didn't + # classify the failure (e.g. legacy vendor 429 with no header hints). + self.rate_limit_type: Optional[str] = ( + rate_limit_type.value + if isinstance(rate_limit_type, RateLimitType) + else rate_limit_type + ) + # Headers explicitly attached to the error (e.g. retry-after, + # rate_limit_type, reset_at). Preserved across the proxy boundary so + # clients can react appropriately. + # + # IMPORTANT: we deliberately do NOT auto-populate self.headers from + # response.headers when only `response` is provided. A vendor 429 can + # set arbitrary response headers (Set-Cookie, CORS overrides, …); if + # those leaked into e.headers and a downstream proxy serializer + # forwarded them to the client, a malicious upstream could inject + # browser-interpreted headers for the proxy origin. Vendor response + # headers stay reachable on `e.response.headers` for callers that + # explicitly want them; only the proxy-supplied `headers=` kwarg + # makes it onto `self.headers`. _response_headers = ( getattr(response, "headers", None) if response is not None else None ) + self.headers: Optional[Dict[str, str]] = ( + {k: str(v) for k, v in headers.items()} if headers else None + ) + # Mirrors FastAPI HTTPException.detail so the same instance can be + # serialized through both the ProxyException and HTTPException paths. + self.detail = detail if detail is not None else self.message self.response = httpx.Response( status_code=429, headers=_response_headers, @@ -843,11 +987,24 @@ def __init__( class BudgetExceededError(Exception): def __init__( - self, current_cost: float, max_budget: float, message: Optional[str] = None + self, + current_cost: float, + max_budget: float, + message: Optional[str] = None, + llm_provider: Optional[str] = None, ): self.current_cost = current_cost self.max_budget = max_budget self.status_code = 429 + self.llm_provider = llm_provider or "" + # Surface unified rate-limit fields without joining the RateLimitError + # hierarchy so existing `except BudgetExceededError:` handlers keep + # working; custom callbacks reading StandardLoggingPayload pick these + # up via the same `category` / `rate_limit_type` attributes the rest + # of the unified rate-limit error path uses. Stored as plain strings + # to match the normalization RateLimitError.__init__ performs. + self.category: str = RateLimitErrorCategory.LITELLM_RATE_LIMIT.value + self.rate_limit_type: str = RateLimitType.BUDGET.value message = ( message or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}" diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 0ec17bbea5d..390af2cb6e6 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -37,6 +37,8 @@ VirtualKeyEvent, WebhookEvent, ) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository from litellm.types.integrations.slack_alerting import * from ..email_templates.templates import * @@ -1231,7 +1233,7 @@ async def send_key_created_or_user_invited_email( and recipient_user_id is not None and prisma_client is not None ): - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": recipient_user_id} ) @@ -1263,7 +1265,7 @@ async def send_key_created_or_user_invited_email( team_id = webhook_event.team_id team_name = "Default Team" if team_id is not None and prisma_client is not None: - team_row = await prisma_client.db.litellm_teamtable.find_unique( + team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) if team_row is not None: diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index c6ae7d9e82b..8899089500d 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -72,8 +72,13 @@ def initialize_from_proxy_config( compression_params: CompressionInterceptionConfig = {} if "compression_interception_params" in litellm_settings: compression_params = litellm_settings["compression_interception_params"] - elif "compression_interception" in callback_specific_params: - compression_params = callback_specific_params["compression_interception"] + elif "compression_interception" in callback_specific_params and isinstance( + callback_specific_params["compression_interception"], dict + ): + compression_params = cast( + CompressionInterceptionConfig, + callback_specific_params["compression_interception"], + ) return CompressionInterceptionLogger.from_config_yaml(compression_params) async def async_pre_call_deployment_hook( diff --git a/litellm/integrations/email_alerting.py b/litellm/integrations/email_alerting.py index b45b9aa7f5c..b721dc50464 100644 --- a/litellm/integrations/email_alerting.py +++ b/litellm/integrations/email_alerting.py @@ -7,6 +7,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.proxy._types import WebhookEvent +from litellm.repositories.team_repository import TeamRepository # we use this for the email header, please send a test email if you change this. verify it looks good on email LITELLM_LOGO_URL = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" @@ -24,7 +25,7 @@ async def get_all_team_member_emails(team_id: Optional[str] = None) -> list: if prisma_client is None: raise Exception("Not connected to DB!") - team_row = await prisma_client.db.litellm_teamtable.find_unique( + team_row = await TeamRepository(prisma_client).table.find_unique( where={ "team_id": team_id, } diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 298254670eb..3ae3f6b53ac 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -80,11 +80,15 @@ async def get_usage_data( vt.team_id, vt.key_alias as api_key_alias, tt.team_alias, - ut.user_email as user_email + ut.user_email as user_email, + COALESCE(vt.organization_id, tt.organization_id) as organization_id, + ot.organization_alias as organization_alias FROM "LiteLLM_DailyUserSpend" dus LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id + LEFT JOIN "LiteLLM_OrganizationTable" ot + ON ot.organization_id = COALESCE(vt.organization_id, tt.organization_id) {where_clause} ORDER BY dus.date DESC, dus.created_at DESC {limit_clause} diff --git a/litellm/integrations/focus/destinations/__init__.py b/litellm/integrations/focus/destinations/__init__.py index 775d3a259d2..e0cd90c1d61 100644 --- a/litellm/integrations/focus/destinations/__init__.py +++ b/litellm/integrations/focus/destinations/__init__.py @@ -2,12 +2,14 @@ from .base import FocusDestination, FocusTimeWindow from .factory import FocusDestinationFactory +from .gcs_destination import FocusGCSDestination from .s3_destination import FocusS3Destination from .vantage_destination import FocusVantageDestination __all__ = [ "FocusDestination", "FocusDestinationFactory", + "FocusGCSDestination", "FocusTimeWindow", "FocusS3Destination", "FocusVantageDestination", diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py index 706e10624ce..7ce21d4040a 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -6,6 +6,7 @@ from typing import Any, Dict, Optional from .base import FocusDestination +from .gcs_destination import FocusGCSDestination from .s3_destination import FocusS3Destination from .vantage_destination import FocusVantageDestination @@ -29,6 +30,8 @@ def create( return FocusS3Destination(prefix=prefix, config=normalized_config) if provider_lower == "vantage": return FocusVantageDestination(prefix=prefix, config=normalized_config) + if provider_lower == "gcs": + return FocusGCSDestination(prefix=prefix, config=normalized_config) raise NotImplementedError( f"Provider '{provider}' not supported for Focus export" ) @@ -72,6 +75,18 @@ def _resolve_config( "VANTAGE_INTEGRATION_TOKEN must be provided for Vantage exports" ) return {k: v for k, v in resolved.items() if v is not None} + if provider == "gcs": + resolved = { + "bucket_name": overrides.get("bucket_name") + or os.getenv("FOCUS_GCS_BUCKET_NAME"), + "service_account_json": overrides.get("service_account_json") + or os.getenv("FOCUS_GCS_PATH_SERVICE_ACCOUNT"), + } + if not resolved.get("bucket_name"): + raise ValueError( + "FOCUS_GCS_BUCKET_NAME must be provided for GCS exports" + ) + return {k: v for k, v in resolved.items() if v is not None} raise NotImplementedError( f"Provider '{provider}' not supported for Focus export configuration" ) diff --git a/litellm/integrations/focus/destinations/gcs_destination.py b/litellm/integrations/focus/destinations/gcs_destination.py new file mode 100644 index 00000000000..b04c16c9d32 --- /dev/null +++ b/litellm/integrations/focus/destinations/gcs_destination.py @@ -0,0 +1,74 @@ +"""GCS destination for Focus export — reuses GCSBucketBase auth and httpx client.""" + +from __future__ import annotations + +from datetime import timezone +from typing import Any, Optional + +from litellm._logging import verbose_logger +from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase +from litellm.litellm_core_utils.cloud_storage_security import ( + encode_gcs_object_name_for_url, +) + +from .base import FocusDestination, FocusTimeWindow + + +class FocusGCSDestination(GCSBucketBase, FocusDestination): + """Upload serialized Focus exports to GCS using the GCS JSON API.""" + + def __init__( + self, + *, + prefix: str, + config: Optional[dict[str, Any]] = None, + ) -> None: + config = config or {} + bucket_name = config.get("bucket_name") + if not bucket_name: + raise ValueError("bucket_name must be provided for GCS destination") + super().__init__(bucket_name=bucket_name) + service_account_json = config.get("service_account_json") + if service_account_json is not None: + self.path_service_account_json = service_account_json + self.prefix = prefix.rstrip("/") + + async def deliver( + self, + *, + content: bytes, + time_window: FocusTimeWindow, + filename: str, + ) -> None: + object_name = self._build_object_key(time_window=time_window, filename=filename) + headers = await self.construct_request_headers( + service_account_json=self.path_service_account_json + ) + headers["Content-Type"] = "application/octet-stream" + encoded_name = encode_gcs_object_name_for_url(object_name) + url = ( + f"https://storage.googleapis.com/upload/storage/v1/b/" + f"{self.BUCKET_NAME}/o?uploadType=media&name={encoded_name}" + ) + response = await self.async_httpx_client.post( + url=url, headers=headers, data=content + ) + if response.status_code != 200: + raise RuntimeError( + f"GCS upload failed: status={response.status_code} body={response.text}" + ) + verbose_logger.debug( + "Focus GCS: uploaded %d bytes to gs://%s/%s", + len(content), + self.BUCKET_NAME, + object_name, + ) + + def _build_object_key(self, *, time_window: FocusTimeWindow, filename: str) -> str: + start_utc = time_window.start_time.astimezone(timezone.utc) + date_component = f"date={start_utc.strftime('%Y-%m-%d')}" + parts = [self.prefix, date_component] + if time_window.frequency == "hourly": + parts.append(f"hour={start_utc.strftime('%H')}") + key_prefix = "/".join(filter(None, parts)) + return f"{key_prefix}/{filename}" if key_prefix else filename diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index 8496b7ec159..a17df29b912 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -12,6 +12,8 @@ _TAG_KEYS = ( "team_id", "team_alias", + "organization_id", + "organization_alias", "user_id", "user_email", "api_key_alias", diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index 8fef90c24e0..f9ff7e8c7a1 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -26,6 +26,7 @@ get_async_httpx_client, httpxSpecialProvider, ) +from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus GALILEO_CLOUD_API_BASE_URL = "https://api.galileo.ai" # Cap the in-memory buffer so persistent flush failures (e.g. Galileo @@ -89,6 +90,52 @@ def _is_configured(self) -> bool: return bool(self.api_key) return bool(self.username and self.password) + async def async_health_check(self) -> IntegrationHealthCheckStatus: + try: + if not self.project_id: + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message="GALILEO_PROJECT_ID environment variable not set", + ) + + if not self.base_url: + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message="GALILEO_BASE_URL environment variable not set", + ) + + if not self.use_v2_api and (not self.username or not self.password): + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message=( + "GALILEO_API_KEY or GALILEO_USERNAME and GALILEO_PASSWORD " + "environment variables must be set" + ), + ) + + if not await self._ensure_headers(): + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message="Galileo authentication failed", + ) + + response = await self.async_httpx_handler.get( + url=f"{self.base_url}/current_user", + headers=self.headers, + ) + if response.status_code >= 400: + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message=(f"Galileo API returned HTTP {response.status_code}"), + ) + + return IntegrationHealthCheckStatus(status="healthy", error_message=None) + except Exception as e: + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message=f"Galileo health check failed: {str(e)}", + ) + async def async_set_galileo_headers(self) -> None: galileo_login_response = await self.async_httpx_handler.post( url=f"{self.base_url}/login", @@ -399,9 +446,9 @@ def _build_prompt(kwargs: Dict[str, Any]) -> Dict[str, Any]: return prompt @staticmethod - def _serialize_galileo_output(value: Any) -> Optional[str]: + def _serialize_galileo_output(value: Any) -> str: if value is None: - return None + return "" if isinstance(value, str): return value @@ -460,11 +507,11 @@ def _get_galileo_input_output_content( response_obj: Any, level: str = "DEFAULT", status_message: Optional[str] = None, - ) -> Tuple[str, Optional[str], Any]: + ) -> Tuple[str, str, Any]: """ Mirror Langfuse _get_langfuse_input_output_content for Galileo ingest. - Returns (input_text, output_text, messages_for_span). output_text None skips ingest. + Returns (input_text, output_text, messages_for_span). """ call_type = kwargs.get("call_type") prompt = self._build_prompt(kwargs) @@ -477,10 +524,11 @@ def _get_galileo_input_output_content( return self._prompt_to_input_text(prompt), status_message, prompt if response_obj is not None and ( - call_type == "embedding" + call_type in ("embedding", "aembedding") or isinstance(response_obj, litellm.EmbeddingResponse) ): - return self._prompt_to_input_text(prompt), None, prompt + # Match Langfuse OTEL: log embeddings without serializing vectors. + return self._prompt_to_input_text(prompt), "embedding-output", prompt if response_obj is not None and isinstance(response_obj, litellm.ModelResponse): output = self._get_chat_content_for_galileo(response_obj) @@ -549,7 +597,7 @@ def _get_galileo_input_output_content( ): input_val = kwargs.get("input") return ( - self._serialize_galileo_output(input_val) or "", + self._serialize_galileo_output(input_val), self._serialize_galileo_output(response_obj), input_val, ) @@ -574,11 +622,11 @@ def _get_galileo_input_output_content( kwargs.get("messages") or [], ) - return self._prompt_to_input_text(prompt), None, kwargs.get("messages") or [] + return self._prompt_to_input_text(prompt), "", kwargs.get("messages") or [] def get_output_str_from_response( self, response_obj: Any, kwargs: Dict[str, Any] - ) -> Optional[str]: + ) -> str: _, output_text, _ = self._get_galileo_input_output_content( kwargs=kwargs, response_obj=response_obj ) @@ -659,11 +707,6 @@ async def _async_log_success_event_impl( input_text, output_text, messages = self._get_galileo_input_output_content( kwargs=kwargs, response_obj=response_obj ) - if output_text is None: - verbose_logger.debug( - "Galileo Logger: skipping %s — no text output to log", _call_type - ) - return raw_start = slo.get("startTime") raw_end = slo.get("endTime") diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 648fe671140..2119527a8e5 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -24,14 +24,18 @@ import litellm from litellm._logging import print_verbose, verbose_logger -from litellm.integrations.custom_logger import CustomLogger -from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import ( - BoundedPrometheusSeriesTracker, +from litellm.exceptions import ( + validate_rate_limit_category, + validate_rate_limit_type, ) +from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus_helpers import ( PrometheusLabelFactoryContext, _get_cached_end_user_id_for_cost_tracking, ) +from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import ( + BoundedPrometheusSeriesTracker, +) from litellm.litellm_core_utils.core_helpers import ( get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, @@ -42,6 +46,9 @@ LiteLLM_UserTable, UserAPIKeyAuth, ) +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository from litellm.types.integrations.prometheus import * from litellm.types.integrations.prometheus import ( _sanitize_prometheus_label_name, @@ -78,6 +85,20 @@ def __init__( # noqa: PLR0915 # Always initialize label_filters, even for non-premium users self.label_filters = self._parse_prometheus_config() + # Cache resolved label sets per metric. Several entries in + # ``PrometheusMetricLabels.get_labels`` read module-level toggles + # (e.g. ``litellm.prometheus_emit_stream_label``, + # ``litellm.prometheus_emit_rate_limit_labels``) that can be + # changed at runtime. Prometheus counters/gauges/histograms are + # created with a *fixed* ``labelnames`` set; if a runtime call + # to ``get_labels_for_metric`` returned a different set, the + # subsequent ``counter.labels(**_labels)`` would raise a + # ``ValueError`` from the prometheus client. Snapshotting at + # logger init time pins the label set for the lifetime of the + # logger so toggling these flags only takes effect after a + # restart, keeping init-time and runtime label sets in sync. + self._cached_metric_labels: Dict[str, List[str]] = {} + _custom_buckets = litellm.prometheus_latency_buckets self.latency_buckets = ( tuple(_custom_buckets) @@ -1033,13 +1054,27 @@ def get_labels_for_metric( self, metric_name: DEFINED_PROMETHEUS_METRICS ) -> List[str]: """ - Get the labels for a metric, filtered if configured + Get the labels for a metric, filtered if configured. + + The result is cached on the instance so the label set used to + construct each Prometheus metric at ``__init__`` time stays in lock + step with the label set passed to ``counter.labels(...)`` at + runtime, even if the underlying module-level toggles consulted by + :meth:`PrometheusMetricLabels.get_labels` (e.g. + ``litellm.prometheus_emit_rate_limit_labels``, + ``litellm.prometheus_emit_stream_label``) are flipped after the + logger has been created. """ + cached = self._cached_metric_labels.get(metric_name) + if cached is not None: + return cached + # Get default labels for this metric from PrometheusMetricLabels default_labels = PrometheusMetricLabels.get_labels(metric_name) # If no label filtering is configured for this metric, use default labels if metric_name not in self.label_filters: + self._cached_metric_labels[metric_name] = default_labels return default_labels # Get configured labels for this metric @@ -1050,6 +1085,7 @@ def get_labels_for_metric( label for label in default_labels if label in configured_labels ] + self._cached_metric_labels[metric_name] = filtered_labels return filtered_labels def _track_end_user_metric_series( @@ -2029,14 +2065,8 @@ async def async_post_call_failure_hook( Proxy level tracking - failed client side requests - labelnames=[ - "end_user", - "hashed_api_key", - "api_key_alias", - REQUESTED_MODEL, - "team", - "team_alias", - ] + EXCEPTION_LABELS, + See :attr:`PrometheusMetricLabels.litellm_proxy_failed_requests_metric` + for the authoritative list of labels emitted on this metric. """ from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, @@ -2059,6 +2089,9 @@ async def async_post_call_failure_hook( model_id = _metadata.get("model_info", {}).get("id") or request_data.get( "model_info", {} ).get("id") + rate_limit_category, rate_limit_type = self._extract_rate_limit_labels( + original_exception + ) enum_values = UserAPIKeyLabelValues( end_user=user_api_key_dict.end_user_id, user=user_api_key_dict.user_id, @@ -2073,6 +2106,8 @@ async def async_post_call_failure_hook( status_code=str(status_code), exception_status=str(status_code), exception_class=self._get_exception_class_name(original_exception), + rate_limit_category=rate_limit_category, + rate_limit_type=rate_limit_type, tags=_tags, route=user_api_key_dict.request_route, client_ip=_metadata.get("requester_ip_address"), @@ -2843,6 +2878,33 @@ def record_check_batch_cost_error(self, error_type: str): @staticmethod def _get_exception_class_name(exception: Exception) -> str: + # Some exception types pin the ``exception_class`` label to a legacy + # value for back-compat with existing dashboards (e.g. proxy-side 429s + # keep reporting as "HTTPException"). Honor that opt-in marker before + # deriving the label from the runtime class name. Reading it via + # ``getattr`` keeps this core integrations module free of a transitive + # ``fastapi`` dependency. + legacy_class_name = getattr(exception, "prometheus_exception_class_name", None) + if isinstance(legacy_class_name, str) and legacy_class_name: + return legacy_class_name + + # Same back-compat reasoning for ``BudgetExceededError``: the unified + # rate-limit error work attached ``.llm_provider`` to budget errors + # too (so callbacks reading ``StandardLoggingPayload`` get provider + # attribution). Without this short-circuit, the provider prefix below + # would silently flip the label from "BudgetExceededError" to e.g. + # "Openai.BudgetExceededError" and break dashboards keyed on the + # original value. + try: + from litellm.exceptions import BudgetExceededError + except ImportError: + BudgetExceededError = None # type: ignore[assignment,misc] + + if BudgetExceededError is not None and isinstance( + exception, BudgetExceededError + ): + return "BudgetExceededError" + exception_class_name = "" if hasattr(exception, "llm_provider"): exception_class_name = getattr(exception, "llm_provider") or "" @@ -2857,6 +2919,27 @@ def _get_exception_class_name(exception: Exception) -> str: exception_class_name += exception.__class__.__name__ return exception_class_name + @staticmethod + def _extract_rate_limit_labels( + exception: Optional[Exception], + ) -> Tuple[Optional[str], Optional[str]]: + """ + Pull the unified ``category`` / ``rate_limit_type`` fields off any + exception that declares them (``litellm.RateLimitError`` and bare- + Exception subclasses like ``BudgetExceededError``). + + Values are validated against the :class:`RateLimitErrorCategory` / + :class:`RateLimitType` enums so unrelated third-party exceptions that + happen to declare ``.category`` / ``.rate_limit_type`` string attributes + can't leak garbage into Prometheus label cardinality. + """ + if exception is None: + return None, None + return ( + validate_rate_limit_category(getattr(exception, "category", None)), + validate_rate_limit_type(getattr(exception, "rate_limit_type", None)), + ) + async def log_success_fallback_event( self, original_model_group: str, kwargs: dict, original_exception: Exception ): @@ -3198,12 +3281,12 @@ async def fetch_users( page_size: int, page: int ) -> Tuple[List[LiteLLM_UserTable], Optional[int]]: skip = (page - 1) * page_size - users = await prisma_client.db.litellm_usertable.find_many( + users = await UserRepository(prisma_client).table.find_many( skip=skip, take=page_size, order={"created_at": "desc"}, ) - total_count = await prisma_client.db.litellm_usertable.count() + total_count = await UserRepository(prisma_client).table.count() return users, total_count await self._initialize_budget_metrics( @@ -3226,13 +3309,13 @@ async def _initialize_org_budget_metrics(self): async def fetch_orgs(page_size: int, page: int) -> Tuple[list, Optional[int]]: skip = (page - 1) * page_size - orgs = await prisma_client.db.litellm_organizationtable.find_many( + orgs = await OrganizationRepository(prisma_client).table.find_many( skip=skip, take=page_size, order={"created_at": "desc"}, include={"litellm_budget_table": True}, ) - total_count = await prisma_client.db.litellm_organizationtable.count() + total_count = await OrganizationRepository(prisma_client).table.count() return orgs, total_count await self._initialize_budget_metrics( @@ -3300,14 +3383,14 @@ async def _initialize_user_and_team_count_metrics(self): try: # Get total user count - total_users = await prisma_client.db.litellm_usertable.count() + total_users = await UserRepository(prisma_client).table.count() self.litellm_total_users_metric.set(total_users) verbose_logger.debug( f"Prometheus: set litellm_total_users to {total_users}" ) # Get total team count - total_teams = await prisma_client.db.litellm_teamtable.count() + total_teams = await TeamRepository(prisma_client).table.count() self.litellm_teams_count_metric.set(total_teams) verbose_logger.debug( f"Prometheus: set litellm_teams_count to {total_teams}" diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 37528e7dcd5..79f9b16bba0 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1339,8 +1339,13 @@ def initialize_from_proxy_config( websearch_params: WebSearchInterceptionConfig = {} if "websearch_interception_params" in litellm_settings: websearch_params = litellm_settings["websearch_interception_params"] - elif "websearch_interception" in callback_specific_params: - websearch_params = callback_specific_params["websearch_interception"] + elif "websearch_interception" in callback_specific_params and isinstance( + callback_specific_params["websearch_interception"], dict + ): + websearch_params = cast( + WebSearchInterceptionConfig, + callback_specific_params["websearch_interception"], + ) # Use classmethod to initialize from config return WebSearchInterceptionLogger.from_config_yaml(websearch_params) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 3776d276912..eb01359cdc0 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -37,7 +37,7 @@ def get_litellm_gateway_api_key( """ Get the stored CLI API key for use with LiteLLM SDK. - This function reads the token file created by `litellm-proxy login` + This function reads the token file created by `lite login` and returns the API key for use in Python scripts. Args: diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index b32803b5dfc..6e655b03fed 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -32,8 +32,10 @@ "aws_sts_endpoint", "aws_external_id", "aws_bedrock_runtime_endpoint", + "aws_bedrock_project_id", "tpm", "rpm", + "use_xai_oauth", } ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f20b66790c4..dbfcf55d75d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -37,6 +37,10 @@ turn_off_message_logging, ) from litellm._logging import _is_debugging_on, _redact_string, verbose_logger +from litellm.exceptions import ( + validate_rate_limit_category, + validate_rate_limit_type, +) from litellm._uuid import uuid from litellm.batches.batch_utils import _handle_completed_batch from litellm.caching.caching import DualCache, InMemoryCache @@ -5318,12 +5322,27 @@ def get_error_information( else str(original_exception) ) + # Duck-typed read so bare-Exception subclasses like + # `litellm.BudgetExceededError` can participate without joining the + # RateLimitError hierarchy (which would break `except BudgetExceededError`). + # Validated against the enum value sets so a third-party exception that + # happens to declare a `.category` or `.rate_limit_type` string attribute + # can't leak garbage into the payload or Prometheus label cardinality. + rate_limit_category = validate_rate_limit_category( + getattr(original_exception, "category", None) + ) + rate_limit_type = validate_rate_limit_type( + getattr(original_exception, "rate_limit_type", None) + ) + return StandardLoggingPayloadErrorInformation( error_code=error_status, error_class=error_class, llm_provider=_llm_provider_in_exception, traceback=traceback_info, error_message=error_message if original_exception else "", + error_rate_limit_category=rate_limit_category, + error_rate_limit_type=rate_limit_type, ) @staticmethod diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 8da66d4600d..413ddb71bf8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -6,6 +6,7 @@ import litellm from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS +from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests from litellm.types.llms.openai import ( FileSearchTool, ResponsesAPIResponse, @@ -339,8 +340,7 @@ def response_object_includes_web_search_call( # and _handle_web_search_cost() is never called. if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True return False @@ -352,8 +352,7 @@ def response_object_includes_web_search_call( elif usage is not None: if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True elif ( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index f39c942f90f..93049adf75a 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1,7 +1,7 @@ # What is this? ## Helper utilities for cost_per_token() -from typing import Literal, Optional, Tuple, TypedDict, cast +from typing import Any, Literal, Optional, Tuple, TypedDict, cast import litellm from litellm._logging import verbose_logger @@ -42,6 +42,26 @@ def _get_token_detail_value(details: object, key: str) -> Optional[int]: return value if isinstance(value, int) else None +def _get_web_search_requests(server_tool_use: Any) -> Optional[int]: + """ + Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value + that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance, + or any other object supporting attribute access. + + Returns ``None`` when the value cannot be resolved — callers can + distinguish "absent" from "zero" using ``is None``. + + See https://github.com/BerriAI/litellm/issues/26153 — ``stream_chunk_builder`` + historically left this as a plain ``dict``, which broke direct attribute + access in cost calculation. + """ + if server_tool_use is None: + return None + if isinstance(server_tool_use, dict): + return server_tool_use.get("web_search_requests") + return getattr(server_tool_use, "web_search_requests", None) + + def _is_above_128k(tokens: float) -> bool: if tokens > 128000: return True diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 2547fd4d8c6..4e5b53a13d7 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -633,11 +633,6 @@ def convert_to_model_response_object( # noqa: PLR0915 thinking_blocks = choice["message"]["thinking_blocks"] provider_specific_fields["thinking_blocks"] = thinking_blocks - if reasoning_content: - provider_specific_fields["reasoning_content"] = ( - reasoning_content - ) - message = Message( content=content, role=choice["message"]["role"] or "assistant", diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 1460dbaf0a9..b09f2bb130e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3653,17 +3653,13 @@ def stringify_json_tool_call_content(messages: List) -> List: from litellm.types.llms.bedrock import DocumentBlock as BedrockDocumentBlock from litellm.types.llms.bedrock import ImageBlock as BedrockImageBlock from litellm.types.llms.bedrock import SourceBlock as BedrockSourceBlock +from litellm.types.llms.bedrock import BedrockToolSpec from litellm.types.llms.bedrock import ToolBlock as BedrockToolBlock -from litellm.types.llms.bedrock import ( - ToolInputSchemaBlock as BedrockToolInputSchemaBlock, -) -from litellm.types.llms.bedrock import ToolJsonSchemaBlock as BedrockToolJsonSchemaBlock from litellm.types.llms.bedrock import SearchResultBlock from litellm.types.llms.bedrock import ToolResultBlock as BedrockToolResultBlock from litellm.types.llms.bedrock import ( ToolResultContentBlock as BedrockToolResultContentBlock, ) -from litellm.types.llms.bedrock import ToolSpecBlock as BedrockToolSpecBlock from litellm.types.llms.bedrock import ToolUseBlock as BedrockToolUseBlock from litellm.types.llms.bedrock import VideoBlock as BedrockVideoBlock @@ -4294,6 +4290,49 @@ def _deduplicate_bedrock_tool_content( return _deduplicate_bedrock_content_blocks(tool_content, "toolResult") +def _rename_duplicate_bedrock_document_names( + contents: List[BedrockMessageBlock], +) -> List[BedrockMessageBlock]: + """ + Rename duplicate document names across all messages in a Bedrock request. + + Document names are derived from a content hash, so the same file appearing + in multiple conversation turns produces identical names and Bedrock rejects + the request with "Messages can not contain duplicate document names". The + first occurrence keeps its original name so prompt-cache prefixes stay + stable; later occurrences get a deterministic positional suffix + (``_2``, ``_3``, ...), bumped further if the suffixed name already + belongs to another document (e.g. an organic name ending in ``_2``). + """ + used_names: Set[str] = set() + for message in contents: + for block in message.get("content") or []: + document = block.get("document") + if isinstance(document, dict) and document.get("name"): + used_names.add(document["name"]) + + name_counts: Dict[str, int] = {} + for message in contents: + for block in message.get("content") or []: + document = block.get("document") + if not isinstance(document, dict): + continue + name = document.get("name") + if not name: + continue + count = name_counts.get(name, 0) + 1 + name_counts[name] = count + if count > 1: + suffix = count + new_name = f"{name}_{suffix}" + while new_name in used_names: + suffix += 1 + new_name = f"{name}_{suffix}" + used_names.add(new_name) + document["name"] = new_name + return contents + + def _sort_bedrock_assistant_content_blocks( blocks: List[BedrockContentBlock], ) -> List[BedrockContentBlock]: @@ -4942,7 +4981,7 @@ async def _bedrock_converse_messages_pt_async( # noqa: PLR0915 llm_provider=llm_provider, ) - return contents + return _rename_duplicate_bedrock_document_names(contents) @staticmethod def translate_thinking_blocks_to_reasoning_content_blocks( @@ -5364,7 +5403,7 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 llm_provider=llm_provider, ) - return contents + return _rename_duplicate_bedrock_document_names(contents) def make_valid_bedrock_tool_name(input_tool_name: str) -> str: @@ -5496,6 +5535,7 @@ def _bedrock_tools_pt( ] """ from litellm.llms.bedrock.common_utils import ( + get_bedrock_base_model, normalize_json_schema_custom_types_to_object, ) from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs @@ -5503,6 +5543,11 @@ def _bedrock_tools_pt( _valid_json_schema_root_types = frozenset( ("array", "boolean", "integer", "null", "number", "object", "string") ) + # Only Claude on Bedrock honours strict tool schemas; other families + # (Nova, Llama, GPT-OSS) reject the strict field outright. + supports_strict_tools = bool( + model and get_bedrock_base_model(model).startswith("anthropic") + ) tool_block_list: List[BedrockToolBlock] = [] for tool_idx, tool in enumerate(tools): # Check if tool is already a BedrockToolBlock (e.g., systemTool for Nova grounding) @@ -5548,17 +5593,16 @@ def _bedrock_tools_pt( normalize_json_schema_custom_types_to_object(parameters) if parameters.get("type") not in _valid_json_schema_root_types: parameters["type"] = "object" - tool_input_schema = BedrockToolInputSchemaBlock( - json=BedrockToolJsonSchemaBlock( - type=parameters["type"], - properties=parameters.get("properties", {}), - required=parameters.get("required", []), - ) - ) - tool_spec = BedrockToolSpecBlock( - inputSchema=tool_input_schema, name=name, description=description + tool_block = cast( + BedrockToolBlock, + BedrockToolSpec( + name=name, + description=description, + parameters=parameters, + strict=tool.get("function", {}).get("strict", None), + supports_strict_tools=supports_strict_tools, + ), ) - tool_block = BedrockToolBlock(toolSpec=tool_spec) tool_block_list.append(tool_block) ## ADD CACHE POINT TOOL BLOCK ## diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 772f058d9bb..4b7f0e22198 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -144,7 +144,7 @@ def _should_store_message( return True return False - def store_message(self, message: Union[str, bytes, OpenAIRealtimeEvents]): + def store_message(self, message: Union[str, bytes, dict, OpenAIRealtimeEvents]): """Store message in list""" if isinstance(message, bytes): message = message.decode("utf-8") @@ -154,22 +154,20 @@ def store_message(self, message: Union[str, bytes, OpenAIRealtimeEvents]): else: message_obj = cast(Dict[str, Any], json.loads(cast(str, message))) self._collect_tool_calls_from_response_done(cast(dict, message_obj)) + if not self._should_store_message(message_obj): + return try: event_type = message_obj.get("type", "") if event_type in self._SESSION_EVENT_TYPES: - typed_obj = OpenAIRealtimeStreamSessionEvents(**message_obj) # type: ignore + typed_obj: OpenAIRealtimeEvents = OpenAIRealtimeStreamSessionEvents(**message_obj) # type: ignore else: - # Use the base object as a safe catch-all for all other event types - # (both beta and GA), so unknown/new event names never raise here. + # Catch-all base object so unknown/new event names never raise. typed_obj = OpenAIRealtimeStreamResponseBaseObject(**message_obj) # type: ignore except Exception as e: verbose_logger.debug(f"Error parsing message for logging: {e}") - # Don't re-raise — a parse failure must not drop or delay the message - if self._should_store_message(message_obj): - self.messages.append(message_obj) # type: ignore[arg-type] + self.messages.append(message_obj) # type: ignore[arg-type] return - if self._should_store_message(typed_obj): - self.messages.append(typed_obj) + self.messages.append(typed_obj) def _collect_user_input_from_client_event(self, message: Union[str, dict]) -> None: """Extract user text content from client WebSocket events for spend logging.""" @@ -358,8 +356,7 @@ async def _flush_pending_messages_until_setup(self) -> bool: for msg in self._pending_messages_until_setup ) verbose_logger.debug( - "Failed to flush buffered client message after setup: %s " - "(%d buffered message(s) retained)", + "Failed to flush buffered client message after setup: %s (%d buffered message(s) retained)", e, len(unsent), ) @@ -376,8 +373,7 @@ async def _send_event_to_client(self, event: Any, event_str: str) -> bool: return True except Exception as e: verbose_logger.warning( - "Failed to translate %s to beta protocol, forwarding " - "untranslated event to client: %s", + "Failed to translate %s to beta protocol, forwarding untranslated event to client: %s", event.get("type"), e, ) @@ -705,48 +701,48 @@ async def _handle_provider_config_message(self, raw_response) -> None: self.store_message(event_str) await self._send_event_to_client(event, event_str) - async def _handle_raw_backend_message(self, raw_response) -> bool: + @staticmethod + def _parse_backend_event(raw_response: str) -> Optional[dict]: + """Parse a backend frame once. Returns None for non-JSON or non-object frames.""" + try: + event = json.loads(raw_response) + except (json.JSONDecodeError, TypeError): + return None + return event if isinstance(event, dict) else None + + async def _handle_raw_backend_message( + self, event_obj: dict, raw_response: str + ) -> bool: """Process a backend message without provider_config (raw path). Returns True if the caller should skip the default store+forward (i.e. continue the loop). """ - try: - event_obj = json.loads(raw_response) + event_type = event_obj.get("type") - # For audio/VAD guardrail path: once the session is ready, tell the backend - # not to auto-respond after VAD detects end-of-speech. We send the - # session.created to the client FIRST so the client is always in sync, then - # inject the session.update so a potential error from the backend doesn't - # arrive before the client sees session.created. - if ( - event_obj.get("type") == "session.created" - and self._has_audio_transcription_guardrails() - ): - self.store_message(raw_response) - await self.websocket.send_text(raw_response) - await self._send_to_backend(self._make_disable_auto_response_message()) - return True + # Send session.created to the client FIRST so it stays in sync, then inject + # the disable-auto-response session.update; otherwise a backend error could + # reach the client before it sees session.created. + if ( + event_type == "session.created" + and self._has_audio_transcription_guardrails() + ): + self.store_message(event_obj) + await self.websocket.send_text(raw_response) + await self._send_to_backend(self._make_disable_auto_response_message()) + return True - if ( - event_obj.get("type") - == "conversation.item.input_audio_transcription.completed" - ): - transcript = event_obj.get("transcript", "") - self._collect_user_input_from_backend_event(event_obj) - ## LOGGING — must happen before continue below - self.store_message(raw_response) - # Forward transcript to client so user sees what they said - await self.websocket.send_text(raw_response) - blocked = await self.run_realtime_guardrails( - transcript, - item_id=event_obj.get("item_id"), - ) - if not blocked: - # Clean — trigger LLM response - await self._send_to_backend(json.dumps({"type": "response.create"})) - return True - except (json.JSONDecodeError, AttributeError): - pass + if event_type == "conversation.item.input_audio_transcription.completed": + transcript = event_obj.get("transcript", "") + self._collect_user_input_from_backend_event(event_obj) + self.store_message(event_obj) + await self.websocket.send_text(raw_response) + blocked = await self.run_realtime_guardrails( + transcript, + item_id=event_obj.get("item_id"), + ) + if not blocked: + await self._send_to_backend(json.dumps({"type": "response.create"})) + return True return False async def backend_to_client_send_messages(self): @@ -779,25 +775,25 @@ async def backend_to_client_send_messages(self): ) continue else: - handled = await self._handle_raw_backend_message(raw_response) - if handled: + event = self._parse_backend_event(raw_response) + if event is None: + await self.websocket.send_text(raw_response) + continue + + if await self._handle_raw_backend_message(event, raw_response): continue - ## LOGGING - self.store_message(raw_response) - - # If the client opted into beta protocol, translate GA event - # names/shapes back to the beta equivalents before forwarding. - if self._client_wants_beta: - try: - event_dict = json.loads(raw_response) - translated = self._translate_event_to_beta(event_dict) - if translated is None: - continue # drop GA-only events (e.g. conversation.item.done) - await self.websocket.send_text(json.dumps(translated)) - except Exception: - await self.websocket.send_text(raw_response) - else: + self.store_message(event) + + if not self._client_wants_beta: await self.websocket.send_text(raw_response) + continue + + translated = self._translate_event_to_beta(event) + if translated is None: + continue + await self.websocket.send_text( + raw_response if translated is event else json.dumps(translated) + ) except websockets.exceptions.ConnectionClosed as e: # type: ignore verbose_logger.exception( @@ -927,41 +923,43 @@ def _remap_beta_session_to_ga(session: dict) -> dict: def _translate_event_to_beta(event: dict) -> Optional[dict]: """Translate a single GA event dict to its beta equivalent. - Returns None if the event should be dropped entirely (e.g. the GA-only - conversation.item.done has no beta counterpart). - Returns the (possibly mutated copy of the) event otherwise. + Returns None when the event must be dropped (the GA-only + conversation.item.done has no beta counterpart). Returns the original + event object unchanged when no translation applies, so the caller can + forward the raw frame without re-serializing; otherwise returns a + translated copy. """ event_type = event.get("type", "") - # conversation.item.done has no beta equivalent — the client already - # received conversation.item.created (translated from .added). if event_type == "conversation.item.done": return None - # Shallow-copy so we don't mutate the stored message - translated = dict(event) - - # Rename the type field - if event_type in RealTimeStreaming._GA_TO_BETA_EVENT_TYPES: - translated["type"] = RealTimeStreaming._GA_TO_BETA_EVENT_TYPES[event_type] + renamed_type = RealTimeStreaming._GA_TO_BETA_EVENT_TYPES.get(event_type) + has_item = isinstance(event.get("item"), dict) + response = event.get("response") + has_response_output = isinstance(response, dict) and isinstance( + response.get("output"), list + ) + if renamed_type is None and not has_item and not has_response_output: + return event - # Fix content block types inside items (response.done output list, - # conversation.item.created item content, etc.) - if "item" in translated and isinstance(translated["item"], dict): + translated = dict(event) + if renamed_type is not None: + translated["type"] = renamed_type + if has_item: translated["item"] = RealTimeStreaming._translate_item_content_types( dict(translated["item"]) ) - if "response" in translated and isinstance(translated["response"], dict): + if has_response_output: resp = dict(translated["response"]) - if "output" in resp and isinstance(resp["output"], list): - resp["output"] = [ - ( - RealTimeStreaming._translate_item_content_types(dict(o)) - if isinstance(o, dict) - else o - ) - for o in resp["output"] - ] + resp["output"] = [ + ( + RealTimeStreaming._translate_item_content_types(dict(o)) + if isinstance(o, dict) + else o + ) + for o in resp["output"] + ] translated["response"] = resp return translated diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index dbc9cabdc7a..763596336a0 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -17,6 +17,10 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) +from litellm.llms.vertex_ai.common_utils import ( + redact_vertex_ai_metadata_from_litellm_params, + redact_vertex_ai_metadata_from_logged_object, +) from litellm.secret_managers.main import str_to_bool from litellm.types.utils import StandardCallbackDynamicParams @@ -119,10 +123,12 @@ def _redact_standard_logging_object(model_call_details: dict): # ResponsesAPIResponse format - redact content in output items if isinstance(response.get("output"), list): _redact_responses_api_output_dict(response["output"], redacted_str) + redact_vertex_ai_metadata_from_logged_object(response) elif isinstance(response, dict) and "choices" in response: # ModelResponse dict format - redact content in choices if isinstance(response.get("choices"), list): _redact_model_response_dict_choices(response["choices"], redacted_str) + redact_vertex_ai_metadata_from_logged_object(response) elif isinstance(response, str): standard_logging_object["response"] = redacted_str else: @@ -164,6 +170,7 @@ def perform_redaction(model_call_details: dict, result): model_call_details["prompt"] = "" model_call_details["input"] = "" _redact_standard_logging_object(model_call_details) + redact_vertex_ai_metadata_from_litellm_params(model_call_details) # Redact streaming response if ( @@ -174,6 +181,7 @@ def perform_redaction(model_call_details: dict, result): if hasattr(_streaming_response, "choices"): for choice in _streaming_response.choices: _redact_choice_content(choice) + redact_vertex_ai_metadata_from_logged_object(_streaming_response) elif hasattr(_streaming_response, "output"): _redact_responses_api_output(_streaming_response.output) # Redact reasoning field in ResponsesAPIResponse @@ -200,12 +208,14 @@ def perform_redaction(model_call_details: dict, result): if hasattr(_result, "choices") and _result.choices is not None: for choice in _result.choices: _redact_choice_content(choice) + redact_vertex_ai_metadata_from_logged_object(_result) elif isinstance(_result, dict) and "choices" in _result: # Handle dict representation of ModelResponse (e.g., from model_dump()) if _result.get("choices") is not None: _redact_model_response_dict_choices( _result["choices"], "redacted-by-litellm" ) + redact_vertex_ai_metadata_from_logged_object(_result) elif isinstance(_result, dict) and "output" in _result: if isinstance(_result.get("output"), list): _redact_responses_api_output_dict( diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fe7c62c3842..b495b183ec0 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -20,6 +20,7 @@ ServerToolUse, Usage, ) +from litellm._logging import verbose_logger from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: @@ -79,6 +80,54 @@ def update_model_response_with_hidden_params( model_response._hidden_params = chunk.get("_hidden_params", {}) return model_response + @staticmethod + def apply_provider_assembled_streaming_metadata( + response: ModelResponse, + chunks: List[Any], + logging_obj: Optional[Any] = None, + ) -> None: + if not chunks: + return + + model = getattr(response, "model", None) + if not model: + return + + custom_llm_provider = None + if logging_obj is not None: + custom_llm_provider = logging_obj.model_call_details.get( + "custom_llm_provider" + ) + + try: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider, + ) + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + if custom_llm_provider: + provider = LlmProviders(custom_llm_provider) + else: + _, provider_str, _, _ = get_llm_provider(model) + provider = LlmProviders(provider_str) + + provider_config = ProviderConfigManager.get_provider_chat_config( + model=model, + provider=provider, + ) + if provider_config is not None: + provider_config.apply_assembled_streaming_response_metadata( + response=response, + chunks=chunks, + ) + except Exception as e: + verbose_logger.debug( + "apply_provider_assembled_streaming_metadata failed for model=%s: %s", + model, + e, + ) + @staticmethod def _get_chunk_id(chunks: List[Dict[str, Any]]) -> str: """ @@ -588,7 +637,18 @@ def _calculate_usage_per_chunk( hasattr(usage_chunk, "server_tool_use") and usage_chunk.server_tool_use is not None ): - server_tool_use = usage_chunk.server_tool_use + # Coerce dict to ServerToolUse so downstream cost-calc code + # (which accesses .web_search_requests as an attribute) + # doesn't raise AttributeError. Some providers / streaming + # paths leave server_tool_use as a plain dict on the chunk. + if isinstance(usage_chunk.server_tool_use, dict): + server_tool_use = ServerToolUse(**usage_chunk.server_tool_use) + elif isinstance(usage_chunk.server_tool_use, ServerToolUse): + server_tool_use = usage_chunk.server_tool_use + else: + server_tool_use = ServerToolUse.model_validate( + usage_chunk.server_tool_use + ) if ( usage_chunk_dict["prompt_tokens_details"] is not None and getattr( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 7949e150c23..9ecd0df0cb8 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1455,10 +1455,15 @@ def map_openai_params( # noqa: PLR0915 _value = self._map_stop_sequences(value) if _value is not None: optional_params["stop_sequences"] = _value - elif param == "temperature": - optional_params["temperature"] = value - elif param == "top_p": - optional_params["top_p"] = value + elif param == "temperature" or param == "top_p": + AnthropicConfig._apply_sampling_param( + optional_params=optional_params, + model=model, + param=param, + value=value, + drop_params=drop_params, + output_key=param, + ) elif param == "response_format" and isinstance(value, dict): if any( substring in model @@ -1607,6 +1612,15 @@ def _create_json_tool_call_for_response_format( ) return _tool + def should_strip_billing_metadata(self) -> bool: + """ + Whether to drop x-anthropic-billing-header system blocks before sending upstream. + + The first-party Anthropic API uses these blocks for Claude Code attribution, so the + base config keeps them. Providers that reject them (e.g. Bedrock) override this to True. + """ + return False + def translate_system_message( self, messages: List[AllMessageValues] ) -> List[AnthropicSystemMessageContent]: @@ -1614,7 +1628,7 @@ def translate_system_message( Translate system message to anthropic format. Removes system message from the original list and returns a new list of anthropic system message content. - Filters out system messages containing x-anthropic-billing-header metadata. + When should_strip_billing_metadata() is True, x-anthropic-billing-header system blocks are dropped. """ system_prompt_indices = [] anthropic_system_message_list: List[AnthropicSystemMessageContent] = [] @@ -1626,10 +1640,9 @@ def translate_system_message( # Skip empty text blocks - Anthropic API raises errors for empty text if not system_message_block["content"]: continue - # Skip system messages containing x-anthropic-billing-header metadata - if system_message_block["content"].startswith( - "x-anthropic-billing-header:" - ): + if self.should_strip_billing_metadata() and system_message_block[ + "content" + ].startswith("x-anthropic-billing-header:"): continue anthropic_system_message_content = AnthropicSystemMessageContent( type="text", @@ -1648,9 +1661,9 @@ def translate_system_message( text_value = _content.get("text") if _content.get("type") == "text" and not text_value: continue - # Skip system messages containing x-anthropic-billing-header metadata if ( - _content.get("type") == "text" + self.should_strip_billing_metadata() + and _content.get("type") == "text" and text_value and text_value.startswith("x-anthropic-billing-header:") ): @@ -1967,6 +1980,20 @@ def transform_request( optional_params.pop("is_vertex_request", None) optional_params.pop("client_metadata", None) + # ``top_k`` is a provider-specific kwarg that bypasses + # ``map_openai_params``; gate it here, the single boundary shared by + # the direct Anthropic, Bedrock invoke, Vertex, and Azure paths. + top_k = optional_params.pop("top_k", None) + if top_k is not None: + AnthropicConfig._apply_sampling_param( + optional_params=optional_params, + model=model, + param="top_k", + value=top_k, + drop_params=litellm_params.get("drop_params") is True, + output_key="top_k", + ) + data = { "model": model, "messages": anthropic_messages, diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 3f002d73cbc..5741513903c 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -272,23 +272,68 @@ def _is_claude_4_7_model(model: str) -> bool: ) @staticmethod - def _supports_model_capability(model: str, key: str) -> bool: - """Check a boolean capability ``key`` in the model map. + def _supports_sampling_params(model: str) -> bool: + """Claude 4.7+ (Opus 4.7/4.8, Fable 5) removed sampling params: the API + rejects ``top_p``, ``top_k``, and any ``temperature`` other than 1 with + a 400 ("`temperature` is deprecated for this model"). + + Driven by the ``supports_sampling_params`` flag in the model map; the + name check remains only as a fallback for provider-routed ids whose + map entries predate the flag.""" + flag = AnthropicModelInfo._get_model_capability( + model, "supports_sampling_params" + ) + if flag is not None: + return flag + model_lower = model.lower() + return not any( + v in model_lower + for v in ( + "fable", + "opus-4-7", + "opus_4_7", + "opus-4.7", + "opus_4.7", + "opus-4-8", + "opus_4_8", + "opus-4.8", + "opus_4.8", + ) + ) - Strips bedrock/vertex prefixes so a provider-routed Claude still - resolves to the Anthropic model-map entry. - """ - from litellm.utils import _supports_factory + @staticmethod + def _apply_sampling_param( + optional_params: dict, + model: str, + param: str, + value: Any, + drop_params: bool, + output_key: str, + ) -> None: + """Forward ``temperature``/``top_p``/``top_k`` to + ``optional_params[output_key]`` unless the model removed sampling + params, in which case drop the param (with drop_params) or raise a + clean client-side 400.""" + if AnthropicModelInfo._supports_sampling_params(model) or ( + param == "temperature" and value == 1 + ): + optional_params[output_key] = value + elif not (litellm.drop_params or drop_params): + supported_hint = ( + "Only temperature=1 is supported. " if param == "temperature" else "" + ) + raise litellm.utils.UnsupportedParamsError( + message=( + f"{model} does not support {param}={value}. {supported_hint}" + "To drop unsupported params, set `litellm.drop_params = True`." + ), + status_code=400, + ) - try: - if _supports_factory( - model=model, - custom_llm_provider="anthropic", - key=key, - ): - return True - except Exception: - pass + @staticmethod + def _model_map_lookup_candidates(model: str) -> List[str]: + """Model-map keys to try for ``model``, stripping bedrock/vertex + prefixes so a provider-routed Claude still resolves to its entry.""" candidates = [model] for prefix in ( "bedrock/converse/", @@ -307,15 +352,40 @@ def _supports_model_capability(model: str, key: str) -> bool: candidates.append(f"bedrock/{base}") except Exception: pass + return candidates + + @staticmethod + def _get_model_capability(model: str, key: str) -> Optional[bool]: + """Read boolean capability ``key`` from the model map, or None when + no entry declares it.""" try: - for cand in candidates: - if cand in litellm.model_cost and ( - litellm.model_cost[cand].get(key) is True - ): - return True + for cand in AnthropicModelInfo._model_map_lookup_candidates(model): + value = litellm.model_cost.get(cand, {}).get(key) + if isinstance(value, bool): + return value except Exception: pass - return False + return None + + @staticmethod + def _supports_model_capability(model: str, key: str) -> bool: + """Check a boolean capability ``key`` in the model map. + + Strips bedrock/vertex prefixes so a provider-routed Claude still + resolves to the Anthropic model-map entry. + """ + from litellm.utils import _supports_factory + + try: + if _supports_factory( + model=model, + custom_llm_provider="anthropic", + key=key, + ): + return True + except Exception: + pass + return AnthropicModelInfo._get_model_capability(model, key) is True @staticmethod def _is_adaptive_thinking_model(model: str) -> bool: diff --git a/litellm/llms/anthropic/cost_calculation.py b/litellm/llms/anthropic/cost_calculation.py index 3882d8f978c..6a031498dae 100644 --- a/litellm/llms/anthropic/cost_calculation.py +++ b/litellm/llms/anthropic/cost_calculation.py @@ -7,6 +7,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( _get_token_base_cost, + _get_web_search_requests, _parse_prompt_tokens_details, calculate_cache_writing_cost, generic_cost_per_token, @@ -110,11 +111,12 @@ def get_cost_for_anthropic_web_search( if model_info is None: return 0.0 - if ( - usage is None - or usage.server_tool_use is None - or usage.server_tool_use.web_search_requests is None - ): + if usage is None: + return 0.0 + web_search_requests = _get_web_search_requests( + getattr(usage, "server_tool_use", None) + ) + if web_search_requests is None: return 0.0 ## Get the cost per web search request @@ -128,5 +130,5 @@ def get_cost_for_anthropic_web_search( return 0.0 ## Calculate the total cost - total_cost = cost_per_web_search_request * usage.server_tool_use.web_search_requests + total_cost = cost_per_web_search_request * web_search_requests return total_cost diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index bacb9f8ddf6..8c20f4c430e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -1,5 +1,6 @@ # What is this? ## Translates OpenAI call to Anthropic `/v1/messages` format +import copy import json import traceback from collections import deque @@ -29,6 +30,98 @@ from litellm.types.utils import ModelResponseStream +class _CombinedChunkSplitter: + """ + Splits a streaming chunk that carries BOTH response content and a + ``finish_reason`` into two chunks: a content-only chunk followed by a + finish-only chunk. + + ``AnthropicStreamWrapper`` (via ``translate_streaming_openai_response_to_anthropic``) + assumes content and ``finish_reason`` never arrive in the same chunk — true for + real provider streams, but false for fake-streamed providers (e.g. Vertex AI + Gemma ``:predict``) where ``MockResponseIterator`` collapses the entire response + into a single chunk. Without this split the assumption causes all content to be + silently dropped (only the ``message_delta`` stop event is emitted). + + Supports both sync and async iteration, since ``AnthropicStreamWrapper`` exposes + both ``__next__`` and ``__anext__``. An instance is single-mode: callers must + iterate it either synchronously or asynchronously, never both — the two modes + hold independent iterator references on the upstream stream and mixing them + would advance them out of sync. + """ + + def __init__(self, completion_stream: Any): + self._stream = completion_stream + self._sync_iter: Optional[Iterator[Any]] = None + self._async_iter: Optional[AsyncIterator[Any]] = None + self._buffer: deque = deque() + + @staticmethod + def _is_combined(chunk: Any) -> bool: + """True if ``chunk`` carries response content AND a finish_reason.""" + choices = getattr(chunk, "choices", None) + if not choices: + return False + choice = choices[0] + if getattr(choice, "finish_reason", None) is None: + return False + delta = getattr(choice, "delta", None) + if delta is None: + return False + return bool( + getattr(delta, "content", None) + or getattr(delta, "tool_calls", None) + or getattr(delta, "reasoning_content", None) + or getattr(delta, "thinking_blocks", None) + ) + + @staticmethod + def _split(chunk: Any) -> List[Any]: + """Return ``[chunk]``, or ``[content_chunk, finish_chunk]`` if combined.""" + if not _CombinedChunkSplitter._is_combined(chunk): + return [chunk] + + # Content chunk: keep the delta payload, clear the finish_reason. + content_chunk = copy.deepcopy(chunk) + content_chunk.choices[0].finish_reason = None + + # Finish chunk: keep finish_reason (and usage), clear the delta payload. + finish_chunk = copy.deepcopy(chunk) + finish_delta = finish_chunk.choices[0].delta + finish_delta.content = None + if hasattr(finish_delta, "tool_calls"): + finish_delta.tool_calls = None + if hasattr(finish_delta, "reasoning_content"): + finish_delta.reasoning_content = None + if hasattr(finish_delta, "thinking_blocks"): + finish_delta.thinking_blocks = None + return [content_chunk, finish_chunk] + + def __iter__(self) -> "Iterator[Any]": + return self + + def __next__(self) -> Any: + if self._buffer: + return self._buffer.popleft() + if self._sync_iter is None: + self._sync_iter = iter(self._stream) + chunk = next(self._sync_iter) # propagates StopIteration when exhausted + self._buffer.extend(self._split(chunk)) + return self._buffer.popleft() + + def __aiter__(self) -> "AsyncIterator[Any]": + return self + + async def __anext__(self) -> Any: + if self._buffer: + return self._buffer.popleft() + if self._async_iter is None: + self._async_iter = self._stream.__aiter__() + chunk = await self._async_iter.__anext__() # propagates StopAsyncIteration + self._buffer.extend(self._split(chunk)) + return self._buffer.popleft() + + class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): """ - first chunk return 'message_start' @@ -62,7 +155,10 @@ def __init__( compaction_block: Optional[CompactionBlock] = None, iterations_usage: Optional[List[UsageIteration]] = None, ): - super().__init__(completion_stream) + # Wrap the upstream stream so chunks that carry both content and a + # finish_reason (fake-streamed providers) are split into two — see + # _CombinedChunkSplitter. + super().__init__(_CombinedChunkSplitter(completion_stream)) self.model = model # Mapping of truncated tool names to original names (for OpenAI's 64-char limit) self.tool_name_mapping = tool_name_mapping or {} diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 3a2c09f2183..07e8270b496 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -84,6 +84,15 @@ def _process_content_list(content: list) -> None: if isinstance(content, list): _process_content_list(content) + def should_strip_billing_metadata(self) -> bool: + """ + Whether to drop x-anthropic-billing-header system blocks before sending upstream. + + The first-party Anthropic API uses these blocks for Claude Code attribution, so the + base config keeps them. Providers that reject them override this to True. + """ + return False + @staticmethod def _filter_billing_headers_from_system(system_param): """ @@ -286,14 +295,12 @@ def transform_anthropic_messages_request( optional_params=anthropic_messages_optional_request_params, ) - # Filter out x-anthropic-billing-header from system messages system_param = anthropic_messages_optional_request_params.get("system") - if system_param is not None: + if self.should_strip_billing_metadata() and system_param is not None: filtered_system = self._filter_billing_headers_from_system(system_param) if filtered_system is not None and len(filtered_system) > 0: anthropic_messages_optional_request_params["system"] = filtered_system else: - # Remove system parameter if all content was filtered out anthropic_messages_optional_request_params.pop("system", None) # Transform context_management from OpenAI format to Anthropic format if needed diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 94c5200be64..5f1362e259f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -155,10 +155,24 @@ def _process_event(self, event: Any) -> None: # noqa: PLR0915 event.get("delta", "") if isinstance(event, dict) else "" ) block_idx = ( - self._item_id_to_block_index.get(item_id, self._current_block_index) + self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index ) + if block_idx < 0: + # Some providers (e.g. LMStudio) skip response.output_item.added, + # so no text block is open yet; synthesize content_block_start + # instead of emitting a delta with index -1 + block_idx = self._next_block_index() + if item_id: + self._item_id_to_block_index[item_id] = block_idx + self._chunk_queue.append( + { + "type": "content_block_start", + "index": block_idx, + "content_block": {"type": "text", "text": ""}, + } + ) self._chunk_queue.append( { "type": "content_block_delta", diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 734b8ecef16..56cf035d0f7 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -43,7 +43,10 @@ process_azure_headers, select_azure_base_url_or_endpoint, ) -from .image_generation import get_azure_image_generation_config +from .image_generation import ( + AzureFoundryMAIImageGenerationConfig, + get_azure_image_generation_config, +) from .image_generation.http_utils import azure_deployment_image_generation_json_body @@ -1097,10 +1100,14 @@ def make_sync_azure_httpx_request( ) def create_azure_base_url( - self, azure_client_params: dict, model: Optional[str] + self, + azure_client_params: dict, + model: Optional[str], + base_model: Optional[str] = None, ) -> str: from litellm.llms.azure_ai.image_generation import ( AzureFoundryFluxImageGenerationConfig, + AzureFoundryMAIImageGenerationConfig, ) api_base: str = azure_client_params.get( @@ -1112,6 +1119,12 @@ def create_azure_base_url( if model is None: model = "" + if AzureFoundryMAIImageGenerationConfig.is_mai_model(base_model or model): + return AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( + api_base=api_base, + api_version=api_version, + ) + # Handle FLUX 2 models on Azure AI which use a different URL pattern # e.g., /providers/blackforestlabs/v1/flux-2-pro instead of /openai/deployments/{model}/images/generations if AzureFoundryFluxImageGenerationConfig.is_flux2_model(model): @@ -1153,10 +1166,10 @@ async def aimage_generation( if api_base.endswith("/"): api_base = api_base.rstrip("/") api_version: str = azure_client_params.get("api_version", "") - # Use the deployment name (model) for URL construction, not the base_model from data img_gen_api_base = self.create_azure_base_url( azure_client_params=azure_client_params, model=model or data.get("model", ""), + base_model=data.get("model", ""), ) ## LOGGING @@ -1285,9 +1298,10 @@ def image_generation( if aimg_generation is True: return self.aimage_generation(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_key=api_key, client=client, azure_client_params=azure_client_params, timeout=timeout, headers=headers, model=model) # type: ignore - # Use the deployment name (model) for URL construction, not the base_model from data img_gen_api_base = self.create_azure_base_url( - azure_client_params=azure_client_params, model=model + azure_client_params=azure_client_params, + model=model, + base_model=base_model, ) ## LOGGING @@ -1309,6 +1323,21 @@ def image_generation( data=data, headers=headers, ) + provider_config = get_azure_image_generation_config( + data.get("model", "dall-e-2") + ) + if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig): + return provider_config.transform_image_generation_response( + model=data.get("model", "dall-e-2"), + raw_response=httpx_response, + model_response=model_response or ImageResponse(), + logging_obj=logging_obj, + request_data=data, + optional_params=data, + litellm_params=data, + encoding=litellm.encoding, + ) + response = httpx_response.json() ## LOGGING diff --git a/litellm/llms/azure/image_generation/__init__.py b/litellm/llms/azure/image_generation/__init__.py index f60e446f0c4..64636bc689d 100644 --- a/litellm/llms/azure/image_generation/__init__.py +++ b/litellm/llms/azure/image_generation/__init__.py @@ -1,4 +1,5 @@ from litellm._logging import verbose_logger +from litellm.llms.azure_ai.image_generation import AzureFoundryMAIImageGenerationConfig from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -24,6 +25,8 @@ def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig: return AzureDallE2ImageGenerationConfig() elif "dalle3" in model: return AzureDallE3ImageGenerationConfig() + elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + return AzureFoundryMAIImageGenerationConfig() else: verbose_logger.debug( f"Using AzureGPTImageGenerationConfig for model: {model}. This follows the gpt-image model format." diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index a81218ab76a..59b6ee2b424 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -21,6 +21,9 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): and Azure endpoint format. """ + def should_strip_billing_metadata(self) -> bool: + return True + def validate_anthropic_messages_environment( self, headers: dict, diff --git a/litellm/llms/azure_ai/anthropic/transformation.py b/litellm/llms/azure_ai/anthropic/transformation.py index e176a4d860e..367ca75c196 100644 --- a/litellm/llms/azure_ai/anthropic/transformation.py +++ b/litellm/llms/azure_ai/anthropic/transformation.py @@ -40,6 +40,9 @@ class AzureAnthropicConfig(AnthropicConfig): def custom_llm_provider(self) -> Optional[str]: return "azure_ai" + def should_strip_billing_metadata(self) -> bool: + return True + def validate_environment( self, headers: dict, diff --git a/litellm/llms/azure_ai/image_edit/__init__.py b/litellm/llms/azure_ai/image_edit/__init__.py index e3acd610446..42ece6d19ec 100644 --- a/litellm/llms/azure_ai/image_edit/__init__.py +++ b/litellm/llms/azure_ai/image_edit/__init__.py @@ -1,21 +1,33 @@ from litellm.llms.azure_ai.image_generation.flux_transformation import ( AzureFoundryFluxImageGenerationConfig, ) +from litellm.llms.azure_ai.image_generation.mai_transformation import ( + AzureFoundryMAIImageGenerationConfig, +) from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from .flux2_transformation import AzureFoundryFlux2ImageEditConfig +from .mai_transformation import AzureFoundryMAIImageEditConfig from .transformation import AzureFoundryFluxImageEditConfig -__all__ = ["AzureFoundryFluxImageEditConfig", "AzureFoundryFlux2ImageEditConfig"] +__all__ = [ + "AzureFoundryFluxImageEditConfig", + "AzureFoundryFlux2ImageEditConfig", + "AzureFoundryMAIImageEditConfig", +] def get_azure_ai_image_edit_config(model: str) -> BaseImageEditConfig: """ Get the appropriate image edit config for an Azure AI model. + - MAI models use /mai/v1/images/edits with multipart form data and size - FLUX 2 models use JSON with base64 image - FLUX 1 models use multipart/form-data """ + if AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + return AzureFoundryMAIImageEditConfig() + # Check if it's a FLUX 2 model if AzureFoundryFluxImageGenerationConfig.is_flux2_model(model): return AzureFoundryFlux2ImageEditConfig() diff --git a/litellm/llms/azure_ai/image_edit/mai_transformation.py b/litellm/llms/azure_ai/image_edit/mai_transformation.py new file mode 100644 index 00000000000..75bfc913a8f --- /dev/null +++ b/litellm/llms/azure_ai/image_edit/mai_transformation.py @@ -0,0 +1,199 @@ +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast + +import httpx +from httpx._types import RequestFiles + +from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo +from litellm.llms.azure_ai.image_generation.mai_transformation import ( + AzureFoundryMAIImageGenerationConfig, +) +from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.llms.openai import FileTypes +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ImageResponse +from litellm.utils import convert_to_model_response_object + +if TYPE_CHECKING: + from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj + + +class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig): + """Azure AI Foundry MAI image editing (e.g. MAI-Image-2.5).""" + + DEFAULT_SIZE = "1024x1024" + + def get_supported_openai_params(self, model: str) -> list: + return ["prompt", "image", "model", "n", "size"] + + def map_openai_params( + self, + image_edit_optional_params: ImageEditOptionalRequestParams, + model: str, + drop_params: bool, + ) -> Dict: + optional_params: Dict[str, Any] = {} + supported_params = self.get_supported_openai_params(model) + + for key, value in dict(image_edit_optional_params).items(): + if value is None or key in optional_params: + continue + + if key in supported_params: + if key == "size" and value: + size_param = cast(str, value) + self._validate_size_param(size_param) + optional_params[key] = size_param + else: + optional_params[key] = value + elif not drop_params: + raise ValueError( + f"Parameter {key} is not supported for model {model}. " + f"Supported parameters are {supported_params}. " + f"Set drop_params=True to drop unsupported parameters." + ) + + if "size" not in optional_params: + optional_params["size"] = self.DEFAULT_SIZE + + return optional_params + + def _validate_size_param(self, size: str) -> None: + known_sizes = { + "1024x1024", + "1792x1024", + "1024x1792", + "512x512", + "256x256", + } + + if size in known_sizes: + return + + if "x" in size: + try: + tuple(map(int, size.lower().split("x", 1))) + return + except ValueError: + raise ValueError( + f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')." + ) + + raise ValueError( + f"Unsupported size value: '{size}'. " + f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." + ) + + def validate_environment( + self, + headers: dict, + model: str, + api_key: Optional[str] = None, + litellm_params: Optional[dict] = None, + api_base: Optional[str] = None, + ) -> dict: + api_key = AzureFoundryModelInfo.get_api_key(api_key) + + if not api_key: + raise ValueError( + f"Azure AI API key is required for model {model}. " + "Set AZURE_AI_API_KEY environment variable or pass api_key parameter." + ) + + headers.update({"api-key": api_key}) + return headers + + def get_complete_url( + self, + model: str, + api_base: Optional[str], + litellm_params: dict, + ) -> str: + api_base = AzureFoundryModelInfo.get_api_base(api_base) + + if api_base is None: + raise ValueError( + "Azure AI API base is required. Set AZURE_AI_API_BASE environment variable or pass api_base parameter." + ) + + api_version = ( + litellm_params.get("api_version") + or get_secret_str("AZURE_AI_API_VERSION") + or "preview" + ) + + return AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url( + api_base=api_base, + api_version=api_version, + ) + + def transform_image_edit_request( + self, + model: str, + prompt: Optional[str], + image: Optional[FileTypes], + image_edit_optional_request_params: Dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> Tuple[Dict, RequestFiles]: + request_params = { + "model": model, + **image_edit_optional_request_params, + } + if prompt is not None: + request_params["prompt"] = prompt + + data_without_files = { + key: value + for key, value in request_params.items() + if key not in ["image", "mask"] + } + files_list: List[Tuple[str, Any]] = [] + + if image is not None: + image_list = [image] if not isinstance(image, list) else image + for _image in image_list: + if _image is not None: + self._add_image_to_files( + files_list=files_list, + image=_image, + field_name="image", + ) + break + + return data_without_files, files_list + + def transform_image_edit_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> ImageResponse: + try: + response = raw_response.json() + except Exception: + raise OpenAIError( + message=raw_response.text, status_code=raw_response.status_code + ) + + if "usage" in response: + response["usage"] = ( + AzureFoundryMAIImageGenerationConfig.normalize_mai_image_usage( + response.get("usage") + ) + ) + + logging_obj.post_call( + input="", + api_key="", + additional_args={"complete_input_dict": {}}, + original_response=response, + ) + + return convert_to_model_response_object( + response_object=response, + model_response_object=ImageResponse(), + response_type="image_generation", + ) diff --git a/litellm/llms/azure_ai/image_generation/__init__.py b/litellm/llms/azure_ai/image_generation/__init__.py index cebab3de16e..70821d5d764 100644 --- a/litellm/llms/azure_ai/image_generation/__init__.py +++ b/litellm/llms/azure_ai/image_generation/__init__.py @@ -7,12 +7,14 @@ from .dall_e_3_transformation import AzureFoundryDallE3ImageGenerationConfig from .flux_transformation import AzureFoundryFluxImageGenerationConfig from .gpt_transformation import AzureFoundryGPTImageGenerationConfig +from .mai_transformation import AzureFoundryMAIImageGenerationConfig __all__ = [ "AzureFoundryFluxImageGenerationConfig", "AzureFoundryGPTImageGenerationConfig", "AzureFoundryDallE2ImageGenerationConfig", "AzureFoundryDallE3ImageGenerationConfig", + "AzureFoundryMAIImageGenerationConfig", ] @@ -24,6 +26,8 @@ def get_azure_ai_image_generation_config(model: str) -> BaseImageGenerationConfi return AzureFoundryDallE2ImageGenerationConfig() elif "dalle3" in model: return AzureFoundryDallE3ImageGenerationConfig() + elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + return AzureFoundryMAIImageGenerationConfig() elif "flux" in model: return AzureFoundryFluxImageGenerationConfig() else: diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index b67de9cb70d..f8c876bb5be 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -1,6 +1,9 @@ from typing import Any import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + calculate_image_response_cost_from_usage, +) from litellm.types.utils import ImageResponse @@ -9,19 +12,28 @@ def cost_calculator( image_response: Any, ) -> float: """ - Recraft image generation cost calculator + Azure AI image generation cost calculator """ _model_info = litellm.get_model_info( model=model, custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, ) - output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 - num_images: int = 0 + if isinstance(image_response, ImageResponse): + token_based_cost = calculate_image_response_cost_from_usage( + model=model, + image_response=image_response, + custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, + ) + if token_based_cost is not None: + return token_based_cost + + output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 + num_images: int = 0 if image_response.data: num_images = len(image_response.data) return output_cost_per_image * num_images - else: - raise ValueError( - f"image_response must be of type ImageResponse got type={type(image_response)}" - ) + + raise ValueError( + f"image_response must be of type ImageResponse got type={type(image_response)}" + ) diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py new file mode 100644 index 00000000000..071ca9d9895 --- /dev/null +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -0,0 +1,236 @@ +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import httpx + +from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, +) +from litellm.llms.openai.common_utils import OpenAIError +from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams +from litellm.types.utils import ImageResponse +from litellm.utils import convert_to_model_response_object + +if TYPE_CHECKING: + from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj + + +class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): + """Azure AI Foundry MAI image generation (e.g. MAI-Image-2.5).""" + + DEFAULT_WIDTH = 1024 + DEFAULT_HEIGHT = 1024 + + @staticmethod + def get_mai_image_generation_url( + api_base: Optional[str], + api_version: Optional[str], + ) -> str: + if api_base is None: + raise ValueError("api_base is required for Azure AI MAI image generation") + + api_version = api_version or "preview" + path, separator, query = api_base.partition("?") + path = path.rstrip("/") + + if "/mai/" in path: + prefix, _, _ = path.partition("/images/") + path = f"{prefix}/images/generations" + else: + path = f"{path}/mai/v1/images/generations" + + if separator: + return f"{path}?{query}" + return f"{path}?api-version={api_version}" + + @staticmethod + def get_mai_image_edit_url( + api_base: Optional[str], + api_version: Optional[str], + ) -> str: + if api_base is None: + raise ValueError("api_base is required for Azure AI MAI image editing") + + api_version = api_version or "preview" + path, separator, query = api_base.partition("?") + path = path.rstrip("/") + + if "/mai/" in path: + prefix, _, _ = path.partition("/images/") + path = f"{prefix}/images/edits" + else: + path = f"{path}/mai/v1/images/edits" + + if separator: + return f"{path}?{query}" + return f"{path}?api-version={api_version}" + + @staticmethod + def is_mai_model(model: str) -> bool: + model_normalized = model.lower().replace("-", "").replace("_", "") + return "maiimage" in model_normalized + + @staticmethod + def normalize_mai_image_usage(usage: Optional[Dict[str, Any]]) -> Dict[str, Any]: + """Map Azure MAI usage fields to OpenAI ImageUsage schema.""" + if usage is None: + return { + "input_tokens": 0, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "output_tokens": 0, + "total_tokens": 0, + } + + normalized_usage = dict(usage) + input_tokens_details = normalized_usage.get("input_tokens_details") + if not isinstance(input_tokens_details, dict): + input_tokens_details = {} + + text_tokens = normalized_usage.get("num_input_text_tokens") + if text_tokens is None: + text_tokens = input_tokens_details.get("text_tokens") + if text_tokens is None: + text_tokens = normalized_usage.get("input_tokens", 0) or 0 + + image_tokens = normalized_usage.get("num_input_image_tokens") + if image_tokens is None: + image_tokens = input_tokens_details.get("image_tokens") + if image_tokens is None: + image_tokens = 0 + + output_tokens = normalized_usage.get("output_tokens") + if output_tokens is None: + output_tokens = normalized_usage.get("num_output_tokens") + if output_tokens is None: + output_tokens = normalized_usage.get("output_image_tokens") + if output_tokens is None: + output_tokens = 0 + + input_tokens = normalized_usage.get("input_tokens") + if input_tokens is None: + input_tokens = text_tokens + image_tokens + + total_tokens = normalized_usage.get("total_tokens") + if total_tokens is None: + total_tokens = input_tokens + output_tokens + + normalized_usage.update( + { + "input_tokens": input_tokens, + "input_tokens_details": { + "image_tokens": image_tokens, + "text_tokens": text_tokens, + }, + "output_tokens": output_tokens, + "total_tokens": total_tokens, + } + ) + return normalized_usage + + def get_supported_openai_params( + self, model: str + ) -> List[OpenAIImageGenerationOptionalParams]: + return ["n", "size"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + supported_params = self.get_supported_openai_params(model) + + for k, v in non_default_params.items(): + if k in optional_params: + continue + + if k in supported_params: + if k == "size" and v: + self._map_size_param(v, optional_params) + else: + optional_params[k] = v + elif k in ("width", "height"): + optional_params[k] = v + elif not drop_params: + raise ValueError( + f"Parameter {k} is not supported for model {model}. " + f"Supported parameters are {supported_params} and width/height. " + f"Set drop_params=True to drop unsupported parameters." + ) + + if "width" not in optional_params: + optional_params["width"] = self.DEFAULT_WIDTH + if "height" not in optional_params: + optional_params["height"] = self.DEFAULT_HEIGHT + + optional_params.pop("size", None) + return optional_params + + def _map_size_param(self, size: str, optional_params: dict) -> None: + size_mapping = { + "1024x1024": (1024, 1024), + "1792x1024": (1792, 1024), + "1024x1792": (1024, 1792), + "512x512": (512, 512), + "256x256": (256, 256), + } + + if size in size_mapping: + width, height = size_mapping[size] + optional_params["width"] = width + optional_params["height"] = height + elif "x" in size: + try: + width, height = map(int, size.lower().split("x")) + optional_params["width"] = width + optional_params["height"] = height + except ValueError: + raise ValueError( + f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')." + ) + else: + raise ValueError( + f"Unsupported size value: '{size}'. " + f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." + ) + + def transform_image_generation_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ImageResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: dict, + optional_params: dict, + litellm_params: dict, + encoding: Any, + api_key: Optional[str] = None, + json_mode: Optional[bool] = None, + ) -> ImageResponse: + try: + response = raw_response.json() + except Exception: + raise OpenAIError( + message=raw_response.text, status_code=raw_response.status_code + ) + + if "usage" in response: + response["usage"] = self.normalize_mai_image_usage(response.get("usage")) + + logging_obj.post_call( + input=request_data.get("prompt", ""), + api_key=api_key, + additional_args={"complete_input_dict": request_data}, + original_response=response, + ) + + image_response: ImageResponse = convert_to_model_response_object( + response_object=response, + model_response_object=model_response, + response_type="image_generation", + ) + + width = optional_params.get("width", self.DEFAULT_WIDTH) + height = optional_params.get("height", self.DEFAULT_HEIGHT) + image_response.size = f"{width}x{height}" # type: ignore[assignment] + return image_response diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index cf1fd6f786e..bf1bfd06537 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -50,6 +50,11 @@ def convert_model_response_to_streaming( model=model_response.model, choices=streaming_choices, ) + # Carry usage onto the streaming chunk so fake-streamed responses + # (e.g. Vertex AI Gemma :predict) still report token counts. + usage = getattr(model_response, "usage", None) + if usage is not None: + setattr(processed_chunk, "usage", usage) return processed_chunk except Exception as e: raise ValueError( diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 5f35a58ce1f..8f9d5cad7c4 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -442,6 +442,14 @@ def post_stream_processing(self, stream: Any) -> Any: """Hook for providers to post-process streaming responses. Default: pass-through.""" return stream + def apply_assembled_streaming_response_metadata( + self, + response: "ModelResponse", + chunks: List[Any], + ) -> None: + """Hook for providers to merge chunk metadata into assembled streaming responses.""" + return None + def calculate_additional_costs( self, model: str, prompt_tokens: int, completion_tokens: int ) -> Optional[dict]: diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 853eb282758..407d5ad8146 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -62,6 +62,26 @@ def supports_native_file_search(self) -> bool: """ return False + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: dict, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + stream: Optional[bool] = None, + fake_stream: Optional[bool] = None, + ) -> Tuple[dict, Optional[bytes]]: + """Sign the request after the body is finalized. + + Default is a no-op (returns headers unchanged, no signed body). Providers + whose endpoint requires request signing (e.g. Bedrock Mantle SigV4) + override this and return the signed body bytes so the handler sends those + exact bytes. + """ + return headers, None + @abstractmethod def get_supported_openai_params(self, model: str) -> list: pass diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 90dfa13e938..b5e5e4de6fc 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -920,10 +920,15 @@ def map_openai_params( continue value = [value] optional_params["stopSequences"] = value - if param == "temperature": - optional_params["temperature"] = value - if param == "top_p": - optional_params["topP"] = value + if param == "temperature" or param == "top_p": + AnthropicConfig._apply_sampling_param( + optional_params=optional_params, + model=model, + param=param, + value=value, + drop_params=drop_params, + output_key="topP" if param == "top_p" else param, + ) if param == "tools" and isinstance(value, list): self._apply_tool_call_transformation( tools=cast(List[OpenAIChatCompletionToolParam], value), @@ -1221,7 +1226,9 @@ def _transform_inference_params(self, inference_params: dict) -> InferenceConfig inference_params["topK"] = inference_params.pop("top_k") return InferenceConfig(**inference_params) - def _handle_top_k_value(self, model: str, inference_params: dict) -> dict: + def _handle_top_k_value( + self, model: str, inference_params: dict, drop_params: bool = False + ) -> dict: base_model = BedrockModelInfo.get_base_model(model) val_top_k = None @@ -1230,16 +1237,25 @@ def _handle_top_k_value(self, model: str, inference_params: dict) -> dict: elif "top_k" in inference_params: val_top_k = inference_params.pop("top_k") - if val_top_k: + if val_top_k is not None: if base_model.startswith("anthropic"): - return {"top_k": val_top_k} + top_k_params: dict = {} + AnthropicConfig._apply_sampling_param( + optional_params=top_k_params, + model=model, + param="top_k", + value=val_top_k, + drop_params=drop_params, + output_key="top_k", + ) + return top_k_params if base_model.startswith("amazon.nova"): return {"inferenceConfig": {"topK": val_top_k}} return {} def _prepare_request_params( - self, optional_params: dict, model: str + self, optional_params: dict, model: str, drop_params: bool = False ) -> Tuple[dict, dict, dict, Optional[OutputConfigBlock]]: """Prepare and separate request parameters.""" # Consume the internal ``_output_config_normalized`` marker set by @@ -1338,7 +1354,7 @@ def _prepare_request_params( # Only set the topK value in for models that support it additional_request_params.update( - self._handle_top_k_value(model, inference_params) + self._handle_top_k_value(model, inference_params, drop_params) ) # Filter out internal/MCP-related parameters that shouldn't be sent to the API @@ -1572,6 +1588,7 @@ def _transform_request_helper( optional_params: dict, messages: Optional[List[AllMessageValues]] = None, headers: Optional[dict] = None, + drop_params: bool = False, ) -> CommonRequestObject: ## VALIDATE REQUEST """ @@ -1618,7 +1635,7 @@ def _transform_request_helper( additional_request_params, request_metadata, output_config, - ) = self._prepare_request_params(optional_params, model) + ) = self._prepare_request_params(optional_params, model, drop_params) original_tools = inference_params.pop("tools", []) @@ -1649,12 +1666,14 @@ def _transform_request_helper( bedrock_tool_config["toolChoice"] = tool_choice_values data: CommonRequestObject = { - "additionalModelRequestFields": additional_request_params, - "system": system_content_blocks, "inferenceConfig": self._transform_inference_params( inference_params=inference_params ), } + if additional_request_params: + data["additionalModelRequestFields"] = additional_request_params + if system_content_blocks: + data["system"] = system_content_blocks # Handle all config blocks for config_name, config_class in self.get_config_blocks().items(): @@ -1699,6 +1718,7 @@ async def _async_transform_request( optional_params=optional_params, messages=messages, headers=headers, + drop_params=litellm_params.get("drop_params") is True, ) bedrock_messages = ( @@ -1756,6 +1776,7 @@ def _transform_request( optional_params=optional_params, messages=messages, headers=headers, + drop_params=litellm_params.get("drop_params") is True, ) ## TRANSFORMATION ## diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index a13336b6c88..4887cbd23be 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -60,6 +60,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): def custom_llm_provider(self) -> Optional[str]: return "bedrock" + def should_strip_billing_metadata(self) -> bool: + return True + def get_supported_openai_params(self, model: str) -> List[str]: return AnthropicConfig.get_supported_openai_params(self, model) diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index ef0199031af..cbed2232be5 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -48,6 +48,30 @@ def get_complete_url( region = self._get_aws_region_name(optional_params=optional_params, model=model) return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + headers = super().validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + project_id = litellm_params.get("aws_bedrock_project_id") + if project_id: + headers["anthropic-workspace"] = project_id + return headers + def transform_request( self, model: str, diff --git a/litellm/llms/bedrock/claude_platform/transformation.py b/litellm/llms/bedrock/claude_platform/transformation.py index 0167c457c96..c20dc63444f 100644 --- a/litellm/llms/bedrock/claude_platform/transformation.py +++ b/litellm/llms/bedrock/claude_platform/transformation.py @@ -17,6 +17,9 @@ class BedrockClaudePlatformConfig(BedrockClaudePlatformMixin, AnthropicConfig): def custom_llm_provider(self) -> Optional[str]: return "bedrock" + def should_strip_billing_metadata(self) -> bool: + return True + def validate_environment( self, headers: dict, diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index c967fd334bc..bdef3349e00 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -11,6 +11,11 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import get_bedrock_base_model +# Placeholder satisfying the Anthropic InvokeModel schema's required +# max_tokens field; CountTokens only counts input, so it has no effect +# on any generation. +DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS = 1024 + class BedrockCountTokensConfig(BaseAWSLLM): """ @@ -32,8 +37,20 @@ def _detect_input_type(self, request_data: Dict[str, Any]) -> str: Returns: 'converse' or 'invokeModel' """ - # If the request has messages in the expected Anthropic format, use converse - if "messages" in request_data and isinstance(request_data["messages"], list): + messages = request_data.get("messages") + if isinstance(messages, list): + # Anthropic content blocks carry a "type" key ({"type": "text", ...}); + # Converse blocks don't ({"text": ...}, {"toolUse": ...}). Converse + # rejects Anthropic-shape blocks, so route those to invokeModel, + # which forwards the body verbatim. + for message in messages: + if not isinstance(message, dict): + continue + content = message.get("content") + if isinstance(content, list) and any( + isinstance(block, dict) and "type" in block for block in content + ): + return "invokeModel" return "converse" # For raw text or other formats, use invokeModel @@ -68,7 +85,7 @@ def transform_anthropic_to_bedrock_count_tokens( { "input": { "invokeModel": { - "body": "{...raw model input...}" + "body": "" } } } @@ -168,13 +185,24 @@ def _transform_to_invoke_model_format( self, request_data: Dict[str, Any] ) -> Dict[str, Any]: """Transform to InvokeModel input format.""" + import base64 import json # For InvokeModel, we need to provide the raw body that would be sent to the model # Remove the 'model' field from the body as it's not part of the model input body_data = {k: v for k, v in request_data.items() if k != "model"} - return {"input": {"invokeModel": {"body": json.dumps(body_data)}}} + if "messages" in body_data: + # Bedrock validates the body against the model's InvokeModel schema; + # Anthropic Messages bodies require these fields. + body_data.setdefault("anthropic_version", "bedrock-2023-05-31") + body_data.setdefault( + "max_tokens", DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS + ) + + # The CountTokens API expects invokeModel.body as a base64-encoded blob + encoded_body = base64.b64encode(json.dumps(body_data).encode()).decode() + return {"input": {"invokeModel": {"body": encoded_body}}} def get_bedrock_count_tokens_endpoint( self, diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index a78f696a057..900d9aa97d8 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -6,7 +6,7 @@ stripping that are specific to the bedrock-mantle endpoint. """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeMessagesConfig, @@ -45,6 +45,30 @@ def get_complete_url( region = self._get_aws_region_name(optional_params=optional_params, model=model) return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + def validate_anthropic_messages_environment( + self, + headers: dict, + model: str, + messages: List[Any], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> Tuple[dict, Optional[str]]: + headers, api_base = super().validate_anthropic_messages_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + project_id = litellm_params.get("aws_bedrock_project_id") + if project_id: + headers["anthropic-workspace"] = project_id + return headers, api_base + def transform_anthropic_messages_request( self, model: str, diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 81a56030a5c..ad37a1990d3 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -8,11 +8,12 @@ or region-aware key via BEDROCK_MANTLE_{REGION}_API_KEY. """ -from typing import Iterator, AsyncIterator, Any, Optional, Tuple, Union +from typing import Iterator, AsyncIterator, Any, List, Optional, Tuple, Union import litellm from litellm._logging import verbose_logger from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues from ...openai_like.chat.transformation import OpenAILikeChatConfig @@ -48,6 +49,30 @@ def _get_openai_compatible_provider_info( dynamic_api_key = api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") return api_base, dynamic_api_key + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + headers = super().validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + project_id = litellm_params.get("aws_bedrock_project_id") + if project_id: + headers["OpenAI-Project"] = project_id + return headers + def get_supported_openai_params(self, model: str) -> list: base_params = super().get_supported_openai_params(model) try: diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index b63fd0ecdb1..29248e1ca50 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -1,17 +1,31 @@ """ Amazon Bedrock Mantle - Responses API backend. -gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses` -path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI +Mantle serves Responses on two upstream paths: gpt frontier models (gpt-5.5 / +gpt-5.4) on `/openai/v1/responses`, and everything else that supports Responses +(e.g. gpt-oss) on the standard `/v1/responses`. The gate picks the path per +model and injects it via `use_openai_path`. Payloads and SSE follow the OpenAI Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides -only the endpoint URL and Bearer authentication. +only the endpoint URL and authentication. -Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the -standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4. +Auth: Bearer token (BEDROCK_MANTLE_API_KEY or the standard +AWS_BEARER_TOKEN_BEDROCK, or litellm_params.api_key) when present; otherwise +AWS SigV4 (service name "bedrock") using the standard credential chain (IAM +role / access key / profile / web identity), signed via the shared +BaseAWSLLM._sign_request after the request body is finalized. """ -from typing import Optional +import re +from typing import Optional, Tuple +from botocore.exceptions import ( + CredentialRetrievalError, + NoCredentialsError, + PartialCredentialsError, + ProfileNotFound, +) + +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -29,22 +43,49 @@ "/v1", ) +# Standard Mantle host: https://bedrock-mantle..api.aws (group 1 = region). +_MANTLE_HOST_RE = re.compile( + r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE +) + class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): + def __init__( + self, + aws_signer: Optional[BaseAWSLLM] = None, + use_openai_path: bool = True, + ): + super().__init__() + self._aws_signer = aws_signer or BaseAWSLLM() + self.use_openai_path = use_openai_path + @property def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK_MANTLE + @staticmethod + def _resolve_region(params: dict) -> str: + region = params.get("aws_region_name") + if region: + return region + base = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE") + if base: + match = _MANTLE_HOST_RE.match(base.rstrip("/")) + if match: + return match.group(1) + return ( + get_secret_str("BEDROCK_MANTLE_REGION") + or get_secret_str("AWS_REGION_NAME") + or get_secret_str("AWS_REGION") + or BEDROCK_MANTLE_DEFAULT_REGION + ) + def get_complete_url( self, api_base: Optional[str], litellm_params: dict, ) -> str: - region = ( - get_secret_str("BEDROCK_MANTLE_REGION") - or get_secret_str("AWS_REGION") - or BEDROCK_MANTLE_DEFAULT_REGION - ) + region = self._resolve_region({**litellm_params, "api_base": api_base}) base = ( api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") @@ -55,7 +96,13 @@ def get_complete_url( if base.endswith(suffix): base = base[: -len(suffix)] break - return f"{base}/openai/v1/responses" + # For the standard Mantle host (including the default-region base that + # responses/main.py auto-injects into litellm_params.api_base), pin to the + # single resolved region so aws_region_name wins; preserve custom proxy hosts. + if _MANTLE_HOST_RE.match(base): + base = f"https://bedrock-mantle.{region}.api.aws" + path = "/openai/v1/responses" if self.use_openai_path else "/v1/responses" + return f"{base}{path}" def validate_environment( self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] @@ -66,12 +113,10 @@ def validate_environment( or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") ) - if not api_key: - raise ValueError( - "Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY " - "(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key." - ) - headers["Authorization"] = f"Bearer {api_key}" + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + if litellm_params.aws_bedrock_project_id: + headers["OpenAI-Project"] = litellm_params.aws_bedrock_project_id return headers def supports_native_file_search(self) -> bool: @@ -79,3 +124,58 @@ def supports_native_file_search(self) -> bool: def supports_native_websocket(self) -> bool: return False + + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: dict, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + stream: Optional[bool] = None, + fake_stream: Optional[bool] = None, + ) -> Tuple[dict, Optional[bytes]]: + bearer = ( + api_key + or get_secret_str("BEDROCK_MANTLE_API_KEY") + or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + ) + if not bearer: + # SigV4 path. Pin the credential-scope region to the region of the actual + # signing URL (api_base, already region-resolved by get_complete_url) so the + # SigV4 scope and the URL host can never disagree. Resolve from api_base first, + # then fall back to the regular precedence. Also drop any caller Authorization + # so _sign_request's restore-original-Authorization step cannot override the + # SigV4 header. + optional_params = { + **optional_params, + "aws_region_name": self._resolve_region( + {**optional_params, "api_base": api_base} + ), + } + headers = {k: v for k, v in headers.items() if k.lower() != "authorization"} + try: + return self._aws_signer._sign_request( + service_name="bedrock", + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=api_base, + api_key=bearer, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + except ( + NoCredentialsError, + PartialCredentialsError, + ProfileNotFound, + CredentialRetrievalError, + ) as e: + raise ValueError( + "Bedrock Mantle auth failed: no Bearer token and no usable AWS " + "credentials. Set BEDROCK_MANTLE_API_KEY (or AWS_BEARER_TOKEN_BEDROCK) " + "or pass api_key for Bearer auth, or provide AWS credentials " + "(IAM role / access key / profile / web identity) for SigV4." + ) from e diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 31c772510ba..25424feaeb4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2318,6 +2318,31 @@ def response_api_handler( # but never included in the outbound provider payload. request_context["litellm_params"] = dict(litellm_params) + is_stream_request = bool(stream) + if is_stream_request and fake_stream is True: + stream, data = self._prepare_fake_stream_request( + stream=stream, + data=data, + fake_stream=fake_stream, + ) + + # Sign after the body is final (post-transform/normalize/extra_body and post + # fake-stream prep) so signed bytes match what we send. No-op for providers + # that inherit the default sign_request. + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=api_base, + api_key=litellm_params.api_key, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -2330,22 +2355,14 @@ def response_api_handler( ) try: - if stream: - # For streaming, use stream=True in the request - if fake_stream is True: - stream, data = self._prepare_fake_stream_request( - stream=stream, - data=data, - fake_stream=fake_stream, - ) - + if is_stream_request: response = sync_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, + **body_kwargs, ) if fake_stream is True: return MockResponsesAPIStreamingIterator( @@ -2370,13 +2387,12 @@ def response_api_handler( call_type=CallTypes.responses.value, ) else: - # For non-streaming requests response = sync_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), + **body_kwargs, ) except Exception as e: raise self._handle_error( @@ -2464,6 +2480,28 @@ async def async_response_api_handler( # but never included in the outbound provider payload. request_context["litellm_params"] = dict(litellm_params) + is_stream_request = bool(stream) + if is_stream_request and fake_stream is True: + stream, data = self._prepare_fake_stream_request( + stream=stream, + data=data, + fake_stream=fake_stream, + ) + + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=api_base, + api_key=litellm_params.api_key, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -2476,22 +2514,14 @@ async def async_response_api_handler( ) try: - if stream: - # For streaming, we need to use stream=True in the request - if fake_stream is True: - stream, data = self._prepare_fake_stream_request( - stream=stream, - data=data, - fake_stream=fake_stream, - ) - + if is_stream_request: response = await async_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, + **body_kwargs, ) if fake_stream is True: @@ -2518,13 +2548,12 @@ async def async_response_api_handler( call_type=CallTypes.responses.value, ) else: - # For non-streaming, proceed as before response = await async_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), + **body_kwargs, ) except Exception as e: @@ -4005,6 +4034,18 @@ def compact_response_api_handler( ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=url, + api_key=litellm_params.api_key, + model=model, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -4018,7 +4059,7 @@ def compact_response_api_handler( try: response = sync_httpx_client.post( - url=url, headers=headers, json=data, timeout=timeout + url=url, headers=headers, timeout=timeout, **body_kwargs ) except Exception as e: @@ -4088,6 +4129,18 @@ async def async_compact_response_api_handler( ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=url, + api_key=litellm_params.api_key, + model=model, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -4101,7 +4154,7 @@ async def async_compact_response_api_handler( try: response = await async_httpx_client.post( - url=url, headers=headers, json=data, timeout=timeout + url=url, headers=headers, timeout=timeout, **body_kwargs ) except Exception as e: diff --git a/litellm/llms/databricks/streaming_utils.py b/litellm/llms/databricks/streaming_utils.py index eebe3182881..7a7330227d6 100644 --- a/litellm/llms/databricks/streaming_utils.py +++ b/litellm/llms/databricks/streaming_utils.py @@ -25,6 +25,28 @@ def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: finish_reason = "" usage: Optional[ChatCompletionUsageBlock] = None + # Usage-only final chunk (OpenAI ``stream_options.include_usage``) + # arrives with an empty ``choices`` list — return usage without + # indexing ``choices[0]``. + if len(processed_chunk.choices) == 0: + final_usage = getattr(processed_chunk, "usage", None) + return GenericStreamingChunk( + text="", + tool_use=None, + is_finished=False, + finish_reason="", + usage=( + ChatCompletionUsageBlock( + prompt_tokens=final_usage.prompt_tokens or 0, + completion_tokens=final_usage.completion_tokens or 0, + total_tokens=final_usage.total_tokens or 0, + ) + if final_usage is not None + else None + ), + index=0, + ) + if processed_chunk.choices[0].delta.content is not None: # type: ignore text = processed_chunk.choices[0].delta.content # type: ignore diff --git a/litellm/llms/deepseek/messages/transformation.py b/litellm/llms/deepseek/messages/transformation.py index ad60478960e..63b736ffd1d 100644 --- a/litellm/llms/deepseek/messages/transformation.py +++ b/litellm/llms/deepseek/messages/transformation.py @@ -26,6 +26,9 @@ class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): def custom_llm_provider(self) -> Optional[str]: return "deepseek" + def should_strip_billing_metadata(self) -> bool: + return True + @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: return api_key or get_secret_str("DEEPSEEK_API_KEY") or litellm.api_key diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index 0929f95cf43..3406538c774 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -2,7 +2,7 @@ GitHub Copilot Responses API Configuration. This module provides the configuration for GitHub Copilot's Responses API, -which is required for models like gpt-5.1-codex that only support the /responses endpoint. +which is required for models like gpt-5.3-codex that only support the /responses endpoint. Implementation based on analysis of the copilot-api project by caozhiyuan: https://github.com/caozhiyuan/copilot-api @@ -12,6 +12,7 @@ import os +import litellm from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.exceptions import AuthenticationError @@ -22,6 +23,7 @@ ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +from litellm.utils import _cached_get_model_info_helper from ..authenticator import Authenticator from ..common_utils import ( @@ -38,6 +40,47 @@ LiteLLMLoggingObj = Any +def github_copilot_supports_responses_api(model: str) -> bool: + """ + Gate native /v1/responses dispatch per github_copilot model. + + Resolution (first match wins): mode "responses" -> True; mode "chat" -> + False (opt-out wins for dual-endpoint models); "/v1/responses" in + supported_endpoints -> True; else False. Unknown model -> False (the bridge + always works since every Copilot model supports /chat/completions). + + Reads merged model info (per-deployment model_info applied via the router's + register_model, which also clears the cache used here). + """ + try: + info = _cached_get_model_info_helper( + model=model, custom_llm_provider="github_copilot" + ) + except Exception as e: + verbose_logger.debug( + "github_copilot_supports_responses_api: get_model_info failed " + "for %s: %s", + model, + e, + ) + return False + + mode = info.get("mode") + if mode == "responses": + return True + if mode == "chat": + return False + + # supported_endpoints is dropped by ModelInfoBase; read it from the raw + # model_cost entry via the resolved key. + key = info.get("key") + raw_info = litellm.model_cost.get(key) if isinstance(key, str) else None + endpoints = ( + raw_info.get("supported_endpoints") if isinstance(raw_info, dict) else None + ) + return isinstance(endpoints, list) and "/v1/responses" in endpoints + + class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): """ Configuration for GitHub Copilot's Responses API. diff --git a/litellm/llms/litellm_proxy/skills/README.md b/litellm/llms/litellm_proxy/skills/README.md index 1dfeff1a42c..a896aa1166e 100644 --- a/litellm/llms/litellm_proxy/skills/README.md +++ b/litellm/llms/litellm_proxy/skills/README.md @@ -18,7 +18,7 @@ flowchart TB F[Request with container.skills] --> G[SkillsInjectionHook] G --> H{skill_id prefix?} - H -->|"litellm:skill_abc"| I[Fetch from LiteLLM DB] + H -->|"litellm_skill_abc"| I[Fetch from LiteLLM DB] H -->|"skill_xyz" no prefix| J[Pass to Anthropic as native skill] I --> K{Model provider?} @@ -57,7 +57,7 @@ sequenceDiagram Note over LiteLLM,PreHook: PRE-CALL HOOK LiteLLM->>PreHook: Intercept request - PreHook->>PreHook: Fetch skill from DB (litellm:skill_id) + PreHook->>PreHook: Fetch skill from DB (litellm_skill_id) PreHook->>PreHook: Extract SKILL.md from ZIP PreHook->>PreHook: Inject SKILL.md into system prompt PreHook->>PreHook: Add litellm_code_execution tool @@ -105,7 +105,7 @@ response = await litellm.acompletion( model="gpt-4o-mini", messages=[{"role": "user", "content": "Create a bouncing ball GIF"}], container={ - "skills": [{"type": "custom", "skill_id": "litellm:skill_abc123"}] + "skills": [{"type": "custom", "skill_id": "litellm_skill_abc123"}] }, ) @@ -261,7 +261,7 @@ response = litellm.completion( messages=[{"role": "user", "content": "Analyze this data..."}], container={ "skills": [ - {"type": "custom", "skill_id": "litellm:skill_abc123"} # litellm: prefix + {"type": "custom", "skill_id": "litellm_skill_abc123"} # litellm_skill_ prefix ] } ) @@ -277,7 +277,7 @@ response = litellm.completion( "messages": [{"role": "user", "content": "Help me analyze data"}], "container": { "skills": [ - {"type": "custom", "skill_id": "litellm:skill_abc123"} + {"type": "custom", "skill_id": "litellm_skill_abc123"} ] } } @@ -287,7 +287,7 @@ response = litellm.completion( The hook (`litellm/proxy/hooks/litellm_skills/main.py`) intercepts the request: -1. **Detects `litellm:` prefix** → Fetches skill from database +1. **Detects `litellm_skill_` prefix** → Fetches skill from database 2. **Checks model provider** → Bedrock is not Anthropic 3. **Extracts SKILL.md** from stored ZIP file 4. **Converts skill to tool** + **Injects content into system prompt** @@ -361,8 +361,8 @@ model LiteLLM_SkillsTable { | Create skill on Anthropic | `anthropic` | N/A | Forward to Anthropic API | | Create skill in LiteLLM DB | `litellm_proxy` | N/A | Store in database | | Use Anthropic native skill | N/A | `skill_xyz` | Pass to Anthropic container.skills | -| Use LiteLLM skill on Anthropic | N/A | `litellm:skill_abc` | Convert to tools | -| Use LiteLLM skill on Bedrock/OpenAI | N/A | `litellm:skill_abc` | Convert to tools + inject SKILL.md | +| Use LiteLLM skill on Anthropic | N/A | `litellm_skill_abc` | Convert to tools | +| Use LiteLLM skill on Bedrock/OpenAI | N/A | `litellm_skill_abc` | Convert to tools + inject SKILL.md | ## Testing diff --git a/litellm/llms/litellm_proxy/skills/constants.py b/litellm/llms/litellm_proxy/skills/constants.py index a8c2697fcee..0c60a60842a 100644 --- a/litellm/llms/litellm_proxy/skills/constants.py +++ b/litellm/llms/litellm_proxy/skills/constants.py @@ -4,6 +4,10 @@ Centralized constants for skills processing, code execution, and sandbox configuration. """ +LITELLM_SKILL_ID_PREFIX: str = "litellm_skill_" +"""Prefix for DB-backed skill IDs. The model-facing tool name is the skill ID +with hyphens/spaces replaced by underscores, which leaves this prefix intact.""" + # Code execution loop settings DEFAULT_MAX_ITERATIONS: int = 10 """Maximum number of iterations for the automatic code execution loop.""" diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 37aabd8b477..9138b9a712f 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -10,6 +10,7 @@ from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache +from litellm.llms.litellm_proxy.skills.constants import LITELLM_SKILL_ID_PREFIX from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth from litellm.proxy.common_utils.resource_ownership import ( get_primary_resource_owner_scope, @@ -17,6 +18,7 @@ is_proxy_admin, user_can_access_resource_owner, ) +from litellm.repositories.table_repositories import SkillsRepository # Skills are looked up on every chat completion that has skills enabled # (`SkillsInjectionHook` calls ``fetch_skill_from_db``). 60s LRU/TTL cache @@ -67,7 +69,7 @@ async def create_skill( ) -> LiteLLM_SkillsTable: prisma_client = await LiteLLMSkillsHandler._get_prisma_client() - skill_id = f"litellm_skill_{uuid.uuid4()}" + skill_id = f"{LITELLM_SKILL_ID_PREFIX}{uuid.uuid4()}" owner = get_primary_resource_owner_scope(user_api_key_dict) or user_id if owner is None: # Identity-less callers (no user_id / team_id / org_id / @@ -107,7 +109,7 @@ async def create_skill( f"LiteLLMSkillsHandler: Creating skill {skill_id} with title={data.display_title}" ) - new_skill = await prisma_client.db.litellm_skillstable.create(data=skill_data) + new_skill = await SkillsRepository(prisma_client).table.create(data=skill_data) return _prisma_skill_to_litellm(new_skill) @staticmethod @@ -133,7 +135,7 @@ async def list_skills( return [] find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}} - skills = await prisma_client.db.litellm_skillstable.find_many( + skills = await SkillsRepository(prisma_client).table.find_many( **find_many_kwargs ) return [_prisma_skill_to_litellm(s) for s in skills] @@ -150,7 +152,7 @@ async def _load_skill(skill_id: str) -> Optional[Any]: return cached prisma_client = await LiteLLMSkillsHandler._get_prisma_client() - skill = await prisma_client.db.litellm_skillstable.find_unique( + skill = await SkillsRepository(prisma_client).table.find_unique( where={"skill_id": skill_id} ) _SKILL_CACHE.set_cache( @@ -189,7 +191,7 @@ async def delete_skill( ): raise ValueError(f"Skill not found: {skill_id}") - await prisma_client.db.litellm_skillstable.delete(where={"skill_id": skill_id}) + await SkillsRepository(prisma_client).table.delete(where={"skill_id": skill_id}) _SKILL_CACHE.set_cache(skill_id, _NEGATIVE_SKILL_SENTINEL) return {"id": skill_id, "type": "skill_deleted"} diff --git a/litellm/llms/minimax/messages/transformation.py b/litellm/llms/minimax/messages/transformation.py index 3190a5f5412..57cfcbf0621 100644 --- a/litellm/llms/minimax/messages/transformation.py +++ b/litellm/llms/minimax/messages/transformation.py @@ -28,6 +28,9 @@ class MinimaxMessagesConfig(AnthropicMessagesConfig): def custom_llm_provider(self) -> Optional[str]: return "minimax" + def should_strip_billing_metadata(self) -> bool: + return True + @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: """ diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index 1641615126e..63d39151254 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -49,6 +49,8 @@ def completion( headers: Optional[dict] = None, ): try: + if headers: + optional_params = {**optional_params, "extra_headers": headers} if headers is None: headers = self.validate_environment(api_key=api_key) if model is None or messages is None: diff --git a/litellm/llms/openai_like/dynamic_config.py b/litellm/llms/openai_like/dynamic_config.py index fac453447fa..9ed9734edae 100644 --- a/litellm/llms/openai_like/dynamic_config.py +++ b/litellm/llms/openai_like/dynamic_config.py @@ -187,6 +187,7 @@ def create_responses_config_class(provider: SimpleProviderConfig): from litellm.llms.openai_like.responses.transformation import ( OpenAILikeResponsesConfig, ) + from litellm.types.llms.openai import ResponseInputParam from litellm.types.router import GenericLiteLLMParams class JSONProviderResponsesConfig(OpenAILikeResponsesConfig): @@ -223,5 +224,23 @@ def get_complete_url( api_base = api_base.rstrip("/") return f"{api_base}/responses" + def transform_responses_api_request( + self, + model: str, + input: Union[str, ResponseInputParam], + response_api_optional_request_params: dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> dict: + if provider.special_handling.get("force_store_false"): + response_api_optional_request_params["store"] = False + return super().transform_responses_api_request( + model=model, + input=input, + response_api_optional_request_params=response_api_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + _responses_config_cache[provider.slug] = JSONProviderResponsesConfig return JSONProviderResponsesConfig diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 49b3801c82f..13d22488838 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -132,5 +132,14 @@ "param_mappings": { "max_completion_tokens": "max_tokens" } + }, + "parasail": { + "base_url": "https://api.parasail.io/v1", + "api_key_env": "PARASAIL_API_KEY", + "api_base_env": "PARASAIL_API_BASE", + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "special_handling": { + "force_store_false": true + } } } diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index e6e39651109..85c23d8603c 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -12,7 +12,11 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues -from litellm.types.llms.vertex_ai import PartType, Schema +from litellm.types.llms.vertex_ai import ( + VERTEX_AI_PROVIDER_METADATA_FIELDS, + PartType, + Schema, +) from litellm.types.utils import TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages @@ -27,6 +31,47 @@ def __init__( super().__init__(message=message, status_code=status_code, headers=headers) +def redact_vertex_ai_metadata_from_logged_object(obj: Any) -> None: + if isinstance(obj, dict): + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + if field in obj: + obj[field] = [] + hidden_params = obj.get("_hidden_params") + if isinstance(hidden_params, dict): + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + return + + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + if hasattr(obj, field): + setattr(obj, field, []) + hidden_params = getattr(obj, "_hidden_params", None) + if isinstance(hidden_params, dict): + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + + +def redact_vertex_ai_metadata_from_litellm_params(model_call_details: dict) -> None: + """ + success_handler() merges response._hidden_params into + litellm_params.metadata['hidden_params'] before redaction runs, so the Vertex + metadata must be scrubbed from that copy too. + """ + litellm_params = model_call_details.get("litellm_params") + if not isinstance(litellm_params, dict): + return + + for metadata_key in ("metadata", "litellm_metadata"): + metadata = litellm_params.get(metadata_key) + if not isinstance(metadata, dict): + continue + hidden_params = metadata.get("hidden_params") + if not isinstance(hidden_params, dict): + continue + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + + def vertex_request_labels_from_litellm_params( litellm_params: Optional[dict], ) -> Optional[Dict[str, str]]: diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index e9f08f403f9..103801a1e8d 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -19,7 +19,7 @@ VertexAICachedContentResponseObject, ) -from ..common_utils import VertexAIError +from ..common_utils import VertexAIError, get_vertex_base_url from ..vertex_llm_base import VertexBase from .transformation import ( separate_cached_messages, @@ -69,17 +69,13 @@ def _get_token_and_url_context_caching( elif custom_llm_provider == "vertex_ai": auth_header = vertex_auth_header endpoint = "cachedContents" - if vertex_location == "global": - url = f"https://aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" - else: - url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" + base_url = get_vertex_base_url(vertex_location) + url = f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" else: auth_header = vertex_auth_header endpoint = "cachedContents" - if vertex_location == "global": - url = f"https://aiplatform.googleapis.com/v1beta1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" - else: - url = f"https://{vertex_location}-aiplatform.googleapis.com/v1beta1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" + base_url = get_vertex_base_url(vertex_location) + url = f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" return self._check_custom_proxy( api_base=api_base, 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 5cd02293f14..430a789d2a0 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 @@ -63,6 +63,7 @@ OpenAIChatCompletionFinishReason, ) from litellm.types.llms.vertex_ai import ( + VERTEX_AI_PROVIDER_METADATA_FIELDS, VERTEX_CREDENTIALS_TYPES, Candidates, ContentType, @@ -1111,6 +1112,7 @@ def _map_audio_params(self, value: dict) -> dict: { "voice": "alloy", "format": "mp3", + "language_code": "en-US", } Expected output: @@ -1119,7 +1121,8 @@ def _map_audio_params(self, value: dict) -> dict: prebuiltVoiceConfig: { voiceName: "alloy", } - } + }, + languageCode: "en-US", } """ from litellm.types.llms.vertex_ai import ( @@ -1145,6 +1148,9 @@ def _map_audio_params(self, value: dict) -> dict: voice_config: VoiceConfig = {"prebuiltVoiceConfig": prebuilt_voice_config} speech_config["voiceConfig"] = voice_config + if "language_code" in value: + speech_config["languageCode"] = value["language_code"] + return cast(dict, speech_config) @staticmethod @@ -2253,6 +2259,71 @@ def _extract_candidate_metadata( citation_metadata, ) + @staticmethod + def _get_stream_chunk_attr(chunk: Any, field_name: str) -> Any: + if isinstance(chunk, dict): + value = chunk.get(field_name) + if value is not None: + return value + model_extra = chunk.get("model_extra") + if isinstance(model_extra, dict): + value = model_extra.get(field_name) + if value is not None: + return value + hidden_params = chunk.get("_hidden_params") + if isinstance(hidden_params, dict): + return hidden_params.get(field_name) + return None + return getattr(chunk, field_name, None) + + @staticmethod + def _set_stream_metadata_on_response( + model_response: Any, + grounding_metadata: List[dict], + url_context_metadata: List[dict], + safety_ratings: List[dict], + citation_metadata: List[dict], + ) -> None: + setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore + if grounding_metadata: + model_response._hidden_params["vertex_ai_grounding_metadata"] = ( + grounding_metadata + ) + setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore + if url_context_metadata: + model_response._hidden_params["vertex_ai_url_context_metadata"] = ( + url_context_metadata + ) + setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore + setattr(model_response, "vertex_ai_safety_results", safety_ratings) # type: ignore + if safety_ratings: + model_response._hidden_params["vertex_ai_safety_ratings"] = safety_ratings + model_response._hidden_params["vertex_ai_safety_results"] = safety_ratings + setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore + if citation_metadata: + model_response._hidden_params["vertex_ai_citation_metadata"] = ( + citation_metadata + ) + + def apply_assembled_streaming_response_metadata( + self, + response: ModelResponse, + chunks: List[Any], + ) -> None: + for field_name in VERTEX_AI_PROVIDER_METADATA_FIELDS: + merged: List[Any] = [] + for chunk in chunks: + value = VertexGeminiConfig._get_stream_chunk_attr(chunk, field_name) + if not value: + continue + if isinstance(value, list): + merged.extend(value) + else: + merged.append(value) + if merged: + setattr(response, field_name, merged) + response._hidden_params[field_name] = merged + @staticmethod def _convert_grounding_metadata_to_annotations( grounding_metadata: List[dict], @@ -3385,10 +3456,13 @@ def _apply_stream_candidates( if choice.finish_reason == "stop": choice.finish_reason = "tool_calls" - setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore - setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore - setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore - setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore + VertexGeminiConfig._set_stream_metadata_on_response( + model_response, + grounding_metadata, + url_context_metadata, + safety_ratings, + citation_metadata, + ) return ( grounding_metadata, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 1e92754857b..8a92e7ec4a5 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -17,6 +17,9 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, VertexBase): + def should_strip_billing_metadata(self) -> bool: + return True + def validate_anthropic_messages_environment( self, headers: dict, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index c852909d475..ae8bdc55443 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -52,6 +52,9 @@ class VertexAIAnthropicConfig(AnthropicConfig): def custom_llm_provider(self) -> Optional[str]: return "vertex_ai" + def should_strip_billing_metadata(self) -> bool: + return True + def _add_context_management_beta_headers( self, beta_set: set, context_management: dict ) -> None: diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index c06928516ef..8019bb67991 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -5,6 +5,7 @@ import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE +from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, strip_name_from_messages, @@ -39,6 +40,72 @@ def _get_openai_compatible_provider_info( dynamic_api_key = XAIModelInfo.get_api_key(api_key) return api_base, dynamic_api_key + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + from litellm.llms.xai.oauth import ( + XAIOAuthAuthenticator, + XAIOAuthError, + should_use_xai_oauth, + ) + + dynamic_api_key = XAIModelInfo.get_api_key(api_key) + if should_use_xai_oauth(litellm_params) and not dynamic_api_key: + try: + headers["Authorization"] = ( + f"Bearer {XAIOAuthAuthenticator().get_access_token()}" + ) + except XAIOAuthError as exc: + raise AuthenticationError( + model=model, + llm_provider=self.custom_llm_provider or "xai", + message=str(exc), + ) from exc + if "content-type" not in headers and "Content-Type" not in headers: + headers["Content-Type"] = "application/json" + return headers + + return super().validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=dynamic_api_key, + api_base=api_base, + ) + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + from litellm.llms.xai.oauth import XAIOAuthAuthenticator, should_use_xai_oauth + + dynamic_api_key = XAIModelInfo.get_api_key(api_key) + if should_use_xai_oauth(litellm_params) and not dynamic_api_key: + api_base = XAIOAuthAuthenticator().get_api_base() + + return super().get_complete_url( + api_base=api_base, + api_key=dynamic_api_key, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + stream=stream, + ) + def get_supported_openai_params(self, model: str) -> list: base_openai_params = [ "logit_bias", diff --git a/litellm/llms/xai/oauth.py b/litellm/llms/xai/oauth.py new file mode 100644 index 00000000000..30c717b7ca0 --- /dev/null +++ b/litellm/llms/xai/oauth.py @@ -0,0 +1,421 @@ +import base64 +import hashlib +import json +import os +import secrets +import sys +import threading +import time +import uuid +import webbrowser +from http.server import BaseHTTPRequestHandler, HTTPServer +from typing import Any, Dict, Optional, Tuple, Union +from urllib.parse import parse_qs, urlencode, urlparse + +import httpx + +from litellm._logging import verbose_logger +from litellm.constants import XAI_API_BASE +from litellm.llms.custom_httpx.http_handler import HTTPHandler, _get_httpx_client +from litellm.secret_managers.main import get_secret_str + +XAI_OAUTH_ISSUER = "https://auth.x.ai" +XAI_OAUTH_DISCOVERY_URL = f"{XAI_OAUTH_ISSUER}/.well-known/openid-configuration" +XAI_OAUTH_CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828" +XAI_OAUTH_SCOPE = "openid profile email offline_access grok-cli:access api:access" +XAI_OAUTH_REDIRECT_HOST = "127.0.0.1" +XAI_OAUTH_REDIRECT_PORT = 56121 +XAI_OAUTH_REDIRECT_PATH = "/callback" +XAI_OAUTH_EXPIRY_SKEW_SECONDS = 120 +XAI_OAUTH_CALLBACK_TIMEOUT_SECONDS = 180 +_XAI_OAUTH_REFRESH_LOCK = threading.Lock() + + +class XAIOAuthError(Exception): + pass + + +class XAIOAuthLoginRequiredError(XAIOAuthError): + pass + + +class _CallbackHandler(BaseHTTPRequestHandler): + server: "_CallbackServer" + + def do_GET(self) -> None: + parsed = urlparse(self.path) + if parsed.path != XAI_OAUTH_REDIRECT_PATH: + self.send_response(404) + self.end_headers() + return + + params = parse_qs(parsed.query) + result = { + "code": params.get("code", [None])[0], + "state": params.get("state", [None])[0], + "error": params.get("error", [None])[0], + "error_description": params.get("error_description", [None])[0], + } + self.server.callback_result = result + + if result["state"] != self.server.expected_state: + self.send_response(400) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.end_headers() + self.wfile.write( + b"

xAI authorization state mismatch.

" + ) + return + + self.send_response(200) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.end_headers() + body = ( + b"

xAI authorization failed.

You can close this tab." + if result["error"] + else b"

xAI authorization received.

You can close this tab." + ) + self.wfile.write(body) + + def log_message(self, format: str, *args: Any) -> None: + return + + +class _CallbackServer(HTTPServer): + expected_state: str + callback_result: Optional[Dict[str, Optional[str]]] + + +class XAIOAuthAuthenticator: + def __init__( + self, http_client: Optional[Union[httpx.Client, HTTPHandler]] = None + ) -> None: + self.token_dir = get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser( + "~/.config/litellm/xai_oauth" + ) + self.auth_file = os.path.join( + self.token_dir, get_secret_str("XAI_OAUTH_AUTH_FILE") or "auth.json" + ) + self.http_client = http_client + + def get_api_base(self) -> str: + return ( + get_secret_str("XAI_OAUTH_API_BASE") + or get_secret_str("XAI_API_BASE") + or XAI_API_BASE + ) + + def get_access_token(self) -> str: + auth_data = self._read_auth_file() + if not auth_data: + raise XAIOAuthLoginRequiredError( + "xAI OAuth login required. Run `litellm xai-oauth login`." + ) + + access_token = auth_data.get("access_token") + if access_token and not self._is_expired(auth_data): + return access_token + + refresh_token = auth_data.get("refresh_token") + if not refresh_token: + raise XAIOAuthLoginRequiredError( + "xAI OAuth refresh token missing. Run `litellm xai-oauth login`." + ) + + with _XAI_OAUTH_REFRESH_LOCK: + locked_auth_data = self._read_auth_file() or auth_data + access_token = locked_auth_data.get("access_token") + if access_token and not self._is_expired(locked_auth_data): + return access_token + + refreshed = self._refresh_tokens(locked_auth_data) + return refreshed["access_token"] + + def login(self, force: bool = False, no_browser: bool = False) -> Dict[str, Any]: + existing = self._read_auth_file() + if existing and not force and existing.get("access_token"): + if not self._is_expired(existing): + return existing + if existing.get("refresh_token"): + try: + return self._refresh_tokens(existing) + except XAIOAuthError: + pass + + discovery = self._discover() + verifier, challenge = self._pkce_pair() + state = uuid.uuid4().hex + nonce = uuid.uuid4().hex + server, redirect_uri = self._start_callback_server(state) + authorize_url = self._build_authorize_url( + authorization_endpoint=discovery["authorization_endpoint"], + redirect_uri=redirect_uri, + challenge=challenge, + state=state, + nonce=nonce, + ) + + if no_browser or not webbrowser.open(authorize_url): + sys.stdout.write( + f"Open this URL to authenticate with xAI:\n{authorize_url}\n" + ) + sys.stdout.flush() + + result = self._wait_for_callback(server) + if result.get("state") != state: + raise XAIOAuthError("xAI OAuth state mismatch") + if result.get("error"): + description = result.get("error_description") or result["error"] + raise XAIOAuthError(f"xAI authorization failed: {description}") + code = result.get("code") + if not code: + raise XAIOAuthError("xAI authorization failed: no code returned") + + token_payload = self._exchange_token( + discovery["token_endpoint"], + { + "grant_type": "authorization_code", + "code": code, + "redirect_uri": redirect_uri, + "client_id": XAI_OAUTH_CLIENT_ID, + "code_verifier": verifier, + }, + ) + auth_data = self._build_auth_record(token_payload, discovery["token_endpoint"]) + self._write_auth_file(auth_data) + return auth_data + + def _client(self) -> Union[httpx.Client, HTTPHandler]: + return self.http_client or _get_httpx_client() + + def _ensure_token_dir(self) -> None: + os.makedirs(self.token_dir, mode=0o700, exist_ok=True) + try: + os.chmod(self.token_dir, 0o700) + except OSError: + verbose_logger.debug("Could not chmod xAI OAuth token directory") + + def _read_auth_file(self) -> Optional[Dict[str, Any]]: + try: + with open(self.auth_file, "r") as f: + data = json.load(f) + return data if isinstance(data, dict) else None + except (IOError, json.JSONDecodeError): + return None + + def _write_auth_file(self, data: Dict[str, Any]) -> None: + self._ensure_token_dir() + tmp_file = os.path.join( + self.token_dir, + f".{os.path.basename(self.auth_file)}.{uuid.uuid4().hex}.tmp", + ) + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL + if hasattr(os, "O_NOFOLLOW"): + flags |= os.O_NOFOLLOW + fd = os.open(tmp_file, flags, 0o600) + try: + with os.fdopen(fd, "w") as f: + json.dump(data, f) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_file, self.auth_file) + try: + os.chmod(self.auth_file, 0o600) + except OSError: + verbose_logger.debug("Could not chmod xAI OAuth auth file") + except Exception: + try: + os.close(fd) + except OSError: + pass + try: + os.unlink(tmp_file) + except OSError: + pass + raise + + def _is_expired(self, auth_data: Dict[str, Any]) -> bool: + expires_at = auth_data.get("expires_at") + if expires_at is None: + return True + try: + return time.time() >= float(expires_at) - XAI_OAUTH_EXPIRY_SKEW_SECONDS + except (TypeError, ValueError): + return True + + def _discover(self) -> Dict[str, str]: + try: + response = self._client().get( + XAI_OAUTH_DISCOVERY_URL, headers={"Accept": "application/json"} + ) + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise XAIOAuthError( + f"xAI OAuth discovery request failed: {exc.response.status_code} {exc.response.text}" + ) from exc + try: + data = response.json() + except ValueError as exc: + raise XAIOAuthError( + "xAI OAuth discovery response was not valid JSON" + ) from exc + authorization_endpoint = data.get("authorization_endpoint") + token_endpoint = data.get("token_endpoint") + if not authorization_endpoint or not token_endpoint: + raise XAIOAuthError("xAI OAuth discovery missing endpoints") + return { + "authorization_endpoint": self._validate_xai_endpoint( + authorization_endpoint + ), + "token_endpoint": self._validate_xai_endpoint(token_endpoint), + } + + def _validate_xai_endpoint(self, url: str) -> str: + parsed = urlparse(url) + host = (parsed.hostname or "").lower() + if parsed.scheme != "https" or (host != "x.ai" and not host.endswith(".x.ai")): + raise XAIOAuthError( + f"xAI OAuth discovery returned unexpected endpoint: {url}" + ) + return url + + def _pkce_pair(self) -> Tuple[str, str]: + verifier = ( + base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode() + ) + challenge = ( + base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()) + .rstrip(b"=") + .decode() + ) + return verifier, challenge + + def _start_callback_server(self, state: str) -> Tuple[_CallbackServer, str]: + last_error: Optional[OSError] = None + for port in (XAI_OAUTH_REDIRECT_PORT, 0): + try: + server = _CallbackServer( + (XAI_OAUTH_REDIRECT_HOST, port), _CallbackHandler + ) + server.expected_state = state + server.callback_result = None + actual_port = server.server_address[1] + redirect_uri = f"http://{XAI_OAUTH_REDIRECT_HOST}:{actual_port}{XAI_OAUTH_REDIRECT_PATH}" + return server, redirect_uri + except OSError as exc: + last_error = exc + raise XAIOAuthError(f"Could not start xAI OAuth callback server: {last_error}") + + def _build_authorize_url( + self, + authorization_endpoint: str, + redirect_uri: str, + challenge: str, + state: str, + nonce: str, + ) -> str: + params = { + "response_type": "code", + "client_id": XAI_OAUTH_CLIENT_ID, + "redirect_uri": redirect_uri, + "scope": XAI_OAUTH_SCOPE, + "code_challenge": challenge, + "code_challenge_method": "S256", + "state": state, + "nonce": nonce, + } + return f"{authorization_endpoint}?{urlencode(params)}" + + def _wait_for_callback(self, server: _CallbackServer) -> Dict[str, Optional[str]]: + server.timeout = 1 + deadline = time.time() + XAI_OAUTH_CALLBACK_TIMEOUT_SECONDS + try: + while time.time() < deadline: + server.handle_request() + if server.callback_result is not None: + return server.callback_result + finally: + server.server_close() + raise XAIOAuthError("Timed out waiting for xAI OAuth callback") + + def _exchange_token( + self, token_endpoint: str, data: Dict[str, str] + ) -> Dict[str, Any]: + try: + response = self._client().post( + token_endpoint, + headers={ + "Accept": "application/json", + "Content-Type": "application/x-www-form-urlencoded", + }, + data=data, + ) + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise XAIOAuthError( + f"xAI OAuth token request failed: {exc.response.status_code} {exc.response.text}" + ) from exc + try: + body = response.json() + except ValueError as exc: + raise XAIOAuthError("xAI OAuth token response was not valid JSON") from exc + if not isinstance(body, dict): + raise XAIOAuthError("xAI OAuth token response was not an object") + return body + + def _build_auth_record( + self, + token_payload: Dict[str, Any], + token_endpoint: str, + fallback_refresh_token: Optional[str] = None, + ) -> Dict[str, Any]: + access_token = token_payload.get("access_token") + refresh_token = token_payload.get("refresh_token") or fallback_refresh_token + if not access_token: + raise XAIOAuthError("xAI OAuth token response missing access_token") + if not refresh_token: + raise XAIOAuthError("xAI OAuth token response missing refresh_token") + expires_in = token_payload.get("expires_in") or 3600 + try: + expires_at = int(time.time() + int(expires_in)) + except (TypeError, ValueError): + expires_at = int(time.time() + 3600) + return { + "access_token": access_token, + "refresh_token": refresh_token, + "id_token": token_payload.get("id_token"), + "token_type": token_payload.get("token_type") or "Bearer", + "token_endpoint": token_endpoint, + "expires_at": expires_at, + } + + def _refresh_tokens(self, auth_data: Dict[str, Any]) -> Dict[str, Any]: + token_endpoint = auth_data.get("token_endpoint") + if not token_endpoint: + token_endpoint = self._discover()["token_endpoint"] + token_endpoint = self._validate_xai_endpoint(token_endpoint) + refresh_token = auth_data.get("refresh_token") + if not refresh_token: + raise XAIOAuthLoginRequiredError( + "xAI OAuth refresh token missing. Run `litellm xai-oauth login`." + ) + + token_payload = self._exchange_token( + token_endpoint, + { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": XAI_OAUTH_CLIENT_ID, + }, + ) + refreshed = self._build_auth_record( + token_payload, + token_endpoint, + fallback_refresh_token=refresh_token, + ) + self._write_auth_file(refreshed) + return refreshed + + +def should_use_xai_oauth(litellm_params: Optional[Dict[str, Any]]) -> bool: + return bool((litellm_params or {}).get("use_xai_oauth")) diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 55805ddaede..f81e860a8ce 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -3,6 +3,7 @@ import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE +from litellm.exceptions import AuthenticationError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.llms.xai.common_utils import XAIModelInfo from litellm.secret_managers.main import get_secret_str @@ -220,10 +221,27 @@ def validate_environment( litellm_params.api_key, legacy_generic_before_env=True ) + if not api_key: + from litellm.llms.xai.oauth import ( + XAIOAuthAuthenticator, + XAIOAuthError, + should_use_xai_oauth, + ) + + if should_use_xai_oauth(litellm_params.model_dump()): + try: + api_key = XAIOAuthAuthenticator().get_access_token() + except XAIOAuthError as exc: + raise AuthenticationError( + model=model, + llm_provider=self.custom_llm_provider.value, + message=str(exc), + ) from exc + if not api_key: raise ValueError( "XAI API key is required. Set api_key, litellm.xai_key, " - "litellm.api_key, or XAI_API_KEY." + "litellm.api_key, XAI_API_KEY, or use_xai_oauth=True." ) headers.update( @@ -244,12 +262,20 @@ def get_complete_url( Returns: str: The full URL for the XAI /responses endpoint """ - api_base = ( - api_base - or litellm.api_base - or get_secret_str("XAI_API_BASE") - or XAI_API_BASE + from litellm.llms.xai.oauth import XAIOAuthAuthenticator, should_use_xai_oauth + + api_key = XAIModelInfo.get_api_key( + litellm_params.get("api_key"), legacy_generic_before_env=True ) + if should_use_xai_oauth(litellm_params) and not api_key: + api_base = XAIOAuthAuthenticator().get_api_base() + else: + api_base = ( + api_base + or litellm.api_base + or get_secret_str("XAI_API_BASE") + or XAI_API_BASE + ) # Remove trailing slashes api_base = api_base.rstrip("/") diff --git a/litellm/main.py b/litellm/main.py index 64891e2def9..02609217ddb 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1638,6 +1638,8 @@ def completion( # type: ignore # noqa: PLR0915 litellm_request_debug=kwargs.get("litellm_request_debug", False), tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), + use_xai_oauth=kwargs.get("use_xai_oauth", False), + aws_bedrock_project_id=kwargs.get("aws_bedrock_project_id"), ) cast(LiteLLMLoggingObj, logging).update_environment_variables( model=model, @@ -2134,9 +2136,6 @@ def completion( # type: ignore # noqa: PLR0915 headers = headers or litellm.headers - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - ## LOAD CONFIG - if set config = litellm.OpenAITextCompletionConfig.get_config() for k, v in config.items(): @@ -2162,6 +2161,7 @@ def completion( # type: ignore # noqa: PLR0915 _response = openai_text_completions.completion( model=model, messages=messages, + headers=headers, model_response=model_response, print_verbose=print_verbose, api_key=api_key, @@ -7761,6 +7761,9 @@ def stream_chunk_builder( # noqa: PLR0915 "cost", logging_obj._response_cost_calculator(result=response), ) + processor.apply_provider_assembled_streaming_metadata( + response, chunks, logging_obj + ) return response tool_call_chunks = [ @@ -7940,6 +7943,9 @@ def stream_chunk_builder( # noqa: PLR0915 usage, "cost", logging_obj._response_cost_calculator(result=response) ) + processor.apply_provider_assembled_streaming_metadata( + response, chunks, logging_obj + ) return response except Exception as e: verbose_logger.exception( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 397f96fdb1e..aab0e4264d0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1156,6 +1156,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1202,6 +1203,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1233,6 +1235,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1264,6 +1267,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1295,6 +1299,139 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "global.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "us.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "eu.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1327,6 +1464,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1359,6 +1497,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1391,6 +1530,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1423,6 +1563,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1455,6 +1596,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1485,6 +1627,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2208,6 +2351,37 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "azure_ai/claude-fable-5": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2237,6 +2411,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -6889,6 +7064,43 @@ "/v1/images/generations" ] }, + "azure_ai/MAI-Image-2.5": { + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.05, + "output_cost_per_image_token": 4.7e-05, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/new-mai-models-in-microsoft-foundry-across-text-image-voice-and-speech/4524632", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "azure_ai/MAI-Image-2.5-Flash": { + "input_cost_per_image_token": 1.75e-06, + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.0338, + "output_cost_per_image_token": 3.3e-05, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/new-mai-models-in-microsoft-foundry-across-text-image-voice-and-speech/4524632", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "azure_ai/MAI-Image-2e": { + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.02, + "output_cost_per_image_token": 1.95e-05, + "source": "https://aka.ms/mai-image-2e-foundryblog", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "azure_ai/Llama-3.2-11B-Vision-Instruct": { "input_cost_per_token": 3.7e-07, "litellm_provider": "azure_ai", @@ -10133,6 +10345,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10167,6 +10380,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10177,6 +10391,40 @@ }, "supports_output_config": true }, + "claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true + }, "claude-opus-4-8": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -10201,6 +10449,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -14286,10 +14535,10 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://fireworks.ai/models/fireworks/glm-5p1", - "supports_function_calling": false, + "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": false, - "supports_tool_choice": false + "supports_response_schema": true, + "supports_tool_choice": true }, "fireworks_ai/accounts/fireworks/models/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, @@ -14567,10 +14816,10 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://fireworks.ai/models/fireworks/glm-5p1", - "supports_function_calling": false, + "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": false, - "supports_tool_choice": false + "supports_response_schema": true, + "supports_tool_choice": true }, "fireworks_ai/kimi-k2p5": { "cache_read_input_token_cost": 1e-07, @@ -24143,9 +24392,12 @@ "max_output_tokens": 8192 }, "minimax/MiniMax-M3": { - "input_cost_per_token": 6e-07, - "output_cost_per_token": 2.4e-06, - "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 3e-07, + "input_cost_per_token_above_512k_tokens": 6e-07, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_512k_tokens": 2.4e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_512k_tokens": 1.2e-07, "litellm_provider": "minimax", "mode": "chat", "supports_function_calling": true, @@ -24154,7 +24406,7 @@ "supports_reasoning": true, "supports_system_messages": true, "supports_vision": true, - "max_input_tokens": 512000, + "max_input_tokens": 1000000, "max_output_tokens": 128000 }, "mistral.devstral-2-123b": { @@ -33967,6 +34219,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -33995,6 +34248,67 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5@default": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34024,6 +34338,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34053,6 +34368,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -41333,6 +41649,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "use_openai_responses_path": true, "supported_endpoints": ["/v1/responses"], "supported_modalities": ["text", "image"], "supported_output_modalities": ["text"], @@ -41352,6 +41669,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "use_openai_responses_path": true, "supported_endpoints": ["/v1/responses"], "supported_modalities": ["text", "image"], "supported_output_modalities": ["text"], @@ -41653,5 +41971,164 @@ "/v1/audio/transcriptions" ], "supports_audio_input": true + }, + "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 6e-07, + "output_cost_per_token": 3.6e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 4.5e-07, + "output_cost_per_token": 1.8e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3.6-27B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3.2e-07, + "output_cost_per_token": 3.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 202752, + "max_output_tokens": 202752, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/deepseek-ai/DeepSeek-V4-Flash": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/moonshotai/Kimi-K2.6": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 9.6e-07, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/MiniMaxAI/MiniMax-M2.5": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 196608, + "max_output_tokens": 196608, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/google/gemma-4-31B-it": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 5.6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-120b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-20b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" } } \ No newline at end of file diff --git a/litellm/models/__init__.py b/litellm/models/__init__.py new file mode 100644 index 00000000000..7e2d2c0ed9d --- /dev/null +++ b/litellm/models/__init__.py @@ -0,0 +1,66 @@ +""" +Domain models for LiteLLM backend. +""" + +from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.models.budget import ( + LiteLLM_BudgetTable, + LiteLLM_BudgetTableFull, + LiteLLM_TeamMemberTable, +) +from litellm.models.config import LiteLLM_Config +from litellm.models.credentials import ( + CreateCredentialItem, + CredentialBase, + CredentialItem, +) +from litellm.models.end_user import LiteLLM_EndUserTable +from litellm.models.managed_files import ( + LiteLLM_ManagedFileTable, + LiteLLM_ManagedObjectTable, + LiteLLM_ManagedVectorStoresTable, + LiteLLM_ManagedVectorStoreTable, +) +from litellm.models.mcp_server import LiteLLM_MCPServerTable +from litellm.models.model import LiteLLM_ProxyModelTable +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.models.organization import LiteLLM_OrganizationTable +from litellm.models.organization_membership import LiteLLM_OrganizationMembershipTable +from litellm.models.project import LiteLLM_ProjectTable +from litellm.models.skills import LiteLLM_SkillsTable +from litellm.models.spend_logs import LiteLLM_ErrorLogs, LiteLLM_SpendLogs +from litellm.models.tag import LiteLLM_TagTable +from litellm.models.team import LiteLLM_TeamTable +from litellm.models.team_membership import LiteLLM_TeamMembership +from litellm.models.user import LiteLLM_UserTable +from litellm.models.verification_token import LiteLLM_VerificationToken + +__all__ = [ + "LiteLLM_AccessGroupTable", + "LiteLLM_BudgetTable", + "LiteLLM_BudgetTableFull", + "LiteLLM_TeamMemberTable", + "LiteLLM_Config", + "CredentialBase", + "CredentialItem", + "CreateCredentialItem", + "LiteLLM_EndUserTable", + "LiteLLM_ManagedFileTable", + "LiteLLM_ManagedObjectTable", + "LiteLLM_ManagedVectorStoreTable", + "LiteLLM_ManagedVectorStoresTable", + "LiteLLM_MCPServerTable", + "LiteLLM_ProxyModelTable", + "LiteLLM_ObjectPermissionTable", + "LiteLLM_OrganizationTable", + "LiteLLM_OrganizationMembershipTable", + "LiteLLM_ProjectTable", + "LiteLLM_SkillsTable", + "LiteLLM_ErrorLogs", + "LiteLLM_SpendLogs", + "LiteLLM_TagTable", + "LiteLLM_TeamTable", + "LiteLLM_TeamMembership", + "LiteLLM_UserTable", + "LiteLLM_VerificationToken", +] diff --git a/litellm/models/access_group.py b/litellm/models/access_group.py new file mode 100644 index 00000000000..682e779e531 --- /dev/null +++ b/litellm/models/access_group.py @@ -0,0 +1,26 @@ +""" +Access group table model. + +Canonical definition for ``litellm_accessgrouptable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import List, Optional + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_AccessGroupTable(LiteLLMPydanticObjectBase): + access_group_id: str + access_group_name: str + description: Optional[str] = None + access_model_names: List[str] = [] + access_mcp_server_ids: List[str] = [] + access_agent_ids: List[str] = [] + assigned_team_ids: List[str] = [] + assigned_key_ids: List[str] = [] + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None diff --git a/litellm/models/base.py b/litellm/models/base.py new file mode 100644 index 00000000000..01981297bd5 --- /dev/null +++ b/litellm/models/base.py @@ -0,0 +1,38 @@ +""" +Base model class for domain models. +""" + +from datetime import datetime +from typing import Any, Dict, Optional + +from pydantic import BaseModel, ConfigDict + + +class DomainModel(BaseModel): + """Base class for all domain models.""" + + model_config = ConfigDict( + from_attributes=True, + protected_namespaces=(), + extra="ignore", + ) + + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + + @classmethod + def from_db_record(cls, record: Any) -> "DomainModel": + """Create a domain model from a database record.""" + if record is None: + raise ValueError("Cannot create domain model from None record") + if isinstance(record, dict): + return cls(**record) + if hasattr(record, "model_dump") and callable(record.model_dump): + return cls(**record.model_dump()) + if hasattr(record, "dict") and callable(record.dict): + return cls(**record.dict()) + return cls(**dict(record)) + + def to_db_dict(self, exclude_unset: bool = False) -> Dict[str, Any]: + """Convert domain model to a dictionary for database operations.""" + return self.model_dump(exclude_none=True, exclude_unset=exclude_unset) diff --git a/litellm/models/budget.py b/litellm/models/budget.py new file mode 100644 index 00000000000..e7dfe2f8fbc --- /dev/null +++ b/litellm/models/budget.py @@ -0,0 +1,56 @@ +""" +Budget table model. + +Canonical definition for ``litellm_budgettable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import List, Optional + +from pydantic import ConfigDict + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): + """Represents user-controllable params for a LiteLLM_BudgetTable record. + + Budget-write paths use `model_fields.keys()` on this class as an allowlist + for user input. Keep server-managed fields (e.g. `budget_reset_at`) on + `LiteLLM_BudgetTableFull` so they aren't user-settable. + """ + + budget_id: Optional[str] = None + soft_budget: Optional[float] = None + max_budget: Optional[float] = None + max_parallel_requests: Optional[int] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + model_max_budget: Optional[dict] = None + budget_duration: Optional[str] = None + allowed_models: Optional[List[str]] = ( + None # per-member model scope; empty = inherit team models + ) + + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable): + """LiteLLM_BudgetTable + server-managed fields returned on API responses.""" + + budget_reset_at: Optional[datetime] = None + created_at: datetime + + +class LiteLLM_TeamMemberTable(LiteLLM_BudgetTable): + """ + Used to track spend of a user_id within a team_id + """ + + spend: Optional[float] = None + user_id: Optional[str] = None + team_id: Optional[str] = None + budget_id: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/models/config.py b/litellm/models/config.py new file mode 100644 index 00000000000..99b5c5692fd --- /dev/null +++ b/litellm/models/config.py @@ -0,0 +1,15 @@ +""" +Config table model. + +Canonical definition for ``litellm_config``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from typing import Dict + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_Config(LiteLLMPydanticObjectBase): + param_name: str + param_value: Dict diff --git a/litellm/models/credentials.py b/litellm/models/credentials.py new file mode 100644 index 00000000000..b74ea055d21 --- /dev/null +++ b/litellm/models/credentials.py @@ -0,0 +1,31 @@ +""" +Credential table models. + +These are the canonical credential types for the proxy. They live in the model +layer; ``litellm.types.utils`` re-exports them for backwards compatibility. +""" + +from typing import Optional + +from pydantic import BaseModel, model_validator + + +class CredentialBase(BaseModel): + credential_name: str + credential_info: dict + + +class CredentialItem(CredentialBase): + credential_values: dict + + +class CreateCredentialItem(CredentialBase): + credential_values: Optional[dict] = None + model_id: Optional[str] = None + + @model_validator(mode="before") + @classmethod + def check_credential_params(cls, values): + if not values.get("credential_values") and not values.get("model_id"): + raise ValueError("Either credential_values or model_id must be set") + return values diff --git a/litellm/models/end_user.py b/litellm/models/end_user.py new file mode 100644 index 00000000000..15fd03ec2ca --- /dev/null +++ b/litellm/models/end_user.py @@ -0,0 +1,35 @@ +""" +End-user table model. + +Canonical definition for ``litellm_endusertable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from typing import Literal, Optional + +from pydantic import ConfigDict, model_validator + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase): + user_id: str + blocked: bool + alias: Optional[str] = None + spend: float = 0.0 + allowed_model_region: Optional[Literal["eu", "us"]] = None + default_model: Optional[str] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + object_permission_id: Optional[str] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("spend") is None: + values.update({"spend": 0.0}) + return values + + model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/models/managed_files.py b/litellm/models/managed_files.py new file mode 100644 index 00000000000..24154768860 --- /dev/null +++ b/litellm/models/managed_files.py @@ -0,0 +1,62 @@ +""" +Managed file, object, and vector store table models. + +Canonical definitions for the ``litellm_managed*`` tables. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Any, Dict, List, Literal, Optional, Union + +from litellm.types.llms.base import LiteLLMPydanticObjectBase +from litellm.types.llms.openai import OpenAIFileObject, ResponsesAPIResponse +from litellm.types.utils import LiteLLMBatch, LiteLLMFineTuningJob + + +class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): + unified_file_id: str + file_object: Optional[OpenAIFileObject] = None + model_mappings: Dict[str, str] + flat_model_file_ids: List[str] + created_by: Optional[str] = None + team_id: Optional[str] = None + updated_by: Optional[str] = None + storage_backend: Optional[str] = None + storage_url: Optional[str] = None + + +class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): + unified_object_id: str + model_object_id: str + file_purpose: Literal["batch", "fine-tune", "response", "container"] + file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse] + created_by: Optional[str] = None + team_id: Optional[str] = None + + +class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): + """Table for managing vector stores with target_model_names support.""" + + unified_resource_id: str + resource_object: Optional[Any] = None + model_mappings: Dict[str, str] + flat_model_resource_ids: List[str] + created_by: Optional[str] = None + team_id: Optional[str] = None + updated_by: Optional[str] = None + storage_backend: Optional[str] = None + storage_url: Optional[str] = None + + +class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase): + vector_store_id: str + custom_llm_provider: str + vector_store_name: Optional[str] + vector_store_description: Optional[str] + vector_store_metadata: Optional[Dict[str, Any]] + created_at: Optional[datetime] + updated_at: Optional[datetime] + litellm_credential_name: Optional[str] + litellm_params: Optional[Dict[str, Any]] + team_id: Optional[str] + user_id: Optional[str] diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py new file mode 100644 index 00000000000..3d03eff6df8 --- /dev/null +++ b/litellm/models/mcp_server.py @@ -0,0 +1,103 @@ +""" +MCP server table model. + +Canonical definition for ``litellm_mcpservertable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +import enum +from datetime import datetime +from typing import Dict, List, Literal, Optional + +from pydantic import Field + +from litellm.types.llms.base import LiteLLMPydanticObjectBase +from litellm.types.mcp import MCPAuthType, MCPCredentials, MCPTransportType +from litellm.types.mcp_server.mcp_server_manager import MCPInfo + + +class MCPEnvVarScope(str, enum.Enum): + """Scope for an MCP server environment variable. + + - ``global``: value is provided by the admin and used for all users. + - ``user``: each user must provide their own value via the per-user + env-var endpoint. The admin-supplied ``value`` is treated as a + placeholder/hint and is not used at request time. + """ + + global_ = "global" + user = "user" + + +class MCPEnvVar(LiteLLMPydanticObjectBase): + """One environment variable for an MCP server. + + Variables can be interpolated into ``static_headers`` using ``${NAME}`` + syntax. ``scope=global`` values are stored on the server. ``scope=user`` + values are stored per-user in ``LiteLLM_MCPUserEnvVars`` and supplied by + each user. + """ + + name: str + value: str = "" + scope: MCPEnvVarScope = MCPEnvVarScope.global_ + description: Optional[str] = None + + +class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): + """Represents a LiteLLM_MCPServerTable record""" + + server_id: str + server_name: Optional[str] = None + alias: Optional[str] = None + description: Optional[str] = None + url: Optional[str] = None + spec_path: Optional[str] = None + transport: MCPTransportType + auth_type: Optional[MCPAuthType] = None + credentials: Optional[MCPCredentials] = None + instructions: Optional[str] = None + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None + teams: List[Dict[str, Optional[str]]] = Field(default_factory=list) + mcp_access_groups: List[str] = Field(default_factory=list) + allowed_tools: List[str] = Field(default_factory=list) + tool_name_to_display_name: Optional[Dict[str, str]] = None + tool_name_to_description: Optional[Dict[str, str]] = None + extra_headers: List[str] = Field(default_factory=list) + mcp_info: Optional[MCPInfo] = None + static_headers: Optional[Dict[str, str]] = None + env_vars: Optional[List[MCPEnvVar]] = None + status: Optional[Literal["healthy", "unhealthy", "unknown"]] = Field( + default="unknown", + description="Health status: 'healthy', 'unhealthy', 'unknown'", + ) + last_health_check: Optional[datetime] = None + health_check_error: Optional[str] = None + command: Optional[str] = None + args: List[str] = Field(default_factory=list) + env: Dict[str, str] = Field(default_factory=dict) + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None + oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None + allow_all_keys: bool = False + available_on_public_internet: bool = True + delegate_auth_to_upstream: bool = False + oauth_passthrough: bool = False + is_byok: bool = False + byok_description: List[str] = Field(default_factory=list) + byok_api_key_help_url: Optional[str] = None + has_user_credential: Optional[bool] = None + source_url: Optional[str] = None + timeout: Optional[float] = None + approval_status: Optional[str] = Field( + default="active", + description="Approval status: 'pending_review', 'active', 'rejected'", + ) + submitted_by: Optional[str] = None + submitted_at: Optional[datetime] = None + reviewed_at: Optional[datetime] = None + review_notes: Optional[str] = None diff --git a/litellm/models/model.py b/litellm/models/model.py new file mode 100644 index 00000000000..7657e4d30f8 --- /dev/null +++ b/litellm/models/model.py @@ -0,0 +1,59 @@ +""" +Proxy model table model. + +Canonical definition for ``litellm_proxymodeltable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +import json +from datetime import datetime +from typing import Optional + +from pydantic import ConfigDict, model_validator + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase): + model_id: str + model_name: str + litellm_params: dict + model_info: Optional[dict] = None + blocked: bool = False + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="before") + @classmethod + def check_potential_json_str(cls, values): + if isinstance(values.get("litellm_params"), str): + try: + values["litellm_params"] = json.loads(values["litellm_params"]) + except json.JSONDecodeError: + pass + if isinstance(values.get("model_info"), str): + try: + values["model_info"] = json.loads(values["model_info"]) + except json.JSONDecodeError: + pass + return values + + @property + def is_blocked(self) -> bool: + return self.blocked + + @property + def team_id(self) -> Optional[str]: + if self.model_info: + return self.model_info.get("team_id") + return None + + @property + def team_public_model_name(self) -> Optional[str]: + if self.model_info: + return self.model_info.get("team_public_model_name") + return None diff --git a/litellm/models/object_permission.py b/litellm/models/object_permission.py new file mode 100644 index 00000000000..6c0d100046c --- /dev/null +++ b/litellm/models/object_permission.py @@ -0,0 +1,26 @@ +""" +Object permission table model. + +Canonical definition for ``litellm_objectpermissiontable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from typing import Dict, List, Optional + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): + """Represents a LiteLLM_ObjectPermissionTable record""" + + object_permission_id: str + mcp_servers: Optional[List[str]] = [] + mcp_access_groups: Optional[List[str]] = [] + mcp_tool_permissions: Optional[Dict[str, List[str]]] = None + vector_stores: Optional[List[str]] = [] + agents: Optional[List[str]] = [] + agent_access_groups: Optional[List[str]] = [] + models: Optional[List[str]] = [] + mcp_toolsets: Optional[List[str]] = None + blocked_tools: Optional[List[str]] = [] + search_tools: Optional[List[str]] = [] diff --git a/litellm/models/organization.py b/litellm/models/organization.py new file mode 100644 index 00000000000..8b2d95c3e09 --- /dev/null +++ b/litellm/models/organization.py @@ -0,0 +1,31 @@ +""" +Organization table model. + +Canonical definition for ``litellm_organizationtable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from typing import List, Optional + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.models.user import LiteLLM_UserTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_OrganizationTable(LiteLLMPydanticObjectBase): + """Represents user-controllable params for a LiteLLM_OrganizationTable record""" + + organization_id: Optional[str] = None + organization_alias: Optional[str] = None + budget_id: str + spend: float = 0.0 + metadata: Optional[dict] = None + models: List[str] = [] + model_spend: Optional[dict] = {} + created_by: str + updated_by: str + users: Optional[List[LiteLLM_UserTable]] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + object_permission_id: Optional[str] = None diff --git a/litellm/models/organization_membership.py b/litellm/models/organization_membership.py new file mode 100644 index 00000000000..9957c0c21af --- /dev/null +++ b/litellm/models/organization_membership.py @@ -0,0 +1,40 @@ +""" +Organization membership table model. + +Canonical definition for ``litellm_organizationmembership``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Any, Optional + +from pydantic import ConfigDict, model_validator + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): + """Tracks which organizations a user belongs to and their spend within it.""" + + user_id: str + organization_id: str + user_role: Optional[str] = None + spend: float = 0.0 + budget_id: Optional[str] = None + created_at: datetime + updated_at: datetime + user: Optional[Any] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + user_email: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="after") + def populate_user_email(self) -> "LiteLLM_OrganizationMembershipTable": + if self.user_email is None and self.user is not None: + if isinstance(self.user, dict): + self.user_email = self.user.get("user_email") + else: + self.user_email = getattr(self.user, "user_email", None) + return self diff --git a/litellm/models/project.py b/litellm/models/project.py new file mode 100644 index 00000000000..083c7ee3cc5 --- /dev/null +++ b/litellm/models/project.py @@ -0,0 +1,41 @@ +""" +Project table model. + +Canonical definition for ``litellm_projecttable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import List, Optional + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_ProjectTable(LiteLLMPydanticObjectBase): + """Database model representation for project""" + + project_id: str + project_alias: Optional[str] = None + description: Optional[str] = None + team_id: Optional[str] = None + budget_id: Optional[str] = None + metadata: Optional[dict] = None + models: List[str] = [] + spend: float = 0.0 + model_spend: Optional[dict] = None + model_rpm_limit: Optional[dict] = None + model_tpm_limit: Optional[dict] = None + blocked: bool = False + object_permission_id: Optional[str] = None + created_by: Optional[str] = None + updated_by: Optional[str] = None + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + + @property + def is_blocked(self) -> bool: + return self.blocked diff --git a/litellm/models/skills.py b/litellm/models/skills.py new file mode 100644 index 00000000000..62091c0ca01 --- /dev/null +++ b/litellm/models/skills.py @@ -0,0 +1,30 @@ +""" +Skills table model. + +Canonical definition for ``litellm_skillstable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Any, Dict, Optional + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_SkillsTable(LiteLLMPydanticObjectBase): + """Represents a LiteLLM_SkillsTable record""" + + skill_id: str + display_title: Optional[str] = None + description: Optional[str] = None + instructions: Optional[str] = None + source: str = "custom" + latest_version: Optional[str] = None + file_content: Optional[bytes] = None + file_name: Optional[str] = None + file_type: Optional[str] = None + metadata: Optional[Dict[str, Any]] = None + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None diff --git a/litellm/models/spend_logs.py b/litellm/models/spend_logs.py new file mode 100644 index 00000000000..96bd328c3ca --- /dev/null +++ b/litellm/models/spend_logs.py @@ -0,0 +1,50 @@ +""" +Spend and error log table models. + +Canonical definitions for ``litellm_spendlogs`` and ``litellm_errorlogs``. +Re-exported from ``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Optional, Union + +from pydantic import Json + +from litellm._uuid import uuid +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_SpendLogs(LiteLLMPydanticObjectBase): + request_id: str + api_key: str + model: Optional[str] = "" + api_base: Optional[str] = "" + call_type: str + spend: Optional[float] = 0.0 + total_tokens: Optional[int] = 0 + prompt_tokens: Optional[int] = 0 + completion_tokens: Optional[int] = 0 + startTime: Union[str, datetime, None] + endTime: Union[str, datetime, None] + user: Optional[str] = "" + metadata: Optional[Json] = {} + cache_hit: Optional[str] = "False" + cache_key: Optional[str] = None + request_tags: Optional[Json] = None + requester_ip_address: Optional[str] = None + messages: Optional[Union[str, list, dict]] + response: Optional[Union[str, list, dict]] + + +class LiteLLM_ErrorLogs(LiteLLMPydanticObjectBase): + request_id: Optional[str] = str(uuid.uuid4()) + api_base: Optional[str] = "" + model_group: Optional[str] = "" + litellm_model_name: Optional[str] = "" + model_id: Optional[str] = "" + request_kwargs: Optional[dict] = {} + exception_type: Optional[str] = "" + status_code: Optional[str] = "" + exception_string: Optional[str] = "" + startTime: Union[str, datetime, None] + endTime: Union[str, datetime, None] diff --git a/litellm/models/tag.py b/litellm/models/tag.py new file mode 100644 index 00000000000..02d8f58916d --- /dev/null +++ b/litellm/models/tag.py @@ -0,0 +1,36 @@ +""" +Tag table model. + +Canonical definition for ``litellm_tagtable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import List, Optional + +from pydantic import model_validator + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_TagTable(LiteLLMPydanticObjectBase): + tag_name: str + description: Optional[str] = None + models: List[str] = [] + model_info: Optional[dict] = None + spend: float = 0.0 + budget_id: Optional[str] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("spend") is None: + values.update({"spend": 0.0}) + if values.get("models") is None: + values.update({"models": []}) + return values diff --git a/litellm/models/team.py b/litellm/models/team.py new file mode 100644 index 00000000000..aa0798955f2 --- /dev/null +++ b/litellm/models/team.py @@ -0,0 +1,154 @@ +""" +Team table models. + +Canonical definitions for ``litellm_teamtable`` (plus the shared Member and +budget-window value types and the team-model alias table). Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +import json +from datetime import datetime +from typing import List, Literal, Optional, Union + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class MemberBase(LiteLLMPydanticObjectBase): + user_id: Optional[str] = Field( + default=None, + description="The unique ID of the user to add. Either user_id or user_email must be provided", + ) + user_email: Optional[str] = Field( + default=None, + description="The email address of the user to add. Either user_id or user_email must be provided", + ) + + @model_validator(mode="before") + @classmethod + def check_user_info(cls, values): + if not isinstance(values, dict): + raise ValueError("input needs to be a dictionary") + if values.get("user_id") is None and values.get("user_email") is None: + raise ValueError("Either user id or user email must be provided") + return values + + +class Member(MemberBase): + role: Literal["admin", "user"] = Field( + description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member" + ) + + +class BudgetLimitEntry(LiteLLMPydanticObjectBase): + """A single budget window with its own limit and independent reset schedule.""" + + budget_duration: str + max_budget: float + reset_at: Optional[datetime] = None + + +class LiteLLM_ModelTable(LiteLLMPydanticObjectBase): + id: Optional[int] = None + model_aliases: Optional[Union[str, dict]] = None + created_by: str + updated_by: str + team: Optional["LiteLLM_TeamTable"] = None + + model_config = ConfigDict(protected_namespaces=()) + + +class TeamBase(LiteLLMPydanticObjectBase): + team_alias: Optional[str] = None + team_id: Optional[str] = None + organization_id: Optional[str] = None + admins: list = [] + members: list = [] + members_with_roles: List[Member] = [] + team_member_permissions: Optional[List[str]] = None + metadata: Optional[dict] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + max_budget: Optional[float] = None + soft_budget: Optional[float] = None + budget_duration: Optional[str] = None + budget_limits: Optional[List[BudgetLimitEntry]] = None + models: list = [] + blocked: bool = False + router_settings: Optional[dict] = None + access_group_ids: Optional[List[str]] = None + default_team_member_models: Optional[List[str]] = None + + +class LiteLLM_TeamTable(TeamBase): + team_id: str # type: ignore + spend: Optional[float] = None + max_parallel_requests: Optional[int] = None + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + model_id: Optional[int] = None + model_spend: Optional[dict] = {} + model_max_budget: Optional[dict] = {} + policies: Optional[List[str]] = None + allow_team_guardrail_config: Optional[bool] = False + litellm_model_table: Optional[LiteLLM_ModelTable] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + object_permission_id: Optional[str] = None + updated_at: Optional[datetime] = None + created_at: Optional[datetime] = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + dict_fields = [ + "metadata", + "aliases", + "config", + "permissions", + "model_max_budget", + "model_aliases", + "router_settings", + "budget_limits", + ] + + if isinstance(values, BaseModel): + values = values.model_dump() + + if ( + isinstance(values.get("members_with_roles"), dict) + and not values["members_with_roles"] + ): + values["members_with_roles"] = [] + + for field in dict_fields: + value = values.get(field) + if value is not None and isinstance(value, str): + try: + values[field] = json.loads(value) + except json.JSONDecodeError: + raise ValueError(f"Field {field} should be a valid dictionary") + + return values + + +class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable): + last_refreshed_at: Optional[float] = None + + +class LiteLLM_DeletedTeamTable(LiteLLM_TeamTable): + """Audit record for deleted teams; mirrors the team plus deletion metadata.""" + + id: Optional[str] = None + deleted_at: Optional[datetime] = None + deleted_by: Optional[str] = None + deleted_by_api_key: Optional[str] = None + litellm_changed_by: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + +LiteLLM_ModelTable.model_rebuild() diff --git a/litellm/models/team_membership.py b/litellm/models/team_membership.py new file mode 100644 index 00000000000..d0a1308ce7c --- /dev/null +++ b/litellm/models/team_membership.py @@ -0,0 +1,32 @@ +""" +Team membership table model. + +Canonical definition for ``litellm_teammembership``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from typing import Optional, Union + +from litellm.models.budget import LiteLLM_BudgetTable, LiteLLM_BudgetTableFull +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): + user_id: str + team_id: str + budget_id: Optional[str] = None + spend: Optional[float] = 0.0 + total_spend: Optional[float] = 0.0 + litellm_budget_table: Optional[ + Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable] + ] = None + + def safe_get_team_member_rpm_limit(self) -> Optional[int]: + if self.litellm_budget_table is not None: + return self.litellm_budget_table.rpm_limit + return None + + def safe_get_team_member_tpm_limit(self) -> Optional[int]: + if self.litellm_budget_table is not None: + return self.litellm_budget_table.tpm_limit + return None diff --git a/litellm/models/user.py b/litellm/models/user.py new file mode 100644 index 00000000000..cd7e9db4aec --- /dev/null +++ b/litellm/models/user.py @@ -0,0 +1,70 @@ +""" +User table model. + +Canonical definition for ``litellm_usertable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Dict, List, Optional + +from pydantic import ConfigDict, Field, model_validator + +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.models.organization_membership import ( + LiteLLM_OrganizationMembershipTable, +) +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_UserTable(LiteLLMPydanticObjectBase): + user_id: str + user_alias: Optional[str] = None + team_id: Optional[str] = None + sso_user_id: Optional[str] = None + organization_id: Optional[str] = None + object_permission_id: Optional[str] = None + password: Optional[str] = Field(default=None, exclude=True) + teams: List[str] = [] + user_role: Optional[str] = None + max_budget: Optional[float] = None + spend: float = 0.0 + user_email: Optional[str] = None + models: list = [] + metadata: Optional[dict] = None + max_parallel_requests: Optional[int] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + allowed_cache_controls: List[str] = [] + policies: List[str] = [] + model_spend: Optional[Dict] = {} + model_max_budget: Optional[Dict] = {} + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + organization_memberships: Optional[List[LiteLLM_OrganizationMembershipTable]] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("spend") is None: + values.update({"spend": 0.0}) + if values.get("models") is None: + values.update({"models": []}) + if values.get("teams") is None: + values.update({"teams": []}) + return values + + def is_over_budget(self) -> bool: + if self.max_budget is None: + return False + return self.spend >= self.max_budget + + def has_model_access(self, model_name: str) -> bool: + if not self.models: + return True + return model_name in self.models diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py new file mode 100644 index 00000000000..8bddd1c1619 --- /dev/null +++ b/litellm/models/verification_token.py @@ -0,0 +1,74 @@ +""" +Verification token table model. + +Canonical definition for ``litellm_verificationtoken``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Dict, List, Optional, Union + +from pydantic import ConfigDict + +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): + token: Optional[str] = None + key_name: Optional[str] = None + key_alias: Optional[str] = None + spend: float = 0.0 + max_budget: Optional[float] = None + expires: Optional[Union[str, datetime]] = None + models: List = [] + aliases: Dict = {} + config: Dict = {} + user_id: Optional[str] = None + team_id: Optional[str] = None + agent_id: Optional[str] = None + project_id: Optional[str] = None + max_parallel_requests: Optional[int] = None + metadata: Dict = {} + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + allowed_cache_controls: Optional[list] = [] + allowed_routes: Optional[list] = [] + permissions: Dict = {} + model_spend: Dict = {} + model_max_budget: Dict = {} + soft_budget_cooldown: bool = False + blocked: Optional[bool] = None + litellm_budget_table: Optional[dict] = None + budget_id: Optional[str] = None + org_id: Optional[str] = None # org id for a given key + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None + last_active: Optional[datetime] = None + object_permission_id: Optional[str] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + access_group_ids: Optional[List[str]] = None + rotation_count: Optional[int] = 0 + auto_rotate: Optional[bool] = False + rotation_interval: Optional[str] = None + last_rotation_at: Optional[datetime] = None + key_rotation_at: Optional[datetime] = None + router_settings: Optional[dict] = None + budget_limits: Optional[List[dict]] = None + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken): + """Audit record for deleted keys; mirrors the token plus deletion metadata.""" + + id: Optional[str] = None + deleted_at: Optional[datetime] = None + deleted_by: Optional[str] = None + deleted_by_api_key: Optional[str] = None + litellm_changed_by: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) 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 863e6acd41e..dcf7660d002 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 @@ -14,8 +14,12 @@ SpecialHeaders, UserAPIKeyAuth, ) -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import ( + AgentsRepository, + MCPServerRepository, +) def _parse_mcp_server_names_from_path( @@ -1445,7 +1449,7 @@ async def _get_agent_object_permission( return None if object_permission_id is None: - agent_row = await prisma_client.db.litellm_agentstable.find_unique( + agent_row = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id}, ) object_permission_id = ( @@ -1600,7 +1604,7 @@ async def _get_db_server_ids_for_access_groups( server_ids: Set[str] = set() if access_groups and prisma_client is not None: try: - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many( + mcp_servers = await MCPServerRepository(prisma_client).table.find_many( where={"mcp_access_groups": {"hasSome": access_groups}} ) for server in mcp_servers: diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 0ba0181200f..8edb831a9df 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import ( LiteLLM_MCPServerTable, LiteLLM_ObjectPermissionTable, @@ -25,8 +26,16 @@ decrypt_value_helper, encrypt_value_helper, ) -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy.utils import PrismaClient +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.table_repositories import ( + MCPServerRepository, + MCPUserCredentialsRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPCredentials @@ -354,7 +363,7 @@ async def get_all_mcp_servers( where: Dict[str, Any] = {} if approval_status is not None: where["approval_status"] = approval_status - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many( + mcp_servers = await MCPServerRepository(prisma_client).table.find_many( where=where if where else {} ) @@ -380,7 +389,9 @@ async def get_mcp_server( """ Returns the matching mcp server from the db iff exists """ - mcp_server = await prisma_client.db.litellm_mcpservertable.find_unique( + mcp_server: Optional[LiteLLM_MCPServerTable] = await MCPServerRepository( + prisma_client + ).table.find_unique( where={ "server_id": server_id, } @@ -398,12 +409,12 @@ async def get_mcp_servers( """ Returns the matching mcp servers from the db with the server_ids """ - _mcp_servers: List[LiteLLM_MCPServerTable] = ( - await prisma_client.db.litellm_mcpservertable.find_many( - where={ - "server_id": {"in": server_ids}, - } - ) + _mcp_servers: List[LiteLLM_MCPServerTable] = await MCPServerRepository( + prisma_client + ).table.find_many( + where={ + "server_id": {"in": server_ids}, + } ) final_mcp_servers: List[LiteLLM_MCPServerTable] = [] for _mcp_server in _mcp_servers: @@ -420,15 +431,15 @@ async def get_mcp_servers_by_verificationtoken( """ Returns the mcp servers from the db for the verification token """ - verification_token_record: LiteLLM_TeamTable = ( - await prisma_client.db.litellm_verificationtoken.find_unique( - where={ - "token": token, - }, - include={ - "object_permission": True, - }, - ) + verification_token_record: LiteLLM_TeamTable = await VerificationTokenRepository( + prisma_client + ).table.find_unique( + where={ + "token": token, + }, + include={ + "object_permission": True, + }, ) mcp_servers: Optional[List[str]] = [] @@ -446,15 +457,15 @@ async def get_mcp_servers_by_team( """ Returns the mcp servers from the db for the team id """ - team_record: LiteLLM_TeamTable = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={ - "team_id": team_id, - }, - include={ - "object_permission": True, - }, - ) + team_record: LiteLLM_TeamTable = await TeamRepository( + prisma_client + ).table.find_unique( + where={ + "team_id": team_id, + }, + include={ + "object_permission": True, + }, ) mcp_servers: Optional[List[str]] = [] @@ -505,16 +516,16 @@ async def get_objectpermissions_for_mcp_server( """ Get all the object permissions records and the associated team and verficiationtoken records that have access to the mcp server """ - object_permission_records = ( - await prisma_client.db.litellm_objectpermissiontable.find_many( - where={ - "mcp_servers": {"has": mcp_server_id}, - }, - include={ - "teams": True, - "verification_tokens": True, - }, - ) + object_permission_records = await ObjectPermissionRepository( + prisma_client + ).table.find_many( + where={ + "mcp_servers": {"has": mcp_server_id}, + }, + include={ + "teams": True, + "verification_tokens": True, + }, ) return object_permission_records @@ -526,7 +537,7 @@ async def get_virtualkeys_for_mcp_server( """ Get all the virtual keys that have access to the mcp server """ - virtual_keys = await prisma_client.db.litellm_verificationtoken.find_many( + virtual_keys = await VerificationTokenRepository(prisma_client).table.find_many( where={ "mcp_servers": {"has": server_id}, }, @@ -557,30 +568,35 @@ async def delete_mcp_server( """ Delete the mcp server from the db by server_id - The server-row delete is the commit point. Per-user env var rows have no FK - cascade, so they are cleaned up afterwards on a best-effort basis: a transient - failure there leaves only orphaned rows pointing at a now-missing server and - must not turn a successful delete into a caller-visible error. + The server-row delete is the commit point. Per-user credential and env var + rows have no FK cascade, so they are cleaned up afterwards on a best-effort + basis: a transient failure there leaves only orphaned rows pointing at a + now-missing server and must not turn a successful delete into a + caller-visible error. Each table is cleaned independently so a failure on one + still attempts the other. Returns the deleted mcp server record if it exists, otherwise None """ - deleted_server = await prisma_client.db.litellm_mcpservertable.delete( + deleted_server = await MCPServerRepository(prisma_client).table.delete( where={ "server_id": server_id, }, ) if deleted_server is not None: - try: - await prisma_client.db.litellm_mcpuserenvvars.delete_many( - where={"server_id": server_id} - ) - except Exception as e: - verbose_proxy_logger.warning( - "MCP server %s deleted but per-user env var cleanup failed; " - "orphaned rows can be removed on a later delete: %s", - server_id, - e, - ) + for model, label in ( + (prisma_client.db.litellm_mcpusercredentials, "credential"), + (prisma_client.db.litellm_mcpuserenvvars, "env var"), + ): + try: + await model.delete_many(where={"server_id": server_id}) + except Exception as e: + verbose_proxy_logger.warning( + "MCP server %s deleted but per-user %s cleanup failed; " + "orphaned rows can be removed on a later delete: %s", + server_id, + label, + e, + ) return deleted_server @@ -600,7 +616,7 @@ async def create_mcp_server( data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by - new_mcp_server = await prisma_client.db.litellm_mcpservertable.create( + new_mcp_server = await MCPServerRepository(prisma_client).table.create( data=data_dict # type: ignore ) @@ -635,7 +651,7 @@ async def update_mcp_server( "credentials" in data_dict and data_dict["credentials"] is not None ) if data.auth_type or has_credentials: - existing = await prisma_client.db.litellm_mcpservertable.find_unique( + existing = await MCPServerRepository(prisma_client).table.find_unique( where={"server_id": data.server_id} ) @@ -678,7 +694,7 @@ async def update_mcp_server( # Add audit fields data_dict["updated_by"] = touched_by - updated_mcp_server = await prisma_client.db.litellm_mcpservertable.update( + updated_mcp_server = await MCPServerRepository(prisma_client).table.update( where={"server_id": data.server_id}, data=data_dict # type: ignore ) @@ -691,7 +707,7 @@ async def rotate_mcp_server_credentials_master_key( ): from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many() + mcp_servers = await MCPServerRepository(prisma_client).table.find_many() updated = 0 for mcp_server in mcp_servers: @@ -719,7 +735,7 @@ async def rotate_mcp_server_credentials_master_key( continue update_data["updated_by"] = touched_by - await prisma_client.db.litellm_mcpservertable.update( + await MCPServerRepository(prisma_client).table.update( where={"server_id": mcp_server.server_id}, data=update_data, ) @@ -781,7 +797,7 @@ async def rotate_mcp_user_credentials_master_key( under the new master key. Rows that are unreadable under both paths are logged and skipped so one corrupt row does not abort the rotation. """ - rows = await prisma_client.db.litellm_mcpusercredentials.find_many() + rows = await MCPUserCredentialsRepository(prisma_client).table.find_many() rotated = 0 skipped = 0 for row in rows: @@ -798,7 +814,7 @@ async def rotate_mcp_user_credentials_master_key( re_encrypted = encrypt_value_helper( plaintext, new_encryption_key=new_master_key ) - await prisma_client.db.litellm_mcpusercredentials.update( + await MCPUserCredentialsRepository(prisma_client).table.update( where={ "user_id_server_id": { "user_id": row.user_id, @@ -873,7 +889,7 @@ async def store_user_credential( """Store a user credential for a BYOK MCP server.""" encoded = encrypt_value_helper(credential) - await prisma_client.db.litellm_mcpusercredentials.upsert( + await MCPUserCredentialsRepository(prisma_client).table.upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ "create": { @@ -893,7 +909,7 @@ async def get_user_credential( ) -> Optional[str]: """Return credential for a user+server pair, or None.""" - row = await prisma_client.db.litellm_mcpusercredentials.find_unique( + row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if row is None: @@ -907,7 +923,7 @@ async def has_user_credential( server_id: str, ) -> bool: """Return True if the user has a stored credential for this server.""" - row = await prisma_client.db.litellm_mcpusercredentials.find_unique( + row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) return row is not None @@ -919,7 +935,7 @@ async def delete_user_credential( server_id: str, ) -> None: """Delete the user's stored credential for a BYOK MCP server.""" - await prisma_client.db.litellm_mcpusercredentials.delete( + await MCPUserCredentialsRepository(prisma_client).table.delete( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) @@ -966,7 +982,7 @@ async def store_user_oauth_credential( # Skip the guard when the caller knows the row is already an OAuth2 credential # (e.g. during token refresh), saving an extra DB round-trip. if not skip_byok_guard: - existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( + existing = await MCPUserCredentialsRepository(prisma_client).table.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if ( @@ -984,7 +1000,7 @@ async def store_user_oauth_credential( ) encoded = encrypt_value_helper(json.dumps(payload)) - await prisma_client.db.litellm_mcpusercredentials.upsert( + await MCPUserCredentialsRepository(prisma_client).table.upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ "create": { @@ -1025,7 +1041,7 @@ async def get_user_oauth_credential( ) -> Optional[Dict[str, Any]]: """Return the decoded OAuth2 payload dict for a user+server pair, or None.""" - row = await prisma_client.db.litellm_mcpusercredentials.find_unique( + row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if row is None: @@ -1039,7 +1055,7 @@ async def list_user_oauth_credentials( ) -> List[Dict[str, Any]]: """Return all OAuth2 credential payloads for a user, tagged with server_id.""" - rows = await prisma_client.db.litellm_mcpusercredentials.find_many( + rows = await MCPUserCredentialsRepository(prisma_client).table.find_many( where={"user_id": user_id} ) results: List[Dict[str, Any]] = [] @@ -1212,7 +1228,7 @@ async def approve_mcp_server( ) -> LiteLLM_MCPServerTable: """Set approval_status=active and record reviewed_at.""" now = datetime.now(timezone.utc) - updated = await prisma_client.db.litellm_mcpservertable.update( + updated = await MCPServerRepository(prisma_client).table.update( where={"server_id": server_id}, data={ "approval_status": MCPApprovalStatus.active, @@ -1240,7 +1256,7 @@ async def reject_mcp_server( } if review_notes is not None: data["review_notes"] = review_notes - updated = await prisma_client.db.litellm_mcpservertable.update( + updated = await MCPServerRepository(prisma_client).table.update( where={"server_id": server_id}, data=data, ) @@ -1257,7 +1273,7 @@ async def get_mcp_submissions( along with a summary count breakdown by approval_status. Mirrors get_guardrail_submissions() from guardrail_endpoints.py. """ - rows = await prisma_client.db.litellm_mcpservertable.find_many( + rows = await MCPServerRepository(prisma_client).table.find_many( where={"submitted_at": {"not": None}}, order={"submitted_at": "desc"}, take=500, # safety cap; paginate if needed in a future iteration diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ed374635fea..3beddd2c435 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -512,12 +512,13 @@ async def exchange_token_with_server( result = { "access_token": access_token, "token_type": token_response.get("token_type", "Bearer"), - "expires_in": token_response.get("expires_in", 3600), } - if "refresh_token" in token_response and token_response["refresh_token"]: + if token_response.get("expires_in") is not None: + result["expires_in"] = token_response["expires_in"] + if token_response.get("refresh_token"): result["refresh_token"] = token_response["refresh_token"] - if "scope" in token_response and token_response["scope"]: + if token_response.get("scope"): result["scope"] = token_response["scope"] # RFC 6749 §5.1: token responses must not be cached. diff --git a/litellm/proxy/_experimental/mcp_server/mcp_context.py b/litellm/proxy/_experimental/mcp_server/mcp_context.py index a60138dd340..51918509441 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_context.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_context.py @@ -19,3 +19,9 @@ _mcp_gateway_initialize_instructions: ContextVar[Optional[str]] = ContextVar( "_mcp_gateway_initialize_instructions", default=None ) + +# Per-request scoped server name; set in MCP HTTP/SSE handlers when the path +# identifies exactly one upstream server. Never populated from client-supplied headers. +_mcp_gateway_server_name: ContextVar[Optional[str]] = ContextVar( + "_mcp_gateway_server_name", default=None +) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7048f5bf7c4..73935beeb3a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -42,8 +42,8 @@ MCP_TOOL_LISTING_TIMEOUT, ) from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException -from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth +from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, @@ -85,6 +85,7 @@ from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.utils import ProxyLogging +from litellm.repositories.table_repositories import MCPServerRepository from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPAuth, MCPStdioConfig from litellm.types.mcp_server.mcp_server_manager import ( @@ -353,6 +354,52 @@ def _deserialize_json_list(data: Any) -> Optional[List[Dict[str, Any]]]: ] +def _normalize_mcp_server_cost_info(mcp_info: MCPInfo) -> None: + """Coerce ``mcp_server_cost_info`` numeric fields to ``float`` at ingest. + + YAML 1.1 parses scientific notation without a decimal point (e.g. + ``7e-05``) as a string, and ``MCPServerCostInfo`` is a TypedDict with no + runtime validation, so string-typed costs flow through to the UI and + crash its ``.toFixed`` formatting. Values that cannot be coerced are + dropped with a warning instead of failing the server load. + """ + cost_info = mcp_info.get("mcp_server_cost_info") + if not isinstance(cost_info, dict): + return + + server_name = mcp_info.get("server_name") + normalized = dict(cost_info) + + default_cost = normalized.get("default_cost_per_query") + if default_cost is not None: + try: + normalized["default_cost_per_query"] = float(default_cost) + except (TypeError, ValueError): + verbose_logger.warning( + "MCP server '%s' has non-numeric default_cost_per_query %r; ignoring it", + server_name, + default_cost, + ) + del normalized["default_cost_per_query"] + + tool_costs = normalized.get("tool_name_to_cost_per_query") + if isinstance(tool_costs, dict): + normalized_tool_costs = {} + for tool_name, cost in tool_costs.items(): + try: + normalized_tool_costs[tool_name] = float(cost) + except (TypeError, ValueError): + verbose_logger.warning( + "MCP server '%s' has non-numeric cost %r for tool '%s'; ignoring it", + server_name, + cost, + tool_name, + ) + normalized["tool_name_to_cost_per_query"] = normalized_tool_costs + + mcp_info["mcp_server_cost_info"] = normalized + + def _create_sampling_callback(user_api_key_auth: Optional[Any] = None): """ Create a sampling callback for MCP ClientSession. @@ -620,6 +667,7 @@ async def load_servers_from_config( mcp_info["server_name"] = server_name if "description" not in mcp_info and server_config.get("description"): mcp_info["description"] = server_config.get("description") + _normalize_mcp_server_cost_info(mcp_info) # Use alias for name if present, else server_name alias = server_config.get("alias", None) @@ -1090,6 +1138,7 @@ async def build_mcp_server_from_table( mcp_info["server_name"] = mcp_server.server_name or mcp_server.server_id if "description" not in mcp_info and mcp_server.description: mcp_info["description"] = mcp_server.description + _normalize_mcp_server_cost_info(mcp_info) auth_type = cast(MCPAuthType, mcp_server.auth_type) server_url = mcp_server.url @@ -3817,7 +3866,7 @@ async def reload_servers_from_database(self): # Pending/rejected servers are excluded at the DB level so we never load them. from litellm.proxy._experimental.mcp_server.db import LiteLLM_MCPServerTable - raw_rows = await prisma_client.db.litellm_mcpservertable.find_many( + raw_rows = await MCPServerRepository(prisma_client).table.find_many( where={ "OR": [ {"approval_status": None}, diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 725f7a335bc..2149f079a3d 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -386,8 +386,15 @@ async def _get_tools_for_single_server( raw_headers: Optional[Dict[str, str]] = None, user_api_key_auth: Optional[UserAPIKeyAuth] = None, extra_headers: Optional[Dict[str, str]] = None, + apply_tool_filters: bool = True, ): - """Helper function to get tools for a single server.""" + """Helper function to get tools for a single server. + + When ``apply_tool_filters`` is False the raw server catalog is returned + without the allowed_tools/disallowed_tools gate or the per-key tool + permissions. This is the admin-only configuration view; every runtime + path keeps the default True so callable tools stay filtered. + """ tools = await global_mcp_server_manager._get_tools_from_server( server=server, mcp_auth_header=server_auth_header, @@ -397,6 +404,9 @@ async def _get_tools_for_single_server( user_api_key_auth=user_api_key_auth, ) + if not apply_tool_filters: + return _create_tool_response_objects(tools, server.mcp_info) + # Always apply allowed_tools/disallowed_tools so the blacklist is # enforced even when no allowlist is set (matches the SSE/HTTP path). tools = filter_tools_by_allowed_tools(tools, server) @@ -463,6 +473,7 @@ async def _list_tools_for_single_server( mcp_auth_header: Optional[str], raw_headers_from_request: dict, user_api_key_dict: UserAPIKeyAuth, + apply_tool_filters: bool = True, ) -> dict: """Handle tool listing for a single server_id request.""" # Resolve a server name to its UUID if needed @@ -527,6 +538,7 @@ async def _list_tools_for_single_server( raw_headers_from_request, user_api_key_dict, extra_headers=user_oauth_extra_headers, + apply_tool_filters=apply_tool_filters, ) except MCPUpstreamAuthError: # Surface the upstream 401/403 to the caller so it can emit the @@ -552,6 +564,14 @@ async def list_tool_rest_api( server_id: Optional[str] = Query( None, description="The server id to list tools for" ), + include_disabled_tools: bool = Query( + False, + description=( + "Admin only. Return the full server tool catalog without the " + "allowed_tools filter or per-key tool permissions, so the MCP " + "settings UI can configure the allowlist. Ignored for non-admins." + ), + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> dict: """ @@ -579,6 +599,13 @@ async def list_tool_rest_api( ) try: + # The full catalog (allowlist filter skipped) is admin-only so the + # REST endpoint can't be used to enumerate deliberately-disabled tools. + apply_tool_filters = not ( + include_disabled_tools + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + ) + # Extract auth headers from request headers = request.headers raw_headers_from_request = dict(headers) @@ -620,6 +647,7 @@ async def list_tool_rest_api( mcp_auth_header=mcp_auth_header, raw_headers_from_request=raw_headers_from_request, user_api_key_dict=user_api_key_dict, + apply_tool_filters=apply_tool_filters, ) else: if not allowed_server_ids: @@ -677,6 +705,7 @@ async def list_tool_rest_api( raw_headers_from_request, user_api_key_dict, extra_headers=user_oauth_extra_headers, + apply_tool_filters=apply_tool_filters, ) list_tools_result.extend(tools_result) except Exception as e: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 0477a5d3244..731493b1337 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -47,6 +47,7 @@ from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_active_toolset_id, _mcp_gateway_initialize_instructions, + _mcp_gateway_server_name, ) from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug from litellm.proxy._experimental.mcp_server.utils import ( @@ -323,10 +324,14 @@ def _gateway_create_initialization_options( notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, ) + updates: Dict[str, Any] = {} merged = _mcp_gateway_initialize_instructions.get() if merged is not None: - return opts.model_copy(update={"instructions": merged}) - return opts + updates["instructions"] = merged + scoped_server_name = _mcp_gateway_server_name.get() + if scoped_server_name is not None: + updates["server_name"] = scoped_server_name + return opts.model_copy(update=updates) if updates else opts ######################################################## ############ Initialize the MCP Server ################# @@ -1544,6 +1549,7 @@ async def _gateway_initialize_instructions_request_scope( user_api_key_auth: Optional[UserAPIKeyAuth], mcp_servers: Optional[List[str]], client_ip: Optional[str], + scoped_server_endpoint: bool = False, ) -> AsyncIterator[None]: allowed = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, @@ -1565,11 +1571,22 @@ async def _gateway_initialize_instructions_request_scope( return_exceptions=True, ) merged = _merge_gateway_initialize_instructions(allowed_mcp_servers=allowed) - tok = _mcp_gateway_initialize_instructions.set(merged) + scoped_server_name = None + if scoped_server_endpoint and len(allowed) == 1: + scoped_server = allowed[0] + scoped_server_name = ( + scoped_server.alias + or scoped_server.server_name + or scoped_server.name + or scoped_server.server_id + ) + instructions_token = _mcp_gateway_initialize_instructions.set(merged) + server_name_token = _mcp_gateway_server_name.set(scoped_server_name) try: yield finally: - _mcp_gateway_initialize_instructions.reset(tok) + _mcp_gateway_initialize_instructions.reset(instructions_token) + _mcp_gateway_server_name.reset(server_name_token) async def _get_tools_from_mcp_servers( # noqa: PLR0915 user_api_key_auth: Optional[UserAPIKeyAuth], @@ -3620,6 +3637,7 @@ async def handle_streamable_http_mcp( # noqa: PLR0915 oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) + scoped_server_endpoint = len(_get_mcp_servers_in_path(path) or []) == 1 # Extract client IP for MCP access control _client_ip = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope)) @@ -3896,6 +3914,7 @@ async def _dispatch() -> None: user_api_key_auth, mcp_servers, _client_ip, + scoped_server_endpoint=scoped_server_endpoint, ): await target_manager.handle_request(scope, receive, local_send) if use_stateful and session_id and scope.get("method") == "DELETE": @@ -3980,6 +3999,7 @@ async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) + scoped_server_endpoint = len(_get_mcp_servers_in_path(path) or []) == 1 # Extract client IP for MCP access control _sse_client_ip = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope)) @@ -4052,6 +4072,7 @@ async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: user_api_key_auth, mcp_servers, _sse_client_ip, + scoped_server_endpoint=scoped_server_endpoint, ): await sse_session_manager.handle_request(scope, receive, send) except MCPUpstreamAuthError as e: diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index 08ac7dbd33b..a996131653f 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -4,6 +4,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import MCPToolsetRepository from litellm.types.mcp_server.mcp_toolset import ( MCPToolset, NewMCPToolsetRequest, @@ -30,7 +31,7 @@ async def create_mcp_toolset( data_dict["tools"] = json.dumps(data_dict.get("tools", [])) data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by - row = await prisma_client.db.litellm_mcptoolsettable.create(data=data_dict) + row = await MCPToolsetRepository(prisma_client).table.create(data=data_dict) return _toolset_from_row(row) @@ -38,7 +39,7 @@ async def get_mcp_toolset( prisma_client: PrismaClient, toolset_id: str, ) -> Optional[MCPToolset]: - row = await prisma_client.db.litellm_mcptoolsettable.find_unique( + row = await MCPToolsetRepository(prisma_client).table.find_unique( where={"toolset_id": toolset_id} ) if row is None: @@ -54,7 +55,7 @@ async def list_mcp_toolsets( where = {} if toolset_ids is not None: where = {"toolset_id": {"in": toolset_ids}} - rows = await prisma_client.db.litellm_mcptoolsettable.find_many(where=where) + 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( @@ -69,7 +70,7 @@ async def get_mcp_toolset_by_name( prisma_client: PrismaClient, toolset_name: str, ) -> Optional[MCPToolset]: - row = await prisma_client.db.litellm_mcptoolsettable.find_first( + row = await MCPToolsetRepository(prisma_client).table.find_first( where={"toolset_name": toolset_name} ) if row is None: @@ -87,7 +88,7 @@ async def update_mcp_toolset( data_dict["tools"] = json.dumps(data_dict["tools"]) data_dict["updated_by"] = touched_by try: - row = await prisma_client.db.litellm_mcptoolsettable.update( + row = await MCPToolsetRepository(prisma_client).table.update( where={"toolset_id": data.toolset_id}, data=data_dict, ) @@ -105,7 +106,7 @@ async def delete_mcp_toolset( toolset_id: str, ) -> Optional[MCPToolset]: try: - row = await prisma_client.db.litellm_mcptoolsettable.delete( + row = await MCPToolsetRepository(prisma_client).table.delete( where={"toolset_id": toolset_id} ) except Exception as e: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e5d70933063..1b594e20d32 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -23,7 +23,6 @@ from litellm.types.integrations.slack_alerting import AlertType from litellm.types.llms.openai import ( AllMessageValues, - OpenAIFileObject, ResponsesAPIResponse, ) from litellm.types.mcp import ( @@ -41,8 +40,6 @@ EmbeddingResponse, GenericBudgetConfigType, ImageResponse, - LiteLLMBatch, - LiteLLMFineTuningJob, LiteLLMPydanticObjectBase, ModelResponse, ProviderField, @@ -1014,12 +1011,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): search_tools: Optional[List[str]] = None -class BudgetLimitEntry(LiteLLMPydanticObjectBase): - """A single budget window with its own limit and independent reset schedule.""" - - budget_duration: str # e.g. "24h", "7d", "30d" - max_budget: float # max spend in USD for this window - reset_at: Optional[datetime] = None # populated at creation/reset time +from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402 class GenerateRequestBase(LiteLLMPydanticObjectBase): @@ -1217,40 +1209,10 @@ def validate_at_least_one(cls, values): return values -class LiteLLM_ModelTable(LiteLLMPydanticObjectBase): - id: Optional[int] = None - model_aliases: Optional[Union[str, dict]] = None # json dump the dict - created_by: str - updated_by: str - team: Optional["LiteLLM_TeamTable"] = None - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase): - model_id: str - model_name: str - litellm_params: dict - model_info: dict - created_at: Optional[datetime] = None - created_by: str - updated_at: Optional[datetime] = None - updated_by: str - - @model_validator(mode="before") - @classmethod - def check_potential_json_str(cls, values): - if isinstance(values.get("litellm_params"), str): - try: - values["litellm_params"] = json.loads(values["litellm_params"]) - except json.JSONDecodeError: - pass - if isinstance(values.get("model_info"), str): - try: - values["model_info"] = json.loads(values["model_info"]) - except json.JSONDecodeError: - pass - return values +from litellm.models.model import ( # noqa: E402 + LiteLLM_ProxyModelTable as LiteLLM_ProxyModelTable, +) +from litellm.models.team import LiteLLM_ModelTable as LiteLLM_ModelTable # noqa: E402 # MCP Types @@ -1265,32 +1227,12 @@ class MCPApprovalStatus(str, enum.Enum): rejected = "rejected" -class MCPEnvVarScope(str, enum.Enum): - """Scope for an MCP server environment variable. - - - ``global``: value is provided by the admin and used for all users. - - ``user``: each user must provide their own value via the per-user - env-var endpoint. The admin-supplied ``value`` is treated as a - placeholder/hint and is not used at request time. - """ - - global_ = "global" - user = "user" - - -class MCPEnvVar(LiteLLMPydanticObjectBase): - """One environment variable for an MCP server. - - Variables can be interpolated into ``static_headers`` using ``${NAME}`` - syntax. ``scope=global`` values are stored on the server. ``scope=user`` - values are stored per-user in ``LiteLLM_MCPUserEnvVars`` and supplied by - each user. - """ - - name: str - value: str = "" - scope: MCPEnvVarScope = MCPEnvVarScope.global_ - description: Optional[str] = None +from litellm.models.mcp_server import ( # noqa: E402 + MCPEnvVar as MCPEnvVar, +) +from litellm.models.mcp_server import ( # noqa: E402 + MCPEnvVarScope as MCPEnvVarScope, +) # MCP Proxy Request Types @@ -1443,66 +1385,9 @@ def validate_transport_fields(cls, values): return values -class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): - """Represents a LiteLLM_MCPServerTable record""" - - server_id: str - server_name: Optional[str] = None - alias: Optional[str] = None - description: Optional[str] = None - url: Optional[str] = None - spec_path: Optional[str] = None - transport: MCPTransportType - auth_type: Optional[MCPAuthType] = None - credentials: Optional[MCPCredentials] = None - instructions: Optional[str] = None - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - updated_by: Optional[str] = None - teams: List[Dict[str, Optional[str]]] = Field(default_factory=list) - mcp_access_groups: List[str] = Field(default_factory=list) - allowed_tools: List[str] = Field(default_factory=list) - tool_name_to_display_name: Optional[Dict[str, str]] = None - tool_name_to_description: Optional[Dict[str, str]] = None - extra_headers: List[str] = Field(default_factory=list) - mcp_info: Optional[MCPInfo] = None - static_headers: Optional[Dict[str, str]] = None - env_vars: Optional[List[MCPEnvVar]] = None - # Health check status - status: Optional[Literal["healthy", "unhealthy", "unknown"]] = Field( - default="unknown", - description="Health status: 'healthy', 'unhealthy', 'unknown'", - ) - last_health_check: Optional[datetime] = None - health_check_error: Optional[str] = None - # Stdio-specific fields - command: Optional[str] = None - args: List[str] = Field(default_factory=list) - env: Dict[str, str] = Field(default_factory=dict) - authorization_url: Optional[str] = None - token_url: Optional[str] = None - registration_url: Optional[str] = None - oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None - allow_all_keys: bool = False - available_on_public_internet: bool = True - delegate_auth_to_upstream: bool = False - oauth_passthrough: bool = False - is_byok: bool = False - byok_description: List[str] = Field(default_factory=list) - byok_api_key_help_url: Optional[str] = None - has_user_credential: Optional[bool] = None - source_url: Optional[str] = None - timeout: Optional[float] = None - # BYOM submission fields - approval_status: Optional[str] = Field( - default="active", - description="Approval status: 'pending_review', 'active', 'rejected'", - ) - submitted_by: Optional[str] = None - submitted_at: Optional[datetime] = None - reviewed_at: Optional[datetime] = None - review_notes: Optional[str] = None +from litellm.models.mcp_server import ( # noqa: E402 + LiteLLM_MCPServerTable as LiteLLM_MCPServerTable, +) class MakeMCPServersPublicRequest(LiteLLMPydanticObjectBase): @@ -1622,23 +1507,9 @@ class UpdateSkillRequest(LiteLLMPydanticObjectBase): metadata: Optional[Dict[str, Any]] = None -class LiteLLM_SkillsTable(LiteLLMPydanticObjectBase): - """Represents a LiteLLM_SkillsTable record""" - - skill_id: str - display_title: Optional[str] = None - description: Optional[str] = None - instructions: Optional[str] = None - source: str = "custom" - latest_version: Optional[str] = None - file_content: Optional[bytes] = None # Binary content of skill files (zip) - file_name: Optional[str] = None # Original filename - file_type: Optional[str] = None # MIME type - metadata: Optional[Dict[str, Any]] = None - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - updated_by: Optional[str] = None +from litellm.models.skills import ( # noqa: E402 + LiteLLM_SkillsTable as LiteLLM_SkillsTable, +) class ListSkillsRequest(LiteLLMPydanticObjectBase): @@ -1839,33 +1710,8 @@ class DeleteCustomerRequest(LiteLLMPydanticObjectBase): user_ids: List[str] -class MemberBase(LiteLLMPydanticObjectBase): - user_id: Optional[str] = Field( - default=None, - description="The unique ID of the user to add. Either user_id or user_email must be provided", - ) - user_email: Optional[str] = Field( - default=None, - description="The email address of the user to add. Either user_id or user_email must be provided", - ) - - @model_validator(mode="before") - @classmethod - def check_user_info(cls, values): - if not isinstance(values, dict): - raise ValueError("input needs to be a dictionary") - if values.get("user_id") is None and values.get("user_email") is None: - raise ValueError("Either user id or user email must be provided") - return values - - -class Member(MemberBase): - role: Literal[ - "admin", - "user", - ] = Field( - description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member" - ) +from litellm.models.team import Member as Member # noqa: E402 +from litellm.models.team import MemberBase as MemberBase # noqa: E402 class OrgMember(MemberBase): @@ -1876,33 +1722,7 @@ class OrgMember(MemberBase): ] -class TeamBase(LiteLLMPydanticObjectBase): - team_alias: Optional[str] = None - team_id: Optional[str] = None - organization_id: Optional[str] = None - admins: list = [] - members: list = [] - members_with_roles: List[Member] = [] - team_member_permissions: Optional[List[str]] = None - metadata: Optional[dict] = None - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - - # Budget fields - max_budget: Optional[float] = None - soft_budget: Optional[float] = None - budget_duration: Optional[str] = None - budget_limits: Optional[List[BudgetLimitEntry]] = ( - None # multiple concurrent budget windows - ) - - models: list = [] - blocked: bool = False - router_settings: Optional[dict] = None - access_group_ids: Optional[List[str]] = None - default_team_member_models: Optional[List[str]] = ( - None # default allowed_models seeded onto new team members - ) +from litellm.models.team import TeamBase as TeamBase # noqa: E402 class NewTeamRequest(TeamBase): @@ -2100,147 +1920,31 @@ def validate_callback_vars(cls, values): return values -class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): - """Represents a LiteLLM_ObjectPermissionTable record""" - - object_permission_id: str - mcp_servers: Optional[List[str]] = [] - mcp_access_groups: Optional[List[str]] = [] - mcp_tool_permissions: Optional[Dict[str, List[str]]] = None - """ - Mapping - server_id -> list of tools - - Enforces allowed tools for a specific key/team/organization - { - "1234567890": ["tool_name_1", "tool_name_2"] - } - """ - - vector_stores: Optional[List[str]] = [] - agents: Optional[List[str]] = [] - agent_access_groups: Optional[List[str]] = [] - mcp_toolsets: Optional[List[str]] = None - blocked_tools: Optional[List[str]] = [] - search_tools: Optional[List[str]] = [] - - -class LiteLLM_TeamTable(TeamBase): - team_id: str # type: ignore - spend: Optional[float] = None - max_parallel_requests: Optional[int] = None - budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None - model_id: Optional[int] = None - litellm_model_table: Optional[LiteLLM_ModelTable] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - updated_at: Optional[datetime] = None - created_at: Optional[datetime] = None - - ######################################################### - # Object Permission - MCP, Vector Stores etc. - ######################################################### - object_permission_id: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - dict_fields = [ - "metadata", - "aliases", - "config", - "permissions", - "model_max_budget", - "model_aliases", - "router_settings", - "budget_limits", - ] - - if isinstance(values, BaseModel): - values = values.model_dump() - - if ( - isinstance(values.get("members_with_roles"), dict) - and not values["members_with_roles"] - ): - values["members_with_roles"] = [] - - for field in dict_fields: - value = values.get(field) - if value is not None and isinstance(value, str): - try: - values[field] = json.loads(value) - except json.JSONDecodeError: - raise ValueError(f"Field {field} should be a valid dictionary") - - return values - - -class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable): - last_refreshed_at: Optional[float] = None - - -class LiteLLM_DeletedTeamTable(LiteLLM_TeamTable): - """ - Recording of deleted teams for audit purposes. Mirrors LiteLLM_TeamTable - plus metadata captured at deletion time. - """ - - id: Optional[str] = None - deleted_at: Optional[datetime] = None - deleted_by: Optional[str] = None - deleted_by_api_key: Optional[str] = None - litellm_changed_by: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) +from litellm.models.object_permission import ( # noqa: E402 + LiteLLM_ObjectPermissionTable as LiteLLM_ObjectPermissionTable, +) +from litellm.models.team import ( # noqa: E402 + LiteLLM_DeletedTeamTable as LiteLLM_DeletedTeamTable, +) +from litellm.models.team import LiteLLM_TeamTable as LiteLLM_TeamTable # noqa: E402 +from litellm.models.team import ( # noqa: E402 + LiteLLM_TeamTableCachedObj as LiteLLM_TeamTableCachedObj, +) class TeamRequest(LiteLLMPydanticObjectBase): teams: List[str] -class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): - """Represents user-controllable params for a LiteLLM_BudgetTable record. - - Budget-write paths use `model_fields.keys()` on this class as an allowlist - for user input. Keep server-managed fields (e.g. `budget_reset_at`) on - `LiteLLM_BudgetTableFull` so they aren't user-settable. - """ - - budget_id: Optional[str] = None - soft_budget: Optional[float] = None - max_budget: Optional[float] = None - max_parallel_requests: Optional[int] = None - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - model_max_budget: Optional[dict] = None - budget_duration: Optional[str] = None - allowed_models: Optional[List[str]] = ( - None # per-member model scope; empty = inherit team models - ) - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable): - """LiteLLM_BudgetTable + server-managed fields returned on API responses.""" - - budget_reset_at: Optional[datetime] = None - created_at: datetime - - -class LiteLLM_TeamMemberTable(LiteLLM_BudgetTable): - """ - Used to track spend of a user_id within a team_id - """ - - spend: Optional[float] = None - user_id: Optional[str] = None - team_id: Optional[str] = None - budget_id: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) +from litellm.models.budget import ( # noqa: E402 + LiteLLM_BudgetTable as LiteLLM_BudgetTable, +) +from litellm.models.budget import ( # noqa: E402 + LiteLLM_BudgetTableFull as LiteLLM_BudgetTableFull, +) +from litellm.models.budget import ( # noqa: E402 + LiteLLM_TeamMemberTable as LiteLLM_TeamMemberTable, +) class NewOrganizationRequest(LiteLLM_BudgetTable): @@ -2473,6 +2177,17 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "`statement_cache_size`). Keys here override any default LiteLLM sets." ), ) + database_disable_prepared_statements: Optional[bool] = Field( + None, + description=( + "Disable server-side prepared statements by setting Prisma's " + "`pgbouncer=true` URL param. Use this for pgbouncer transaction-pooling " + "deployments, or to prevent the 'cached plan must not change result " + "type' error that pooled connections hit during rolling schema " + "migrations. An explicit `pgbouncer` in `database_extra_connection_params` " + "takes precedence." + ), + ) database_type: Optional[Literal["dynamo_db"]] = Field( None, description="to use dynamodb instead of postgres db" ) @@ -2609,6 +2324,24 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).", ) + disable_budget_reservation: Optional[bool] = Field( + None, + description=( + "If True, disables the optimistic per-request budget reservation " + "introduced in v1.84.0. " + "WARNING: This weakens hard budget enforcement. Without the reservation, " + "a burst of concurrent requests from a single key can each pass the " + "read-time spend check before any of them is charged, allowing a " + "configured budget to be exceeded under high concurrency. " + "Budgets are still evaluated on every request at read time, so " + "an already-exhausted budget is still rejected. " + "Enable only if your deployment is experiencing phantom " + "BudgetExceededError responses caused by leaked reservations " + "(see GitHub issue #27639). " + "A proxy-level WARNING is logged on every request while this flag " + "is active as a reminder that hard enforcement is relaxed." + ), + ) class ConfigYAML(LiteLLMPydanticObjectBase): @@ -2637,66 +2370,12 @@ class ConfigYAML(LiteLLMPydanticObjectBase): model_config = ConfigDict(protected_namespaces=()) -class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): - token: Optional[str] = None - key_name: Optional[str] = None - key_alias: Optional[str] = None - spend: float = 0.0 - max_budget: Optional[float] = None - expires: Optional[Union[str, datetime]] = None - models: List = [] - aliases: Dict = {} - config: Dict = {} - user_id: Optional[str] = None - team_id: Optional[str] = None - agent_id: Optional[str] = None - project_id: Optional[str] = None - max_parallel_requests: Optional[int] = None - metadata: Dict = {} - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None - allowed_cache_controls: Optional[list] = [] - allowed_routes: Optional[list] = [] - permissions: Dict = {} - model_spend: Dict = {} - model_max_budget: Dict = {} - soft_budget_cooldown: bool = False - blocked: Optional[bool] = None - litellm_budget_table: Optional[dict] = None - org_id: Optional[str] = None # org id for a given key - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - updated_by: Optional[str] = None - last_active: Optional[datetime] = None - object_permission_id: Optional[str] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - access_group_ids: Optional[List[str]] = None - rotation_count: Optional[int] = 0 # Number of times key has been rotated - auto_rotate: Optional[bool] = False # Whether this key should be auto-rotated - rotation_interval: Optional[str] = None # How often to rotate (e.g., "30d", "90d") - last_rotation_at: Optional[datetime] = None # When this key was last rotated - key_rotation_at: Optional[datetime] = None # When this key should next be rotated - router_settings: Optional[dict] = None - budget_limits: Optional[List[dict]] = None # multiple concurrent budget windows - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken): - """ - Recording of deleted keys for audit purposes. Mirrors LiteLLM_VerificationToken - plus metadata captured at deletion time. - """ - - id: Optional[str] = None - deleted_at: Optional[datetime] = None - deleted_by: Optional[str] = None - deleted_by_api_key: Optional[str] = None - litellm_changed_by: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) +from litellm.models.verification_token import ( # noqa: E402 + LiteLLM_DeletedVerificationToken as LiteLLM_DeletedVerificationToken, +) +from litellm.models.verification_token import ( # noqa: E402 + LiteLLM_VerificationToken as LiteLLM_VerificationToken, +) class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): @@ -2935,39 +2614,10 @@ class UserInfoV2Response(LiteLLMPydanticObjectBase): teams: List[str] = [] # Just team IDs, not full team objects -class LiteLLM_Config(LiteLLMPydanticObjectBase): - param_name: str - param_value: Dict - - -class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): - """ - This is the table that track what organizations a user belongs to and users spend within the organization - """ - - user_id: str - organization_id: str - user_role: Optional[str] = None - spend: float = 0.0 - budget_id: Optional[str] = None - created_at: datetime - updated_at: datetime - user: Optional[Any] = ( - None # You might want to replace 'Any' with a more specific type if available - ) - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - user_email: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) - - @model_validator(mode="after") - def populate_user_email(self) -> "LiteLLM_OrganizationMembershipTable": - if self.user_email is None and self.user is not None: - if isinstance(self.user, dict): - self.user_email = self.user.get("user_email") - else: - self.user_email = getattr(self.user, "user_email", None) - return self +from litellm.models.config import LiteLLM_Config as LiteLLM_Config # noqa: E402 +from litellm.models.organization_membership import ( # noqa: E402 + LiteLLM_OrganizationMembershipTable as LiteLLM_OrganizationMembershipTable, +) class LiteLLM_OrganizationTableUpdate(LiteLLM_BudgetTable): @@ -2997,61 +2647,10 @@ def set_model_info(cls, values): return values -class LiteLLM_UserTable(LiteLLMPydanticObjectBase): - user_id: str - max_budget: Optional[float] = None - spend: float = 0.0 - model_max_budget: Optional[Dict] = {} - model_spend: Optional[Dict] = {} - user_email: Optional[str] = None - user_alias: Optional[str] = None - models: list = [] - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - user_role: Optional[str] = None - organization_memberships: Optional[List[LiteLLM_OrganizationMembershipTable]] = None - teams: List[str] = [] - sso_user_id: Optional[str] = None - budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None - metadata: Optional[dict] = None - created_at: Optional[datetime] = None - updated_at: Optional[datetime] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("spend") is None: - values.update({"spend": 0.0}) - if values.get("models") is None: - values.update({"models": []}) - if values.get("teams") is None: - values.update({"teams": []}) - return values - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_OrganizationTable(LiteLLMPydanticObjectBase): - """Represents user-controllable params for a LiteLLM_OrganizationTable record""" - - organization_id: Optional[str] = None - organization_alias: Optional[str] = None - budget_id: str - spend: float = 0.0 - metadata: Optional[dict] = None - models: List[str] - created_by: str - updated_by: str - users: Optional[List[LiteLLM_UserTable]] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - - ######################################################### - # Object Permission - MCP, Vector Stores etc. - ######################################################### - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - object_permission_id: Optional[str] = None +from litellm.models.organization import ( # noqa: E402 + LiteLLM_OrganizationTable as LiteLLM_OrganizationTable, +) +from litellm.models.user import LiteLLM_UserTable as LiteLLM_UserTable # noqa: E402 class LiteLLM_OrganizationTableWithMembers(LiteLLM_OrganizationTable): @@ -3160,28 +2759,9 @@ class DeleteProjectRequest(LiteLLMPydanticObjectBase): project_ids: List[str] -class LiteLLM_ProjectTable(LiteLLMPydanticObjectBase): - """Database model representation for project""" - - project_id: str - project_alias: Optional[str] = None - description: Optional[str] = None - team_id: Optional[str] = None - budget_id: Optional[str] = None - metadata: Optional[dict] = None - models: List[str] = [] - spend: float = 0.0 - model_spend: Optional[dict] = None - model_rpm_limit: Optional[dict] = None - model_tpm_limit: Optional[dict] = None - blocked: bool = False - object_permission_id: Optional[str] = None - created_by: str - updated_by: str - created_at: Optional[datetime] = None - updated_at: Optional[datetime] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None +from litellm.models.project import ( # noqa: E402 + LiteLLM_ProjectTable as LiteLLM_ProjectTable, +) class NewProjectResponse(LiteLLM_ProjectTable): @@ -3207,101 +2787,19 @@ class LiteLLM_UserTableWithKeyCount(LiteLLM_UserTable): key_count: int = 0 -class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase): - user_id: str - blocked: bool - alias: Optional[str] = None - spend: float = 0.0 - allowed_model_region: Optional[AllowedModelRegion] = None - default_model: Optional[str] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - object_permission_id: Optional[str] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("spend") is None: - values.update({"spend": 0.0}) - return values - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_TagTable(LiteLLMPydanticObjectBase): - tag_name: str - description: Optional[str] = None - models: List[str] = [] - model_info: Optional[dict] = None - spend: float = 0.0 - budget_id: Optional[str] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("spend") is None: - values.update({"spend": 0.0}) - if values.get("models") is None: - values.update({"models": []}) - return values - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_AccessGroupTable(LiteLLMPydanticObjectBase): - access_group_id: str - access_group_name: str - description: Optional[str] = None - access_model_names: List[str] = [] - access_mcp_server_ids: List[str] = [] - access_agent_ids: List[str] = [] - assigned_team_ids: List[str] = [] - assigned_key_ids: List[str] = [] - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - updated_by: Optional[str] = None - - -class LiteLLM_SpendLogs(LiteLLMPydanticObjectBase): - request_id: str - api_key: str - model: Optional[str] = "" - api_base: Optional[str] = "" - call_type: str - spend: Optional[float] = 0.0 - total_tokens: Optional[int] = 0 - prompt_tokens: Optional[int] = 0 - completion_tokens: Optional[int] = 0 - startTime: Union[str, datetime, None] - endTime: Union[str, datetime, None] - user: Optional[str] = "" - metadata: Optional[Json] = {} - cache_hit: Optional[str] = "False" - cache_key: Optional[str] = None - request_tags: Optional[Json] = None - requester_ip_address: Optional[str] = None - messages: Optional[Union[str, list, dict]] - response: Optional[Union[str, list, dict]] - - -class LiteLLM_ErrorLogs(LiteLLMPydanticObjectBase): - request_id: Optional[str] = str(uuid.uuid4()) - api_base: Optional[str] = "" - model_group: Optional[str] = "" - litellm_model_name: Optional[str] = "" - model_id: Optional[str] = "" - request_kwargs: Optional[dict] = {} - exception_type: Optional[str] = "" - status_code: Optional[str] = "" - exception_string: Optional[str] = "" - startTime: Union[str, datetime, None] - endTime: Union[str, datetime, None] - +from litellm.models.access_group import ( # noqa: E402 + LiteLLM_AccessGroupTable as LiteLLM_AccessGroupTable, +) +from litellm.models.end_user import ( # noqa: E402 + LiteLLM_EndUserTable as LiteLLM_EndUserTable, +) +from litellm.models.spend_logs import ( # noqa: E402 + LiteLLM_ErrorLogs as LiteLLM_ErrorLogs, +) +from litellm.models.spend_logs import ( # noqa: E402 + LiteLLM_SpendLogs as LiteLLM_SpendLogs, +) +from litellm.models.tag import LiteLLM_TagTable as LiteLLM_TagTable # noqa: E402 AUDIT_ACTIONS = Literal[ "created", "updated", "deleted", "blocked", "unblocked", "rotated" @@ -3982,29 +3480,9 @@ class CreatePassThroughEndpoint(LiteLLMPydanticObjectBase): headers: dict -class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): - user_id: str - team_id: str - budget_id: Optional[str] = None - spend: Optional[float] = 0.0 - total_spend: Optional[float] = 0.0 - # Union so Pydantic picks Full when data has server-managed fields - # (/team/info) and Base when callers/tests construct with only - # user-settable fields. - litellm_budget_table: Optional[ - Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable] - ] = None - - def safe_get_team_member_rpm_limit(self) -> Optional[int]: - if self.litellm_budget_table is not None: - return self.litellm_budget_table.rpm_limit - return None - - def safe_get_team_member_tpm_limit(self) -> Optional[int]: - if self.litellm_budget_table is not None: - return self.litellm_budget_table.tpm_limit - return None - +from litellm.models.team_membership import ( # noqa: E402 + LiteLLM_TeamMembership as LiteLLM_TeamMembership, +) #### Organization / Team Member Requests #### @@ -4125,6 +3603,10 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): rpm_limit: Optional[int] = Field( default=None, description="Requests per minute limit for this team member" ) + budget_duration: Optional[str] = Field( + default=None, + description="Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets.", + ) allowed_models: Optional[List[str]] = Field( default=None, description="List of models this team member can access. Pass an empty list to remove per-member model restrictions.", @@ -4136,6 +3618,7 @@ class TeamMemberUpdateResponse(MemberUpdateResponse): max_budget_in_team: Optional[float] = None tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None + budget_duration: Optional[str] = None allowed_models: Optional[List[str]] = None @@ -4363,6 +3846,7 @@ class UserManagementEndpointParamDocStringEnums(str, enum.Enum): EmbeddingResponse, VideoObject, StandardPassThroughResponseObject, + ResponsesAPIResponse, ] @@ -4724,6 +4208,16 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): default=None, description="Optional claim-based routing overrides for JWT-shaped tokens. Matching rules route requests to oauth2 before default JWT flow.", ) + team_claim_fallback: bool = Field( + default=False, + description=( + "If True, when a configured team_id_jwt_field / team_ids_jwt_field " + "claim is present but does not resolve to any known team, defer to " + "the single-team DB fallback (caller's only team membership) " + "instead of raising. Default False preserves strict claim-based " + "authorization." + ), + ) issuers: Optional[List[JWTIssuerConfig]] = Field( default=None, description="Optional issuer-bound JWT validation rules. When a token's `iss` matches a configured issuer, validation uses that issuer's JWKS, audience, and claim mappings. Tokens with an unlisted `iss` fall back to the global JWT_AUDIENCE/JWT_ISSUER validation path — this is additive routing, not an allow-list.", @@ -4917,39 +4411,18 @@ class ToolDiscoveryQueueItem(TypedDict, total=False): user_agent: Optional[str] # HTTP User-Agent of the caller -class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): - unified_file_id: str - file_object: Optional[OpenAIFileObject] = None - model_mappings: Dict[str, str] - flat_model_file_ids: List[str] - created_by: Optional[str] = None - team_id: Optional[str] = None - updated_by: Optional[str] = None - storage_backend: Optional[str] = None - storage_url: Optional[str] = None - - -class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): - unified_object_id: str - model_object_id: str - file_purpose: Literal["batch", "fine-tune", "response", "container"] - file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse] - created_by: Optional[str] = None - team_id: Optional[str] = None - - -class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): - """Table for managing vector stores with target_model_names support.""" - - unified_resource_id: str - resource_object: Optional[Any] = None # VectorStoreCreateResponse - model_mappings: Dict[str, str] - flat_model_resource_ids: List[str] - created_by: Optional[str] = None - team_id: Optional[str] = None - updated_by: Optional[str] = None - storage_backend: Optional[str] = None - storage_url: Optional[str] = None +from litellm.models.managed_files import ( # noqa: E402 + LiteLLM_ManagedFileTable as LiteLLM_ManagedFileTable, +) +from litellm.models.managed_files import ( # noqa: E402 + LiteLLM_ManagedObjectTable as LiteLLM_ManagedObjectTable, +) +from litellm.models.managed_files import ( # noqa: E402 + LiteLLM_ManagedVectorStoresTable as LiteLLM_ManagedVectorStoresTable, +) +from litellm.models.managed_files import ( # noqa: E402 + LiteLLM_ManagedVectorStoreTable as LiteLLM_ManagedVectorStoreTable, +) class EnterpriseLicenseData(TypedDict, total=False): @@ -4960,20 +4433,6 @@ class EnterpriseLicenseData(TypedDict, total=False): max_teams: int -class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase): - vector_store_id: str - custom_llm_provider: str - vector_store_name: Optional[str] - vector_store_description: Optional[str] - vector_store_metadata: Optional[Dict[str, Any]] - created_at: Optional[datetime] - updated_at: Optional[datetime] - litellm_credential_name: Optional[str] - litellm_params: Optional[Dict[str, Any]] - team_id: Optional[str] - user_id: Optional[str] - - class ResponseLiteLLM_ManagedVectorStore(TypedDict, total=False): vector_store: LiteLLM_ManagedVectorStoresTable diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 13a2dd9f040..11fd01e2369 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -9,6 +9,7 @@ handle_update_object_permission_common, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import AgentsRepository from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -174,7 +175,7 @@ async def add_agent_to_db( create_data[rate_field] = _val # Create agent in DB - created_agent = await prisma_client.db.litellm_agentstable.create( + created_agent = await AgentsRepository(prisma_client).table.create( data=create_data, include={"object_permission": True}, ) @@ -200,7 +201,7 @@ async def delete_agent_from_db( Delete an agent from the database """ try: - deleted_agent = await prisma_client.db.litellm_agentstable.delete( + deleted_agent = await AgentsRepository(prisma_client).table.delete( where={"agent_id": agent_id} ) return dict(deleted_agent) @@ -229,7 +230,7 @@ async def patch_agent_in_db( The patched agent """ try: - existing_agent = await prisma_client.db.litellm_agentstable.find_unique( + existing_agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if existing_agent is not None: @@ -282,7 +283,7 @@ async def patch_agent_in_db( if object_permission_id is not None: update_data["object_permission_id"] = object_permission_id # Patch agent in DB - patched_agent = await prisma_client.db.litellm_agentstable.update( + patched_agent = await AgentsRepository(prisma_client).table.update( where={"agent_id": agent_id}, data={ **update_data, @@ -368,9 +369,9 @@ async def update_agent_in_db( update_data[rate_field] = _val if agent.get("object_permission") is not None: - existing_agent = await prisma_client.db.litellm_agentstable.find_unique( - where={"agent_id": agent_id} - ) + existing_agent = await AgentsRepository( + prisma_client + ).table.find_unique(where={"agent_id": agent_id}) existing_object_permission_id = ( existing_agent.object_permission_id if existing_agent is not None @@ -386,7 +387,7 @@ async def update_agent_in_db( update_data["object_permission_id"] = object_permission_id # Update agent in DB - updated_agent = await prisma_client.db.litellm_agentstable.update( + updated_agent = await AgentsRepository(prisma_client).table.update( where={"agent_id": agent_id}, data=update_data, include={"object_permission": True}, @@ -414,7 +415,7 @@ async def get_all_agents_from_db( Get all agents from the database """ try: - agents_from_db = await prisma_client.db.litellm_agentstable.find_many( + agents_from_db = await AgentsRepository(prisma_client).table.find_many( order={"created_at": "desc"}, include={"object_permission": True}, ) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 42cf31e1e2b..2577615fc8e 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -9,11 +9,12 @@ from litellm._logging import verbose_logger from litellm.proxy._types import ( + UI_TEAM_ID, LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, - UI_TEAM_ID, UserAPIKeyAuth, ) +from litellm.repositories.table_repositories import AgentsRepository class AgentRequestHandler: @@ -298,7 +299,7 @@ async def _get_db_agent_ids_for_access_groups( agent_ids: Set[str] = set() if access_groups and prisma_client is not None: try: - agents = await prisma_client.db.litellm_agentstable.find_many( + agents = await AgentsRepository(prisma_client).table.find_many( where={"agent_access_groups": {"hasSome": access_groups}} ) for agent in agents: diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 19dbfe33d32..d19008856bd 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -17,6 +17,7 @@ import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.a2a.agent_card import merge_agent_card @@ -30,7 +31,6 @@ MakeAgentsPublicRequest, PatchAgentRequest, ) -from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, @@ -219,7 +219,7 @@ async def get_agents( if prisma_client is not None: agent_ids = [agent.agent_id for agent in returned_agents] if agent_ids: - db_agents = await prisma_client.db.litellm_agentstable.find_many( + db_agents = await AgentsRepository(prisma_client).table.find_many( where={"agent_id": {"in": agent_ids}}, ) spend_map = {a.agent_id: a.spend for a in db_agents} @@ -301,6 +301,7 @@ async def get_agents( from litellm.proxy.agent_endpoints.agent_registry import ( global_agent_registry as AGENT_REGISTRY, ) +from litellm.repositories.table_repositories import AgentsRepository @router.post( @@ -471,7 +472,7 @@ async def get_agent_by_id( try: agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: - agent_row = await prisma_client.db.litellm_agentstable.find_unique( + agent_row = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id}, include={"object_permission": True}, ) @@ -489,7 +490,7 @@ async def get_agent_by_id( agent = AgentResponse(**agent_dict) # type: ignore else: # Agent found in memory — refresh spend from DB - db_row = await prisma_client.db.litellm_agentstable.find_unique( + db_row = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if db_row is not None: @@ -570,7 +571,7 @@ async def update_agent( try: # Check if agent exists - existing_agent = await prisma_client.db.litellm_agentstable.find_unique( + existing_agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if existing_agent is not None: @@ -678,7 +679,7 @@ async def patch_agent( try: # Check if agent exists - existing_agent = await prisma_client.db.litellm_agentstable.find_unique( + existing_agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if existing_agent is not None: @@ -769,7 +770,7 @@ async def delete_agent( try: # Check if agent exists - existing_agent = await prisma_client.db.litellm_agentstable.find_unique( + existing_agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if existing_agent is not None: @@ -859,7 +860,7 @@ async def make_agent_public( agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: # check if agent exists in DB - agent = await prisma_client.db.litellm_agentstable.find_unique( + agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if agent is not None: @@ -982,7 +983,7 @@ async def make_agents_public( agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: # check if agent exists in DB - agent = await prisma_client.db.litellm_agentstable.find_unique( + agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if agent is not None: @@ -1082,7 +1083,7 @@ async def get_agent_daily_activity( if user_api_key_dict.user_id is None: permitted_agent_ids = [] else: - owned_records = await prisma_client.db.litellm_agentstable.find_many( + owned_records = await AgentsRepository(prisma_client).table.find_many( where={"created_by": user_api_key_dict.user_id} ) permitted_agent_ids = [a.agent_id for a in owned_records] @@ -1118,7 +1119,7 @@ async def get_agent_daily_activity( if agent_ids_list: where_condition["agent_id"] = {"in": list(agent_ids_list)} - agent_records = await prisma_client.db.litellm_agentstable.find_many( + agent_records = await AgentsRepository(prisma_client).table.find_many( where=where_condition ) agent_metadata = { 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 20b1659fa1a..dd7350e13ce 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 @@ -26,6 +26,7 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import ClaudeCodePluginRepository from litellm.types.proxy.claude_code_endpoints import ( ListPluginsResponse, PluginListItem, @@ -71,7 +72,7 @@ async def get_marketplace(): try: prisma_client = await _get_prisma_client() - plugins = await prisma_client.db.litellm_claudecodeplugintable.find_many( + plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many( where={"enabled": True} ) @@ -268,12 +269,12 @@ async def register_plugin( manifest["namespace"] = request.namespace # Check if plugin exists - existing = await prisma_client.db.litellm_claudecodeplugintable.find_unique( + existing = await ClaudeCodePluginRepository(prisma_client).table.find_unique( where={"name": request.name} ) if existing: - plugin = await prisma_client.db.litellm_claudecodeplugintable.update( + plugin = await ClaudeCodePluginRepository(prisma_client).table.update( where={"name": request.name}, data={ "version": request.version, @@ -285,7 +286,7 @@ async def register_plugin( ) action = "updated" else: - plugin = await prisma_client.db.litellm_claudecodeplugintable.create( + plugin = await ClaudeCodePluginRepository(prisma_client).table.create( data={ "name": request.name, "version": request.version, @@ -348,7 +349,7 @@ async def list_plugins( prisma_client = await _get_prisma_client() where = {"enabled": True} if enabled_only else {} - plugins = await prisma_client.db.litellm_claudecodeplugintable.find_many( + plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many( where=where ) @@ -415,7 +416,7 @@ async def get_plugin( try: prisma_client = await _get_prisma_client() - plugin = await prisma_client.db.litellm_claudecodeplugintable.find_unique( + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( where={"name": plugin_name} ) @@ -471,7 +472,7 @@ async def enable_plugin( try: prisma_client = await _get_prisma_client() - plugin = await prisma_client.db.litellm_claudecodeplugintable.find_unique( + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( where={"name": plugin_name} ) if not plugin: @@ -480,7 +481,7 @@ async def enable_plugin( detail={"error": f"Plugin '{plugin_name}' not found"}, ) - await prisma_client.db.litellm_claudecodeplugintable.update( + await ClaudeCodePluginRepository(prisma_client).table.update( where={"name": plugin_name}, data={"enabled": True, "updated_at": datetime.now(timezone.utc)}, ) @@ -516,7 +517,7 @@ async def disable_plugin( try: prisma_client = await _get_prisma_client() - plugin = await prisma_client.db.litellm_claudecodeplugintable.find_unique( + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( where={"name": plugin_name} ) if not plugin: @@ -525,7 +526,7 @@ async def disable_plugin( detail={"error": f"Plugin '{plugin_name}' not found"}, ) - await prisma_client.db.litellm_claudecodeplugintable.update( + await ClaudeCodePluginRepository(prisma_client).table.update( where={"name": plugin_name}, data={"enabled": False, "updated_at": datetime.now(timezone.utc)}, ) @@ -561,7 +562,7 @@ async def delete_plugin( try: prisma_client = await _get_prisma_client() - plugin = await prisma_client.db.litellm_claudecodeplugintable.find_unique( + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( where={"name": plugin_name} ) if not plugin: @@ -570,7 +571,7 @@ async def delete_plugin( detail={"error": f"Plugin '{plugin_name}' not found"}, ) - await prisma_client.db.litellm_claudecodeplugintable.delete( + await ClaudeCodePluginRepository(prisma_client).table.delete( where={"name": plugin_name} ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 94ae3f5eacc..6eae9d0d475 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -61,19 +61,33 @@ UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, _safe_get_request_query_params, ) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names, ) -from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.project_repository import ProjectRepository +from litellm.repositories.table_repositories import ( + AccessGroupRepository, + EndUserRepository, + JWTKeyMappingRepository, + ManagedVectorStoresRepository, + TagRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository from litellm.router import Router from litellm.utils import get_utc_datetime @@ -957,7 +971,7 @@ async def get_default_end_user_budget( # Fetch from database try: - budget_record = await prisma_client.db.litellm_budgettable.find_unique( + budget_record = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": litellm.max_end_user_budget_id} ) @@ -1016,7 +1030,7 @@ async def get_team_member_default_budget( return LiteLLM_BudgetTable(**cached_budget) try: - budget_record = await prisma_client.db.litellm_budgettable.find_unique( + budget_record = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": budget_id} ) @@ -1175,7 +1189,7 @@ async def get_end_user_object( # Fetch from database try: - response = await prisma_client.db.litellm_endusertable.find_unique( + response = await EndUserRepository(prisma_client).table.find_unique( where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True}, ) @@ -1375,7 +1389,7 @@ async def get_tag_objects_batch( # Batch fetch uncached tags from DB in one query if uncached_tags: try: - db_tags = await prisma_client.db.litellm_tagtable.find_many( + db_tags = await TagRepository(prisma_client).table.find_many( where={"tag_name": {"in": uncached_tags}}, include={"litellm_budget_table": True}, ) @@ -1469,7 +1483,7 @@ async def get_team_membership( # else, check db try: - response = await prisma_client.db.litellm_teammembership.find_unique( + response = await TeamMembershipRepository(prisma_client).table.find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}, ) @@ -1622,7 +1636,7 @@ async def _get_fuzzy_user_object( response = None if sso_user_id is not None: - response = await prisma_client.db.litellm_usertable.find_unique( + response = await UserRepository(prisma_client).table.find_unique( where={"sso_user_id": sso_user_id}, include={"organization_memberships": True}, ) @@ -1630,14 +1644,14 @@ async def _get_fuzzy_user_object( if response is None and user_email is not None: # Use case-insensitive query to handle emails with different casing # This matches the pattern used in _check_duplicate_user_email - response = await prisma_client.db.litellm_usertable.find_first( + response = await UserRepository(prisma_client).table.find_first( where={"user_email": {"equals": user_email, "mode": "insensitive"}}, include={"organization_memberships": True}, ) if response is not None and sso_user_id is not None: # update sso_user_id asyncio.create_task( # background task to update user with sso id - prisma_client.db.litellm_usertable.update( + UserRepository(prisma_client).table.update( where={"user_id": response.user_id}, data={"sso_user_id": sso_user_id}, ) @@ -1687,7 +1701,7 @@ async def get_user_object( ) if should_check_db: - response = await prisma_client.db.litellm_usertable.find_unique( + response = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id}, include={"organization_memberships": True} ) @@ -1711,7 +1725,7 @@ async def get_user_object( if litellm.default_internal_user_params is not None: new_user_params.update(litellm.default_internal_user_params) - response = await prisma_client.db.litellm_usertable.create( + response = await UserRepository(prisma_client).table.create( data=new_user_params, include={"organization_memberships": True}, ) @@ -1860,7 +1874,7 @@ async def _delete_cache_key_object( async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: Optional[bool] = None ): - response = await prisma_client.db.litellm_teamtable.find_unique( + response = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) @@ -1882,7 +1896,7 @@ async def _get_team_db_check( async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient): - return await prisma_client.db.litellm_teamtable.find_unique( + return await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) @@ -2111,7 +2125,7 @@ async def get_access_object( # Not in cache - fetch from DB try: - response = await prisma_client.db.litellm_accessgrouptable.find_unique( + response = await AccessGroupRepository(prisma_client).table.find_unique( where={"access_group_id": access_group_id} ) @@ -2193,7 +2207,7 @@ async def get_team_object_by_alias( # Query database by team_alias try: - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"team_alias": team_alias} ) @@ -2301,7 +2315,7 @@ async def get_org_object_by_alias( # Query database by organization_alias try: - orgs = await prisma_client.db.litellm_organizationtable.find_many( + orgs = await OrganizationRepository(prisma_client).table.find_many( where={"organization_alias": org_alias} ) @@ -2526,7 +2540,7 @@ async def get_jwt_key_mapping_object( Returns the hashed token (str) if a matching active mapping is found, else None. """ - mapping = await prisma_client.db.litellm_jwtkeymapping.find_first( + mapping = await JWTKeyMappingRepository(prisma_client).table.find_first( where={ "jwt_claim_name": jwt_claim_name, "jwt_claim_value": jwt_claim_value, @@ -2659,7 +2673,7 @@ async def get_object_permission( # else, check db try: - response = await prisma_client.db.litellm_objectpermissiontable.find_unique( + response = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": object_permission_id} ) @@ -2715,7 +2729,7 @@ async def get_managed_vector_store_rows_by_uuids( if not cache_misses: return result - rows = await prisma_client.db.litellm_managedvectorstorestable.find_many( + rows = await ManagedVectorStoresRepository(prisma_client).table.find_many( where={"vector_store_id": {"in": cache_misses}}, take=len(cache_misses), ) @@ -2790,7 +2804,7 @@ async def get_org_object( if include_budget_table: query_kwargs["include"] = {"litellm_budget_table": True} - response = await prisma_client.db.litellm_organizationtable.find_unique( + response = await OrganizationRepository(prisma_client).table.find_unique( **query_kwargs ) @@ -3502,10 +3516,13 @@ async def _virtual_key_max_budget_check( if valid_token.max_budget is not None: from litellm.proxy.proxy_server import get_current_spend + fallback_spend = valid_token.spend or 0.0 + counter_key = f"spend:key:{valid_token.token}" + # Read spend from cross-pod counter (Redis-first) or cached object (fallback) spend = await get_current_spend( - counter_key=f"spend:key:{valid_token.token}", - fallback_spend=valid_token.spend or 0.0, + counter_key=counter_key, + fallback_spend=fallback_spend, ) #################################### @@ -4180,7 +4197,7 @@ async def get_project_object( return deserialized_project # Fetch from DB - project_row = await prisma_client.db.litellm_projecttable.find_unique( + project_row = await ProjectRepository(prisma_client).table.find_unique( where={"project_id": project_id}, include={"litellm_budget_table": True}, ) @@ -4480,10 +4497,10 @@ async def vector_store_access_check( ######################################################### # Check if the key can access the vector store if valid_token is not None and valid_token.object_permission_id is not None: - key_object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": valid_token.object_permission_id}, - ) + key_object_permission = await ObjectPermissionRepository( + prisma_client + ).table.find_unique( + where={"object_permission_id": valid_token.object_permission_id}, ) if key_object_permission is not None: _can_object_call_vector_stores( @@ -4494,10 +4511,10 @@ async def vector_store_access_check( # Check if the team can access the vector store if team_object is not None and team_object.object_permission_id is not None: - team_object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": team_object.object_permission_id}, - ) + team_object_permission = await ObjectPermissionRepository( + prisma_client + ).table.find_unique( + where={"object_permission_id": team_object.object_permission_id}, ) if team_object_permission is not None: _can_object_call_vector_stores( diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index e06ac760237..83f18173182 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -126,6 +126,20 @@ async def _handle_authentication_error( model=request_data.get("model"), ) + # Budget checks live in tenant-scoped helpers (key / team / org / tag) + # that don't see the request model, so the BudgetExceededError they + # raise carries `llm_provider=""`. Resolve it here off `request_data` + # so custom-callback consumers reading StandardLoggingPayload get + # the same `llm_provider` attribution as for RPM/TPM 429s. + if isinstance(e, litellm.BudgetExceededError) and not e.llm_provider: + from litellm.proxy.hooks.rate_limiter_utils import ( + resolve_llm_provider_for_rate_limit, + ) + + _, e.llm_provider = resolve_llm_provider_for_rate_limit( + request_data.get("model") + ) + # Allow callbacks to transform the error response transformed_exception = await proxy_logging_obj.post_call_failure_hook( request_data=request_data, @@ -154,6 +168,16 @@ async def _handle_authentication_error( ) elif isinstance(e, ProxyException): raise e + if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + raise ProxyException( + message=( + "Service Unavailable, the authentication database is " + "temporarily unreachable. Please retry shortly." + ), + type=ProxyErrorTypes.no_db_connection, + param="None", + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) raise ProxyException( message="Authentication Error, " + str(e), type=ProxyErrorTypes.auth_error, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 71cf5197dec..c868d3d22b2 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -271,6 +271,11 @@ def _build_banned_observability_params() -> FrozenSet[str]: # tokens) to the attacker's host, or coerces the proxy into # authenticating against the attacker's host with admin secrets. "aws_bedrock_runtime_endpoint", + # Bedrock project/workspace association. Deployments pin this to + # enforce a data-retention policy, so a caller-supplied value would + # re-route the request's retention and accounting to any project + # reachable with the deployment's shared AWS credentials. + "aws_bedrock_project_id", # Provider-specific endpoint overrides that flow into the outbound # request via ``optional_params``. Same threat as ``api_base``: # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 6d3d49b71ec..fd6ff2ada7f 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -14,11 +14,11 @@ import re from typing import Any, List, Literal, Optional, Set, Tuple, Union, cast +import jwt from cryptography import x509 from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization from fastapi import HTTPException, status -import jwt from jwt.api_jwk import PyJWK from litellm._logging import verbose_proxy_logger @@ -50,6 +50,7 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.user_repository import UserRepository from .auth_checks import ( _allowed_routes_check, @@ -1298,15 +1299,29 @@ async def find_and_validate_specific_team_id( # First try to get team by team_id if individual_team_id: - team_object = await get_team_object( - team_id=individual_team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert, - ) - return individual_team_id, team_object + try: + team_object = await get_team_object( + team_id=individual_team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert, + ) + return individual_team_id, team_object + except HTTPException as e: + if ( + e.status_code != 404 + or not jwt_handler.litellm_jwtauth.team_claim_fallback + ): + raise + # Claim doesn't map to a known team — defer to fallback. + verbose_proxy_logger.debug( + "JWT team_id claim '%s' did not resolve to a team: %s", + individual_team_id, + e.detail, + ) + return None, None # If no team_id found, try to resolve via team_alias_jwt_field team_alias = jwt_handler.get_team_alias( @@ -1430,6 +1445,7 @@ async def find_team_with_model_access( ) return None, None + any_claim_team_resolved = False for team_id in team_ids: try: team_object = await get_team_object( @@ -1440,6 +1456,9 @@ async def find_team_with_model_access( proxy_logging_obj=proxy_logging_obj, ) + if team_object is not None: + any_claim_team_resolved = True + if team_object and team_object.models is not None: team_models = team_object.models if isinstance(team_models, list) and ( @@ -1477,12 +1496,17 @@ async def find_team_with_model_access( if denied_auth_enforced_pass_through_route: JWTAuthManager._raise_team_passthrough_route_denial(route=route) - if requested_model: + if requested_model and ( + any_claim_team_resolved + or not jwt_handler.litellm_jwtauth.team_claim_fallback + ): + # Claim resolved but no model access, or fallback disabled — deny. raise HTTPException( status_code=403, detail=f"No team has access to the requested model: {requested_model}. Checked teams={team_ids}. Check `/models` to see all available models.", ) + # No claim team resolved and fallback enabled — defer to fallback. return None, None @staticmethod @@ -1790,7 +1814,7 @@ async def sync_user_role_and_teams( # Update user role new_role = jwt_handler.map_jwt_role_to_litellm_role(jwt_valid_token) if new_role and user_object.user_role != new_role.value: - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": user_object.user_id}, data={"user_role": new_role.value}, ) diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 34085d5685a..d0818b95363 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -34,6 +34,7 @@ hash_password, verify_password, ) +from litellm.repositories.user_repository import UserRepository from litellm.secret_managers.main import get_secret_bool from litellm.types.proxy.ui_sso import ReturnedUITokenObject @@ -45,7 +46,7 @@ async def _rehash_password_if_needed(user_id: str, password: str, stored: str) - from litellm.proxy.proxy_server import prisma_client if prisma_client is not None: - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": user_id}, data={"password": hash_password(password)}, ) @@ -151,7 +152,7 @@ async def authenticate_user( # noqa: PLR0915 if prisma_client is not None: _user_row = cast( Optional[LiteLLM_UserTable], - await prisma_client.db.litellm_usertable.find_first( + await UserRepository(prisma_client).table.find_first( where={"user_email": {"equals": username, "mode": "insensitive"}} ), ) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index d9600e1a4b4..00f276dc970 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -6,6 +6,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth +from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.router import Router from litellm.router_utils.fallback_event_handlers import get_fallback_model_group from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params @@ -86,7 +87,7 @@ async def get_mcp_server_ids( # Make a direct SQL query to get just the mcp_servers try: - result = await prisma_client.db.litellm_objectpermissiontable.find_unique( + result = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": user_api_key_dict.object_permission_id}, ) if result and result.mcp_servers: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a5501fefa4e..666c01562b5 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -21,9 +21,9 @@ import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.integrations.otel.runtime import phase_span, seed_request_identity -from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * @@ -65,7 +65,6 @@ from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -73,12 +72,14 @@ populate_request_with_path_params, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.utils import ( PrismaClient, ProxyLogging, normalize_route_for_root_path, ) +from litellm.repositories.table_repositories import TeamMembershipRepository from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes @@ -1797,7 +1798,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 _team_id = valid_token.team_id if _user_id is not None and _team_id is not None: - _db_member = await prisma_client.db.litellm_teammembership.find_first( + _db_member = await TeamMembershipRepository( + prisma_client + ).table.find_first( where={ "user_id": _user_id, "team_id": _team_id, @@ -2422,6 +2425,7 @@ async def _run_centralized_common_checks( # noqa: PLR0915 user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, skip_budget_checks=skip_budget_checks, + general_settings=general_settings, ) @@ -2442,12 +2446,23 @@ async def _reserve_budget_after_common_checks( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, skip_budget_checks: bool, + general_settings: dict, end_user_id: Optional[str] = None, end_user_object: Optional[LiteLLM_EndUserTable] = None, ) -> None: user_api_key_auth_obj.budget_reservation = None if skip_budget_checks: return + if general_settings.get("disable_budget_reservation") is True: + verbose_proxy_logger.warning( + "disable_budget_reservation is enabled: skipping optimistic budget " + "reservation. Budget enforcement is read-time only — concurrent " + "requests can each pass the spend check before their cost is recorded, " + "so a configured budget may be briefly exceeded under high concurrency. " + "Set disable_budget_reservation to False or remove it to restore " + "hard per-request budget enforcement." + ) + return from litellm.proxy.spend_tracking.budget_reservation import ( reserve_budget_for_request, diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index 9fbc6f2197d..c2ce28884c7 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -338,9 +338,9 @@ sequenceDiagram The CLI provides three authentication commands: -- **`litellm-proxy login`** - Start SSO authentication flow -- **`litellm-proxy logout`** - Clear stored authentication token -- **`litellm-proxy whoami`** - Show current authentication status +- **`lite login`** - Start SSO authentication flow +- **`lite logout`** - Clear stored authentication token +- **`lite whoami`** - Show current authentication status ### Authentication Flow Steps @@ -382,14 +382,14 @@ Once authenticated, the CLI will automatically use the stored token for all requ ```bash # Login -litellm-proxy login +lite login # Use CLI without specifying API key -litellm-proxy models list +lite models list # Check authentication status -litellm-proxy whoami +lite whoami # Logout -litellm-proxy logout +lite logout ``` diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 6ef837cb521..333e2029e46 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -22,11 +22,11 @@ The CLI can be configured using environment variables or command-line options: Example: ```bash -litellm-proxy version +lite version # or -litellm-proxy --version +lite --version # or -litellm-proxy -v +lite -v ``` ## Commands @@ -40,7 +40,7 @@ The CLI provides several commands for managing models on your LiteLLM proxy serv View all available models: ```bash -litellm-proxy models list [--format table|json] +lite models list [--format table|json] ``` Options: @@ -52,7 +52,7 @@ Options: Get detailed information about all models: ```bash -litellm-proxy models info [options] +lite models info [options] ``` Options: @@ -75,7 +75,7 @@ Default columns: `public_model`, `upstream_model`, `updated_at` Add a new model to the proxy: ```bash -litellm-proxy models add [options] +lite models add [options] ``` Options: @@ -86,7 +86,7 @@ Options: Example: ```bash -litellm-proxy models add gpt-4 -p api_key=sk-123 -p api_base=https://api.openai.com -i description="GPT-4 model" +lite models add gpt-4 -p api_key=sk-123 -p api_base=https://api.openai.com -i description="GPT-4 model" ``` #### Get Model Info @@ -94,7 +94,7 @@ litellm-proxy models add gpt-4 -p api_key=sk-123 -p api_base=https://api.openai. Get information about a specific model: ```bash -litellm-proxy models get [--id MODEL_ID] [--name MODEL_NAME] +lite models get [--id MODEL_ID] [--name MODEL_NAME] ``` Options: @@ -107,7 +107,7 @@ Options: Delete a model from the proxy: ```bash -litellm-proxy models delete +lite models delete ``` #### Update Model @@ -115,7 +115,7 @@ litellm-proxy models delete Update an existing model's configuration: ```bash -litellm-proxy models update [options] +lite models update [options] ``` Options: @@ -128,7 +128,7 @@ Options: Import models from a YAML file: ```bash -litellm-proxy models import models.yaml +lite models import models.yaml ``` Options: @@ -142,31 +142,31 @@ Examples: 1. Import all models from a YAML file: ```bash -litellm-proxy models import models.yaml +lite models import models.yaml ``` 2. Dry run (show what would be imported): ```bash -litellm-proxy models import models.yaml --dry-run +lite models import models.yaml --dry-run ``` 3. Only import models where the model name contains 'gpt': ```bash -litellm-proxy models import models.yaml --only-models-matching-regex gpt +lite models import models.yaml --only-models-matching-regex gpt ``` 4. Only import models with access group containing 'beta': ```bash -litellm-proxy models import models.yaml --only-access-groups-matching-regex beta +lite models import models.yaml --only-access-groups-matching-regex beta ``` 5. Combine both filters: ```bash -litellm-proxy models import models.yaml --only-models-matching-regex gpt --only-access-groups-matching-regex beta +lite models import models.yaml --only-models-matching-regex gpt --only-access-groups-matching-regex beta ``` ### Credentials Management @@ -178,7 +178,7 @@ The CLI provides commands for managing credentials on your LiteLLM proxy server: View all available credentials: ```bash -litellm-proxy credentials list [--format table|json] +lite credentials list [--format table|json] ``` Options: @@ -194,7 +194,7 @@ The table format displays: Create a new credential: ```bash -litellm-proxy credentials create --info --values +lite credentials create --info --values ``` Options: @@ -205,7 +205,7 @@ Options: Example: ```bash -litellm-proxy credentials create azure-cred \ +lite credentials create azure-cred \ --info '{"custom_llm_provider": "azure"}' \ --values '{"api_key": "sk-123", "api_base": "https://example.azure.openai.com"}' ``` @@ -215,7 +215,7 @@ litellm-proxy credentials create azure-cred \ Get information about a specific credential: ```bash -litellm-proxy credentials get +lite credentials get ``` #### Delete Credential @@ -223,7 +223,7 @@ litellm-proxy credentials get Delete a credential: ```bash -litellm-proxy credentials delete +lite credentials delete ``` ### Keys Management @@ -235,7 +235,7 @@ The CLI provides commands for managing API keys on your LiteLLM proxy server: View all API keys: ```bash -litellm-proxy keys list [--format table|json] [options] +lite keys list [--format table|json] [options] ``` Options: @@ -256,7 +256,7 @@ Options: Generate a new API key: ```bash -litellm-proxy keys generate [options] +lite keys generate [options] ``` Options: @@ -274,7 +274,7 @@ Options: Example: ```bash -litellm-proxy keys generate --models gpt-4,gpt-3.5-turbo --spend 100 --duration 24h --key-alias my-key --team-id team123 +lite keys generate --models gpt-4,gpt-3.5-turbo --spend 100 --duration 24h --key-alias my-key --team-id team123 ``` #### Delete Keys @@ -282,7 +282,7 @@ litellm-proxy keys generate --models gpt-4,gpt-3.5-turbo --spend 100 --duration Delete API keys by key or alias: ```bash -litellm-proxy keys delete [--keys ] [--key-aliases ] +lite keys delete [--keys ] [--key-aliases ] ``` Options: @@ -293,7 +293,7 @@ Options: Example: ```bash -litellm-proxy keys delete --keys sk-key1,sk-key2 --key-aliases alias1,alias2 +lite keys delete --keys sk-key1,sk-key2 --key-aliases alias1,alias2 ``` #### Get Key Info @@ -301,7 +301,7 @@ litellm-proxy keys delete --keys sk-key1,sk-key2 --key-aliases alias1,alias2 Get information about a specific API key: ```bash -litellm-proxy keys info --key +lite keys info --key ``` Options: @@ -311,7 +311,7 @@ Options: Example: ```bash -litellm-proxy keys info --key sk-key1 +lite keys info --key sk-key1 ``` ### User Management @@ -323,7 +323,7 @@ The CLI provides commands for managing users on your LiteLLM proxy server: View all users: ```bash -litellm-proxy users list +lite users list ``` #### Get User Info @@ -331,7 +331,7 @@ litellm-proxy users list Get information about a specific user: ```bash -litellm-proxy users get --id +lite users get --id ``` #### Create User @@ -339,7 +339,7 @@ litellm-proxy users get --id Create a new user: ```bash -litellm-proxy users create --email user@example.com --role internal_user --alias "Alice" --team team1 --max-budget 100.0 +lite users create --email user@example.com --role internal_user --alias "Alice" --team team1 --max-budget 100.0 ``` #### Delete User @@ -347,7 +347,7 @@ litellm-proxy users create --email user@example.com --role internal_user --alias Delete one or more users by user_id: ```bash -litellm-proxy users delete +lite users delete ``` ### Chat Commands @@ -359,7 +359,7 @@ The CLI provides commands for interacting with chat models through your LiteLLM Create a chat completion: ```bash -litellm-proxy chat completions [options] +lite chat completions [options] ``` Arguments: @@ -379,12 +379,12 @@ Examples: 1. Simple completion: ```bash -litellm-proxy chat completions gpt-4 -m "user:Hello, how are you?" +lite chat completions gpt-4 -m "user:Hello, how are you?" ``` 2. Multi-message conversation: ```bash -litellm-proxy chat completions gpt-4 \ +lite chat completions gpt-4 \ -m "system:You are a helpful assistant" \ -m "user:What's the capital of France?" \ -m "assistant:The capital of France is Paris." \ @@ -393,7 +393,7 @@ litellm-proxy chat completions gpt-4 \ 3. With generation parameters: ```bash -litellm-proxy chat completions gpt-4 \ +lite chat completions gpt-4 \ -m "user:Write a story" \ --temperature 0.7 \ --max-tokens 500 \ @@ -409,7 +409,7 @@ The CLI provides commands for making direct HTTP requests to your LiteLLM proxy Make an HTTP request to any endpoint: ```bash -litellm-proxy http request [options] +lite http request [options] ``` Arguments: @@ -425,19 +425,46 @@ Examples: 1. List models: ```bash -litellm-proxy http request GET /models +lite http request GET /models ``` 2. Create a chat completion: ```bash -litellm-proxy http request POST /chat/completions -j '{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}' +lite http request POST /chat/completions -j '{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}' ``` 3. Test connection with custom headers: ```bash -litellm-proxy http request GET /health/test_connection -H "X-Custom-Header:value" +lite http request GET /health/test_connection -H "X-Custom-Header:value" ``` +### Run a Coding Agent + +Launch a coding agent with all of its LLM traffic routed through your LiteLLM proxy. Each supported agent is its own command, so there is nothing to remember beyond the agent's name: + +```bash +lite claude +lite codex +lite opencode +``` + +Anything you type after the agent name is forwarded to it untouched, so the usual flags keep working: + +```bash +lite claude --resume +lite codex exec "summarize the repo" +``` + +Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process. + +The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). + +Options (these belong to the wrapper, so put them before the agent's own flags): + +- `--skip-verify`: Skip the pre-launch key check (useful offline or with non-standard auth). + +To pin the model, pass the agent's own model flag (for example `lite claude --model my-proxy-model` or `lite codex -m my-proxy-model`), or export the variable the agent reads (`ANTHROPIC_MODEL` / `ANTHROPIC_SMALL_FAST_MODEL` for Claude Code); the wrapper preserves anything you already have set. Whatever model the agent ends up requesting must exist on the proxy, since requests land on the proxy's `/v1/messages` (Anthropic) or `/v1/chat/completions` and `/v1/responses` (OpenAI) endpoints. + ## Environment Variables The CLI respects the following environment variables: @@ -450,37 +477,37 @@ The CLI respects the following environment variables: 1. List all models in table format: ```bash -litellm-proxy models list +lite models list ``` 2. Add a new model with parameters: ```bash -litellm-proxy models add gpt-4 -p api_key=sk-123 -p max_tokens=2048 +lite models add gpt-4 -p api_key=sk-123 -p max_tokens=2048 ``` 3. Get model information in JSON format: ```bash -litellm-proxy models info --format json +lite models info --format json ``` 4. Update model parameters: ```bash -litellm-proxy models update model-123 -p temperature=0.7 -i description="Updated model" +lite models update model-123 -p temperature=0.7 -i description="Updated model" ``` 5. List all credentials in table format: ```bash -litellm-proxy credentials list +lite credentials list ``` 6. Create a new credential for Azure: ```bash -litellm-proxy credentials create azure-prod \ +lite credentials create azure-prod \ --info '{"custom_llm_provider": "azure"}' \ --values '{"api_key": "sk-123", "api_base": "https://prod.azure.openai.com"}' ``` @@ -488,7 +515,7 @@ litellm-proxy credentials create azure-prod \ 7. Make a custom HTTP request: ```bash -litellm-proxy http request POST /chat/completions \ +lite http request POST /chat/completions \ -j '{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}' \ -H "X-Custom-Header:value" ``` @@ -497,29 +524,29 @@ litellm-proxy http request POST /chat/completions \ ```bash # List users -litellm-proxy users list +lite users list # Get user info -litellm-proxy users get --id u1 +lite users get --id u1 # Create a user -litellm-proxy users create --email a@b.com --role internal_user --alias "Alice" --team team1 --max-budget 100.0 +lite users create --email a@b.com --role internal_user --alias "Alice" --team team1 --max-budget 100.0 # Delete users -litellm-proxy users delete u1 u2 +lite users delete u1 u2 ``` 9. Import models from a YAML file (with filters): ```bash # Only import models where the model name contains 'gpt' -litellm-proxy models import models.yaml --only-models-matching-regex gpt +lite models import models.yaml --only-models-matching-regex gpt # Only import models with access group containing 'beta' -litellm-proxy models import models.yaml --only-access-groups-matching-regex beta +lite models import models.yaml --only-access-groups-matching-regex beta # Combine both filters -litellm-proxy models import models.yaml --only-models-matching-regex gpt --only-access-groups-matching-regex beta +lite models import models.yaml --only-models-matching-regex gpt --only-access-groups-matching-regex beta ``` ## Error Handling diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py new file mode 100644 index 00000000000..f39ffb3e864 --- /dev/null +++ b/litellm/proxy/client/cli/commands/agents.py @@ -0,0 +1,303 @@ +import os +import shutil +import sys +from typing import Callable, Dict, FrozenSet, List, Mapping, Optional, Sequence, Tuple + +import click +import requests + +from .auth import get_stored_api_key, login + +ANTHROPIC_BASE_URL_ENV = "ANTHROPIC_BASE_URL" +ANTHROPIC_AUTH_TOKEN_ENV = "ANTHROPIC_AUTH_TOKEN" +ANTHROPIC_API_KEY_ENV = "ANTHROPIC_API_KEY" +OPENAI_BASE_URL_ENV = "OPENAI_BASE_URL" +OPENAI_API_KEY_ENV = "OPENAI_API_KEY" + +PROFILE_ANTHROPIC = "anthropic" +PROFILE_OPENAI = "openai" + +_KNOWN_AGENTS: Dict[str, Tuple[str, FrozenSet[str]]] = { + "claude": ("Claude Code", frozenset({PROFILE_ANTHROPIC})), + "codex": ("Codex", frozenset({PROFILE_OPENAI})), + "opencode": ("OpenCode", frozenset({PROFILE_OPENAI})), +} + +_INSTALL_DOCS: Dict[str, str] = { + "claude": "https://docs.claude.com/en/docs/claude-code/setup", + "codex": "https://developers.openai.com/codex/cli", + "opencode": "https://opencode.ai/docs", +} + +CODEX_PROXY_PROVIDER = "litellm" + + +class AgentRunError(Exception): + """Raised for any user-actionable failure while preparing to run an agent.""" + + +def agent_profile(command: str) -> Tuple[str, FrozenSet[str]]: + """Return the (display name, env profiles) for a wrapped command. + + Known agents map to the API family they speak. Anything else gets both + families so it works regardless of which env vars the tool reads. + """ + base = os.path.basename(command) + if base in _KNOWN_AGENTS: + return _KNOWN_AGENTS[base] + return base, frozenset({PROFILE_ANTHROPIC, PROFILE_OPENAI}) + + +def build_agent_env( + base_env: Mapping[str, str], + base_url: str, + api_key: str, + profiles: FrozenSet[str], +) -> Dict[str, str]: + """Return a copy of base_env wired to route the agent through the proxy. + + Anthropic clients (Claude Code) append /v1/messages to ANTHROPIC_BASE_URL, + so it stays the bare proxy root; OpenAI clients (Codex, OpenCode) expect the + /v1 suffix on OPENAI_BASE_URL. ANTHROPIC_API_KEY is dropped so a stray + Anthropic key cannot win over the bearer token we set. + """ + env = dict(base_env) + root = base_url.rstrip("/") + if PROFILE_ANTHROPIC in profiles: + env[ANTHROPIC_BASE_URL_ENV] = root + env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key + env.pop(ANTHROPIC_API_KEY_ENV, None) + if PROFILE_OPENAI in profiles: + env[OPENAI_BASE_URL_ENV] = root + "/v1" + env[OPENAI_API_KEY_ENV] = api_key + return env + + +def _codex_proxy_args(base_url: str) -> List[str]: + """Codex `-c` overrides that point it at the proxy. + + Codex ignores OPENAI_BASE_URL (it always dials api.openai.com), so the env + profile alone cannot route it. It does honor a custom provider, so define one + inline; supports_websockets=false forces the HTTP/SSE Responses transport + because the proxy does not speak the Responses WebSocket protocol. The key is + read from OPENAI_API_KEY, which build_agent_env already exports. + """ + root = base_url.rstrip("/") + "/v1" + provider = f"model_providers.{CODEX_PROXY_PROVIDER}" + return [ + "-c", + f'model_provider="{CODEX_PROXY_PROVIDER}"', + "-c", + f'{provider}.name="LiteLLM proxy"', + "-c", + f'{provider}.base_url="{root}"', + "-c", + f'{provider}.env_key="{OPENAI_API_KEY_ENV}"', + "-c", + f'{provider}.wire_api="responses"', + "-c", + f"{provider}.supports_websockets=false", + ] + + +_PROXY_ARGS: Dict[str, Callable[[str], List[str]]] = { + "codex": _codex_proxy_args, +} + + +def agent_launch_args(command: str, base_url: str) -> List[str]: + """Extra CLI args an agent needs to actually honor the proxy. + + Claude Code and OpenCode respect the exported env vars, so they get nothing + here; Codex needs its provider pointed via config overrides. + """ + builder = _PROXY_ARGS.get(os.path.basename(command)) + return builder(base_url) if builder else [] + + +def verify_proxy_key( + base_url: str, + api_key: str, + *, + get: Callable[..., requests.Response] = requests.get, +) -> None: + """Probe the proxy with the key so bad creds fail here, not inside the agent. + + Raises AgentRunError when the proxy is unreachable or rejects the key. Other + non-2xx responses are tolerated; the agent's own call is the real test. + """ + url = base_url.rstrip("/") + "/v1/models" + try: + resp = get(url, headers={"Authorization": f"Bearer {api_key}"}, timeout=10) + except requests.RequestException as e: + raise AgentRunError( + f"Could not reach the LiteLLM proxy at {base_url.rstrip('/')}: {e}. " + "Is it running, and is --base-url (or LITELLM_PROXY_URL) correct?" + ) + if resp.status_code in (401, 403): + raise AgentRunError( + f"LiteLLM rejected your key (HTTP {resp.status_code}). " + "Run `lite login` to refresh it, or pass a valid --api-key." + ) + + +def _exec(path: str, args: Sequence[str], env: Mapping[str, str]) -> None: + os.execvpe(path, list(args), dict(env)) + + +def _restore_controlling_terminal() -> None: + """Reattach the controlling terminal to stdin before handing off to the agent. + + Completing the browser SSO login can leave stdin detached from the terminal, + which makes a TUI agent like Claude Code start in non-interactive mode and + exit immediately. Reopening /dev/tty onto fd 0 gives the agent a live + terminal; when stdin is still a tty (no login happened) this is a no-op. + """ + if sys.stdin.isatty(): + return + try: + fd = os.open("/dev/tty", os.O_RDONLY) + except OSError: + return + try: + os.dup2(fd, 0) + finally: + os.close(fd) + + +def run_agent( + base_url: str, + api_key: str, + command: Sequence[str], + *, + skip_verify: bool = False, + base_env: Optional[Mapping[str, str]] = None, + which: Callable[[str], Optional[str]] = shutil.which, + verify: Callable[[str, str], None] = verify_proxy_key, + launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _exec, + reattach_terminal: Optional[Callable[[], None]] = None, +) -> None: + """Validate, wire the environment, and hand off to the agent. + + On success this replaces the current process and never returns. Raises + AgentRunError for missing binaries, an unreachable proxy, or a rejected key. + reattach_terminal, when given, runs just before handoff to restore stdin. + """ + if not command: + raise AgentRunError("Nothing to run.") + + _, profiles = agent_profile(command[0]) + binary = which(command[0]) + if binary is None: + docs = _INSTALL_DOCS.get(os.path.basename(command[0])) + hint = f" Install it first: {docs}" if docs else "" + raise AgentRunError(f"Could not find `{command[0]}` on your PATH.{hint}") + + if not skip_verify: + verify(base_url, api_key) + + env = build_agent_env( + base_env if base_env is not None else os.environ, + base_url, + api_key, + profiles, + ) + extra_args = agent_launch_args(command[0], base_url) + if reattach_terminal is not None: + reattach_terminal() + launcher(binary, [command[0], *extra_args, *command[1:]], env) + + +def _is_interactive() -> bool: + return sys.stdin.isatty() + + +def _resolve_api_key(ctx: click.Context) -> str: + base_url = ctx.obj["base_url"] + api_key = ctx.obj.get("api_key") + if api_key: + return api_key + + if not _is_interactive(): + raise click.ClickException( + "No LiteLLM key found. Set LITELLM_PROXY_API_KEY (or pass --api-key) for " + "non-interactive use, or run `lite login` from a terminal." + ) + + click.echo("No LiteLLM credentials found; starting login...") + ctx.invoke(login) + api_key = get_stored_api_key(expected_base_url=base_url) + if not api_key: + raise click.ClickException( + "Login did not produce an API key; cannot start the agent." + ) + return api_key + + +_SKIP_VERIFY_HELP = "Skip the pre-launch key check against the proxy." + + +def _launch( + ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool +) -> None: + base_url = ctx.obj["base_url"] + started_interactive = _is_interactive() + api_key = _resolve_api_key(ctx) + + display_name, _ = agent_profile(binary) + click.echo( + f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}" + ) + + try: + run_agent( + base_url, + api_key, + [binary, *args], + skip_verify=skip_verify, + reattach_terminal=( + _restore_controlling_terminal if started_interactive else None + ), + ) + except AgentRunError as e: + raise click.ClickException(str(e)) + + +def _make_agent_command(binary: str, display_name: str) -> click.Command: + @click.command( + name=binary, + context_settings={"ignore_unknown_options": True}, + short_help=f"Run {display_name} through your LiteLLM proxy", + ) + @click.option("--skip-verify", is_flag=True, default=False, help=_SKIP_VERIFY_HELP) + @click.argument("args", nargs=-1, type=click.UNPROCESSED) + @click.pass_context + def _command(ctx: click.Context, skip_verify: bool, args: Sequence[str]) -> None: + _launch(ctx, binary, list(args), skip_verify=skip_verify) + + _command.help = ( + f"Run {display_name} routed through your LiteLLM proxy.\n\n" + f"Logs in with LiteLLM if needed, verifies your key against the proxy, " + f"exports the env vars {binary} reads, then hands off. Any arguments are " + f"forwarded to `{binary}`." + ) + return _command + + +def agent_commands() -> List[click.Command]: + """Build one top-level command per known agent, e.g. `lite claude`.""" + return [ + _make_agent_command(binary, name) + for binary, (name, _profiles) in _KNOWN_AGENTS.items() + ] + + +__all__ = [ + "agent_commands", + "run_agent", + "build_agent_env", + "agent_launch_args", + "verify_proxy_key", + "agent_profile", + "AgentRunError", +] diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 447837c35e7..b06d86d5965 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -624,7 +624,7 @@ def whoami(): token_data = load_token() if not token_data: - click.echo("❌ Not authenticated. Run 'litellm-proxy login' to authenticate.") + click.echo("❌ Not authenticated. Run 'lite login' to authenticate.") return click.echo("✅ Authenticated") diff --git a/litellm/proxy/client/cli/commands/chat.py b/litellm/proxy/client/cli/commands/chat.py index a078b766107..696e34c3ecd 100644 --- a/litellm/proxy/client/cli/commands/chat.py +++ b/litellm/proxy/client/cli/commands/chat.py @@ -122,13 +122,13 @@ def chat( Examples: # Chat with a specific model - litellm-proxy chat gpt-4 + lite chat gpt-4 # Chat without specifying model (will show model selection) - litellm-proxy chat + lite chat # Chat with custom settings - litellm-proxy chat gpt-4 --temperature 0.9 --system "You are a helpful coding assistant" + lite chat gpt-4 --temperature 0.9 --system "You are a helpful coding assistant" """ console = Console() diff --git a/litellm/proxy/client/cli/interface.py b/litellm/proxy/client/cli/interface.py index eba693dc18e..a32d60aadd9 100644 --- a/litellm/proxy/client/cli/interface.py +++ b/litellm/proxy/client/cli/interface.py @@ -80,6 +80,8 @@ def styled_prompt(): def show_commands(): """Display available commands.""" + from .commands.agents import agent_commands + commands = [ ("login", "Authenticate with the LiteLLM proxy server"), ("logout", "Clear stored authentication"), @@ -91,6 +93,9 @@ def show_commands(): ("keys", "Manage API keys"), ("teams", "Manage teams and team assignments"), ("users", "Manage users"), + ] + commands += [(c.name, c.get_short_help_str()) for c in agent_commands()] + commands += [ ("version", "Show version information"), ("help", "Show this help message"), ("quit", "Exit the interactive session"), @@ -156,7 +161,7 @@ def execute_command(user_input: str, ctx: click.Context): # Execute the command try: # Create a new argument list for click to parse - sys.argv = ["litellm-proxy"] + [command] + args + sys.argv = ["lite"] + [command] + args # Get the command object and invoke it cmd = cli.commands[command] diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index be55f79c066..b8c483f4b08 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -7,6 +7,7 @@ from litellm._version import version as litellm_version from litellm.proxy.client.health import HealthManagementClient +from .commands.agents import agent_commands from .commands.auth import get_stored_api_key, login, logout, whoami from .commands.chat import chat from .commands.credentials import credentials @@ -112,6 +113,9 @@ def version(ctx: click.Context): cli.add_command(teams) # Add the users command group cli.add_command(users) +# Add a top-level command per coding agent (claude, codex, opencode, ...) +for agent_command in agent_commands(): + cli.add_command(agent_command) if __name__ == "__main__": diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6558543370d..b9a9f3cebb7 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1949,12 +1949,26 @@ async def _handle_llm_api_exception( code=status.HTTP_400_BAD_REQUEST, headers=headers, ) + # Extract status_code from the exception if it carries one. + # Provider exceptions (NotFoundError, BadRequestError, GeminiError, + # VertexAIError, etc.) all have a status_code attribute reflecting + # the upstream API response. Use it to return the correct HTTP code + # instead of defaulting to 500. + _exc_status_code = getattr(e, "status_code", None) + if ( + _exc_status_code is not None + and isinstance(_exc_status_code, int) + and 400 <= _exc_status_code <= 599 + ): + _code = _exc_status_code + else: + _code = status.HTTP_500_INTERNAL_SERVER_ERROR raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), openai_code=getattr(e, "code", None), - code=getattr(e, "status_code", 500), + code=_code, provider_specific_fields=getattr(e, "provider_specific_fields", None), headers=headers, ) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index a65e737f248..c630294c1ec 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -40,8 +40,10 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 premium_user: bool, config_file_path: str, litellm_settings: dict, - callback_specific_params: dict = {}, + callback_specific_params: Optional[dict] = None, ): + if not isinstance(callback_specific_params, dict): + callback_specific_params = {} from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_callback_manager import ( LoggingCallbackManager, @@ -166,7 +168,12 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 ) init_params = {} - if "lakera_prompt_injection" in callback_specific_params: + if ( + "lakera_prompt_injection" in callback_specific_params + and isinstance( + callback_specific_params["lakera_prompt_injection"], dict + ) + ): init_params = callback_specific_params["lakera_prompt_injection"] lakera_moderations_object = lakeraAI_Moderation(**init_params) imported_list.append(lakera_moderations_object) diff --git a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py index 67a24567461..4f3e26ab5fb 100644 --- a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py +++ b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py @@ -8,7 +8,6 @@ from typing import Any, Dict, List, Optional from litellm._logging import verbose_proxy_logger -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, @@ -16,11 +15,15 @@ UI_SESSION_TOKEN_TEAM_ID, ) from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken, UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) class ExpiredUISessionKeyCleanupManager: @@ -147,7 +150,7 @@ async def _find_expired_ui_session_keys(self) -> List[LiteLLM_VerificationToken] Find expired LiteLLM dashboard session keys. """ now = datetime.now(timezone.utc) - return await self.prisma_client.db.litellm_verificationtoken.find_many( + return await VerificationTokenRepository(self.prisma_client).table.find_many( where={ "team_id": UI_SESSION_TOKEN_TEAM_ID, "expires": {"lt": now}, diff --git a/litellm/proxy/common_utils/html_forms/cli_sso_success.py b/litellm/proxy/common_utils/html_forms/cli_sso_success.py index 51f0775d90b..345f3ca5b42 100644 --- a/litellm/proxy/common_utils/html_forms/cli_sso_success.py +++ b/litellm/proxy/common_utils/html_forms/cli_sso_success.py @@ -135,7 +135,7 @@ def render_cli_sso_success_page() -> str: font-size: 14px; }} - .countdown {{ + .status {{ color: #64748b; font-size: 14px; font-weight: 500; @@ -183,23 +183,11 @@ def render_cli_sso_success_page() -> str:

You can now use LiteLLM CLI commands with your authenticated session.

-
This window will close in 3 seconds...
+
You can now close this window and return to your terminal.
- + diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index aaf39a7a19d..d622f612494 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -24,6 +24,12 @@ regenerate_key_fn, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import ( + DeprecatedVerificationTokenRepository, +) +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) class KeyRotationManager: @@ -124,20 +130,20 @@ async def _find_keys_needing_rotation(self) -> List[LiteLLM_VerificationToken]: """ now = datetime.now(timezone.utc) - keys_with_rotation = ( - await self.prisma_client.db.litellm_verificationtoken.find_many( - where={ - "auto_rotate": True, # Only keys marked for auto rotation - "OR": [ - { - "key_rotation_at": None - }, # Keys that need initial rotation time setup - { - "key_rotation_at": {"lte": now} - }, # Keys where rotation time has passed - ], - } - ) + keys_with_rotation = await VerificationTokenRepository( + self.prisma_client + ).table.find_many( + where={ + "auto_rotate": True, # Only keys marked for auto rotation + "OR": [ + { + "key_rotation_at": None + }, # Keys that need initial rotation time setup + { + "key_rotation_at": {"lte": now} + }, # Keys where rotation time has passed + ], + } ) return keys_with_rotation @@ -148,9 +154,9 @@ async def _cleanup_expired_deprecated_keys(self) -> None: """ try: now = datetime.now(timezone.utc) - result = await self.prisma_client.db.litellm_deprecatedverificationtoken.delete_many( - where={"revoke_at": {"lt": now}} - ) + result = await DeprecatedVerificationTokenRepository( + self.prisma_client + ).table.delete_many(where={"revoke_at": {"lt": now}}) if result > 0: verbose_proxy_logger.debug( "Cleaned up %s expired deprecated key(s)", result @@ -206,7 +212,7 @@ async def _rotate_key(self, key: LiteLLM_VerificationToken): # Calculate next rotation time using helper function now = datetime.now(timezone.utc) next_rotation_time = _calculate_key_rotation_time(key.rotation_interval) - await self.prisma_client.db.litellm_verificationtoken.update( + await VerificationTokenRepository(self.prisma_client).table.update( where={"token": response.token_id}, data={ "rotation_count": (key.rotation_count or 0) + 1, diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py new file mode 100644 index 00000000000..24e5c991794 --- /dev/null +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -0,0 +1,196 @@ +""" +ProxyRateLimitError — a unified rate-limit exception used by litellm's +proxy-side hooks. + +Background +---------- +LiteLLM previously surfaced rate-limit conditions through *several* unrelated +exception types: + +* :class:`litellm.exceptions.RateLimitError` — raised by exception mapping when + an upstream LLM provider returns 429. +* :class:`fastapi.HTTPException` (status 429) — raised directly by proxy hooks + such as ``parallel_request_limiter``, ``dynamic_rate_limiter``, + ``batch_rate_limiter``, ``max_budget_limiter``, ``max_iterations_limiter``, + etc. +* :class:`litellm.llms.base_llm.chat.transformation.BaseLLMException` (status + 429) — raised by some provider transports. + +This made it impossible for downstream code (and end users) to express +"is this a rate limit?" with a single ``except`` clause, and impossible to +distinguish *where* the rate limit originated (vendor vs. litellm, batch vs. +chat) without ad-hoc string-matching on the message. + +This module provides a single proxy-side error class that: + +1. Is a subclass of :class:`litellm.exceptions.RateLimitError`, so user code + that catches ``RateLimitError`` works for *every* rate-limit source. +2. Is also a subclass of :class:`fastapi.HTTPException`, so existing proxy + plumbing (``isinstance(e, HTTPException)`` branches in route handlers and + FastAPI's own dispatcher) continues to behave the same way and the + ``retry-after`` / ``rate_limit_type`` / ``reset_at`` headers are preserved + on the wire. +3. Carries a :attr:`category` field (one of + :class:`litellm.exceptions.RateLimitErrorCategory`) so callers can switch on + the rate limit source. +""" + +import json +from typing import Any, Dict, Mapping, Optional, Union + +from fastapi import HTTPException + +from litellm.exceptions import RateLimitError, RateLimitErrorCategory, RateLimitType + + +def map_v3_rate_limit_type( + v3_value: Optional[str], +) -> Optional[RateLimitType]: + """ + Map the v3 rate limiter's internal `status["rate_limit_type"]` strings + onto the public :class:`RateLimitType` enum. + + The v3 limiter uses the literal values ``"requests"``, ``"tokens"``, and + ``"max_parallel_requests"``. We collapse the last one onto + :attr:`RateLimitType.CONCURRENT_REQUESTS` because that's the public name + documented for users and dashboards. Unrecognized values return ``None`` + so the field stays absent rather than carrying garbage downstream. + """ + if v3_value == "tokens": + return RateLimitType.TOKENS + if v3_value == "max_parallel_requests": + return RateLimitType.CONCURRENT_REQUESTS + if v3_value == "requests": + return RateLimitType.REQUESTS + return None + + +def _coerce_message(detail: Any) -> str: + """Best-effort, JSON-friendly stringification of an HTTPException-style detail.""" + if detail is None: + return "" + if isinstance(detail, str): + return detail + if isinstance(detail, Mapping): + for key in ("error", "message"): + if isinstance(detail.get(key), str): + return detail[key] + inner = detail.get(key) + if isinstance(inner, Mapping) and isinstance(inner.get("message"), str): + return inner["message"] + try: + return json.dumps(detail) + except (TypeError, ValueError): + return str(detail) + return str(detail) + + +# NOTE: mypy emits two `[misc]` errors on the class line below because the +# bases declare overlapping attributes with related-but-not-identical +# annotations: +# * `status_code` is `int` on starlette HTTPException but `Literal[429]` on +# openai.RateLimitError (every openai status-error subclass narrows it +# this way and silences pyright with the same convention). +# * `headers` is `Mapping[str, str] | None` on HTTPException; we narrow it +# to `Optional[Dict[str, str]]` on RateLimitError because we always carry +# a stringified dict. +# Both narrowings are intentional and handled at construction time — every +# instance always has status_code == 429 and a Dict-typed headers — so we +# silence the ATTR-overlap check rather than relax the annotations. +class ProxyRateLimitError(HTTPException, RateLimitError): # type: ignore[misc] + """ + A 429 raised by litellm's proxy-side rate limiting hooks. + + This class deliberately inherits from BOTH + :class:`litellm.exceptions.RateLimitError` and :class:`fastapi.HTTPException` + so the same instance can flow through: + + * ``except RateLimitError`` (user / SDK code that wants a category-aware + handler), and + * ``isinstance(e, HTTPException)`` (FastAPI / proxy_server.py route + handlers that need to forward ``status_code``, ``detail`` and + ``headers`` back to the client). + + Downstream code should prefer this class over + ``raise HTTPException(status_code=429, ...)`` for litellm-internal rate + limits. + + Parameters + ---------- + detail: + The structured error payload. Forwarded as ``HTTPException.detail`` so + FastAPI's default exception handler will serialize it verbatim. + headers: + Optional response headers (e.g. ``retry-after``). Values are stringified + to satisfy FastAPI's typing. + category: + One of :class:`RateLimitErrorCategory`. Defaults to + ``LITELLM_RATE_LIMIT`` since this class is only used by litellm's own + proxy-side limiters; pass ``LITELLM_BATCH_RATE_LIMIT`` for the batch + limiter, etc. + model / llm_provider: + Optional context, propagated to the inherited ``RateLimitError`` for + compatibility with logging / standard payload extraction. + """ + + # Prometheus' ``exception_class`` label is pinned to "HTTPException" for + # this type: before the unified class existed, proxy-side 429s surfaced as + # ``fastapi.HTTPException`` and existing dashboards/alerts key off that exact + # value. Distinguishing vendor vs. litellm 429s is now the job of the + # ``rate_limit_category`` / ``rate_limit_type`` labels. + prometheus_exception_class_name = "HTTPException" + + def __init__( + self, + detail: Any, + headers: Optional[Mapping[str, Any]] = None, + category: Union[ + str, RateLimitErrorCategory + ] = RateLimitErrorCategory.LITELLM_RATE_LIMIT, + rate_limit_type: Optional[Union[str, RateLimitType]] = None, + model: Optional[str] = None, + llm_provider: Optional[str] = "litellm_proxy", + ): + # Normalize None → safe defaults so callers (and the resolver helper + # in `rate_limiter_utils`) can pass `None` without producing an + # instance whose `.llm_provider` attribute is `None` — that would + # break Prometheus' `_get_exception_class_name` (it calls + # `.capitalize()` on the provider string). + model = model or "" + llm_provider = llm_provider or "litellm_proxy" + message = _coerce_message(detail) + stringified_headers: Optional[Dict[str, str]] = ( + {k: str(v) for k, v in headers.items()} if headers else None + ) + + # Initialize the FastAPI HTTPException portion first so its attributes + # (status_code, detail, headers) are already on the instance before + # RateLimitError.__init__ runs and possibly overrides them. + HTTPException.__init__( + self, + status_code=429, + detail=detail, + headers=stringified_headers, + ) + + # Now initialize the litellm RateLimitError portion. We deliberately + # pass the structured detail through so RateLimitError preserves it as + # its `.detail` attribute too — keeping both sides of the MRO + # consistent. + RateLimitError.__init__( + self, + message=message, + llm_provider=llm_provider, + model=model, + category=category, + rate_limit_type=rate_limit_type, + headers=stringified_headers, + detail=detail, + ) + # RateLimitError.__init__ overwrites self.headers with its own copy and + # leaves self.status_code at 429 — restore the HTTPException-style + # headers value so downstream code that pulls headers off the + # instance gets back exactly what the limiter passed in. + self.headers = stringified_headers + self.detail = detail + self.status_code = 429 diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 40c8caa49e5..7c1dfe8dc90 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -14,6 +14,16 @@ LiteLLM_VerificationToken, ) from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.table_repositories import ( + EndUserRepository, + TagRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.services import ServiceTypes @@ -159,7 +169,7 @@ async def reset_budget_for_litellm_team_members( """ return await self._cascade_reset_spend_for_budget_link( budgets_to_reset=budgets_to_reset, - table=self.prisma_client.db.litellm_teammembership, + table=TeamMembershipRepository(self.prisma_client).table, counter_key_fn=lambda m: f"spend:team_member:{m.user_id}:{m.team_id}", log_subject="team memberships", cache_key_fn=lambda m: f"{m.team_id}_{m.user_id}", @@ -176,7 +186,7 @@ async def reset_budget_for_keys_linked_to_budgets( """ return await self._cascade_reset_spend_for_budget_link( budgets_to_reset=budgets_to_reset, - table=self.prisma_client.db.litellm_verificationtoken, + table=VerificationTokenRepository(self.prisma_client).table, counter_key_fn=lambda k: f"spend:key:{k.token}", log_subject="keys", extra_where={"budget_duration": None, "spend": {"gt": 0}}, @@ -191,7 +201,7 @@ async def reset_budget_for_orgs_linked_to_budgets( """ return await self._cascade_reset_spend_for_budget_link( budgets_to_reset=budgets_to_reset, - table=self.prisma_client.db.litellm_organizationtable, + table=OrganizationRepository(self.prisma_client).table, counter_key_fn=lambda o: f"spend:org:{o.organization_id}", log_subject="orgs", extra_where={"spend": {"gt": 0}}, @@ -217,7 +227,7 @@ async def reset_budget_for_tags_linked_to_budgets( """ return await self._cascade_reset_spend_for_budget_link( budgets_to_reset=budgets_to_reset, - table=self.prisma_client.db.litellm_tagtable, + table=TagRepository(self.prisma_client).table, counter_key_fn=lambda t: f"spend:tag:{t.tag_name}", log_subject="tags", extra_where={"spend": {"gt": 0}}, @@ -406,7 +416,7 @@ async def _get_endusers_with_no_budget_id( rely on the default budget (litellm.max_end_user_budget_id) applied in-memory during auth checks. """ - rows = await self.prisma_client.db.litellm_endusertable.find_many( + rows = await EndUserRepository(self.prisma_client).table.find_many( where={ "budget_id": None, "spend": {"gt": 0}, @@ -824,7 +834,7 @@ async def reset_budget_windows(self) -> None: ): changed = True if changed: - await self.prisma_client.db.litellm_verificationtoken.update( + await VerificationTokenRepository(self.prisma_client).table.update( where={"token": row["token"]}, data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] ) @@ -852,7 +862,7 @@ async def reset_budget_windows(self) -> None: ): changed = True if changed: - await self.prisma_client.db.litellm_teamtable.update( + await TeamRepository(self.prisma_client).table.update( where={"team_id": row["team_id"]}, data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] ) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index e0015e112e1..8118d53b9f6 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -12,6 +12,7 @@ is_proxy_admin, user_can_access_resource_owner, ) +from litellm.repositories.table_repositories import ManagedObjectRepository from litellm.responses.utils import ResponsesAPIRequestUtils CONTAINER_OBJECT_PURPOSE = "container" @@ -213,7 +214,7 @@ async def record_container_owner( ) return response - table = prisma_client.db.litellm_managedobjecttable + table = ManagedObjectRepository(prisma_client).table existing = await table.find_unique(where={"model_object_id": model_object_id}) if existing is not None: if getattr(existing, "file_purpose", None) != CONTAINER_OBJECT_PURPOSE: @@ -273,7 +274,7 @@ async def _get_container_owner( if prisma_client is None: return None - row = await prisma_client.db.litellm_managedobjecttable.find_first( + row = await ManagedObjectRepository(prisma_client).table.find_first( where={ "model_object_id": model_object_id, "file_purpose": CONTAINER_OBJECT_PURPOSE, @@ -319,7 +320,7 @@ async def _get_stored_container_id( if prisma_client is None: return None - row = await prisma_client.db.litellm_managedobjecttable.find_first( + row = await ManagedObjectRepository(prisma_client).table.find_first( where={ "model_object_id": model_object_id, "file_purpose": CONTAINER_OBJECT_PURPOSE, @@ -411,7 +412,7 @@ async def _get_allowed_container_ids( if prisma_client is None: return set() - rows = await prisma_client.db.litellm_managedobjecttable.find_many( + rows = await ManagedObjectRepository(prisma_client).table.find_many( where={ "file_purpose": CONTAINER_OBJECT_PURPOSE, "created_by": {"in": owner_scopes}, diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 2d05270e2ed..a716857111b 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -14,6 +14,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.utils import handle_exception_on_proxy, jsonify_object +from litellm.repositories.credentials_repository import CredentialsRepository from litellm.types.utils import CreateCredentialItem, CredentialItem router = APIRouter() @@ -96,7 +97,7 @@ async def create_credential( ) credentials_dict = encrypted_credential.model_dump() credentials_dict_jsonified = jsonify_object(credentials_dict) - await prisma_client.db.litellm_credentialstable.create( + await CredentialsRepository(prisma_client).create( data={ **credentials_dict_jsonified, "created_by": user_api_key_dict.user_id, @@ -245,9 +246,7 @@ async def delete_credential( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - await prisma_client.db.litellm_credentialstable.delete( - where={"credential_name": credential_name} - ) + await CredentialsRepository(prisma_client).delete_by_name(credential_name) ## DELETE FROM LITELLM ## litellm.credential_list = [ @@ -326,15 +325,14 @@ async def update_credential( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - db_credential = await prisma_client.db.litellm_credentialstable.find_unique( - where={"credential_name": credential_name}, - ) + credentials_repository = CredentialsRepository(prisma_client) + db_credential = await credentials_repository.find_by_name(credential_name) if db_credential is None: raise HTTPException(status_code=404, detail="Credential not found in DB.") merged_credential = update_db_credential(db_credential, credential) credential_object_jsonified = jsonify_object(merged_credential.model_dump()) - await prisma_client.db.litellm_credentialstable.update( - where={"credential_name": credential_name}, + await credentials_repository.update_by_name( + credential_name, data={ **credential_object_jsonified, "updated_by": user_api_key_dict.user_id, diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index ab9d341aa51..c500e727595 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -109,6 +109,92 @@ def is_database_transport_error(e: Exception) -> bool: return True return False + @staticmethod + def is_prisma_engine_internal_error(e: Exception) -> bool: + """True iff ``e`` is a non-``PrismaError`` exception raised from inside + prisma-client-py's query-engine layer. + + During the instant a DB connection is torn down, the query engine can + return a malformed error payload (``user_facing_error.meta`` is + ``null``). prisma-client-py's ``handle_response_errors`` then crashes + with ``AttributeError: 'NoneType' object has no attribute 'get'`` + before it can raise the proper P1001 "can't reach database server" + error. That AttributeError carries no connection keyword, so it can't + be matched by message; identify it by its ``prisma.engine`` origin + instead. + + Recognized ``PrismaError`` subclasses are excluded: connectivity ones + are already classified by type/keyword above, and data-layer ones + (the DB IS reachable) must stay 401. + """ + import prisma + + if isinstance(e, prisma.errors.PrismaError): + return False + tb = getattr(e, "__traceback__", None) + while tb is not None: + if tb.tb_frame.f_globals.get("__name__", "").startswith("prisma.engine"): + return True + tb = tb.tb_next + return False + + @staticmethod + def is_database_service_unavailable_error(e: Exception) -> bool: + """True iff the exception means the database could not answer at the + infrastructure level (connection refused, socket/interface failure, + timeout) rather than a genuine auth failure (key not found) or a + data-layer error (the DB IS reachable and rejected the data). + + Auth must answer 401 only for a key the DB confirms is invalid. When + the DB itself is unreachable, the request has to surface as 503 so + callers retry instead of treating valid keys as invalid during an + outage. + + Note: prisma-client-py mislabels the P1001 "can't reach database + server" connectivity failure as a ``DataError`` (a data-layer type), + so a type-only check misses real outages. ``is_database_transport_error`` + keyword-matches the connection message and catches that masquerade, + while genuine data errors (no connection keyword) correctly stay 401. + + The Postgres "cached plan must not change result type" error is matched + here, not in ``is_database_transport_error``: it is a transient stale-DB- + state condition (not an invalid key), but the connection is healthy so it + must not trigger a reconnect. + + A non-``PrismaError`` raised from inside the prisma query engine (e.g. + the ``AttributeError`` from ``handle_response_errors`` when the engine + returns a malformed error payload mid-tear-down) is also treated as + unavailable; see ``is_prisma_engine_internal_error``. + """ + import asyncio + + if PrismaDBExceptionHandler.is_database_connection_error(e): + return True + if PrismaDBExceptionHandler.is_database_transport_error(e): + return True + if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e): + return True + if "cached plan must not change result type" in str(e).lower(): + return True + + # OSError already covers ConnectionError and (Py3.3+) TimeoutError. + # asyncio.TimeoutError is a distinct class before Py3.11. + if isinstance(e, (OSError, asyncio.TimeoutError)): + return True + + try: + import asyncpg + except ImportError: + return False + + return isinstance( + e, + ( + asyncpg.exceptions.PostgresConnectionError, + asyncpg.exceptions.InterfaceError, + ), + ) + @staticmethod def handle_db_exception(e: Exception): """ diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index e7c5fa3f72c..2226aeb4b0a 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -20,6 +20,16 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.table_repositories import ( + SpendLogsRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache @@ -83,25 +93,25 @@ async def from_db( try: if counter_key.startswith("spend:key:"): token = counter_key[len("spend:key:") :] - row = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": token} - ) + row = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": token}) elif counter_key.startswith("spend:team_member:"): suffix = counter_key[len("spend:team_member:") :] if ":" not in suffix: return None user_id, team_id = suffix.rsplit(":", 1) - row = await prisma_client.db.litellm_teammembership.find_unique( + row = await TeamMembershipRepository(prisma_client).table.find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} ) elif counter_key.startswith("spend:team:"): team_id = counter_key[len("spend:team:") :] - row = await prisma_client.db.litellm_teamtable.find_unique( + row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) elif counter_key.startswith("spend:user:"): user_id = counter_key[len("spend:user:") :] - row = await prisma_client.db.litellm_usertable.find_unique( + row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) elif counter_key.startswith("spend:end_user:"): @@ -110,7 +120,7 @@ async def from_db( return None elif counter_key.startswith("spend:org:"): org_id = counter_key[len("spend:org:") :] - row = await prisma_client.db.litellm_organizationtable.find_unique( + row = await OrganizationRepository(prisma_client).table.find_unique( where={"organization_id": org_id} ) else: @@ -243,7 +253,7 @@ async def window_from_spend_logs( return None try: - response = await prisma_client.db.litellm_spendlogs.group_by( + response = await SpendLogsRepository(prisma_client).table.group_by( by=[group_field], where=where, # type: ignore[arg-type] sum={"spend": True}, diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 835d76e0ee4..77c06a465f4 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -10,6 +10,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import SpendLogToolIndexRepository def _add_tool_calls_to_set(tool_calls: Any, out: Set[str]) -> None: @@ -141,7 +142,7 @@ async def process_spend_logs_tool_usage( } ) if index_data: - await prisma_client.db.litellm_spendlogtoolindex.create_many( + await SpendLogToolIndexRepository(prisma_client).table.create_many( data=index_data, skip_duplicates=True, ) diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 6b34c974cf4..08bc8944b92 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -11,6 +11,9 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ToolDiscoveryQueueItem +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.table_repositories import ToolRepository from litellm.types.tool_management import ( LiteLLM_ToolTableRow, ToolPolicyOverrideRow, @@ -84,7 +87,7 @@ async def batch_upsert_tools( if not data: return now = datetime.now(timezone.utc) - table = prisma_client.db.litellm_tooltable + table = ToolRepository(prisma_client).table for item in data: tool_name = item.get("tool_name", "") origin = item.get("origin") or "user_defined" @@ -134,7 +137,7 @@ async def list_tools( """Return all tools, optionally filtered by input_policy.""" try: where = {"input_policy": input_policy} if input_policy is not None else {} - rows = await prisma_client.db.litellm_tooltable.find_many( + rows = await ToolRepository(prisma_client).table.find_many( where=where, order={"created_at": "desc"}, ) @@ -150,7 +153,7 @@ async def get_tool( ) -> Optional[LiteLLM_ToolTableRow]: """Return a single tool row by tool_name.""" try: - row = await prisma_client.db.litellm_tooltable.find_unique( + row = await ToolRepository(prisma_client).table.find_unique( where={"tool_name": tool_name}, ) if row is None: @@ -192,7 +195,7 @@ async def update_tool_policy( if output_policy is not None: update_data["output_policy"] = output_policy - await prisma_client.db.litellm_tooltable.upsert( + await ToolRepository(prisma_client).table.upsert( where={"tool_name": tool_name}, data={ "create": create_data, @@ -217,7 +220,7 @@ async def get_tools_by_names( if not tool_names: return {} try: - rows = await prisma_client.db.litellm_tooltable.find_many( + rows = await ToolRepository(prisma_client).table.find_many( where={"tool_name": {"in": tool_names}}, ) return { @@ -244,7 +247,7 @@ async def list_overrides_for_tool( """ out: List[ToolPolicyOverrideRow] = [] try: - perms = await prisma_client.db.litellm_objectpermissiontable.find_many( + perms = await ObjectPermissionRepository(prisma_client).table.find_many( where={"blocked_tools": {"has": tool_name}}, include={ "verification_tokens": True, @@ -307,7 +310,11 @@ def is_initialized(self) -> bool: async def sync_tool_policy_from_db(self, prisma_client: "PrismaClient") -> None: """Load all tool policies and object-permission blocked_tools from DB.""" try: - tools = await prisma_client.db.litellm_tooltable.find_many() + tools = await call_with_db_reconnect_retry( + prisma_client, + lambda: ToolRepository(prisma_client).table.find_many(), + reason="sync_tool_policy_from_db_tools_lookup_failure", + ) self._tool_input_policies = { row.tool_name: getattr(row, "input_policy", "untrusted") or "untrusted" for row in tools @@ -317,7 +324,11 @@ async def sync_tool_policy_from_db(self, prisma_client: "PrismaClient") -> None: for row in tools } - perms = await prisma_client.db.litellm_objectpermissiontable.find_many() + perms = await call_with_db_reconnect_retry( + prisma_client, + lambda: ObjectPermissionRepository(prisma_client).table.find_many(), + reason="sync_tool_policy_from_db_perms_lookup_failure", + ) self._blocked_tools_by_op_id = {} for row in perms: op_id = getattr(row, "object_permission_id", None) @@ -388,7 +399,7 @@ async def add_tool_to_object_permission_blocked( if not object_permission_id or not tool_name: return False try: - row = await prisma_client.db.litellm_objectpermissiontable.find_unique( + row = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": object_permission_id}, ) if row is None: @@ -397,7 +408,7 @@ async def add_tool_to_object_permission_blocked( if tool_name in current: return True current.append(tool_name) - await prisma_client.db.litellm_objectpermissiontable.update( + await ObjectPermissionRepository(prisma_client).table.update( where={"object_permission_id": object_permission_id}, data={"blocked_tools": current}, ) @@ -418,7 +429,7 @@ async def remove_tool_from_object_permission_blocked( if not object_permission_id or not tool_name: return False try: - row = await prisma_client.db.litellm_objectpermissiontable.find_unique( + row = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": object_permission_id}, ) if row is None: @@ -427,7 +438,7 @@ async def remove_tool_from_object_permission_blocked( if tool_name not in current: return False current = [t for t in current if t != tool_name] - await prisma_client.db.litellm_objectpermissiontable.update( + await ObjectPermissionRepository(prisma_client).table.update( where={"object_permission_id": object_permission_id}, data={"blocked_tools": current}, ) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index e0e4bdcf4a4..9f8ea584103 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -13,21 +13,21 @@ from fastapi import APIRouter, Depends, HTTPException, Request from pydantic import BaseModel -from litellm.proxy.common_utils.path_utils import safe_join - from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.proxy.common_utils.path_utils import safe_join from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import ( build_sandbox_globals, compile_sandboxed, ) from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.repositories.table_repositories import GuardrailsRepository from litellm.types.guardrails import ( PII_ENTITY_CATEGORIES_MAP, ApplyGuardrailRequest, @@ -373,7 +373,7 @@ async def create_guardrail( # Configuration error — roll back the DB write so the guardrail isn't orphaned if prisma_client is not None: try: - await prisma_client.db.litellm_guardrailstable.delete( + await GuardrailsRepository(prisma_client).table.delete( where={"guardrail_id": guardrail_id} ) except Exception as rollback_err: @@ -705,7 +705,7 @@ async def register_guardrail( ) try: - existing = await prisma_client.db.litellm_guardrailstable.find_unique( + existing = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_name": request.guardrail_name} ) if existing is not None: @@ -732,7 +732,7 @@ async def register_guardrail( guardrail_info_str = safe_dumps(guardrail_info) try: - created = await prisma_client.db.litellm_guardrailstable.create( + created = await GuardrailsRepository(prisma_client).table.create( data={ "guardrail_name": request.guardrail_name, "litellm_params": litellm_params_str, @@ -874,7 +874,7 @@ async def list_guardrail_submissions( where_clause["team_id"] = {"in": visible_team_ids} # Single query: fetch team guardrails visible to the caller - all_team_rows = await prisma_client.db.litellm_guardrailstable.find_many( + all_team_rows = await GuardrailsRepository(prisma_client).table.find_many( where=where_clause, order={"created_at": "desc"}, ) @@ -945,7 +945,7 @@ async def get_guardrail_submission( is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN try: - row = await prisma_client.db.litellm_guardrailstable.find_unique( + row = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) if row is None: @@ -986,7 +986,7 @@ async def approve_guardrail_submission( raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - row = await prisma_client.db.litellm_guardrailstable.find_unique( + row = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) if row is None: @@ -1000,7 +1000,7 @@ async def approve_guardrail_submission( ) now = datetime.now(timezone.utc) - await prisma_client.db.litellm_guardrailstable.update( + await GuardrailsRepository(prisma_client).table.update( where={"guardrail_id": guardrail_id}, data={"status": "active", "reviewed_at": now, "updated_at": now}, ) @@ -1072,7 +1072,7 @@ async def reject_guardrail_submission( raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - row = await prisma_client.db.litellm_guardrailstable.find_unique( + row = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) if row is None: @@ -1086,7 +1086,7 @@ async def reject_guardrail_submission( ) now = datetime.now(timezone.utc) - await prisma_client.db.litellm_guardrailstable.update( + await GuardrailsRepository(prisma_client).table.update( where={"guardrail_id": guardrail_id}, data={"status": "rejected", "reviewed_at": now, "updated_at": now}, ) @@ -2288,10 +2288,10 @@ async def apply_guardrail( """ import traceback - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.litellm_core_utils.thread_pool_executor import ( executor as thread_pool_executor, ) + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.proxy_server import ( general_settings, proxy_config, diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 14d950ecdf4..248202b644c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -105,6 +105,16 @@ def _extract_text_from_content(content: object) -> str: return "" +def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Optional[dict[str, Any]]: + merged: dict[str, Any] = {} + present = False + for bag in (request_data.get("metadata"), request_data.get("litellm_metadata")): + if isinstance(bag, Mapping): + present = True + merged.update(bag) + return merged if present else None + + class CrowdStrikeAIDRHandler(CustomGuardrail): """ CrowdStrike AIDR AI Guardrail handler to interact with the CrowdStrike AIDR @@ -312,11 +322,27 @@ async def apply_guardrail( event_type = "output" hook_name = "apply_guardrail (response)" - ai_guard_payload = { + ai_guard_payload: dict[str, Any] = { "guard_input": guard_input.model_dump(mode="json"), "event_type": event_type, } + model = inputs.get("model") + if model: + ai_guard_payload["model"] = model + + metadata = _merge_metadata_bags(request_data) + if metadata is not None: + user_id = metadata.get("user_api_key_user_id") + if user_id: + ai_guard_payload["user_id"] = user_id + + extra_info: dict[str, str] = {} + user_email = metadata.get("user_api_key_user_email") + if user_email: + extra_info["user_name"] = user_email + ai_guard_payload["extra_info"] = extra_info + ai_guard_response = await self._call_crowdstrike_aidr_guard( ai_guard_payload, hook_name ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index fc414ab7b54..dd04ae35b3e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -62,6 +62,192 @@ ModelResponseStream, ) +# Anthropic SSE delta types that carry model-generated text, mapped to +# (payload field, needs JSON escaping). partial_json fragments need +# replacements JSON-escaped so the assembled tool input stays valid JSON. +_ANTHROPIC_SSE_DELTA_FIELDS: Dict[str, Tuple[str, bool]] = { + "text_delta": ("text", False), + "thinking_delta": ("thinking", False), +} + +# input_json_delta carries streamed tool-call arguments, which agentic clients +# execute. Restoring PII there hands the original values to tool execution, so +# it is opt-in via `unmask_streamed_tool_calls` rather than unmasked by default. +_ANTHROPIC_SSE_TOOL_INPUT_FIELD: Tuple[str, Tuple[str, bool]] = ( + "input_json_delta", + ("partial_json", True), +) + +# Event types that terminate the current content block / message: any held +# partial-token carry is literal text and must be flushed before them. +# Deliberately excludes "ping" — keepalives interleave mid-block and a +# placeholder may continue right after one. +_ANTHROPIC_SSE_CARRY_FLUSH_EVENTS = { + "content_block_start", + "content_block_stop", + "message_delta", + "message_stop", + "error", +} + + +class _AnthropicSSEUnmasker: + """Incremental PII unmasker for an Anthropic native SSE byte stream. + + feed() accepts arbitrarily split byte chunks and buffers to complete + lines, so a ``data:`` line (or a multi-byte UTF-8 character) split + across chunks is never parsed partially. Placeholder tokens may also be + split across delta *events* (``"……"``): a trailing + fragment that is still a strict prefix of a known token is carried and + prepended to the next delta of the same content block; carries are + flushed as a synthetic delta event when the block ends without + completing one. Call flush() once the stream is exhausted. + + Lives here (not litellm/llms/anthropic/) intentionally: it is private to + the Presidio guardrail's bytes-passthrough path and reads guardrail state + (pii_tokens); nothing provider-side consumes it. + """ + + def __init__( + self, pii_tokens: Dict[str, str], unmask_tool_inputs: bool = False + ) -> None: + self._tokens: Dict[str, str] = dict(pii_tokens) + self._delta_fields: Dict[str, Tuple[str, bool]] = dict( + _ANTHROPIC_SSE_DELTA_FIELDS + ) + if unmask_tool_inputs: + delta_type, field_spec = _ANTHROPIC_SSE_TOOL_INPUT_FIELD + self._delta_fields[delta_type] = field_spec + self._buffer: bytes = b"" + self._carry: str = "" + self._carry_ctx: Optional[Tuple[Any, str]] = None # (index, delta_type) + self._line_ending: str = "\n" + + def feed(self, chunk: bytes) -> bytes: + self._buffer += chunk + out: List[bytes] = [] + while True: + newline_idx = self._buffer.find(b"\n") + if newline_idx == -1: + break + line = self._buffer[: newline_idx + 1] + self._buffer = self._buffer[newline_idx + 1 :] + out.append(self._process_line(line)) + return b"".join(out) + + def flush(self) -> bytes: + """Emit any held carry and unterminated trailing bytes at stream end.""" + out = self._flush_carry() + self._buffer + self._buffer = b"" + return out + + def _process_line(self, line: bytes) -> bytes: + if line.endswith(b"\n"): + self._line_ending = "\r\n" if line.endswith(b"\r\n") else "\n" + if line.startswith(b"event: "): + # Flush a held carry *before* the event line so the synthetic + # delta does not split the next event's `event:`/`data:` framing. + event_name = ( + line[len(b"event: ") :].strip().decode("utf-8", errors="replace") + ) + if event_name in _ANTHROPIC_SSE_CARRY_FLUSH_EVENTS: + return self._flush_carry() + line + return line + if not line.startswith(b"data: "): + return line + + try: + text = line.decode("utf-8") + except UnicodeDecodeError: + return line + stripped = text.rstrip("\r\n") + line_ending = text[len(stripped) :] + raw_json = stripped[len("data: ") :] + if raw_json == "[DONE]": + return line + try: + event = json.loads(raw_json) + except json.JSONDecodeError: + return line + if not isinstance(event, dict): + return line + + event_type = event.get("type") + delta = event.get("delta") + delta_type = delta.get("type") if isinstance(delta, dict) else None + field_spec = self._delta_fields.get(delta_type) if delta_type else None + if ( + event_type == "content_block_delta" + and delta is not None + and field_spec is not None + and isinstance(delta.get(field_spec[0]), str) + ): + field, json_escape = field_spec + prefix = b"" + ctx = (event.get("index"), str(delta_type)) + if self._carry and self._carry_ctx != ctx: + # Safety net: Anthropic blocks are sequential, so a context + # switch without an intervening block stop should not happen. + prefix = self._flush_carry() + delta_text = self._carry + delta[field] + self._carry = "" + self._carry_ctx = None + unmasked = self._replace(delta_text, json_escape) + emit_text, carry = self._split_partial_token(unmasked) + if carry: + self._carry = carry + self._carry_ctx = ctx + delta[field] = emit_text + new_line = "data: " + json.dumps(event, ensure_ascii=False) + line_ending + return prefix + new_line.encode("utf-8") + + if event_type in _ANTHROPIC_SSE_CARRY_FLUSH_EVENTS: + # Fallback for streams without `event:` lines (already a no-op + # when the preceding event line flushed the carry). + return self._flush_carry() + line + return line + + def _replace(self, text: str, json_escape: bool) -> str: + for token, original in self._tokens.items(): + if token in text: + if json_escape: + original = json.dumps(original, ensure_ascii=False)[1:-1] + text = text.replace(token, original) + return text + + def _split_partial_token(self, text: str) -> Tuple[str, str]: + """Split off a trailing fragment that may still grow into a token.""" + idx = text.rfind("<") + if idx == -1: + return text, "" + candidate = text[idx:] + if ">" not in candidate and any( + len(token) > len(candidate) and token.startswith(candidate) + for token in self._tokens + ): + return text[:idx], candidate + return text, "" + + def _flush_carry(self) -> bytes: + if not self._carry: + return b"" + index, delta_type = self._carry_ctx or (0, "text_delta") + field, _ = self._delta_fields[delta_type] + event = { + "type": "content_block_delta", + "index": index, + "delta": {"type": delta_type, field: self._carry}, + } + self._carry = "" + self._carry_ctx = None + le = self._line_ending + return ( + f"event: content_block_delta{le}data: " + + json.dumps(event, ensure_ascii=False) + + le + + le + ).encode("utf-8") + class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_cache = None @@ -75,6 +261,7 @@ def __init__( presidio_analyzer_api_base: Optional[str] = None, presidio_anonymizer_api_base: Optional[str] = None, output_parse_pii: Optional[bool] = False, + unmask_streamed_tool_calls: Optional[bool] = False, apply_to_output: bool = False, presidio_ad_hoc_recognizers: Optional[str] = None, logging_only: Optional[bool] = None, @@ -98,6 +285,7 @@ def __init__( ) # mapping of PII token to original text - only used with Presidio `replace` operation self.mock_redacted_text = mock_redacted_text self.output_parse_pii = output_parse_pii or False + self.unmask_streamed_tool_calls = unmask_streamed_tool_calls or False self.apply_to_output = apply_to_output # When output_parse_pii or apply_to_output is enabled, the guardrail must @@ -1237,15 +1425,39 @@ async def _stream_pii_unmasking( from litellm.main import stream_chunk_builder from litellm.types.utils import ModelResponse + metadata = (request_data.get("metadata") or {}) if request_data else {} + pii_tokens: Dict[str, str] = metadata.get("pii_tokens", {}) + + sse_unmasker: Optional[_AnthropicSSEUnmasker] = None remaining_chunks: List[ModelResponseStream] = [] try: async for chunk in response: if isinstance(chunk, ModelResponseStream): remaining_chunks.append(chunk) elif isinstance(chunk, bytes): - yield chunk # type: ignore[misc] + # bytes chunks only occur on the Anthropic native + # passthrough route (/v1/messages). OpenAI-format streams + # arrive as ModelResponseStream objects and are unmasked + # via the stream_chunk_builder path below — they never + # reach the SSE unmasker. + if pii_tokens: + if sse_unmasker is None: + sse_unmasker = _AnthropicSSEUnmasker( + pii_tokens, + unmask_tool_inputs=self.unmask_streamed_tool_calls, + ) + unmasked_bytes = sse_unmasker.feed(chunk) + if unmasked_bytes: + yield unmasked_bytes # type: ignore[misc] + else: + yield chunk # type: ignore[misc] continue + if sse_unmasker is not None: + trailing_bytes = sse_unmasker.flush() + if trailing_bytes: + yield trailing_bytes # type: ignore[misc] + if not remaining_chunks: return @@ -1275,6 +1487,10 @@ async def _stream_pii_unmasking( except Exception as e: verbose_proxy_logger.error(f"Error in PII streaming processing: {str(e)}") + if sse_unmasker is not None: + trailing_bytes = sse_unmasker.flush() + if trailing_bytes: + yield trailing_bytes # type: ignore[misc] for chunk in remaining_chunks: yield chunk diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index aafcc5f1819..a80bb817890 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -11,12 +11,15 @@ from litellm._uuid import uuid from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy.guardrails.guardrail_hooks.grayswan import GraySwanGuardrail +from litellm.proxy.guardrails.guardrail_hooks.grayswan import ( + GraySwanGuardrail, +) from litellm.proxy.guardrails.guardrail_hooks.grayswan import ( initialize_guardrail as initialize_grayswan, ) from litellm.proxy.types_utils.utils import get_instance_fn from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import GuardrailsRepository from litellm.secret_managers.main import get_secret from litellm.types.guardrails import ( Guardrail, @@ -26,6 +29,9 @@ SupportedGuardrailIntegrations, ) +from .guardrail_hooks.llm_as_a_judge import ( + initialize_guardrail as initialize_llm_as_a_judge, +) from .guardrail_initializers import ( initialize_bedrock, initialize_hide_secrets, @@ -34,9 +40,6 @@ initialize_presidio, initialize_tool_permission, ) -from .guardrail_hooks.llm_as_a_judge import ( - initialize_guardrail as initialize_llm_as_a_judge, -) guardrail_initializer_registry = { SupportedGuardrailIntegrations.BEDROCK.value: initialize_bedrock, @@ -257,7 +260,7 @@ async def add_guardrail_to_db( guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {})) # Create guardrail in DB - created_guardrail = await prisma_client.db.litellm_guardrailstable.create( + created_guardrail = await GuardrailsRepository(prisma_client).table.create( data={ "guardrail_name": guardrail_name, "litellm_params": litellm_params, @@ -283,7 +286,7 @@ async def delete_guardrail_from_db( """ try: # Delete from DB - await prisma_client.db.litellm_guardrailstable.delete( + await GuardrailsRepository(prisma_client).table.delete( where={"guardrail_id": guardrail_id} ) @@ -311,7 +314,7 @@ async def update_guardrail_in_db( guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {})) # Update in DB - updated_guardrail = await prisma_client.db.litellm_guardrailstable.update( + updated_guardrail = await GuardrailsRepository(prisma_client).table.update( where={"guardrail_id": guardrail_id}, data={ "guardrail_name": guardrail_name, @@ -335,11 +338,11 @@ async def get_all_guardrails_from_db( Only rows with status == "active" are returned (pending_review and rejected are excluded). """ try: - guardrails_from_db = ( - await prisma_client.db.litellm_guardrailstable.find_many( - where={"status": "active"}, - order={"created_at": "desc"}, - ) + guardrails_from_db = await GuardrailsRepository( + prisma_client + ).table.find_many( + where={"status": "active"}, + order={"created_at": "desc"}, ) guardrails: List[Guardrail] = [] @@ -357,7 +360,7 @@ async def get_guardrail_by_id_from_db( Get a guardrail by its ID from the database """ try: - guardrail = await prisma_client.db.litellm_guardrailstable.find_unique( + guardrail = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) @@ -375,7 +378,7 @@ async def get_guardrail_by_name_from_db( Get a guardrail by its name from the database """ try: - guardrail = await prisma_client.db.litellm_guardrailstable.find_unique( + guardrail = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_name": guardrail_name} ) diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 529949c6dd8..d8457cf9c86 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -12,6 +12,14 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import ( + DailyGuardrailMetricsRepository, + DailyPolicyMetricsRepository, + GuardrailsRepository, + PolicyRepository, + SpendLogGuardrailIndexRepository, + SpendLogsRepository, +) router = APIRouter() @@ -272,10 +280,10 @@ async def guardrails_usage_overview( try: # Guardrails from DB - guardrails = await prisma_client.db.litellm_guardrailstable.find_many() + guardrails = await GuardrailsRepository(prisma_client).table.find_many() # Daily metrics in range - metrics = await prisma_client.db.litellm_dailyguardrailmetrics.find_many( + metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( where={"date": {"gte": start, "lte": end}} ) @@ -283,9 +291,9 @@ async def guardrails_usage_overview( start_prev = ( datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7) ).strftime("%Y-%m-%d") - metrics_prev = await prisma_client.db.litellm_dailyguardrailmetrics.find_many( - where={"date": {"gte": start_prev, "lt": start}} - ) + metrics_prev = await DailyGuardrailMetricsRepository( + prisma_client + ).table.find_many(where={"date": {"gte": start_prev, "lt": start}}) agg = _aggregate_daily_metrics(metrics, "guardrail_id") prev_agg = _prev_fail_rates(metrics_prev, "guardrail_id") @@ -335,7 +343,7 @@ async def guardrails_usage_detail( end = end_date or now.strftime("%Y-%m-%d") start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") - guardrail = await prisma_client.db.litellm_guardrailstable.find_unique( + guardrail = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) if not guardrail: @@ -349,13 +357,13 @@ async def guardrails_usage_detail( ) metric_ids = [i for i in (logical_id, guardrail_id) if i] - metrics = await prisma_client.db.litellm_dailyguardrailmetrics.find_many( + metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( where={ "guardrail_id": {"in": metric_ids}, "date": {"gte": start, "lte": end}, } ) - metrics_prev = await prisma_client.db.litellm_dailyguardrailmetrics.find_many( + metrics_prev = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( where={ "guardrail_id": {"in": metric_ids}, "date": {"lt": start}, @@ -574,7 +582,7 @@ async def guardrails_usage_logs( # Query by both so we match regardless of which was written. effective_guardrail_ids: List[str] = [guardrail_id] if guardrail_id else [] if guardrail_id: - guardrail = await prisma_client.db.litellm_guardrailstable.find_unique( + guardrail = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) if guardrail: @@ -585,19 +593,23 @@ async def guardrails_usage_logs( where = _build_usage_logs_where( effective_guardrail_ids or None, policy_id, start_date, end_date ) - index_rows = await prisma_client.db.litellm_spendlogguardrailindex.find_many( + index_rows = await SpendLogGuardrailIndexRepository( + prisma_client + ).table.find_many( where=where, order={"start_time": "desc"}, skip=(page - 1) * page_size, take=page_size + 1, ) - total = await prisma_client.db.litellm_spendlogguardrailindex.count(where=where) + total = await SpendLogGuardrailIndexRepository(prisma_client).table.count( + where=where + ) request_ids = [r.request_id for r in index_rows[:page_size]] if not request_ids: return UsageLogsResponse( logs=[], total=total, page=page, page_size=page_size ) - spend_logs = await prisma_client.db.litellm_spendlogs.find_many( + spend_logs = await SpendLogsRepository(prisma_client).table.find_many( where={"request_id": {"in": request_ids}} ) log_by_id = {s.request_id: s for s in spend_logs} @@ -645,11 +657,13 @@ async def policies_usage_overview( start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") try: - policies = await prisma_client.db.litellm_policytable.find_many() - metrics = await prisma_client.db.litellm_dailypolicymetrics.find_many( + policies = await PolicyRepository(prisma_client).table.find_many() + metrics = await DailyPolicyMetricsRepository(prisma_client).table.find_many( where={"date": {"gte": start, "lte": end}} ) - metrics_prev = await prisma_client.db.litellm_dailypolicymetrics.find_many( + metrics_prev = await DailyPolicyMetricsRepository( + prisma_client + ).table.find_many( where={ "date": { "gte": ( diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 8907c9201ad..c55c47ca774 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -10,6 +10,10 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import ( + DailyGuardrailMetricsRepository, + SpendLogGuardrailIndexRepository, +) def _guardrail_status_to_action(status: Optional[str]) -> str: @@ -132,7 +136,7 @@ async def process_spend_logs_guardrail_usage( } ) try: - await prisma_client.db.litellm_spendlogguardrailindex.create_many( + await SpendLogGuardrailIndexRepository(prisma_client).table.create_many( data=index_data, skip_duplicates=True, ) @@ -146,7 +150,7 @@ async def process_spend_logs_guardrail_usage( n = int(agg["requests_evaluated"]) if n == 0: continue - await prisma_client.db.litellm_dailyguardrailmetrics.upsert( + await DailyGuardrailMetricsRepository(prisma_client).table.upsert( where={ "guardrail_id_date": { "guardrail_id": guardrail_id, diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index c109f374993..6ef8bbc4006 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -129,6 +129,7 @@ def get_callback_identifier(callback): "datadog_llm_observability", "generic_api", "arize", + "galileo", "sqs", ], str, @@ -206,6 +207,7 @@ async def health_services_endpoint( # noqa: PLR0915 "datadog_llm_observability", "generic_api", "arize", + "galileo", "sqs", ]: raise HTTPException( @@ -295,6 +297,19 @@ async def health_services_endpoint( # noqa: PLR0915 else "Arize is healthy" ), } + elif service == "galileo": + from litellm.integrations.galileo import GalileoObserve + + galileo_logger = GalileoObserve() + response = await galileo_logger.async_health_check() + return { + "status": response["status"], + "message": ( + response["error_message"] + if response["status"] == "unhealthy" + else "Galileo is healthy" + ), + } elif service == "langfuse": from litellm.integrations.langfuse.langfuse import LangFuseLogger diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 8473b5e77de..5b691beccbf 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -17,7 +17,17 @@ - async_log_success_event() fires on GET /v1/batches/{id} (batch completion) """ -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + NoReturn, + Optional, + Tuple, + Union, +) from fastapi import HTTPException from pydantic import BaseModel @@ -30,6 +40,7 @@ _get_file_content_as_dictionary, _get_models_from_batch_input_file_content, ) +from litellm.exceptions import RateLimitErrorCategory from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ( ProxyErrorTypes, @@ -37,10 +48,11 @@ SpecialModelNames, UserAPIKeyAuth, ) -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, +from litellm.proxy.common_utils.proxy_rate_limit_error import ( + ProxyRateLimitError, + map_v3_rate_limit_type, ) +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -385,8 +397,8 @@ def _raise_rate_limit_error( batch_usage: BatchFileUsage, limit_type: str, requested_model: Optional[str] = None, - ) -> None: - """Raise HTTPException for rate limit exceeded.""" + ) -> NoReturn: + """Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded.""" from datetime import datetime # Find the descriptor for this status @@ -432,14 +444,15 @@ def _raise_rate_limit_error( resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( requested_model ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=detail, headers={ "retry-after": str(window_size), "rate_limit_type": limit_type, "reset_at": reset_time_formatted, }, + category=RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT, + rate_limit_type=map_v3_rate_limit_type(limit_type), model=resolved_model, llm_provider=llm_provider, ) @@ -518,11 +531,17 @@ async def count_input_file_usage( # Check if this is a managed file (base64 encoded unified file ID) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + get_models_from_unified_file_id, ) # Managed files require bypassing the HTTP endpoint (which runs access-check hooks) # and calling the managed files hook directly with the user's credentials. is_managed_file = _is_base64_encoded_unified_file_id(file_id) + target_model_names = ( + get_models_from_unified_file_id(is_managed_file) + if is_managed_file + else [] + ) if is_managed_file and user_api_key_dict is not None: file_content = await self._fetch_managed_file_content( file_id=file_id, @@ -560,6 +579,7 @@ async def count_input_file_usage( await self._enforce_batch_file_model_access( user_api_key_dict=user_api_key_dict, file_content_as_dict=file_content_as_dict, + target_model_names=target_model_names or None, ) input_file_usage = _get_batch_job_input_file_usage( @@ -595,9 +615,13 @@ async def _enforce_batch_file_model_access( self, user_api_key_dict: UserAPIKeyAuth, file_content_as_dict: List[dict], + target_model_names: Optional[List[str]] = None, ) -> None: - """Reject the batch if the caller is not authorized for every - ``body.model`` named inside the JSONL. + """Reject the batch if the caller is not authorized for the upload target. + + For managed files, ``target_model_names`` (from the unified file id) is + the proxy alias the file was uploaded for and is used directly for auth. + For legacy/non-managed files, falls back to ``body.model`` values in the JSONL. Reuses standard auth helpers so the same model access rules the proxy enforces on `/chat/completions` apply here. @@ -614,9 +638,12 @@ async def _enforce_batch_file_model_access( from litellm.proxy.proxy_server import proxy_logging_obj from litellm.proxy.proxy_server import user_api_key_cache - models = _get_models_from_batch_input_file_content(file_content_as_dict) - if not models: - return + if target_model_names: + models = target_model_names + else: + models = _get_models_from_batch_input_file_content(file_content_as_dict) + if not models: + return team_object = None if ( @@ -647,12 +674,7 @@ async def _enforce_batch_file_model_access( llm_model_list = llm_router.model_list if llm_router is not None else None for model in models: - # body.model may be the provider id after replace_model_in_jsonl; map to proxy model_name for auth. model_to_check = model - if llm_router is not None: - proxy_model_name = llm_router.resolve_model_name_from_model_id(model) - if proxy_model_name is not None: - model_to_check = proxy_model_name try: if team_object is not None: try: diff --git a/litellm/proxy/hooks/dynamic_rate_limiter.py b/litellm/proxy/hooks/dynamic_rate_limiter.py index 57cd538507e..b9e2bd12ecf 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter.py @@ -11,9 +11,10 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.exceptions import RateLimitType from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, convert_priority_to_percent, resolve_llm_provider_for_rate_limit, ) @@ -222,8 +223,7 @@ async def async_pre_call_hook( resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( data.get("model") ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail={ "error": "Key={} over available TPM={}. Model TPM={}, Active keys={}".format( user_api_key_dict.api_key, @@ -232,6 +232,7 @@ async def async_pre_call_hook( active_projects, ) }, + rate_limit_type=RateLimitType.TOKENS, model=resolved_model, llm_provider=llm_provider, ) @@ -240,8 +241,7 @@ async def async_pre_call_hook( resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( data.get("model") ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail={ "error": "Key={} over available RPM={}. Model RPM={}, Active keys={}".format( user_api_key_dict.api_key, @@ -250,6 +250,7 @@ async def async_pre_call_hook( active_projects, ) }, + rate_limit_type=RateLimitType.REQUESTS, model=resolved_model, llm_provider=llm_provider, ) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index bfc6e2c2f72..493afe6105a 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -14,13 +14,16 @@ from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ( + ProxyRateLimitError, + map_v3_rate_limit_type, +) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( RateLimitDescriptor, RateLimitDescriptorRateLimitObject, _PROXY_MaxParallelRequestsHandler_v3, ) from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, convert_priority_to_percent, resolve_llm_provider_for_rate_limit, ) @@ -497,8 +500,7 @@ async def _check_rate_limits( continue descriptor_key = status["descriptor_key"] if descriptor_key == "model_saturation_check": - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail={ "error": f"Model capacity reached for {model}. " f"Priority: {priority}, " @@ -512,6 +514,9 @@ async def _check_rate_limits( "rate_limit_type": str(status["rate_limit_type"]), "x-litellm-priority": priority or "default", }, + rate_limit_type=map_v3_rate_limit_type( + status["rate_limit_type"] + ), model=resolved_model, llm_provider=llm_provider, ) @@ -520,8 +525,7 @@ async def _check_rate_limits( f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " f"priority: {priority}" ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail={ "error": f"Priority-based rate limit exceeded. " f"Model: {model}, " @@ -538,6 +542,9 @@ async def _check_rate_limits( "x-litellm-priority": priority or "default", "x-litellm-saturation": f"{saturation:.2%}", }, + rate_limit_type=map_v3_rate_limit_type( + status["rate_limit_type"] + ), model=resolved_model, llm_provider=llm_provider, ) @@ -556,8 +563,7 @@ async def _check_rate_limits( f"Dynamic rate limiter: OVER_LIMIT response with unknown " f"descriptor_key(s) — refusing request. response={atomic_response}" ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail={ "error": "Rate limit exceeded", "descriptor_key": ( @@ -567,6 +573,9 @@ async def _check_rate_limits( str(offending["rate_limit_type"]) if offending else "unknown" ), }, + rate_limit_type=map_v3_rate_limit_type( + offending["rate_limit_type"] if offending else None + ), headers={ "retry-after": str(self.v3_limiter.window_size), "x-litellm-priority": priority or "default", diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index 21e8bbbd308..77ed3493a0c 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -19,7 +19,7 @@ response = await litellm.acompletion( model="gpt-4o-mini", messages=[{"role": "user", "content": "Create a bouncing ball GIF"}], - container={"skills": [{"skill_id": "litellm:skill_abc123"}]}, + container={"skills": [{"skill_id": "litellm_skill_abc123"}]}, ) # Response includes file_ids for generated files """ @@ -31,6 +31,7 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.litellm_proxy.skills.constants import LITELLM_SKILL_ID_PREFIX from litellm.llms.litellm_proxy.skills.prompt_injection import ( SkillPromptInjectionHandler, ) @@ -43,7 +44,7 @@ class SkillsInjectionHook(CustomLogger): Pre/Post-call hook that processes skills from container.skills parameter. Pre-call (async_pre_call_hook): - - Skills with 'litellm:' prefix are fetched from LiteLLM DB + - Skills with 'litellm_skill_' prefix are fetched from LiteLLM DB - For Anthropic models: native skills pass through, LiteLLM skills converted to tools - For non-Anthropic models: LiteLLM skills are converted to tools + execute_code tool @@ -78,7 +79,7 @@ async def async_pre_call_hook( Process skills from container.skills before the LLM call. 1. Check if container.skills exists in request - 2. Separate skills by prefix (litellm: vs native) + 2. Separate skills by prefix (litellm_skill_ vs native) 3. Fetch LiteLLM skills from database 4. For Anthropic: keep native skills in container 5. For non-Anthropic: convert LiteLLM skills to tools, inject content, add execute_code @@ -108,7 +109,7 @@ async def async_pre_call_hook( continue skill_id = skill.get("skill_id", "") - if skill_id.startswith("litellm_"): + if skill_id.startswith(LITELLM_SKILL_ID_PREFIX): # Fetch from LiteLLM DB db_skill = await self._fetch_skill_from_db( skill_id, @@ -287,7 +288,7 @@ async def _fetch_skill_from_db( Fetch a skill from the LiteLLM database. Args: - skill_id: The skill ID (without 'litellm:' prefix) + skill_id: The skill ID (including the 'litellm_skill_' prefix) Returns: LiteLLM_SkillsTable or None if not found @@ -382,10 +383,10 @@ async def async_post_call_success_deployment_hook( has_executable_tool = False for tc in tool_calls: tool_name = tc.get("name", "") - # Execute if it's litellm_code_execution OR a skill tool (skill_xxx) + # Execute if it's litellm_code_execution OR a skill tool (litellm_skill_xxx) if ( tool_name == LiteLLMInternalTools.CODE_EXECUTION.value - or tool_name.startswith("skill_") + or tool_name.startswith(LITELLM_SKILL_ID_PREFIX) ): has_executable_tool = True break @@ -543,7 +544,7 @@ async def _execute_code_loop_messages_api( result = await self._execute_code( code, skill_files, executor, generated_files ) - elif tool_name.startswith("skill_"): + elif tool_name.startswith(LITELLM_SKILL_ID_PREFIX): # Skill tool - execute the skill's code result = await self._execute_skill_tool( tool_name, tool_input, skill_files, executor, generated_files diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 658d7995631..769348a0b88 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -4,11 +4,10 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.exceptions import RateLimitType from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, -) +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit class _PROXY_MaxBudgetLimiter(CustomLogger): @@ -70,9 +69,9 @@ async def async_pre_call_hook( resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( data.get("model") if data else None ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail="Max budget limit reached.", + rate_limit_type=RateLimitType.BUDGET, model=resolved_model, llm_provider=llm_provider, ) diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 0b63465c4a5..20bfeb3a6d5 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -20,11 +20,10 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.exceptions import RateLimitType from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, -) +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache @@ -117,13 +116,13 @@ async def async_pre_call_hook( resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( data.get("model") if data else None ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=( f"Session budget exceeded for session {session_id}. " f"Current spend: ${current_spend:.4f}, " f"max_budget_per_session: ${max_budget:.2f}." ), + rate_limit_type=RateLimitType.BUDGET, model=resolved_model, llm_provider=llm_provider, ) diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index d5bc669c928..525214ff6be 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -16,11 +16,10 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.exceptions import RateLimitType from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, -) +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache @@ -121,12 +120,12 @@ async def async_pre_call_hook( resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( data.get("model") if data else None ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=( f"Max iterations exceeded for session {session_id}. " f"Current count: {current_count}, max_iterations: {max_iterations}." ), + rate_limit_type=RateLimitType.MAX_ITERATIONS, model=resolved_model, llm_provider=llm_provider, ) diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index c6324c3e3a3..874e5aa1939 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -1,26 +1,24 @@ import asyncio import sys from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, List, Literal, NoReturn, Optional, Tuple, Union -from fastapi import HTTPException from pydantic import BaseModel from typing_extensions import TypedDict import litellm -from litellm import DualCache, ModelResponse +from litellm import DualCache, EmbeddingResponse, ModelResponse, TextCompletionResponse from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, UserAPIKeyAuth +from litellm.exceptions import RateLimitType from litellm.proxy.auth.auth_utils import ( get_key_model_rpm_limit, get_key_model_tpm_limit, ) -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, -) +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -75,9 +73,21 @@ async def check_key_in_limits( ) if current is None: if max_parallel_requests == 0 or tpm_limit == 0 or rpm_limit == 0: - # base case - raise self.raise_rate_limit_error( + # base case — at least one dimension is set to 0 (effectively + # disabled). Pick the most specific dimension as the + # rate_limit_type so dashboards can attribute the failure to + # the right cap. Order matters: max_parallel_requests is + # listed first because it's the rarest 0 in practice and the + # most actionable signal. + if max_parallel_requests == 0: + triggered_type = RateLimitType.CONCURRENT_REQUESTS + elif tpm_limit == 0: + triggered_type = RateLimitType.TOKENS + else: + triggered_type = RateLimitType.REQUESTS + self.raise_rate_limit_error( additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current limits: max_parallel_requests: {max_parallel_requests}, tpm_limit: {tpm_limit}, rpm_limit: {rpm_limit}", + rate_limit_type=triggered_type, requested_model=data.get("model") if data else None, ) new_val = { @@ -100,14 +110,23 @@ async def check_key_in_limits( values_to_update_in_cache.append((request_count_api_key, new_val)) else: + # Detect which dimension actually tripped the limit so we can + # surface the right rate_limit_type. Order matches the boolean + # condition above (concurrent → tpm → rpm) — first match wins. + if int(current["current_requests"]) >= max_parallel_requests: + triggered_type = RateLimitType.CONCURRENT_REQUESTS + elif current["current_tpm"] >= tpm_limit: + triggered_type = RateLimitType.TOKENS + else: + triggered_type = RateLimitType.REQUESTS requested_model = data.get("model") if data else None resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( requested_model ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=f"LiteLLM Rate Limit Handler for rate limit type = {rate_limit_type}. {CommonProxyErrors.max_parallel_request_limit_reached.value}. current rpm: {current['current_rpm']}, rpm limit: {rpm_limit}, current tpm: {current['current_tpm']}, tpm limit: {tpm_limit}, current max_parallel_requests: {current['current_requests']}, max_parallel_requests: {max_parallel_requests}", headers={"retry-after": str(self.time_to_next_minute())}, + rate_limit_type=triggered_type, model=resolved_model, llm_provider=llm_provider, ) @@ -135,27 +154,45 @@ def time_to_next_minute(self) -> float: def raise_rate_limit_error( self, additional_details: Optional[str] = None, + rate_limit_type: Optional[RateLimitType] = None, requested_model: Optional[str] = None, - ) -> HTTPException: + ) -> NoReturn: """ - Raise an HTTPException with a 429 status code and a retry-after header. + Raise a 429 with a retry-after header for litellm-proxy parallel-request limits. + + Always raises :class:`ProxyRateLimitError` — never returns. Annotated + ``NoReturn`` so type-checkers know callers after this invocation are + unreachable. The raised exception is both a + :class:`litellm.RateLimitError` (so callers can catch by category) and a + :class:`fastapi.HTTPException` (so the FastAPI dispatcher serializes it + correctly with status 429 and the supplied headers). + + ``rate_limit_type`` defaults to ``CONCURRENT_REQUESTS`` because every + existing internal caller of this helper hits the parallel-request cap + (the global-limit branch in ``async_pre_call_hook`` and the + all-zeros base case in ``check_key_in_limits``). Callers that know + the dimension exactly should pass it explicitly. ``requested_model`` is resolved via :func:`get_llm_provider` so the - raised exception carries ``llm_provider`` for downstream loggers - (Prometheus failure metric, observability callbacks). Falls back to - ``llm_provider="litellm_proxy"`` when the model is missing or - unparseable — see ``resolve_llm_provider_for_rate_limit``. + raised exception carries ``llm_provider`` (and a stripped ``model``) + for downstream loggers (Prometheus failure metric, observability + callbacks). Falls back to ``llm_provider="litellm_proxy"`` when the + model is missing or unparseable — see + :func:`resolve_llm_provider_for_rate_limit`. """ + # additional_details is optional; build the detail with a None-guard + # so callers that pass nothing don't get the literal string "None" + # interpolated into the error message. error_message = "Max parallel request limit reached" if additional_details is not None: error_message = error_message + " " + additional_details resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( requested_model ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=error_message, headers={"retry-after": str(self.time_to_next_minute())}, + rate_limit_type=rate_limit_type or RateLimitType.CONCURRENT_REQUESTS, model=resolved_model, llm_provider=llm_provider, ) @@ -248,7 +285,7 @@ async def async_pre_call_hook( # noqa: PLR0915 current_global_requests = 1 # if above -> raise error if current_global_requests >= global_max_parallel_requests: - return self.raise_rate_limit_error( + self.raise_rate_limit_error( additional_details=f"Hit Global Limit: Limit={global_max_parallel_requests}, current: {current_global_requests}", requested_model=data.get("model") if data else None, ) @@ -533,7 +570,9 @@ async def async_log_success_event( # noqa: PLR0915 total_tokens = 0 - if isinstance(response_obj, ModelResponse): + if isinstance( + response_obj, (ModelResponse, EmbeddingResponse, TextCompletionResponse) + ): total_tokens = response_obj.usage.total_tokens # type: ignore # ------------ @@ -622,7 +661,10 @@ async def async_log_success_event( # noqa: PLR0915 if user_api_key_user_id is not None: total_tokens = 0 - if isinstance(response_obj, ModelResponse): + if isinstance( + response_obj, + (ModelResponse, EmbeddingResponse, TextCompletionResponse), + ): total_tokens = response_obj.usage.total_tokens # type: ignore request_count_api_key = ( @@ -655,7 +697,10 @@ async def async_log_success_event( # noqa: PLR0915 if user_api_key_team_id is not None: total_tokens = 0 - if isinstance(response_obj, ModelResponse): + if isinstance( + response_obj, + (ModelResponse, EmbeddingResponse, TextCompletionResponse), + ): total_tokens = response_obj.usage.total_tokens # type: ignore request_count_api_key = ( @@ -688,7 +733,10 @@ async def async_log_success_event( # noqa: PLR0915 if user_api_key_end_user_id is not None: total_tokens = 0 - if isinstance(response_obj, ModelResponse): + if isinstance( + response_obj, + (ModelResponse, EmbeddingResponse, TextCompletionResponse), + ): total_tokens = response_obj.usage.total_tokens # type: ignore request_count_api_key = ( diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 9fdb146b19d..6b70cea65a3 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -32,13 +32,20 @@ ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import get_model_rate_limit_from_metadata -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, +from litellm.proxy.common_utils.proxy_rate_limit_error import ( + ProxyRateLimitError, + map_v3_rate_limit_type, ) +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject -from litellm.types.utils import CallTypes, ModelResponse, Usage +from litellm.types.utils import ( + CallTypes, + EmbeddingResponse, + ModelResponse, + TextCompletionResponse, + Usage, +) if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -1971,7 +1978,7 @@ def _handle_rate_limit_error( descriptors: List[RateLimitDescriptor], requested_model: Optional[str] = None, ) -> None: - """Handle rate limit exceeded error by raising HTTPException.""" + """Handle rate limit exceeded by raising :class:`ProxyRateLimitError` (a 429).""" for status in response["statuses"]: if status["code"] == "OVER_LIMIT": descriptor_key = status["descriptor_key"] @@ -2005,14 +2012,14 @@ def _handle_rate_limit_error( resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( requested_model ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=detail, headers={ "retry-after": str(self.window_size), "rate_limit_type": str(status["rate_limit_type"]), "reset_at": reset_time_formatted, }, + rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]), model=resolved_model, llm_provider=llm_provider, ) @@ -2735,9 +2742,14 @@ def _build_success_event_pipeline_operations( # Get total tokens from response total_tokens = 0 - # spot fix for /responses api - if isinstance(response_obj, ModelResponse) or isinstance( - response_obj, BaseLiteLLMOpenAIResponseObject + if isinstance( + response_obj, + ( + ModelResponse, + EmbeddingResponse, + TextCompletionResponse, + BaseLiteLLMOpenAIResponseObject, + ), ): _usage = getattr(response_obj, "usage", None) total_tokens = self._get_total_tokens_from_usage( diff --git a/litellm/proxy/hooks/rate_limiter_utils.py b/litellm/proxy/hooks/rate_limiter_utils.py index 0ba3df448e5..07440975476 100644 --- a/litellm/proxy/hooks/rate_limiter_utils.py +++ b/litellm/proxy/hooks/rate_limiter_utils.py @@ -2,13 +2,10 @@ Shared utility functions for rate limiter hooks. """ -from typing import Any, Optional, Tuple, Union - -from fastapi import HTTPException +from typing import Optional, Tuple, Union import litellm from litellm._logging import verbose_proxy_logger -from litellm.exceptions import RateLimitError from litellm.types.router import ModelGroupInfo from litellm.types.utils import PriorityReservationDict @@ -29,11 +26,21 @@ def resolve_llm_provider_for_rate_limit( ``litellm_proxy_failed_requests_metric`` show up with ``exception_class="RateLimitError"`` and no provider attribution. - Wrapped defensively: if ``model`` is missing, malformed, or - ``get_llm_provider`` raises (unknown alias, router-only model, etc.) we - fall back to ``("", "litellm_proxy")`` so we never break the request path - by piling a second exception on top of the rate-limit one we're trying to - raise. + Resolution order: + + 1. ``litellm.get_llm_provider(model)`` — covers raw provider/model + strings the SDK already understands (``"gpt-4o-mini"``, + ``"anthropic/claude-3-5-sonnet"``, ``"bedrock/..."`` etc.). + 2. **Router alias fallback** — nearly every real proxy deployment + routes through a router ``model_name`` alias (e.g. + ``"tpm-locked"`` → ``litellm_params.model: openai/gpt-4o-mini``). + ``get_llm_provider`` doesn't know router aliases, so without this + step every alias call ended up labeled ``"litellm_proxy"``, + defeating the field's purpose for the most common case. + 3. Defensive fallback to ``("", "litellm_proxy")`` — used only when + ``model`` is missing, malformed, or both lookups fail. We never let + a secondary exception escape and mask the rate-limit error we're + trying to surface. """ if not model: return "", PROXY_LLM_PROVIDER_FALLBACK @@ -46,6 +53,9 @@ def resolve_llm_provider_for_rate_limit( custom_llm_provider or PROXY_LLM_PROVIDER_FALLBACK, ) except Exception as e: + alias_resolution = _resolve_provider_from_router_alias(model) + if alias_resolution is not None: + return alias_resolution verbose_proxy_logger.debug( "rate_limiter_utils.resolve_llm_provider_for_rate_limit: " "could not resolve provider for model=%s, falling back to %s. err=%s", @@ -56,50 +66,58 @@ def resolve_llm_provider_for_rate_limit( return model, PROXY_LLM_PROVIDER_FALLBACK -class ProxyHTTPRateLimitError(HTTPException, RateLimitError): # type: ignore[misc] +def _resolve_provider_from_router_alias( + model: str, +) -> Optional[Tuple[str, str]]: """ - HTTPException raised by proxy-side rate-limit hooks that *also* exposes - ``model`` and ``llm_provider`` attributes. - - Why both base classes: - - - The proxy server's exception handler keys off ``HTTPException`` to render - a 429 response, so we must remain an ``HTTPException``. - - Downstream loggers (Prometheus ``async_post_call_failure_hook``, - structured logging, observability callbacks) read ``exception.llm_provider`` - via :meth:`litellm.integrations.prometheus.PrometheusLogger._get_exception_class_name` - and ``isinstance(exc, RateLimitError)`` for category routing. Inheriting - from :class:`litellm.exceptions.RateLimitError` keeps that wiring intact. - - We intentionally do not call ``RateLimitError.__init__`` (which constructs - an httpx.Response) — it isn't needed here and just adds failure surface. - Attribute parity is what downstream consumers rely on. + Resolve a router ``model_name`` alias to ``(underlying_model, provider)`` + by scanning the active router's ``model_list``. + + Returns ``None`` if the router isn't initialized, the alias isn't + registered, the deployment has no usable ``litellm_params.model``, or + any underlying lookup raises. Callers fall through to the defensive + ``litellm_proxy`` fallback in that case — never raising secondary + exceptions out of the rate-limit raise path. """ - - def __init__( - self, - status_code: int, - detail: Any = None, - headers: Optional[dict] = None, - *, - model: str = "", - llm_provider: str = PROXY_LLM_PROVIDER_FALLBACK, - ) -> None: - HTTPException.__init__( - self, status_code=status_code, detail=detail, headers=headers - ) - self.status_code = status_code - self.model = model or "" - self.llm_provider = llm_provider or PROXY_LLM_PROVIDER_FALLBACK - # `message` is what RateLimitError.__str__ would print and what some - # observability callbacks log. Keep it human-readable. - self.message = detail if isinstance(detail, str) else str(detail) - # `RateLimitError.__str__` (resolved via MRO since Starlette's - # HTTPException doesn't define `__str__`) unconditionally reads - # these attributes. Set them so `str(exc)` doesn't raise - # AttributeError from logging/traceback paths. - self.num_retries: Optional[int] = None - self.max_retries: Optional[int] = None + try: + from litellm.proxy.proxy_server import llm_router + except Exception: + return None + if llm_router is None: + return None + try: + model_list = getattr(llm_router, "model_list", None) + if not model_list: + return None + for deployment in model_list: + if not isinstance(deployment, dict): + continue + if deployment.get("model_name") != model: + continue + params = deployment.get("litellm_params") + if not isinstance(params, dict): + continue + underlying_model = params.get("model") + if not isinstance(underlying_model, str) or not underlying_model: + continue + try: + resolved_model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=underlying_model, + ) + except Exception: + continue + if not custom_llm_provider: + continue + # Prefer the underlying provider-qualified model so the failure + # callback / Prometheus label points at the actual deployment, not + # the alias. + return ( + resolved_model or underlying_model, + custom_llm_provider, + ) + return None + except Exception: + return None def convert_priority_to_percent( diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py index 08fa8d4dfad..c22fd1d6579 100644 --- a/litellm/proxy/hooks/user_management_event_hooks.py +++ b/litellm/proxy/hooks/user_management_event_hooks.py @@ -3,7 +3,6 @@ """ import asyncio -from litellm._uuid import uuid from datetime import datetime, timezone from typing import Optional @@ -11,6 +10,7 @@ import litellm from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid from litellm.proxy._types import ( AUDIT_ACTIONS, CommonProxyErrors, @@ -24,6 +24,7 @@ WebhookEvent, ) from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update +from litellm.repositories.user_repository import UserRepository class UserManagementEventHooks: @@ -57,7 +58,7 @@ async def async_user_created_hook( try: if prisma_client is None: raise Exception(CommonProxyErrors.db_not_connected_error.value) - user_row: BaseModel = await prisma_client.db.litellm_usertable.find_first( + user_row: BaseModel = await UserRepository(prisma_client).table.find_first( where={"user_id": response.user_id} ) diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index fe8b7c6fdc9..c217116e45f 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -173,6 +173,16 @@ async def image_generation( ) ) + # Call response headers hook (matches base_process_llm_request behavior) + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + fastapi_response.headers.update(callback_headers) + return response except Exception as e: await proxy_logging_obj.post_call_failure_hook( diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 62a770f46ae..65f7ffc9081 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -19,6 +19,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.utils import get_prisma_client_or_throw +from litellm.repositories.table_repositories import AccessGroupRepository from litellm.types.access_group import ( AccessGroupCreateRequest, AccessGroupResponse, @@ -386,7 +387,7 @@ async def list_access_groups( CommonProxyErrors.db_not_connected_error.value ) - records = await prisma_client.db.litellm_accessgrouptable.find_many( + records = await AccessGroupRepository(prisma_client).table.find_many( order={"created_at": "desc"} ) return [_record_to_response(r) for r in records] @@ -405,7 +406,7 @@ async def get_access_group( CommonProxyErrors.db_not_connected_error.value ) - record = await prisma_client.db.litellm_accessgrouptable.find_unique( + record = await AccessGroupRepository(prisma_client).table.find_unique( where={"access_group_id": access_group_id} ) if record is None: diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 2eda1b30c5d..698155a5c26 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -16,11 +16,12 @@ from fastapi import APIRouter, Depends, HTTPException -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.utils import jsonify_object +from litellm.repositories.budget_repository import BudgetRepository router = APIRouter() @@ -98,7 +99,7 @@ async def new_budget( budget_obj_json = budget_obj.model_dump(exclude_none=True) budget_obj_jsonified = jsonify_object(budget_obj_json) # json dump any dictionaries try: - response = await prisma_client.db.litellm_budgettable.create( + response = await BudgetRepository(prisma_client).table.create( data={ **budget_obj_jsonified, # type: ignore "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -182,7 +183,7 @@ async def update_budget( except ValueError as e: raise HTTPException(status_code=400, detail={"error": str(e)}) - response = await prisma_client.db.litellm_budgettable.update( + response = await BudgetRepository(prisma_client).table.update( where={"budget_id": budget_obj.budget_id}, data={ **budget_obj.model_dump(exclude_unset=True), # type: ignore @@ -217,7 +218,7 @@ async def info_budget(data: BudgetRequest): "error": f"Specify list of budget id's to query. Passed in={data.budgets}" }, ) - response = await prisma_client.db.litellm_budgettable.find_many( + response = await BudgetRepository(prisma_client).table.find_many( where={"budget_id": {"in": data.budgets}}, ) @@ -261,7 +262,7 @@ async def budget_settings( ) ## get budget item from db - db_budget_row = await prisma_client.db.litellm_budgettable.find_first( + db_budget_row = await BudgetRepository(prisma_client).table.find_first( where={"budget_id": budget_id} ) @@ -327,7 +328,7 @@ async def list_budget( }, ) - response = await prisma_client.db.litellm_budgettable.find_many() + response = await BudgetRepository(prisma_client).table.find_many() return response @@ -366,7 +367,7 @@ async def delete_budget( }, ) - response = await prisma_client.db.litellm_budgettable.delete( + response = await BudgetRepository(prisma_client).table.delete( where={"budget_id": data.id} ) diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index 0a26b23beff..b6ddf2d8e07 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -27,6 +27,8 @@ UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry +from litellm.repositories.table_repositories import CacheConfigRepository from litellm.types.management_endpoints import ( CACHE_SETTINGS_FIELDS, REDIS_TYPE_DESCRIPTIONS, @@ -159,8 +161,12 @@ async def init_cache_settings_in_db(prisma_client, proxy_config): import json try: - cache_config = await prisma_client.db.litellm_cacheconfig.find_unique( - where={"id": "cache_config"} + cache_config = await call_with_db_reconnect_retry( + prisma_client, + lambda: CacheConfigRepository(prisma_client).table.find_unique( + where={"id": "cache_config"} + ), + reason="init_cache_settings_in_db_lookup_failure", ) if cache_config is not None and cache_config.cache_settings: # Parse cache settings JSON @@ -274,7 +280,7 @@ async def get_cache_settings( # Try to get cache settings from database current_values = {} if prisma_client is not None: - cache_config = await prisma_client.db.litellm_cacheconfig.find_unique( + cache_config = await CacheConfigRepository(prisma_client).table.find_unique( where={"id": "cache_config"} ) if cache_config is not None and cache_config.cache_settings: @@ -417,7 +423,7 @@ async def update_cache_settings( # Snapshot the prior settings (key set only — values get redacted in # the audit row) so the audit-log entry shows which fields changed. - existing_row = await prisma_client.db.litellm_cacheconfig.find_unique( + existing_row = await CacheConfigRepository(prisma_client).table.find_unique( where={"id": "cache_config"} ) before_settings: Optional[Dict[str, Any]] = None @@ -434,7 +440,7 @@ async def update_cache_settings( ) # Save to database - await prisma_client.db.litellm_cacheconfig.upsert( + await CacheConfigRepository(prisma_client).table.upsert( where={"id": "cache_config"}, data={ "create": { diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index d173cd745ba..92cc2008c73 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -8,6 +8,10 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import DeletedVerificationTokenRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.management_endpoints.common_daily_activity import ( BreakdownMetrics, DailySpendData, @@ -346,7 +350,7 @@ async def get_api_key_metadata( This ensures that key_alias and team_id are preserved in historical activity logs even after a key is deleted or regenerated. """ - key_records = await prisma_client.db.litellm_verificationtoken.find_many( + key_records = await VerificationTokenRepository(prisma_client).table.find_many( where={"token": {"in": list(api_keys)}} ) result = { @@ -357,11 +361,11 @@ async def get_api_key_metadata( missing_keys = api_keys - set(result.keys()) if missing_keys: try: - deleted_key_records = ( - await prisma_client.db.litellm_deletedverificationtoken.find_many( - where={"token": {"in": list(missing_keys)}}, - order={"deleted_at": "desc"}, - ) + deleted_key_records = await DeletedVerificationTokenRepository( + prisma_client + ).table.find_many( + where={"token": {"in": list(missing_keys)}}, + order={"deleted_at": "desc"}, ) # Use the most recent deleted record for each token (ordered by deleted_at desc) for k in deleted_key_records: diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index dc27e87726a..458cba686e6 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, Optional, Union from fastapi import HTTPException, status from pydantic import BaseModel @@ -17,9 +17,13 @@ NewProjectRequest, UpdateProjectRequest, UserAPIKeyAuth, - user_api_key_has_admin_view as _user_has_admin_view, # noqa: F401 re-exported ) +from litellm.proxy._types import ( # noqa: F401 re-exported + user_api_key_has_admin_view as _user_has_admin_view, +) +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.utils import _premium_user_check +from litellm.repositories.team_repository import TeamRepository if TYPE_CHECKING: from litellm.proxy._types import NewProjectRequest, UpdateProjectRequest @@ -204,7 +208,7 @@ async def _user_has_admin_privileges( # Check if user is team admin for any team if user_obj.teams is not None and len(user_obj.teams) > 0: # Get all teams user is in - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": user_obj.teams}} ) @@ -281,7 +285,7 @@ async def _team_admin_can_invite_user( if not target_user_obj.teams or len(target_user_obj.teams) == 0: return False - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": admin_user_obj.teams}} ) admin_team_ids = [ @@ -400,121 +404,127 @@ def _set_object_metadata_field( object_data.metadata[field_name] = value +_TEAM_MEMBER_BUDGET_LIMIT_FIELDS = ( + "max_budget", + "soft_budget", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + "model_max_budget", + "budget_duration", + "allowed_models", +) + + +def _is_set_budget_value(value: Any) -> bool: + if value is None: + return False + if isinstance(value, list) and len(value) == 0: + return False + return True + + +def _has_meaningful_budget_limit(budget_values: Dict[str, Any]) -> bool: + """A budget is meaningful if at least one limit is actually set; an empty + list (no model restriction) and None both count as unset.""" + return any( + _is_set_budget_value(budget_values.get(field)) + for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS + ) + + async def _upsert_budget_and_membership( tx, *, team_id: str, user_id: str, - max_budget: Optional[float], existing_budget_id: Optional[str], user_api_key_dict: UserAPIKeyAuth, - tpm_limit: Optional[int] = None, - rpm_limit: Optional[int] = None, - allowed_models: Optional[List[str]] = None, + budget_patch: Dict[str, Any], team_default_budget_id: Optional[str] = None, ): """ - Helper function to Create/Update or Delete the budget within the team membership - Args: - tx: The transaction object - team_id: The ID of the team - user_id: The ID of the user - max_budget: The maximum budget for the team - existing_budget_id: The ID of the existing budget, if any - user_api_key_dict: User API Key dictionary containing user information - tpm_limit: Tokens per minute limit for the team member - rpm_limit: Requests per minute limit for the team member - allowed_models: Per-member model scope. None = don't change. [] = remove restrictions. Non-empty list = enforce. - team_default_budget_id: The team's shared default member budget id (from - team metadata.team_member_budget_id), if any. When the membership's - existing_budget_id matches this, we clone-on-write so editing one - member's budget does not mutate the shared default (and therefore - every other member who still points at it). - - If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, the user's budget is removed from the team membership. - If any of these values exist, a budget is updated or created and linked to the team membership. + Apply a merge-patch of per-member budget fields to a team membership. + + ``budget_patch`` holds only the budget columns the caller explicitly sent + (RFC 7396 semantics): a value sets the column, ``None`` clears it, and a + column that is absent from the dict is left untouched. Once the patch is + applied, if the budget has no meaningful limit left the member's private + budget is disconnected so they fall back to the team default. + + ``team_default_budget_id`` is the team's shared default member budget id + (from team metadata.team_member_budget_id). When the membership still + points at it, we clone-on-write so editing one member's budget does not + mutate the shared default that every other member points at. """ - if ( - max_budget is None - and tpm_limit is None - and rpm_limit is None - and allowed_models is None - ): - # disconnect the budget since all limits are None - await tx.litellm_teammembership.update( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, - data={"litellm_budget_table": {"disconnect": True}}, - ) + if not budget_patch: return + write_data = dict(budget_patch) + if "budget_duration" in write_data: + duration = write_data["budget_duration"] + write_data["budget_reset_at"] = ( + get_budget_reset_time(budget_duration=duration) + if duration is not None + else None + ) + is_shared_default = ( existing_budget_id is not None and team_default_budget_id is not None and existing_budget_id == team_default_budget_id ) + async def _disconnect(): + await tx.litellm_teammembership.update( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, + data={"litellm_budget_table": {"disconnect": True}}, + ) + if existing_budget_id is not None and not is_shared_default: - # Update the existing budget in-place to preserve fields not being changed. - # Only write fields that the caller explicitly provided (non-None). - update_data: Dict[str, Any] = { - "updated_by": user_api_key_dict.user_id or "", - } - if max_budget is not None: - update_data["max_budget"] = max_budget - if tpm_limit is not None: - update_data["tpm_limit"] = tpm_limit - if rpm_limit is not None: - update_data["rpm_limit"] = rpm_limit - if allowed_models is not None: - update_data["allowed_models"] = allowed_models + existing_budget = await tx.litellm_budgettable.find_unique( + where={"budget_id": existing_budget_id} + ) + merged = existing_budget.model_dump() if existing_budget is not None else {} + merged.update(write_data) + if not _has_meaningful_budget_limit(merged): + await _disconnect() + return await tx.litellm_budgettable.update( where={"budget_id": existing_budget_id}, - data=update_data, + data={"updated_by": user_api_key_dict.user_id or "", **write_data}, ) return - # Either there is no existing budget, OR the membership is still pointing - # at the team's shared default member budget. In both cases we create a - # NEW private budget for this user and (re)link the membership to it. create_data: Dict[str, Any] = { "created_by": user_api_key_dict.user_id or "", "updated_by": user_api_key_dict.user_id or "", } - # If we're forking off the shared default, seed the new row with the - # default's values so fields the caller did not change carry over. if is_shared_default: default_budget_row = await tx.litellm_budgettable.find_unique( where={"budget_id": existing_budget_id} ) if default_budget_row is not None: default_budget_dict = default_budget_row.model_dump() - for field in ( - "max_budget", - "soft_budget", - "max_parallel_requests", - "tpm_limit", - "rpm_limit", - "model_max_budget", - "budget_duration", - "allowed_models", - ): + for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: value = default_budget_dict.get(field) - if value is None: - continue - if isinstance(value, list) and len(value) == 0: - continue - create_data[field] = value - - # Caller-provided values take precedence over the cloned defaults. - if max_budget is not None: - create_data["max_budget"] = max_budget - if tpm_limit is not None: - create_data["tpm_limit"] = tpm_limit - if rpm_limit is not None: - create_data["rpm_limit"] = rpm_limit - if allowed_models is not None: - create_data["allowed_models"] = allowed_models + if _is_set_budget_value(value): + create_data[field] = value + + create_data.update(write_data) + + if create_data.get("budget_duration") is not None: + create_data["budget_reset_at"] = get_budget_reset_time( + budget_duration=create_data["budget_duration"] + ) + else: + create_data.pop("budget_reset_at", None) + + if not _has_meaningful_budget_limit(create_data): + if existing_budget_id is not None: + await _disconnect() + return new_budget = await tx.litellm_budgettable.create( data=create_data, diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index 7f7aa485fb3..97cb5eeddc4 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -30,6 +30,7 @@ UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import ConfigOverridesRepository from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.proxy.management_endpoints.config_overrides import ( ConfigOverrideSettingsResponse, @@ -254,7 +255,7 @@ async def update_hashicorp_vault_config( # Merge ALL fields the user didn't send: try DB first, fall back to env vars. # Omitted field = keep existing; empty string = clear/remove the field. - existing_record = await prisma_client.db.litellm_configoverrides.find_unique( + existing_record = await ConfigOverridesRepository(prisma_client).table.find_unique( where={"config_type": "hashicorp_vault"} ) existing_decrypted: Optional[Dict[str, Any]] = None @@ -321,7 +322,7 @@ async def update_hashicorp_vault_config( # Only persist to DB after successful init encrypted_data = proxy_config._encrypt_env_variables(config_data) config_value = safe_dumps(encrypted_data) - await prisma_client.db.litellm_configoverrides.upsert( + await ConfigOverridesRepository(prisma_client).table.upsert( where={"config_type": "hashicorp_vault"}, data={ "create": { @@ -391,7 +392,7 @@ async def get_hashicorp_vault_config( field_schema = _build_field_schema(HashicorpVaultConfig) # Try to load from DB - db_record = await prisma_client.db.litellm_configoverrides.find_unique( + db_record = await ConfigOverridesRepository(prisma_client).table.find_unique( where={"config_type": "hashicorp_vault"} ) @@ -448,7 +449,7 @@ async def delete_hashicorp_vault_config( # Capture the prior config before delete so the audit-log row can # show *what* was removed (keys only — values get redacted). - existing_record = await prisma_client.db.litellm_configoverrides.find_unique( + existing_record = await ConfigOverridesRepository(prisma_client).table.find_unique( where={"config_type": "hashicorp_vault"} ) before_config: Optional[Dict[str, Any]] = None @@ -463,7 +464,7 @@ async def delete_hashicorp_vault_config( # Delete DB record if it exists — ignore if not found deleted = False try: - await prisma_client.db.litellm_configoverrides.delete( + await ConfigOverridesRepository(prisma_client).table.delete( where={"config_type": "hashicorp_vault"} ) deleted = True diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 1fd8320db20..f1a34bb0ed4 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -17,8 +17,8 @@ from fastapi import APIRouter, Depends, HTTPException, Request import litellm -from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity @@ -27,6 +27,8 @@ handle_update_object_permission_common, ) from litellm.proxy.utils import handle_exception_on_proxy +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.table_repositories import EndUserRepository from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -68,7 +70,7 @@ async def block_user(data: BlockUsers): records = [] if prisma_client is not None: for id in data.user_ids: - record = await prisma_client.db.litellm_endusertable.upsert( + record = await EndUserRepository(prisma_client).table.upsert( where={"user_id": id}, # type: ignore data={ "create": {"user_id": id, "blocked": True}, # type: ignore @@ -337,7 +339,7 @@ async def new_end_user( _new_budget = new_budget_request(data) if _new_budget is not None: try: - budget_record = await prisma_client.db.litellm_budgettable.create( + budget_record = await BudgetRepository(prisma_client).table.create( data={ **_new_budget.model_dump(exclude_unset=True), "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, # type: ignore @@ -373,7 +375,7 @@ async def new_end_user( new_end_user_obj.pop("object_permission", None) ## WRITE TO DB ## - end_user_record = await prisma_client.db.litellm_endusertable.create( + end_user_record = await EndUserRepository(prisma_client).table.create( data=new_end_user_obj, # type: ignore include={"litellm_budget_table": True, "object_permission": True}, ) @@ -446,7 +448,7 @@ async def end_user_info( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - user_info = await prisma_client.db.litellm_endusertable.find_first( + user_info = await EndUserRepository(prisma_client).table.find_first( where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True}, ) @@ -569,7 +571,7 @@ async def update_end_user( non_default_values[k] = v ## Get end user table data ## - end_user_table_data = await prisma_client.db.litellm_endusertable.find_first( + end_user_table_data = await EndUserRepository(prisma_client).table.find_first( where={"user_id": data.user_id}, include={"litellm_budget_table": True} ) @@ -613,17 +615,17 @@ async def update_end_user( if budget_table_data: if end_user_budget_table is None: ## Create new budget ## - budget_table_data_record = ( - await prisma_client.db.litellm_budgettable.create( - data={ - **budget_table_data, - "created_by": user_api_key_dict.user_id - or litellm_proxy_admin_name, - "updated_by": user_api_key_dict.user_id - or litellm_proxy_admin_name, - }, - include={"end_users": True}, - ) + budget_table_data_record = await BudgetRepository( + prisma_client + ).table.create( + data={ + **budget_table_data, + "created_by": user_api_key_dict.user_id + or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id + or litellm_proxy_admin_name, + }, + include={"end_users": True}, ) update_end_user_table_data["budget_id"] = ( @@ -631,11 +633,11 @@ async def update_end_user( ) else: ## Update existing budget ## - budget_table_data_record = ( - await prisma_client.db.litellm_budgettable.update( - where={"budget_id": end_user_budget_table.budget_id}, - data=budget_table_data, - ) + budget_table_data_record = await BudgetRepository( + prisma_client + ).table.update( + where={"budget_id": end_user_budget_table.budget_id}, + data=budget_table_data, ) ## Update user table, with update params + new budget id (if set) ## @@ -652,7 +654,7 @@ async def update_end_user( if data.user_id is not None and len(data.user_id) > 0: update_end_user_table_data["user_id"] = data.user_id # type: ignore verbose_proxy_logger.debug("In update customer, user_id condition block.") - response = await prisma_client.db.litellm_endusertable.update( + response = await EndUserRepository(prisma_client).table.update( where={"user_id": data.user_id}, data=update_end_user_table_data, include={"litellm_budget_table": True, "object_permission": True} # type: ignore ) if response is None: @@ -737,7 +739,7 @@ async def delete_end_user( and len(data.user_ids) > 0 ): # First check if all users exist - existing_users = await prisma_client.db.litellm_endusertable.find_many( + existing_users = await EndUserRepository(prisma_client).table.find_many( where={"user_id": {"in": data.user_ids}} ) existing_user_ids = {user.user_id for user in existing_users} @@ -756,7 +758,7 @@ async def delete_end_user( ) # All users exist, proceed with deletion - response = await prisma_client.db.litellm_endusertable.delete_many( + response = await EndUserRepository(prisma_client).table.delete_many( where={"user_id": {"in": data.user_ids}} ) verbose_proxy_logger.debug( @@ -828,7 +830,7 @@ async def list_end_user( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - response = await prisma_client.db.litellm_endusertable.find_many( + response = await EndUserRepository(prisma_client).table.find_many( include={"litellm_budget_table": True, "object_permission": True} ) @@ -903,7 +905,7 @@ async def get_customer_daily_activity( where_condition = {} if end_user_ids_list: where_condition["user_id"] = {"in": list(end_user_ids_list)} - end_user_aliases = await prisma_client.db.litellm_endusertable.find_many( + end_user_aliases = await EndUserRepository(prisma_client).table.find_many( where=where_condition ) end_user_alias_metadata = {e.user_id: {"alias": e.alias} for e in end_user_aliases} diff --git a/litellm/proxy/management_endpoints/fallback_management_endpoints.py b/litellm/proxy/management_endpoints/fallback_management_endpoints.py index ffb12111d82..1333122c87a 100644 --- a/litellm/proxy/management_endpoints/fallback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/fallback_management_endpoints.py @@ -27,6 +27,7 @@ # fastapi is only required for proxy, not for SDK usage pass +from litellm.repositories.config_repository import ConfigRepository from litellm.types.management_endpoints.router_settings_endpoints import ( FallbackCreateRequest, FallbackDeleteResponse, @@ -157,7 +158,7 @@ async def create_fallback( # Save to database - convert router_settings to JSON string router_settings_json = json.dumps(router_settings) - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "router_settings"}, data={ "create": { @@ -336,7 +337,7 @@ async def delete_fallback( # Save to database - convert router_settings to JSON string router_settings_json = json.dumps(router_settings) - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "router_settings"}, data={ "create": { diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 7b8f0f72e13..b3a5c66e9e1 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -43,6 +43,17 @@ ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.utils import handle_exception_on_proxy, hash_password +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.table_repositories import ( + InvitationLinkRepository, + OrganizationMembershipRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -154,7 +165,7 @@ async def _check_duplicate_user_field( if case_insensitive: where_clause[field_name]["mode"] = "insensitive" - existing_user = await prisma_client.db.litellm_usertable.find_first( + existing_user = await UserRepository(prisma_client).table.find_first( where=where_clause ) @@ -434,7 +445,7 @@ async def new_user( await _check_duplicate_user_email(data.user_email, prisma_client) # Check if license is over limit - total_users = await prisma_client.db.litellm_usertable.count() + total_users = await UserRepository(prisma_client).table.count() if total_users and _license_check.is_over_limit(total_users=total_users): raise HTTPException( status_code=403, @@ -851,7 +862,7 @@ async def _check_user_info_v2_access( # Helper: fetch the target user row (reused across branches) async def _fetch_target_user(): - return await prisma_client.db.litellm_usertable.find_unique( + return await UserRepository(prisma_client).table.find_unique( where={"user_id": target_user_id} ) @@ -866,7 +877,7 @@ async def _fetch_target_user(): # Rule 3: Team admins can look up users in their teams if user_api_key_dict.user_id is not None: # Get caller's teams - caller_user = await prisma_client.db.litellm_usertable.find_unique( + caller_user = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) if caller_user is not None and caller_user.teams: @@ -876,7 +887,7 @@ async def _fetch_target_user(): return None # Get all teams the caller belongs to - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": caller_user.teams}} ) for team in teams: @@ -1165,7 +1176,7 @@ async def _schedule_user_update_audit_log( if prisma_client is None: return try: - updated_user_row = await prisma_client.db.litellm_usertable.find_first( + updated_user_row = await UserRepository(prisma_client).table.find_first( where={"user_id": response["user_id"]} ) if updated_user_row: @@ -1255,11 +1266,11 @@ async def _update_single_user_helper( existing_user_row: Optional[BaseModel] = None if user_request.user_id: - existing_user_row = await prisma_client.db.litellm_usertable.find_first( + existing_user_row = await UserRepository(prisma_client).table.find_first( where={"user_id": user_request.user_id} ) elif user_request.user_email: - existing_user_row = await prisma_client.db.litellm_usertable.find_first( + existing_user_row = await UserRepository(prisma_client).table.find_first( where={"user_email": user_request.user_email} ) @@ -1640,7 +1651,7 @@ async def bulk_user_update( detail="Only proxy admins can update all users at once.", ) # Optimized path for updating all users directly in database - all_users_in_db = await prisma_client.db.litellm_usertable.find_many( + all_users_in_db = await UserRepository(prisma_client).table.find_many( order={"created_at": "desc"} ) @@ -1676,7 +1687,7 @@ async def bulk_user_update( try: # Perform bulk database update - await prisma_client.db.litellm_usertable.update_many( + await UserRepository(prisma_client).table.update_many( where={}, data=non_default_values # Update all users ) @@ -1783,7 +1794,7 @@ async def get_user_key_counts( # Get count for each user_id individually for user_id in user_ids: - count = await prisma_client.db.litellm_verificationtoken.count( + count = await VerificationTokenRepository(prisma_client).table.count( where={ "user_id": user_id, "OR": [ @@ -2056,7 +2067,7 @@ async def get_users( else None ) - users = await prisma_client.db.litellm_usertable.find_many( + users = await UserRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, @@ -2066,7 +2077,9 @@ async def get_users( ) # Get total count of user rows - total_count = await prisma_client.db.litellm_usertable.count(where=where_conditions) + total_count = await UserRepository(prisma_client).table.count( + where=where_conditions + ) # Get key count for each user if users is not None: @@ -2137,14 +2150,14 @@ async def delete_user( from litellm.proxy.management_endpoints.team_endpoints import ( _cleanup_members_with_roles, ) + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + ) from litellm.proxy.proxy_server import ( create_audit_log_for_update, litellm_proxy_admin_name, prisma_client, ) - from litellm.proxy.management_helpers.audit_logs import ( - get_audit_log_changed_by, - ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -2164,7 +2177,7 @@ async def delete_user( caller_admin_org_ids: set = set() if not caller_is_proxy_admin: caller_memberships = ( - await prisma_client.db.litellm_organizationmembership.find_many( + await OrganizationMembershipRepository(prisma_client).table.find_many( where={ "user_id": user_api_key_dict.user_id, "user_role": LitellmUserRoles.ORG_ADMIN.value, @@ -2188,11 +2201,9 @@ async def delete_user( # an N+1 DB call when delete_user is called with a large user_ids list. target_org_ids_by_user: Dict[str, set] = {} if not caller_is_proxy_admin: - all_target_memberships = ( - await prisma_client.db.litellm_organizationmembership.find_many( - where={"user_id": {"in": data.user_ids}} - ) - ) + all_target_memberships = await OrganizationMembershipRepository( + prisma_client + ).table.find_many(where={"user_id": {"in": data.user_ids}}) for m in all_target_memberships: if not m.organization_id: continue @@ -2200,7 +2211,7 @@ async def delete_user( # check that all teams passed exist for user_id in data.user_ids: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) @@ -2254,7 +2265,7 @@ async def delete_user( ) ## CLEANUP MEMBERS_WITH_ROLES - fetch_all_teams = await prisma_client.db.litellm_teamtable.find_many( + fetch_all_teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": user_row.teams}} ) teams_to_update = [] @@ -2277,19 +2288,19 @@ async def delete_user( ## update teams for team in teams_to_update: - await prisma_client.db.litellm_teamtable.update( + await TeamRepository(prisma_client).table.update( where={"team_id": team.team_id}, data={"members_with_roles": team.members_with_roles}, ) # End of Audit logging ## DELETE ASSOCIATED KEYS - await prisma_client.db.litellm_verificationtoken.delete_many( + await VerificationTokenRepository(prisma_client).table.delete_many( where={"user_id": {"in": data.user_ids}} ) ## DELETE ASSOCIATED INVITATION LINKS - await prisma_client.db.litellm_invitationlink.delete_many( + await InvitationLinkRepository(prisma_client).table.delete_many( where={ "OR": [ {"user_id": {"in": data.user_ids}}, @@ -2300,17 +2311,17 @@ async def delete_user( ) ## DELETE ASSOCIATED ORGANIZATION MEMBERSHIPS - await prisma_client.db.litellm_organizationmembership.delete_many( + await OrganizationMembershipRepository(prisma_client).table.delete_many( where={"user_id": {"in": data.user_ids}} ) ## DELETE ASSOCIATED TEAM MEMBERSHIPS - await prisma_client.db.litellm_teammembership.delete_many( + await TeamMembershipRepository(prisma_client).table.delete_many( where={"user_id": {"in": data.user_ids}} ) ## DELETE USERS - deleted_users = await prisma_client.db.litellm_usertable.delete_many( + deleted_users = await UserRepository(prisma_client).table.delete_many( where={"user_id": {"in": data.user_ids}} ) @@ -2340,16 +2351,18 @@ async def add_internal_user_to_organization( try: # Check if organization_id exists - organization_row = await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": organization_id} - ) + organization_row = await OrganizationRepository( + prisma_client + ).table.find_unique(where={"organization_id": organization_id}) if organization_row is None: raise Exception( f"Organization not found, passed organization_id={organization_id}" ) # Create a new organization membership entry - new_membership = await prisma_client.db.litellm_organizationmembership.create( + new_membership = await OrganizationMembershipRepository( + prisma_client + ).table.create( data={ "user_id": user_id, "organization_id": organization_id, @@ -2559,13 +2572,13 @@ async def ui_view_users( } # Query users with pagination and filters - users: Optional[List[BaseModel]] = ( - await prisma_client.db.litellm_usertable.find_many( - where=where_conditions, - skip=skip, - take=page_size, - order={"created_at": "desc"}, - ) + users: Optional[List[BaseModel]] = await UserRepository( + prisma_client + ).table.find_many( + where=where_conditions, + skip=skip, + take=page_size, + order={"created_at": "desc"}, ) if not users: diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 1ee5bfb0226..a5a364c3679 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -11,6 +11,7 @@ ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.repositories.table_repositories import JWTKeyMappingRepository router = APIRouter() @@ -61,7 +62,7 @@ async def create_jwt_key_mapping( if data.description is not None: create_data["description"] = data.description - new_mapping = await prisma_client.db.litellm_jwtkeymapping.create( + new_mapping = await JWTKeyMappingRepository(prisma_client).table.create( data=create_data ) @@ -113,7 +114,7 @@ async def update_jwt_key_mapping( try: # Get old mapping for cache invalidation - old_mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( + old_mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique( where={"id": data.id} ) @@ -123,7 +124,7 @@ async def update_jwt_key_mapping( cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - updated_mapping = await prisma_client.db.litellm_jwtkeymapping.update( + updated_mapping = await JWTKeyMappingRepository(prisma_client).table.update( where={"id": data.id}, data=update_data ) @@ -166,7 +167,7 @@ async def delete_jwt_key_mapping( try: # Get old mapping for cache invalidation - old_mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( + old_mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique( where={"id": data.id} ) @@ -176,7 +177,7 @@ async def delete_jwt_key_mapping( cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - await prisma_client.db.litellm_jwtkeymapping.delete(where={"id": data.id}) + await JWTKeyMappingRepository(prisma_client).table.delete(where={"id": data.id}) return {"status": "success"} except HTTPException: raise @@ -206,12 +207,12 @@ async def list_jwt_key_mappings( try: skip = (page - 1) * size - mappings = await prisma_client.db.litellm_jwtkeymapping.find_many( + mappings = await JWTKeyMappingRepository(prisma_client).table.find_many( skip=skip, take=size, order={"created_at": "desc"}, ) - total_count = await prisma_client.db.litellm_jwtkeymapping.count() + total_count = await JWTKeyMappingRepository(prisma_client).table.count() return { "mappings": [_to_response(m) for m in mappings], "total_count": total_count, @@ -245,7 +246,7 @@ async def info_jwt_key_mapping( raise HTTPException(status_code=500, detail="Database not connected") try: - mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( + mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique( where={"id": id} ) if mapping is None: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 99d0bac88af..eba16c077b0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -28,7 +28,6 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( LENGTH_OF_LITELLM_GENERATED_KEY, LITELLM_PROXY_ADMIN_NAME, @@ -58,6 +57,7 @@ ) from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.common_utils import ( _check_passthrough_routes_caller_permission, @@ -91,6 +91,19 @@ handle_exception_on_proxy, is_valid_api_key, ) +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.credentials_repository import CredentialsRepository +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ( + DeletedVerificationTokenRepository, + DeprecatedVerificationTokenRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.router import Router from litellm.secret_managers.main import get_secret from litellm.types.proxy.management_endpoints.key_management_endpoints import ( @@ -582,7 +595,7 @@ async def validate_team_id_used_in_service_account_request( ) # check if team_id exists in the database - team = await prisma_client.db.litellm_teamtable.find_unique( + team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id}, ) if team is None: @@ -774,7 +787,7 @@ async def _common_key_generation_helper( # noqa: PLR0915 ) new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) - _budget = await prisma_client.db.litellm_budgettable.create( + _budget = await BudgetRepository(prisma_client).table.create( data={ **new_budget, # type: ignore "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -1144,7 +1157,7 @@ async def _check_team_key_limits( # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - keys = await prisma_client.db.litellm_verificationtoken.find_many( + keys = await VerificationTokenRepository(prisma_client).table.find_many( where={"team_id": team_table.team_id}, ) # Exclude the key being updated to avoid double-counting its limits. @@ -1294,7 +1307,7 @@ async def _validate_caller_can_assign_key_org( detail="Cannot assign a key to an organization without a user_id on the caller's token", ) - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id}, include={"organization_memberships": True}, ) @@ -1339,7 +1352,7 @@ async def _check_org_key_limits( # get all organization keys # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - keys = await prisma_client.db.litellm_verificationtoken.find_many( + keys = await VerificationTokenRepository(prisma_client).table.find_many( where={"organization_id": org_table.organization_id}, ) # Exclude the key being updated to avoid double-counting its limits. @@ -1837,11 +1850,10 @@ async def prepare_key_update_data( if "budget_duration" in non_default_values: budget_duration = non_default_values.pop("budget_duration") - if ( - budget_duration - and (isinstance(budget_duration, str)) - and len(budget_duration) > 0 - ): + if budget_duration is None: + non_default_values["budget_duration"] = None + non_default_values["budget_reset_at"] = None + elif isinstance(budget_duration, str) and len(budget_duration) > 0: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time key_reset_at = get_budget_reset_time(budget_duration=budget_duration) @@ -1968,9 +1980,9 @@ async def _get_and_validate_existing_key( hashed_token = _hash_token_if_needed(token=token) - existing_key_row = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_token} - ) + existing_key_row = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": hashed_token}) if existing_key_row is None: raise ProxyException( @@ -2505,7 +2517,7 @@ async def update_key_fn( # noqa: PLR0915 }, ) - data_json: dict = data.model_dump(exclude_unset=True, exclude_none=True) + data_json: dict = data.model_dump(exclude_unset=True) key = data_json.pop("key") # get the row from db @@ -2575,6 +2587,17 @@ async def update_key_fn( # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, ) + if data.spend is not None: + try: + from litellm.proxy.proxy_server import _invalidate_spend_counter + + token_to_invalidate = _hash_token_if_needed(key) + await _invalidate_spend_counter( + counter_key=f"spend:key:{token_to_invalidate}" + ) + except Exception: + pass + asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( data=data, @@ -2869,7 +2892,9 @@ async def bulk_update_team_keys( # `blocked` is Boolean? with no default; `/key/generate` writes NULL. Prisma's `NOT` # excludes NULLs, so explicitly OR `false` with `null` to include them. now = datetime.now(timezone.utc) - existing_keys = await prisma_client.db.litellm_verificationtoken.find_many( + existing_keys = await VerificationTokenRepository( + prisma_client + ).table.find_many( where={ "team_id": data.team_id, "AND": [ @@ -2907,7 +2932,9 @@ async def bulk_update_team_keys( seen_hashes.add(h) requested_tokens.append(k) hashed_key_ids.append(h) - existing_keys = await prisma_client.db.litellm_verificationtoken.find_many( + existing_keys = await VerificationTokenRepository( + prisma_client + ).table.find_many( where={"team_id": data.team_id, "token": {"in": hashed_key_ids}} ) @@ -3232,7 +3259,9 @@ async def info_key_fn_v2( # Resolve key_aliases to tokens so we never pass token=None (unbounded query) tokens_to_query = list(data.keys) if data.keys else [] if data.key_aliases: - alias_rows = await prisma_client.db.litellm_verificationtoken.find_many( + alias_rows = await VerificationTokenRepository( + prisma_client + ).table.find_many( where={"key_alias": {"in": data.key_aliases}}, include={"litellm_budget_table": True}, ) @@ -3311,7 +3340,7 @@ async def info_key_fn( hashed_key: Optional[str] = key if key is not None: hashed_key = _hash_token_if_needed(token=key) - key_info = await prisma_client.db.litellm_verificationtoken.find_unique( + key_info = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed_key}, # type: ignore include={"litellm_budget_table": True}, ) @@ -3851,7 +3880,7 @@ async def delete_verification_tokens( if prisma_client: tokens = [_hash_token_if_needed(token=key) for key in tokens] _keys_being_deleted: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( + await VerificationTokenRepository(prisma_client).table.find_many( where={"token": {"in": tokens}} ) ) @@ -3989,7 +4018,9 @@ async def _save_deleted_verification_token_records( """Save deleted verification token records to the database.""" if not records: return - await prisma_client.db.litellm_deletedverificationtoken.create_many(data=records) + await DeletedVerificationTokenRepository(prisma_client).table.create_many( + data=records + ) async def _persist_deleted_verification_tokens( @@ -4017,9 +4048,9 @@ async def delete_key_aliases( user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str] = None, ) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: - _keys_being_deleted = await prisma_client.db.litellm_verificationtoken.find_many( - where={"key_alias": {"in": key_aliases}} - ) + _keys_being_deleted = await VerificationTokenRepository( + prisma_client + ).table.find_many(where={"key_alias": {"in": key_aliases}}) tokens = [key.token for key in _keys_being_deleted] return await delete_verification_tokens( @@ -4054,9 +4085,7 @@ async def _rotate_master_key( # noqa: PLR0915 from litellm.proxy.proxy_server import proxy_config try: - models: Optional[List] = ( - await prisma_client.db.litellm_proxymodeltable.find_many() - ) + models: Optional[List] = await ModelRepository(prisma_client).table.find_many() except Exception: models = None # 2. process model table @@ -4088,7 +4117,7 @@ async def _rotate_master_key( # noqa: PLR0915 ) # 3. process config table try: - config = await prisma_client.db.litellm_config.find_many() + config = await ConfigRepository(prisma_client).table.find_many() except Exception: config = None @@ -4109,7 +4138,7 @@ async def _rotate_master_key( # noqa: PLR0915 ) if encrypted_env_vars: - await prisma_client.db.litellm_config.update( + await ConfigRepository(prisma_client).table.update( where={"param_name": "environment_variables"}, data={"param_value": prisma.Json(encrypted_env_vars)}, # type: ignore[attr-defined] ) @@ -4148,7 +4177,7 @@ async def _rotate_master_key( # noqa: PLR0915 # 5. process credentials table try: - credentials = await prisma_client.db.litellm_credentialstable.find_many() + credentials = await CredentialsRepository(prisma_client).table.find_many() except Exception: credentials = None if credentials: @@ -4171,7 +4200,7 @@ async def _rotate_master_key( # noqa: PLR0915 _cred_data["credential_info"] = prisma.Json( # type: ignore[attr-defined] _cred_data["credential_info"] ) - await prisma_client.db.litellm_credentialstable.update( + await CredentialsRepository(prisma_client).table.update( where={"credential_name": cred.credential_name}, data={ **_cred_data, @@ -4243,7 +4272,7 @@ async def _insert_deprecated_key( try: revoke_at = datetime.now(timezone.utc) + timedelta(seconds=grace_seconds) - await prisma_client.db.litellm_deprecatedverificationtoken.upsert( + await DeprecatedVerificationTokenRepository(prisma_client).table.upsert( where={"token": old_token_hash}, data={ "create": { @@ -4335,7 +4364,7 @@ async def _execute_virtual_key_regeneration( grace_period=data.grace_period if data else None, ) - updated_token = await prisma_client.db.litellm_verificationtoken.update( + updated_token = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, data=update_data, # type: ignore ) @@ -4530,7 +4559,7 @@ async def regenerate_key_fn( # noqa: PLR0915 else: hashed_api_key = hash_token(key) - _key_in_db = await prisma_client.db.litellm_verificationtoken.find_unique( + _key_in_db = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed_api_key}, ) if _key_in_db is None: @@ -4719,7 +4748,7 @@ async def reset_key_spend_fn( else: hashed_api_key = hash_token(key) - _key_in_db = await prisma_client.db.litellm_verificationtoken.find_unique( + _key_in_db = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed_api_key}, include={"litellm_budget_table": True}, ) @@ -4739,7 +4768,7 @@ async def reset_key_spend_fn( user_api_key_cache=user_api_key_cache, ) - updated_key = await prisma_client.db.litellm_verificationtoken.update( + updated_key = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, data={"spend": reset_to}, ) @@ -4756,6 +4785,13 @@ async def reset_key_spend_fn( proxy_logging_obj=proxy_logging_obj, ) + try: + from litellm.proxy.proxy_server import _invalidate_spend_counter + + await _invalidate_spend_counter(counter_key=f"spend:key:{hashed_api_key}") + except Exception: + pass + max_budget = updated_key.max_budget budget_reset_at = updated_key.budget_reset_at @@ -4792,11 +4828,11 @@ async def validate_key_list_check( param="user_id", code=status.HTTP_403_FORBIDDEN, ) - complete_user_info_db_obj: Optional[BaseModel] = ( - await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_api_key_dict.user_id}, - include={"organization_memberships": True}, - ) + complete_user_info_db_obj: Optional[BaseModel] = await UserRepository( + prisma_client + ).table.find_unique( + where={"user_id": user_api_key_dict.user_id}, + include={"organization_memberships": True}, ) if complete_user_info_db_obj is None: @@ -4846,7 +4882,9 @@ async def validate_key_list_check( if key_hash: try: - key_info = await prisma_client.db.litellm_verificationtoken.find_unique( + key_info = await VerificationTokenRepository( + prisma_client + ).table.find_unique( where={"token": key_hash}, ) except Exception: @@ -4879,11 +4917,9 @@ async def _fetch_user_team_objects( if complete_user_info is None or not complete_user_info.teams: return [] - teams: Optional[List[BaseModel]] = ( - await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"in": complete_user_info.teams}} - ) - ) + teams: Optional[List[BaseModel]] = await TeamRepository( + prisma_client + ).table.find_many(where={"team_id": {"in": complete_user_info.teams}}) if teams is None: return [] @@ -5160,7 +5196,7 @@ async def _apply_non_admin_alias_scope( # Look up the user's teams from the user table user_teams: List[str] = [] if user_api_key_dict.user_id: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) if user_row is not None: @@ -5548,7 +5584,7 @@ async def _list_key_helper( # Fetch keys with pagination if use_deleted_table: - keys = await prisma_client.db.litellm_deletedverificationtoken.find_many( + keys = await DeletedVerificationTokenRepository(prisma_client).table.find_many( where=where, # type: ignore skip=skip, # type: ignore take=size, # type: ignore @@ -5562,7 +5598,7 @@ async def _list_key_helper( ), ) else: - keys = await prisma_client.db.litellm_verificationtoken.find_many( + keys = await VerificationTokenRepository(prisma_client).table.find_many( where=where, # type: ignore skip=skip, # type: ignore take=size, # type: ignore @@ -5581,11 +5617,13 @@ async def _list_key_helper( # Get total count of keys if use_deleted_table: - total_count = await prisma_client.db.litellm_deletedverificationtoken.count( + total_count = await DeletedVerificationTokenRepository( + prisma_client + ).table.count( where=where # type: ignore ) else: - total_count = await prisma_client.db.litellm_verificationtoken.count( + total_count = await VerificationTokenRepository(prisma_client).table.count( where=where # type: ignore ) @@ -5601,7 +5639,7 @@ async def _list_key_helper( created_by_ids = [key.created_by for key in keys if key.created_by] all_ids = list(set(user_ids + created_by_ids)) # Remove duplicates if all_ids: - users = await prisma_client.db.litellm_usertable.find_many( + users = await UserRepository(prisma_client).table.find_many( where={"user_id": {"in": all_ids}} ) user_map = {user.user_id: user for user in users} @@ -5688,7 +5726,7 @@ async def _check_key_admin_access( return # Look up the target key to find its team - target_key_row = await prisma_client.db.litellm_verificationtoken.find_unique( + target_key_row = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed_token} ) if target_key_row is None: @@ -5755,6 +5793,9 @@ async def block_key( Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can block keys. """ + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + ) from litellm.proxy.proxy_server import ( create_audit_log_for_update, hash_token, @@ -5763,9 +5804,6 @@ async def block_key( proxy_logging_obj, user_api_key_cache, ) - from litellm.proxy.management_helpers.audit_logs import ( - get_audit_log_changed_by, - ) if prisma_client is None: raise Exception("{}".format(CommonProxyErrors.db_not_connected_error.value)) @@ -5792,9 +5830,9 @@ async def block_key( ) # Check if the key exists before trying to block it - existing_record = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_token} - ) + existing_record = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": hashed_token}) if existing_record is None: raise ProxyException( message="Key not found.", @@ -5824,7 +5862,7 @@ async def block_key( ) ) - record = await prisma_client.db.litellm_verificationtoken.update( + record = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_token}, data={"blocked": True} # type: ignore ) @@ -5869,6 +5907,9 @@ async def unblock_key( Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can unblock keys. """ + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + ) from litellm.proxy.proxy_server import ( create_audit_log_for_update, hash_token, @@ -5877,9 +5918,6 @@ async def unblock_key( proxy_logging_obj, user_api_key_cache, ) - from litellm.proxy.management_helpers.audit_logs import ( - get_audit_log_changed_by, - ) if prisma_client is None: raise Exception("{}".format(CommonProxyErrors.db_not_connected_error.value)) @@ -5906,9 +5944,9 @@ async def unblock_key( ) # Check if the key exists before trying to unblock it - existing_record = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_token} - ) + existing_record = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": hashed_token}) if existing_record is None: raise ProxyException( message="Key not found.", @@ -5938,7 +5976,7 @@ async def unblock_key( ) ) - record = await prisma_client.db.litellm_verificationtoken.update( + record = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_token}, data={"blocked": False} # type: ignore ) @@ -6202,9 +6240,9 @@ async def _enforce_unique_key_alias( # Exclude the current key from the uniqueness check where_clause["NOT"] = {"token": existing_key_token} - existing_key = await prisma_client.db.litellm_verificationtoken.find_first( - where=where_clause - ) + existing_key = await VerificationTokenRepository( + prisma_client + ).table.find_first(where=where_clause) if existing_key is not None: raise ProxyException( message=f"Key with alias '{key_alias}' already exists. Unique key aliases across all keys are required.", diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f1edcc9c7b2..0df4675b67f 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -21,7 +21,7 @@ import os from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Iterable, List, Literal, Optional +from typing import Any, Dict, Iterable, List, Literal, Optional, Set from fastapi import ( APIRouter, @@ -60,6 +60,10 @@ encrypt_value_helper, ) from litellm.proxy.management_helpers.audit_logs import get_audit_log_changed_by +from litellm.repositories.table_repositories import ( + MCPServerRepository, + MCPUserCredentialsRepository, +) router = APIRouter(prefix="/v1/mcp", tags=["mcp"]) @@ -761,7 +765,7 @@ async def get_mcp_access_groups( # Get from DB if prisma_client is not None: try: - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many() + mcp_servers = await MCPServerRepository(prisma_client).table.find_many() for server in mcp_servers: if ( hasattr(server, "mcp_access_groups") @@ -879,6 +883,32 @@ async def _get_team_scoped_mcp_server_list( return _redact_mcp_credentials_list(servers) + async def _resolve_accessible_mcp_servers( + user_api_key_dict: UserAPIKeyAuth, + ) -> List[LiteLLM_MCPServerTable]: + """The server set the dashboard grid shows (GET /v1/mcp/server, no team + filter), returned unredacted. Callers that surface this to a client must + apply their own redaction; the per-user env-var status endpoint relies on + the raw env_vars and only ever returns is_set booleans, never secrets. + + Sharing this resolution keeps the red "missing user fields" card status + aligned with the cards actually rendered: an admin in view_all mode sees + every server even when their key carries no per-server MCP grant. + """ + if ( + _get_user_mcp_management_mode() == "view_all" + and not _is_restricted_virtual_key_request(user_api_key_dict) + ): + return await global_mcp_server_manager.get_all_mcp_servers_unfiltered() + + aggregated: Dict[str, LiteLLM_MCPServerTable] = {} + for auth_context in await build_effective_auth_contexts(user_api_key_dict): + for server in await global_mcp_server_manager.get_all_allowed_mcp_servers( + user_api_key_auth=auth_context + ): + aggregated.setdefault(server.server_id, server) + return list(aggregated.values()) + @router.get( "/server", description="Returns the mcp server list with associated teams", @@ -950,30 +980,8 @@ async def fetch_all_mcp_servers( sanitized_team_id ) else: - user_mcp_management_mode = _get_user_mcp_management_mode() - - if user_mcp_management_mode == "view_all" and not is_restricted_virtual_key: - servers = ( - await global_mcp_server_manager.get_all_mcp_servers_unfiltered() - ) - redacted_mcp_servers = _redact_mcp_credentials_list(servers) - else: - auth_contexts = await build_effective_auth_contexts(user_api_key_dict) - - aggregated_servers: Dict[str, LiteLLM_MCPServerTable] = {} - for auth_context in auth_contexts: - servers = ( - await global_mcp_server_manager.get_all_allowed_mcp_servers( - user_api_key_auth=auth_context - ) - ) - for server in servers: - if server.server_id not in aggregated_servers: - aggregated_servers[server.server_id] = server - - redacted_mcp_servers = _redact_mcp_credentials_list( - aggregated_servers.values() - ) + servers = await _resolve_accessible_mcp_servers(user_api_key_dict) + redacted_mcp_servers = _redact_mcp_credentials_list(servers) # augment the mcp servers with public status if litellm.public_mcp_servers is not None: @@ -994,10 +1002,10 @@ async def fetch_all_mcp_servers( if getattr(s, "is_byok", False) ] if byok_server_ids: - cred_rows = ( - await _byok_prisma_client.db.litellm_mcpusercredentials.find_many( - where={"user_id": user_id, "server_id": {"in": byok_server_ids}} - ) + cred_rows = await MCPUserCredentialsRepository( + _byok_prisma_client + ).table.find_many( + where={"user_id": user_id, "server_id": {"in": byok_server_ids}} ) cred_set = {r.server_id for r in cred_rows} for server in redacted_mcp_servers: @@ -1714,11 +1722,13 @@ async def _get_cached_temporary_mcp_server_or_404( status_code=status.HTTP_403_FORBIDDEN, detail={"error": f"Access denied to MCP server {server_id}"}, ) - allowed_server_ids = ( - await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_dict + allowed_server_ids: Set[str] = set() + for auth_context in await build_effective_auth_contexts(user_api_key_dict): + allowed_server_ids.update( + await global_mcp_server_manager.get_allowed_mcp_servers( + auth_context + ) ) - ) if server.server_id not in allowed_server_ids: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -2391,9 +2401,7 @@ async def list_mcp_user_env_var_status( user_id = user_api_key_dict.user_id or "" if not user_id: return [] - accessible = await get_all_mcp_servers_for_user( - prisma_client, user_api_key_dict - ) + accessible = await _resolve_accessible_mcp_servers(user_api_key_dict) if not accessible: return [] server_ids = [s.server_id for s in accessible] 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 b05cfef5760..a8551f6333a 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -19,6 +19,7 @@ clear_cache, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.model_repository import ModelRepository from litellm.types.proxy.management_endpoints.model_management_endpoints import ( AccessGroupInfo, DeleteModelGroupResponse, @@ -95,7 +96,7 @@ async def update_deployments_with_access_group( verbose_proxy_logger.debug(f"Updating deployments for model_name: {model_name}") # Get all deployments with this model_name - deployments = await prisma_client.db.litellm_proxymodeltable.find_many( + deployments = await ModelRepository(prisma_client).table.find_many( where={"model_name": model_name} ) @@ -124,7 +125,7 @@ async def update_deployments_with_access_group( # Only update in DB if modified if was_modified: - await prisma_client.db.litellm_proxymodeltable.update( + await ModelRepository(prisma_client).table.update( where={"model_id": deployment.model_id}, data={"model_info": json.dumps(updated_model_info)}, ) @@ -152,7 +153,7 @@ async def update_specific_deployments_with_access_group( models_updated = 0 for model_id in model_ids: verbose_proxy_logger.debug(f"Updating specific deployment model_id: {model_id}") - deployment = await prisma_client.db.litellm_proxymodeltable.find_unique( + deployment = await ModelRepository(prisma_client).table.find_unique( where={"model_id": model_id} ) if deployment is None: @@ -168,7 +169,7 @@ async def update_specific_deployments_with_access_group( access_group=access_group, ) if was_modified: - await prisma_client.db.litellm_proxymodeltable.update( + await ModelRepository(prisma_client).table.update( where={"model_id": model_id}, data={"model_info": json.dumps(updated_model_info)}, ) @@ -215,7 +216,7 @@ async def get_all_access_groups_from_db( Dict[str, AccessGroupInfo]: Dictionary mapping access_group name to info """ # Get all deployments - deployments = await prisma_client.db.litellm_proxymodeltable.find_many() + deployments = await ModelRepository(prisma_client).table.find_many() # Build access group map access_group_map: Dict[str, Dict[str, Any]] = {} @@ -604,7 +605,7 @@ async def update_access_group( try: # Step 1: Remove access group from ALL DB deployments (skip config models) - all_deployments = await prisma_client.db.litellm_proxymodeltable.find_many() + all_deployments = await ModelRepository(prisma_client).table.find_many() for deployment in all_deployments: model_info = deployment.model_info or {} @@ -615,7 +616,7 @@ async def update_access_group( ) if was_modified: - await prisma_client.db.litellm_proxymodeltable.update( + await ModelRepository(prisma_client).table.update( where={"model_id": deployment.model_id}, data={"model_info": json.dumps(updated_model_info)}, ) @@ -722,7 +723,7 @@ async def delete_access_group( try: # Remove access group from all DB deployments (skip config models) - all_deployments = await prisma_client.db.litellm_proxymodeltable.find_many() + all_deployments = await ModelRepository(prisma_client).table.find_many() models_updated = 0 for deployment in all_deployments: @@ -734,7 +735,7 @@ async def delete_access_group( ) if was_modified: - await prisma_client.db.litellm_proxymodeltable.update( + await ModelRepository(prisma_client).table.update( where={"model_id": deployment.model_id}, data={"model_info": json.dumps(updated_model_info)}, ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index e4ecda3fe31..0be476469a6 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -48,6 +48,9 @@ ) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log from litellm.proxy.utils import PrismaClient +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ModelTableRepository +from litellm.repositories.team_repository import TeamRepository from litellm.types.proxy.management_endpoints.model_management_endpoints import ( UpdateUsefulLinksRequest, ) @@ -86,7 +89,7 @@ async def get_db_model( ) -> Optional[Deployment]: db_model = cast( Optional[BaseModel], - await prisma_client.db.litellm_proxymodeltable.find_unique( + await ModelRepository(prisma_client).table.find_unique( where={"model_id": model_id} ), ) @@ -290,7 +293,7 @@ async def patch_model( update_data["updated_at"] = cast(str, get_utc_datetime()) # Perform partial update - updated_model = await prisma_client.db.litellm_proxymodeltable.update( + updated_model = await ModelRepository(prisma_client).table.update( where={"model_id": model_id}, data=update_data, ) @@ -362,7 +365,7 @@ async def _add_model_to_db( if model_params.model_info.id is not None: _data["model_id"] = model_params.model_info.id if should_create_model_in_db: - model_response = await prisma_client.db.litellm_proxymodeltable.create( + model_response = await ModelRepository(prisma_client).table.create( data=_data # type: ignore ) else: @@ -558,7 +561,7 @@ async def _setup_new_team_model_assignment( async def _get_team_deployments( - team_id: str, prisma_client: PrismaClient + team_id: str, prisma_client: PrismaClient, table: Optional[Any] = None ) -> List[LiteLLM_ProxyModelTable]: """ Fetch all deployments for a given team_id from the database. @@ -569,9 +572,13 @@ async def _get_team_deployments( Note: prisma-client-py 0.11.0 does not support JSON path filtering, so we filter by the model_name prefix (team models use "model_name_{team_id}_*") and confirm team_id in model_info with Python-side filtering. + + Pass ``table`` (a transaction's proxy-model table) to run the read inside an + existing transaction. """ prefix = f"model_name_{team_id}_" - response = await prisma_client.db.litellm_proxymodeltable.find_many( + table = table or ModelRepository(prisma_client).table + response = await table.find_many( where={ "model_name": {"startswith": prefix}, } @@ -593,6 +600,42 @@ async def _get_team_deployments( return result +async def delete_team_models( + team_ids: List[str], + prisma_client: PrismaClient, + llm_router: Optional[Any], +) -> List[str]: + """ + Delete every BYOK model owned by the given teams, from the DB and the router. + + The DB rows are removed inside a single transaction, so deletion is atomic + across all team_ids. Each team's rows are deleted by the exact model_ids read + in the same transaction, which keeps the deleted set identical to the set + handed to the router. The router is synced only after the transaction commits, + so a rollback can never leave a deployment live in the router without its row. + + Returns the model_ids that were deleted. + """ + deleted_model_ids: List[str] = [] + async with prisma_client.db.tx() as tx: + for team_id in team_ids: + rows = await _get_team_deployments( + team_id, prisma_client, table=tx.litellm_proxymodeltable + ) + model_ids = [row.model_id for row in rows] + if model_ids: + await tx.litellm_proxymodeltable.delete_many( + where={"model_id": {"in": model_ids}} + ) + deleted_model_ids.extend(model_ids) + + if llm_router is not None: + for model_id in deleted_model_ids: + llm_router.delete_deployment(id=model_id) + + return deleted_model_ids + + async def _get_team_public_model_names( team_id: str, prisma_client: PrismaClient, @@ -828,7 +871,7 @@ async def allow_team_model_action( detail={"error": CommonProxyErrors.not_premium_user.value}, ) - _existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + _existing_team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": model_params.model_info.team_id} ) @@ -857,16 +900,29 @@ async def can_user_make_model_call( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, premium_user: bool, + allow_missing_team: bool = False, ) -> Literal[True]: ## Check team model auth if ( model_params.model_info is not None and model_params.model_info.team_id is not None ): - team_obj_row = await prisma_client.db.litellm_teamtable.find_unique( + team_obj_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": model_params.model_info.team_id} ) if team_obj_row is None: + # The team was deleted. Callers that opt in (e.g. model deletion) may + # act on the orphaned model, but only as a proxy admin -- without the + # team there is no team-admin membership left to verify. + if allow_missing_team: + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return True + raise HTTPException( + status_code=403, + detail={ + "error": "Only a proxy admin can delete a model whose team has been deleted." + }, + ) raise HTTPException( status_code=400, detail={ @@ -937,7 +993,7 @@ async def delete_model( }, ) - model_in_db = await prisma_client.db.litellm_proxymodeltable.find_unique( + model_in_db = await ModelRepository(prisma_client).table.find_unique( where={"model_id": model_info.id} ) if model_in_db is None: @@ -952,6 +1008,7 @@ async def delete_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + allow_missing_team=True, ) # update DB @@ -961,7 +1018,7 @@ async def delete_model( - store keys separately """ # encrypt litellm params # - result = await prisma_client.db.litellm_proxymodeltable.delete( + result = await ModelRepository(prisma_client).table.delete( where={"model_id": model_info.id} ) @@ -1039,7 +1096,7 @@ async def delete_team_model_alias( Returns: - List of team id + model alias pairs that were removed """ - team_model_aliases = await prisma_client.db.litellm_modeltable.find_many( + team_model_aliases = await ModelTableRepository(prisma_client).table.find_many( include={"team": True} ) tasks = [] @@ -1056,7 +1113,7 @@ async def delete_team_model_alias( removed_model_aliases.append((team_model_alias.team.team_id, key)) del model_aliases[key] tasks.append( - prisma_client.db.litellm_modeltable.update( + ModelTableRepository(prisma_client).table.update( where={"id": id}, data={"model_aliases": json.dumps(model_aliases)}, ) @@ -1275,11 +1332,9 @@ async def update_model( if _model_id is None: raise Exception("model_info.id not provided") - _existing_litellm_params = ( - await prisma_client.db.litellm_proxymodeltable.find_unique( - where={"model_id": _model_id} - ) - ) + _existing_litellm_params = await ModelRepository( + prisma_client + ).table.find_unique(where={"model_id": _model_id}) if _existing_litellm_params is None: if ( @@ -1340,7 +1395,7 @@ async def update_model( "litellm_params": json.dumps(merged_dictionary), # type: ignore "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, } - model_response = await prisma_client.db.litellm_proxymodeltable.update( + model_response = await ModelRepository(prisma_client).table.update( where={"model_id": _model_id}, data=_data, # type: ignore ) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 4d4ed53aaa8..99659121b27 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -40,6 +40,15 @@ management_endpoint_wrapper, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.table_repositories import OrganizationMembershipRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -245,7 +254,7 @@ async def new_organization( if user_api_key_dict.user_id is not None: try: - user_object = await prisma_client.db.litellm_usertable.find_unique( + user_object = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) user_object_correct_type = LiteLLM_UserTable(**user_object.model_dump()) @@ -267,7 +276,7 @@ async def new_organization( new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) - _budget = await prisma_client.db.litellm_budgettable.create( + _budget = await BudgetRepository(prisma_client).table.create( data={ **new_budget, # type: ignore "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -323,7 +332,7 @@ async def new_organization( verbose_proxy_logger.info( f"new_organization_row: {json.dumps(new_organization_row, indent=2)}" ) - response = await prisma_client.db.litellm_organizationtable.create( + response = await OrganizationRepository(prisma_client).table.create( data={ **new_organization_row, # type: ignore }, @@ -372,9 +381,9 @@ async def get_organization_daily_activity( # Restrict non-proxy-admins to only organizations where they are org_admin if not _user_has_admin_view(user_api_key_dict): - memberships = await prisma_client.db.litellm_organizationmembership.find_many( - where={"user_id": user_api_key_dict.user_id} - ) + memberships = await OrganizationMembershipRepository( + prisma_client + ).table.find_many(where={"user_id": user_api_key_dict.user_id}) admin_org_ids = [ m.organization_id for m in memberships @@ -400,7 +409,7 @@ async def get_organization_daily_activity( where_condition = {} if org_ids_list: where_condition["organization_id"] = {"in": list(org_ids_list)} - org_aliases = await prisma_client.db.litellm_organizationtable.find_many( + org_aliases = await OrganizationRepository(prisma_client).table.find_many( where=where_condition ) org_alias_metadata = { @@ -439,10 +448,10 @@ async def _set_object_permission( return None if data.object_permission is not None: - created_object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.create( - data=data.object_permission.model_dump(exclude_none=True), - ) + created_object_permission = await ObjectPermissionRepository( + prisma_client + ).table.create( + data=data.object_permission.model_dump(exclude_none=True), ) del data.object_permission return created_object_permission.object_permission_id @@ -525,10 +534,10 @@ async def update_organization( prisma_client=prisma_client, ) - existing_organization_row = ( - await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": data.organization_id}, - ) + existing_organization_row = await OrganizationRepository( + prisma_client + ).table.find_unique( + where={"organization_id": data.organization_id}, ) if existing_organization_row is None: @@ -574,7 +583,7 @@ async def update_organization( for field in LiteLLM_BudgetTable.model_fields.keys(): updated_organization_row.pop(field, None) - response = await prisma_client.db.litellm_organizationtable.update( + response = await OrganizationRepository(prisma_client).table.update( where={"organization_id": data.organization_id}, data=updated_organization_row, include={"members": True, "teams": True, "litellm_budget_table": True}, @@ -644,19 +653,19 @@ async def delete_organization( deleted_orgs = [] for organization_id in data.organization_ids: # delete all teams in the organization - await prisma_client.db.litellm_teamtable.delete_many( + await TeamRepository(prisma_client).table.delete_many( where={"organization_id": organization_id} ) # delete all members in the organization - await prisma_client.db.litellm_organizationmembership.delete_many( + await OrganizationMembershipRepository(prisma_client).table.delete_many( where={"organization_id": organization_id} ) # delete all keys in the organization - await prisma_client.db.litellm_verificationtoken.delete_many( + await VerificationTokenRepository(prisma_client).table.delete_many( where={"organization_id": organization_id} ) # delete the organization - deleted_org = await prisma_client.db.litellm_organizationtable.delete( + deleted_org = await OrganizationRepository(prisma_client).table.delete( where={"organization_id": organization_id}, include={"members": True, "teams": True, "litellm_budget_table": True}, ) @@ -732,17 +741,15 @@ async def list_organization( # if proxy admin or admin viewer - get all orgs (with optional filters) if _user_has_admin_view(user_api_key_dict): - response = await prisma_client.db.litellm_organizationtable.find_many( + response = await OrganizationRepository(prisma_client).table.find_many( where=where_conditions if where_conditions else None, include={"litellm_budget_table": True, "members": True, "teams": True}, ) # if internal user - get orgs they are a member of (with optional filters) else: - org_memberships = ( - await prisma_client.db.litellm_organizationmembership.find_many( - where={"user_id": user_api_key_dict.user_id} - ) - ) + org_memberships = await OrganizationMembershipRepository( + prisma_client + ).table.find_many(where={"user_id": user_api_key_dict.user_id}) membership_org_ids = [ membership.organization_id for membership in org_memberships ] @@ -756,20 +763,20 @@ async def list_organization( response = [] else: where_conditions["organization_id"] = org_id - response = ( - await prisma_client.db.litellm_organizationtable.find_many( - where=where_conditions, - include={ - "litellm_budget_table": True, - "members": True, - "teams": True, - }, - ) + response = await OrganizationRepository( + prisma_client + ).table.find_many( + where=where_conditions, + include={ + "litellm_budget_table": True, + "members": True, + "teams": True, + }, ) else: # Filter by membership and any additional filters where_conditions["organization_id"] = {"in": membership_org_ids} - response = await prisma_client.db.litellm_organizationtable.find_many( + response = await OrganizationRepository(prisma_client).table.find_many( where=where_conditions, include={ "litellm_budget_table": True, @@ -809,20 +816,20 @@ async def info_organization( prisma_client=prisma_client, ) - response: Optional[LiteLLM_OrganizationTableWithMembers] = ( - await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": organization_id}, - include={ - "litellm_budget_table": True, - "members": { - "include": { - "user": True, - } - }, - "teams": True, - "object_permission": True, + response: Optional[ + LiteLLM_OrganizationTableWithMembers + ] = await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": organization_id}, + include={ + "litellm_budget_table": True, + "members": { + "include": { + "user": True, + } }, - ) + "teams": True, + "object_permission": True, + }, ) if response is None: @@ -868,7 +875,7 @@ async def deprecated_info_organization( prisma_client=prisma_client, ) - response = await prisma_client.db.litellm_organizationtable.find_many( + response = await OrganizationRepository(prisma_client).table.find_many( where={"organization_id": {"in": data.organizations}}, include={"litellm_budget_table": True}, ) @@ -945,11 +952,9 @@ async def organization_member_add( ) # Check if organization exists - existing_organization_row = ( - await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": data.organization_id} - ) - ) + existing_organization_row = await OrganizationRepository( + prisma_client + ).table.find_unique(where={"organization_id": data.organization_id}) if existing_organization_row is None: raise HTTPException( status_code=404, @@ -1012,11 +1017,9 @@ async def find_member_if_email( """ try: - existing_user_email_row: BaseModel = ( - await prisma_client.db.litellm_usertable.find_unique( - where={"user_email": user_email} - ) - ) + existing_user_email_row: BaseModel = await UserRepository( + prisma_client + ).table.find_unique(where={"user_email": user_email}) except Exception: raise HTTPException( status_code=400, @@ -1064,11 +1067,9 @@ async def organization_member_update( ) # Check if organization exists - existing_organization_row = ( - await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": data.organization_id} - ) - ) + existing_organization_row = await OrganizationRepository( + prisma_client + ).table.find_unique(where={"organization_id": data.organization_id}) if existing_organization_row is None: raise HTTPException( status_code=400, @@ -1085,15 +1086,15 @@ async def organization_member_update( data.user_id = existing_user_email_row.user_id try: - existing_organization_membership = ( - await prisma_client.db.litellm_organizationmembership.find_unique( - where={ - "user_id_organization_id": { - "user_id": data.user_id, - "organization_id": data.organization_id, - } + existing_organization_membership = await OrganizationMembershipRepository( + prisma_client + ).table.find_unique( + where={ + "user_id_organization_id": { + "user_id": data.user_id, + "organization_id": data.organization_id, } - ) + } ) except Exception as e: raise HTTPException( @@ -1114,7 +1115,7 @@ async def organization_member_update( # org-scoped operations. An org-admin of any org could otherwise # alter a PROXY_ADMIN user's per-org role, which has downstream # effects on admin UI filtering and scope derivation. - target_user_row = await prisma_client.db.litellm_usertable.find_unique( + target_user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": data.user_id} ) if target_user_row is not None and getattr( @@ -1136,7 +1137,7 @@ async def organization_member_update( # Update member role if data.role is not None: - await prisma_client.db.litellm_organizationmembership.update( + await OrganizationMembershipRepository(prisma_client).table.update( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1165,7 +1166,7 @@ async def organization_member_update( ) # update organization membership with new budget_id - await prisma_client.db.litellm_organizationmembership.update( + await OrganizationMembershipRepository(prisma_client).table.update( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1174,16 +1175,16 @@ async def organization_member_update( }, data={"budget_id": budget_id}, ) - final_organization_membership: Optional[BaseModel] = ( - await prisma_client.db.litellm_organizationmembership.find_unique( - where={ - "user_id_organization_id": { - "user_id": data.user_id, - "organization_id": data.organization_id, - } - }, - include={"litellm_budget_table": True}, - ) + final_organization_membership: Optional[ + BaseModel + ] = await OrganizationMembershipRepository(prisma_client).table.find_unique( + where={ + "user_id_organization_id": { + "user_id": data.user_id, + "organization_id": data.organization_id, + } + }, + include={"litellm_budget_table": True}, ) if final_organization_membership is None: @@ -1239,7 +1240,9 @@ async def organization_member_delete( ) data.user_id = existing_user_email_row.user_id - member_to_delete = await prisma_client.db.litellm_organizationmembership.delete( + member_to_delete = await OrganizationMembershipRepository( + prisma_client + ).table.delete( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1273,17 +1276,15 @@ async def add_member_to_organization( existing_user_email_row = None ## Check if user exists in LiteLLM_UserTable - user exists - either the user_id or user_email is in LiteLLM_UserTable if member.user_id is not None: - existing_user_id_row = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": member.user_id} - ) + existing_user_id_row = await UserRepository( + prisma_client + ).table.find_unique(where={"user_id": member.user_id}) if existing_user_id_row is None and member.user_email is not None: try: - existing_user_email_row = ( - await prisma_client.db.litellm_usertable.find_unique( - where={"user_email": member.user_email} - ) - ) + existing_user_email_row = await UserRepository( + prisma_client + ).table.find_unique(where={"user_email": member.user_email}) except Exception as e: raise ValueError( f"Potential NON-Existent or Duplicate user email in DB: Error finding a unique instance of user_email={member.user_email} in LiteLLM_UserTable.: {e}" @@ -1326,14 +1327,14 @@ async def add_member_to_organization( ) # Add user to organization - _organization_membership = ( - await prisma_client.db.litellm_organizationmembership.create( - data={ - "organization_id": organization_id, - "user_id": user_object.user_id, - "user_role": member.role, - } - ) + _organization_membership = await OrganizationMembershipRepository( + prisma_client + ).table.create( + data={ + "organization_id": organization_id, + "user_id": user_object.user_id, + "user_role": member.role, + } ) organization_membership = LiteLLM_OrganizationMembershipTable( **_organization_membership.model_dump() diff --git a/litellm/proxy/management_endpoints/scim/scim_transformations.py b/litellm/proxy/management_endpoints/scim/scim_transformations.py index 28fb87d9b3d..d1e00f87b69 100644 --- a/litellm/proxy/management_endpoints/scim/scim_transformations.py +++ b/litellm/proxy/management_endpoints/scim/scim_transformations.py @@ -6,6 +6,7 @@ Member, NewUserResponse, ) +from litellm.repositories.team_repository import TeamRepository from litellm.types.proxy.management_endpoints.scim_v2 import * @@ -29,7 +30,7 @@ async def transform_litellm_user_to_scim_user( # Get user's teams/groups groups = [] for team_id in user.teams or []: - team = await prisma_client.db.litellm_teamtable.find_unique( + team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) if team: diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 1f20764f837..0798d1a510d 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -22,7 +22,6 @@ import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -41,6 +40,7 @@ ) from litellm.proxy.auth.auth_checks import _delete_cache_key_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.scim.scim_transformations import ( ScimTransformations, @@ -51,6 +51,16 @@ team_member_delete, ) from litellm.proxy.utils import _premium_user_check, handle_exception_on_proxy +from litellm.repositories.table_repositories import ( + InvitationLinkRepository, + OrganizationMembershipRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.management_endpoints.scim_v2 import * @@ -74,7 +84,7 @@ async def handle_existing_user_by_email( if not new_user_request.user_email: return None - existing_user = await prisma_client.db.litellm_usertable.find_first( + existing_user = await UserRepository(prisma_client).table.find_first( where={"user_email": new_user_request.user_email} ) @@ -82,7 +92,7 @@ async def handle_existing_user_by_email( return None # Update the user - updated_user = await prisma_client.db.litellm_usertable.update( + updated_user = await UserRepository(prisma_client).table.update( where={"user_id": existing_user.user_id}, data={ "user_id": new_user_request.user_id, @@ -139,7 +149,7 @@ async def _check_user_exists(user_id: str): """Check if user exists and return user, raise 404 if not found.""" prisma_client = await _get_prisma_client_or_raise_exception() - user = await prisma_client.db.litellm_usertable.find_unique( + user = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) @@ -155,7 +165,7 @@ async def _check_team_exists(team_id: str): """Check if team exists and return team, raise 404 if not found.""" prisma_client = await _get_prisma_client_or_raise_exception() - team = await prisma_client.db.litellm_teamtable.find_unique( + team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) @@ -268,7 +278,7 @@ async def _extract_group_member_ids(group: SCIMGroup) -> GroupMemberExtractionRe ) # Check if user exists - user = await prisma_client.db.litellm_usertable.find_unique( + user = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) @@ -310,7 +320,7 @@ async def _get_team_members_display(member_ids: List[str]) -> List[SCIMMember]: members: List[SCIMMember] = [] for member_id in member_ids: - user = await prisma_client.db.litellm_usertable.find_unique( + user = await UserRepository(prisma_client).table.find_unique( where={"user_id": member_id} ) if user: @@ -367,7 +377,7 @@ async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: # `blocked` is a nullable column with no default, so existing rows # typically hold NULL; treat NULL as "not blocked" since SQL equality # on NULL would otherwise silently skip them. - candidates = await prisma_client.db.litellm_verificationtoken.find_many( + candidates = await VerificationTokenRepository(prisma_client).table.find_many( where={ "user_id": user_id, "OR": [{"blocked": False}, {"blocked": None}], @@ -375,7 +385,7 @@ async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: ) affected_keys = candidates else: - candidates = await prisma_client.db.litellm_verificationtoken.find_many( + candidates = await VerificationTokenRepository(prisma_client).table.find_many( where={"user_id": user_id, "blocked": True}, ) affected_keys = [k for k in candidates if _key_was_scim_blocked(k.metadata)] @@ -395,7 +405,7 @@ async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: for k, v in current_metadata.items() if k != SCIM_BLOCKED_METADATA_KEY } - await prisma_client.db.litellm_verificationtoken.update( + await VerificationTokenRepository(prisma_client).table.update( where={"token": key_row.token}, data={"blocked": blocked, "metadata": safe_dumps(new_metadata)}, ) @@ -423,7 +433,7 @@ async def _delete_rows_referencing_user(prisma_client: Any, *, user_id: str) -> the user delete with an FK constraint violation (e.g. ``LiteLLM_InvitationLink_user_id_fkey``). """ - await prisma_client.db.litellm_invitationlink.delete_many( + await InvitationLinkRepository(prisma_client).table.delete_many( where={ "OR": [ {"user_id": user_id}, @@ -432,10 +442,10 @@ async def _delete_rows_referencing_user(prisma_client: Any, *, user_id: str) -> ] } ) - await prisma_client.db.litellm_organizationmembership.delete_many( + await OrganizationMembershipRepository(prisma_client).table.delete_many( where={"user_id": user_id} ) - await prisma_client.db.litellm_teammembership.delete_many( + await TeamMembershipRepository(prisma_client).table.delete_many( where={"user_id": user_id} ) @@ -897,17 +907,17 @@ async def get_users( where_conditions["user_email"] = filter_value # Get users from database - users: List[LiteLLM_UserTable] = ( - await prisma_client.db.litellm_usertable.find_many( - where=where_conditions, - skip=(startIndex - 1), - take=count, - order={"created_at": "desc"}, - ) + users: List[LiteLLM_UserTable] = await UserRepository( + prisma_client + ).table.find_many( + where=where_conditions, + skip=(startIndex - 1), + take=count, + order={"created_at": "desc"}, ) # Get total count for pagination - total_count = await prisma_client.db.litellm_usertable.count( + total_count = await UserRepository(prisma_client).table.count( where=where_conditions ) @@ -975,7 +985,7 @@ async def create_user( # Check if user already exists if user.userName: - existing_user = await prisma_client.db.litellm_usertable.find_unique( + existing_user = await UserRepository(prisma_client).table.find_unique( where={"user_id": user.userName} ) if existing_user: @@ -1094,7 +1104,7 @@ async def update_user( "metadata": safe_dumps(metadata), } - updated_user = await prisma_client.db.litellm_usertable.update( + updated_user = await UserRepository(prisma_client).table.update( where={"user_id": user_id}, data=update_data, ) @@ -1137,7 +1147,7 @@ async def delete_user( teams = [] if existing_user.teams: for team_id in existing_user.teams: - team = await prisma_client.db.litellm_teamtable.find_unique( + team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) if team: @@ -1148,7 +1158,7 @@ async def delete_user( current_members = team.members or [] if user_id in current_members: new_members = [m for m in current_members if m != user_id] - await prisma_client.db.litellm_teamtable.update( + await TeamRepository(prisma_client).table.update( where={"team_id": team.team_id}, data={"members": new_members} ) @@ -1157,7 +1167,7 @@ async def delete_user( await _delete_rows_referencing_user(prisma_client, user_id=user_id) # Delete user - await prisma_client.db.litellm_usertable.delete(where={"user_id": user_id}) + await UserRepository(prisma_client).table.delete(where={"user_id": user_id}) return Response(status_code=204) except Exception as e: @@ -1413,7 +1423,7 @@ async def patch_user( update_data["metadata"] = safe_dumps(update_data["metadata"]) - updated_user = await prisma_client.db.litellm_usertable.update( + updated_user = await UserRepository(prisma_client).table.update( where={"user_id": user_id}, data=update_data, ) @@ -1465,7 +1475,7 @@ async def get_groups( where_conditions["team_alias"] = team_alias # Get teams from database - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where=where_conditions, skip=(startIndex - 1), take=count, @@ -1473,7 +1483,7 @@ async def get_groups( ) # Get total count for pagination - total_count = await prisma_client.db.litellm_teamtable.count( + total_count = await TeamRepository(prisma_client).table.count( where=where_conditions ) @@ -1561,7 +1571,7 @@ async def create_group( team_id = group.id or group.externalId or str(uuid.uuid4()) # Check if team already exists - existing_team = await prisma_client.db.litellm_teamtable.find_unique( + existing_team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) @@ -1638,7 +1648,7 @@ async def update_group( } # Update team in database - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": group_id}, data=update_data, ) @@ -1683,19 +1693,19 @@ async def delete_group( # For each member, remove this team from their teams list for member_id in existing_team.members or []: - user = await prisma_client.db.litellm_usertable.find_unique( + user = await UserRepository(prisma_client).table.find_unique( where={"user_id": member_id} ) if user: current_teams = user.teams or [] if group_id in current_teams: new_teams = [t for t in current_teams if t != group_id] - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": member_id}, data={"teams": new_teams} ) # Delete team - await prisma_client.db.litellm_teamtable.delete(where={"team_id": group_id}) + await TeamRepository(prisma_client).table.delete(where={"team_id": group_id}) return Response(status_code=204) @@ -1748,7 +1758,7 @@ async def _process_group_patch_operations( detail={"error": "Invalid member: user ID cannot be empty."}, ) - user = await prisma_client.db.litellm_usertable.find_unique( + user = await UserRepository(prisma_client).table.find_unique( where={"user_id": member_id} ) if user: @@ -1805,7 +1815,7 @@ async def _apply_group_patch_updates( update_data["members"] = list(final_members) # Update team in database - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": group_id}, data=update_data, ) @@ -1877,7 +1887,7 @@ async def patch_group( # Refresh team data from database to get the latest state after concurrent updates # This prevents race conditions when multiple PATCH requests come in simultaneously - refreshed_team = await prisma_client.db.litellm_teamtable.find_unique( + refreshed_team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": group_id} ) if refreshed_team: @@ -1894,7 +1904,7 @@ async def patch_group( await _handle_group_membership_changes(group_id, current_members, final_members) # Refresh team one more time to get final state after membership changes - final_team = await prisma_client.db.litellm_teamtable.find_unique( + final_team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": group_id} ) if final_team: diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 49d9b67a28a..f0bb8bdb5ff 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -25,6 +25,14 @@ get_daily_activity, ) from litellm.proxy.management_helpers.utils import handle_budget_for_entity +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ( + DailyTagSpendRepository, + TagRepository, +) +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.tag_management import ( TagConfig, TagDeleteRequest, @@ -56,7 +64,7 @@ async def _get_internal_user_api_keys( if user_id is None: return sorted(user_api_keys) - key_records = await prisma_client.db.litellm_verificationtoken.find_many( + key_records = await VerificationTokenRepository(prisma_client).table.find_many( where={"user_id": user_id}, select={"token": True}, ) @@ -109,7 +117,7 @@ async def _get_tag_daily_activity_api_key_filter( async def _get_model_names(prisma_client, model_ids: list) -> Dict[str, str]: """Helper function to get model names from model IDs""" try: - models = await prisma_client.db.litellm_proxymodeltable.find_many( + models = await ModelRepository(prisma_client).table.find_many( where={"model_id": {"in": model_ids}} ) return {model.model_id: model.model_name for model in models} @@ -189,7 +197,7 @@ async def new_tag( ) try: # Check if tag already exists - existing_tag = await prisma_client.db.litellm_tagtable.find_unique( + existing_tag = await TagRepository(prisma_client).table.find_unique( where={"tag_name": tag.name} ) if existing_tag is not None: @@ -210,7 +218,7 @@ async def new_tag( model_info = await _get_model_names(prisma_client, tag.models or []) # Create new tag in database - new_tag_record = await prisma_client.db.litellm_tagtable.create( + new_tag_record = await TagRepository(prisma_client).table.create( data={ "tag_name": tag.name, "description": tag.description, @@ -267,7 +275,7 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): try: # Get current model from database to preserve encrypted fields - db_model = await prisma_client.db.litellm_proxymodeltable.find_unique( + db_model = await ModelRepository(prisma_client).table.find_unique( where={"model_id": deployment.model_info.id} ) @@ -292,7 +300,7 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): existing_params["tags"].append(tag) # Update database with modified params (keeps encrypted fields encrypted) - await prisma_client.db.litellm_proxymodeltable.update( + await ModelRepository(prisma_client).table.update( where={"model_id": deployment.model_info.id}, data={"litellm_params": json.dumps(existing_params)}, ) @@ -335,7 +343,7 @@ async def update_tag( try: # Check if tag exists - existing_tag = await prisma_client.db.litellm_tagtable.find_unique( + existing_tag = await TagRepository(prisma_client).table.find_unique( where={"tag_name": tag.name} ) if existing_tag is None: @@ -367,7 +375,7 @@ async def update_tag( update_data["budget_id"] = budget_id # Update tag in database - updated_tag_record = await prisma_client.db.litellm_tagtable.update( + updated_tag_record = await TagRepository(prisma_client).table.update( where={"tag_name": tag.name}, data=update_data, ) @@ -414,7 +422,7 @@ async def info_tag( try: # Query tags from database with budget info - tag_records = await prisma_client.db.litellm_tagtable.find_many( + tag_records = await TagRepository(prisma_client).table.find_many( where={"tag_name": {"in": data.names}}, include={"litellm_budget_table": True}, ) @@ -535,7 +543,7 @@ async def list_tags( if start_date is not None and end_date is not None: dynamic_tag_where["date"] = {"gte": start_date, "lte": end_date} - dynamic_tag_rows = await prisma_client.db.litellm_dailytagspend.group_by( + dynamic_tag_rows = await DailyTagSpendRepository(prisma_client).table.group_by( by=["tag"], where=dynamic_tag_where, min={"created_at": True}, @@ -551,7 +559,7 @@ async def list_tags( ) ## QUERY STORED TAGS ## - tag_records = await prisma_client.db.litellm_tagtable.find_many( + tag_records = await TagRepository(prisma_client).table.find_many( where=stored_tag_where, include={"litellm_budget_table": True}, ) @@ -626,14 +634,14 @@ async def delete_tag( try: # Check if tag exists - existing_tag = await prisma_client.db.litellm_tagtable.find_unique( + existing_tag = await TagRepository(prisma_client).table.find_unique( where={"tag_name": data.name} ) if existing_tag is None: raise HTTPException(status_code=404, detail=f"Tag {data.name} not found") # Delete tag from database - await prisma_client.db.litellm_tagtable.delete(where={"tag_name": data.name}) + await TagRepository(prisma_client).table.delete(where={"tag_name": data.name}) return {"message": f"Tag {data.name} deleted successfully"} except Exception as e: diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 63b56425b0e..0c11507697d 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -30,6 +30,7 @@ from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.management_endpoints.team_endpoints import _verify_team_access from litellm.proxy.management_helpers.utils import management_endpoint_wrapper +from litellm.repositories.team_repository import TeamRepository router = APIRouter() @@ -249,7 +250,7 @@ async def add_team_callbacks( team_metadata = encrypt_callback_vars(team_metadata) team_metadata_json = json.dumps(team_metadata) # update team_metadata - new_team_row = await prisma_client.db.litellm_teamtable.update( + new_team_row = await TeamRepository(prisma_client).table.update( where={"team_id": team_id}, data={"metadata": team_metadata_json} # type: ignore ) @@ -353,7 +354,7 @@ async def disable_team_logging( team_metadata_json = json.dumps(team_metadata) # Update team in database - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": team_id}, data={"metadata": team_metadata_json} # type: ignore ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ae7da0d29f2..c894813ada4 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -10,8 +10,8 @@ """ import asyncio -import math import json +import math import traceback from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Tuple, Union, cast @@ -102,6 +102,20 @@ management_endpoint_wrapper, ) from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.table_repositories import ( + AccessGroupRepository, + DeletedTeamRepository, + ModelTableRepository, + OrganizationMembershipRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.router import Router from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, @@ -343,6 +357,42 @@ def _clean_team_member_fields(data_dict: dict) -> None: data_dict.pop("team_member_rpm_limit", None) data_dict.pop("team_member_tpm_limit", None) + @staticmethod + async def clear_team_member_budget_fields( + team_table: LiteLLM_TeamTable, + user_api_key_dict: "UserAPIKeyAuth", + updated_kv: dict, + explicitly_set_fields: set, + ) -> dict: + """Clear explicitly-nulled fields on the team member budget row.""" + from litellm.proxy._types import BudgetNewRequest + from litellm.proxy.management_endpoints.budget_management_endpoints import ( + update_budget, + ) + + if team_table.metadata is None: + team_table.metadata = {} + + team_member_budget_id = team_table.metadata.get("team_member_budget_id") + if team_member_budget_id is not None and isinstance(team_member_budget_id, str): + budget_request = BudgetNewRequest(budget_id=team_member_budget_id) + if "team_member_budget" in explicitly_set_fields: + budget_request.max_budget = None + if "team_member_budget_duration" in explicitly_set_fields: + budget_request.budget_duration = None + budget_request.budget_reset_at = None + if "team_member_rpm_limit" in explicitly_set_fields: + budget_request.rpm_limit = None + if "team_member_tpm_limit" in explicitly_set_fields: + budget_request.tpm_limit = None + await update_budget( + budget_obj=budget_request, + user_api_key_dict=user_api_key_dict, + ) + + TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) + return updated_kv + @staticmethod async def backfill_team_member_budget_entries( team_id: str, @@ -365,9 +415,9 @@ async def backfill_team_member_budget_entries( return # Batch-fetch existing memberships for this team (avoids N+1 queries) - existing_memberships = await prisma_client.db.litellm_teammembership.find_many( - where={"team_id": team_id} - ) + existing_memberships = await TeamMembershipRepository( + prisma_client + ).table.find_many(where={"team_id": team_id}) existing_user_ids = {m.user_id for m in existing_memberships} # Identify members with no existing membership row. @@ -386,7 +436,7 @@ async def backfill_team_member_budget_entries( ) if missing: - await prisma_client.db.litellm_teammembership.create_many( + await TeamMembershipRepository(prisma_client).table.create_many( data=missing, skip_duplicates=True, # safety net against concurrent races ) @@ -400,7 +450,7 @@ async def backfill_team_member_budget_entries( # Heal existing membership rows that predate the team_member_budget # configuration: populate budget_id where it is currently NULL. # Rows with an explicit budget_id (per-member override) are left alone. - updated = await prisma_client.db.litellm_teammembership.update_many( + updated = await TeamMembershipRepository(prisma_client).table.update_many( where={"team_id": team_id, "budget_id": None}, data={"budget_id": team_member_budget_id}, ) @@ -456,7 +506,7 @@ async def get_all_team_memberships( # else: # where_obj = {"user_id": str(user_id), "team_id": {"in": team_id}} - team_memberships = await prisma_client.db.litellm_teammembership.find_many( + team_memberships = await TeamMembershipRepository(prisma_client).table.find_many( where=where_obj, include={"litellm_budget_table": True}, ) @@ -739,7 +789,7 @@ async def _check_org_team_limits( # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"organization_id": org_table.organization_id}, ) @@ -767,14 +817,16 @@ async def _check_user_team_limits( user_api_key_cache: Any, ) -> None: """ - Check user team limits for standalone teams (not org-scoped). + Enforce the caller's personal limits when CREATING a standalone team. - This validates: - - Team budget vs user's max_budget - - Team models vs user's allowed models + This validates the requested team budget / models / tpm / rpm against the + caller's own limits, so a non-admin user cannot mint a brand-new team that + is richer than themselves. - Should only be called for standalone teams (when organization_id is None). - For org-scoped teams, use _check_org_team_limits() instead. + Only used by /team/new for standalone teams (organization_id is None). + /team/update does NOT call this — an existing team's admin is already + authorized via _verify_team_access() and is not gated by their personal + wallet. Org-scoped teams use _check_org_team_limits() instead. """ # Validate team budget against user's max_budget if data.max_budget is not None and user_api_key_dict.user_id is not None: @@ -834,6 +886,45 @@ async def _check_user_team_limits( ) +def _check_team_budget_update_authority( + data: UpdateTeamRequest, + user_api_key_dict: UserAPIKeyAuth, + existing_team_max_budget: Optional[float], +) -> None: + """ + Restrict who can grow a standalone team's spend ceiling on /team/update. + + A team admin (already authorized via _verify_team_access) may keep or lower + the team budget, but only a proxy admin may grow it - by raising max_budget + above the team's current value or by removing the cap (setting it to None). + Setting a finite budget on a team that has no cap is a restriction and is + allowed. Org-scoped teams are governed by _check_org_team_limits(). + """ + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + if existing_team_max_budget is None: + return + + budget_explicitly_set = "max_budget" in ( + getattr(data, "model_fields_set", None) or set() + ) + if budget_explicitly_set and data.max_budget is None: + raise HTTPException( + status_code=403, + detail={ + "error": f"Only a proxy admin can remove a team's max_budget. Team's current max_budget={existing_team_max_budget}." + }, + ) + + if data.max_budget is not None and data.max_budget > existing_team_max_budget: + raise HTTPException( + status_code=403, + detail={ + "error": f"Only a proxy admin can raise a team's max_budget. Team's current max_budget={existing_team_max_budget}, requested={data.max_budget}." + }, + ) + + #### TEAM MANAGEMENT #### @router.post( "/team/new", @@ -931,6 +1022,9 @@ async def new_team( # noqa: PLR0915 ``` """ try: + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + ) from litellm.proxy.proxy_server import ( _license_check, create_audit_log_for_update, @@ -938,9 +1032,6 @@ async def new_team( # noqa: PLR0915 prisma_client, user_api_key_cache, ) - from litellm.proxy.management_helpers.audit_logs import ( - get_audit_log_changed_by, - ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -986,7 +1077,7 @@ async def new_team( # noqa: PLR0915 ) # Check if license is over limit - total_teams = await prisma_client.db.litellm_teamtable.count() + total_teams = await TeamRepository(prisma_client).table.count() if total_teams and _license_check.is_team_count_over_limit( team_count=total_teams ): @@ -1092,7 +1183,7 @@ async def new_team( # noqa: PLR0915 created_by=user_api_key_dict.user_id or litellm_proxy_admin_name, updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, ) - model_dict = await prisma_client.db.litellm_modeltable.create( + model_dict = await ModelTableRepository(prisma_client).table.create( {**litellm_modeltable.json(exclude_none=True)} # type: ignore ) # type: ignore @@ -1195,7 +1286,7 @@ async def new_team( # noqa: PLR0915 db_data=complete_team_data_dict ) - team_row: LiteLLM_TeamTable = await prisma_client.db.litellm_teamtable.create( + team_row: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.create( data=complete_team_data_dict, include={"litellm_model_table": True}, # type: ignore ) @@ -1315,11 +1406,11 @@ async def _update_model_table( updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, ) if model_id is None: - model_dict = await prisma_client.db.litellm_modeltable.create( + model_dict = await ModelTableRepository(prisma_client).table.create( data={**litellm_modeltable.json(exclude_none=True)} # type: ignore ) else: - model_dict = await prisma_client.db.litellm_modeltable.upsert( + model_dict = await ModelTableRepository(prisma_client).table.upsert( where={"id": model_id}, data={ "update": {**litellm_modeltable.json(exclude_none=True)}, # type: ignore @@ -1400,7 +1491,7 @@ async def fetch_and_validate_organization( status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} ) - organization_row = await prisma_client.db.litellm_organizationtable.find_unique( + organization_row = await OrganizationRepository(prisma_client).table.find_unique( where={"organization_id": organization_id}, include={"litellm_budget_table": True, "members": True, "teams": True}, ) @@ -1669,7 +1760,7 @@ async def update_team( # noqa: PLR0915 }, ) - existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + existing_team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) @@ -1738,7 +1829,9 @@ async def update_team( # noqa: PLR0915 ): # Is the caller org_admin of the destination org? caller_memberships = ( - await prisma_client.db.litellm_organizationmembership.find_many( + await OrganizationMembershipRepository( + prisma_client + ).table.find_many( where={ "user_id": user_api_key_dict.user_id, "organization_id": data.organization_id, @@ -1794,21 +1887,14 @@ async def update_team( # noqa: PLR0915 prisma_client=prisma_client, ) - # Check user limits for standalone teams (not org-scoped) - # Skip for PROXY_ADMIN users - if ( - user_api_key_dict.user_role is None - or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN - ): - # Only validate user budget/models for standalone teams - # For org-scoped teams, validation is done by _check_org_team_limits() above - if org_id_to_check is None: - await _check_user_team_limits( - data=data, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) + # Only a proxy admin may grow a standalone team's spend ceiling. + # Org-scoped teams are validated by _check_org_team_limits() above. + if org_id_to_check is None: + _check_team_budget_update_authority( + data=data, + user_api_key_dict=user_api_key_dict, + existing_team_max_budget=existing_team_row.max_budget, + ) updated_kv = data.json(exclude_unset=True) @@ -1822,11 +1908,25 @@ async def update_team( # noqa: PLR0915 # Check budget_duration and budget_reset_at _set_budget_reset_at(data, updated_kv) - if TeamMemberBudgetHandler.should_create_budget( - team_member_budget=data.team_member_budget, - team_member_rpm_limit=data.team_member_rpm_limit, - team_member_tpm_limit=data.team_member_tpm_limit, - team_member_budget_duration=data.team_member_budget_duration, + _team_member_fields_in_request = { + field + for field in [ + "team_member_budget", + "team_member_rpm_limit", + "team_member_tpm_limit", + "team_member_budget_duration", + ] + if field in updated_kv + } + + if ( + _team_member_fields_in_request + and TeamMemberBudgetHandler.should_create_budget( + team_member_budget=data.team_member_budget, + team_member_rpm_limit=data.team_member_rpm_limit, + team_member_tpm_limit=data.team_member_tpm_limit, + team_member_budget_duration=data.team_member_budget_duration, + ) ): updated_kv = await TeamMemberBudgetHandler.upsert_team_member_budget_table( team_table=existing_team_row, @@ -1849,6 +1949,13 @@ async def update_team( # noqa: PLR0915 team_member_budget_id=_backfill_budget_id, prisma_client=prisma_client, ) + elif _team_member_fields_in_request: + updated_kv = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=existing_team_row, + user_api_key_dict=user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields=_team_member_fields_in_request, + ) else: TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) @@ -1885,18 +1992,18 @@ async def update_team( # noqa: PLR0915 updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[LiteLLM_TeamTable] = ( - await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data=updated_kv, - # `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={ - "litellm_model_table": True, - "object_permission": True, - }, # type: ignore - ) + team_row: Optional[LiteLLM_TeamTable] = await TeamRepository( + prisma_client + ).table.update( + where={"team_id": data.team_id}, + data=updated_kv, + # `object_permission` is included so `_refresh_cached_team` + # doesn't write a cached team with the relation nulled out — + # see team_model_add for the full rationale. + include={ + "litellm_model_table": True, + "object_permission": True, + }, # type: ignore ) if team_row is None or team_row.team_id is None: @@ -1937,6 +2044,8 @@ def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None: reset_at = get_budget_reset_time(budget_duration=data.budget_duration) updated_kv["budget_reset_at"] = reset_at + elif "budget_duration" in updated_kv and updated_kv["budget_duration"] is None: + updated_kv["budget_reset_at"] = None if data.budget_limits is not None and len(data.budget_limits) > 0: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time @@ -2306,7 +2415,7 @@ async def _add_team_members_to_team( # ADD MEMBER TO TEAM _db_team_members = [m.model_dump() for m in complete_team_data.members_with_roles] - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore ) @@ -2377,7 +2486,7 @@ async def _validate_and_populate_member_user_info( # Case 2: Only user_email provided - populate user_id from DB if member.user_email is not None and member.user_id is None: - user_by_email = await prisma_client.db.litellm_usertable.find_first( + user_by_email = await UserRepository(prisma_client).table.find_first( where={"user_email": {"equals": member.user_email, "mode": "insensitive"}} ) @@ -2410,7 +2519,7 @@ async def _validate_and_populate_member_user_info( # Case 3: Only user_id provided - populate user_email from DB if user exists if member.user_id is not None and member.user_email is None: - user_by_id = await prisma_client.db.litellm_usertable.find_unique( + user_by_id = await UserRepository(prisma_client).table.find_unique( where={"user_id": member.user_id} ) @@ -2608,7 +2717,7 @@ async def team_member_delete( detail={"error": "Either user_id or user_email needs to be passed in"}, ) - _existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + _existing_team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) @@ -2652,7 +2761,7 @@ async def team_member_delete( _db_new_team_members: List[dict] = [m.model_dump() for m in new_team_members] - _ = await prisma_client.db.litellm_teamtable.update( + _ = await TeamRepository(prisma_client).table.update( where={ "team_id": data.team_id, }, @@ -2666,7 +2775,7 @@ async def team_member_delete( key_val["user_id"] = data.user_id elif data.user_email is not None: key_val["user_email"] = data.user_email - existing_user_rows = await prisma_client.db.litellm_usertable.find_many( + existing_user_rows = await UserRepository(prisma_client).table.find_many( where=key_val # type: ignore ) @@ -2678,7 +2787,7 @@ async def team_member_delete( if data.team_id in existing_user.teams: team_list = existing_user.teams team_list.remove(data.team_id) - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={ "user_id": existing_user.user_id, }, @@ -2695,7 +2804,7 @@ async def team_member_delete( user_ids_to_delete.add(existing_user.user_id) for _uid in user_ids_to_delete: - await prisma_client.db.litellm_teammembership.delete_many( + await TeamMembershipRepository(prisma_client).table.delete_many( where={"team_id": data.team_id, "user_id": _uid} ) @@ -2706,13 +2815,13 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": {"in": list(user_ids_to_delete)}, - "team_id": data.team_id, - } - ) + keys_to_delete: List[ + LiteLLM_VerificationToken + ] = await VerificationTokenRepository(prisma_client).table.find_many( + where={ + "user_id": {"in": list(user_ids_to_delete)}, + "team_id": data.team_id, + } ) if keys_to_delete: @@ -2723,7 +2832,7 @@ async def team_member_delete( litellm_changed_by=None, ) - await prisma_client.db.litellm_verificationtoken.delete_many( + await VerificationTokenRepository(prisma_client).table.delete_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -2733,6 +2842,52 @@ async def team_member_delete( return existing_team_row +_MEMBER_BUDGET_PATCH_FIELDS = { + "max_budget_in_team": "max_budget", + "tpm_limit": "tpm_limit", + "rpm_limit": "rpm_limit", + "budget_duration": "budget_duration", + "allowed_models": "allowed_models", +} + + +def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> Dict[str, Any]: + """Map the budget fields the request actually set (merge-patch: a sent + value updates, an explicit null clears, an absent field is left untouched) + to their budget-table columns.""" + provided = data.model_dump(exclude_unset=True) + return { + column: provided[request_field] + for request_field, column in _MEMBER_BUDGET_PATCH_FIELDS.items() + if request_field in provided + } + + +def _validate_budget_duration(budget_duration: Optional[str]) -> None: + """Reject budget durations that can't be parsed, are non-positive, or + overflow date math, so a bad value can't be persisted and later crash the + budget reset job.""" + if budget_duration is None: + return + + from litellm.litellm_core_utils.duration_parser import duration_in_seconds + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + try: + if duration_in_seconds(budget_duration) <= 0: + raise ValueError("budget_duration must be positive") + get_budget_reset_time(budget_duration=budget_duration) + except (ValueError, OverflowError): + raise HTTPException( + status_code=400, + detail={ + "error": "Invalid budget_duration '{}'. Use a format like '1h', '24h', '7d', or '30d'.".format( + budget_duration + ) + }, + ) + + @router.post( "/team/member_update", tags=["team management"], @@ -2770,7 +2925,9 @@ async def team_member_update( detail={"error": "Either user_id or user_email needs to be passed in"}, ) - _existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + _validate_budget_duration(data.budget_duration) + + _existing_team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) @@ -2843,17 +3000,15 @@ async def team_member_update( team_default_budget_id = raw_default_budget_id ### upsert new budget + budget_patch = _build_member_budget_patch(data) async with prisma_client.db.tx() as tx: await _upsert_budget_and_membership( tx=tx, team_id=data.team_id, user_id=received_user_id, - max_budget=data.max_budget_in_team, existing_budget_id=identified_budget_id, user_api_key_dict=user_api_key_dict, - tpm_limit=data.tpm_limit, - rpm_limit=data.rpm_limit, - allowed_models=data.allowed_models, + budget_patch=budget_patch, team_default_budget_id=team_default_budget_id, ) @@ -2875,7 +3030,7 @@ async def team_member_update( team_table.members_with_roles = team_members _db_team_members: List[dict] = [m.model_dump() for m in team_members] - await prisma_client.db.litellm_teamtable.update( + await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore ) @@ -2887,6 +3042,7 @@ async def team_member_update( max_budget_in_team=data.max_budget_in_team, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, + budget_duration=data.budget_duration, allowed_models=data.allowed_models, ) @@ -3005,7 +3161,7 @@ async def bulk_team_member_add( }, ) # get all users from the database - all_users_in_db = await prisma_client.db.litellm_usertable.find_many( + all_users_in_db = await UserRepository(prisma_client).table.find_many( order={"created_at": "desc"} ) data.members = [ @@ -3106,14 +3262,14 @@ async def delete_team( }' ``` """ + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + ) from litellm.proxy.proxy_server import ( create_audit_log_for_update, litellm_proxy_admin_name, prisma_client, ) - from litellm.proxy.management_helpers.audit_logs import ( - get_audit_log_changed_by, - ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -3125,11 +3281,9 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[BaseModel] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) - ) + team_row_base: Optional[BaseModel] = await TeamRepository( + prisma_client + ).table.find_unique(where={"team_id": team_id}) if team_row_base is None: raise Exception except Exception: @@ -3196,11 +3350,9 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": {"in": data.team_ids}} - ) - ) + keys_to_delete: List[LiteLLM_VerificationToken] = await VerificationTokenRepository( + prisma_client + ).table.find_many(where={"team_id": {"in": data.team_ids}}) if keys_to_delete: await _persist_deleted_verification_tokens( @@ -3212,6 +3364,20 @@ async def delete_team( await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key") + ## DELETE ASSOCIATED BYOK MODELS + # Runs before the team rows are deleted so a mid-flight failure never leaves + # the team gone with its models orphaned. + from litellm.proxy.management_endpoints.model_management_endpoints import ( + delete_team_models, + ) + from litellm.proxy.proxy_server import llm_router + + await delete_team_models( + team_ids=data.team_ids, + prisma_client=prisma_client, + llm_router=llm_router, + ) + # ## DELETE TEAM MEMBERSHIPS for team_row in team_rows: ### get all team members @@ -3291,7 +3457,7 @@ async def _save_deleted_team_records( """Save deleted team records to the database.""" if not records: return - await prisma_client.db.litellm_deletedteamtable.create_many(data=records) + await DeletedTeamRepository(prisma_client).table.create_many(data=records) async def _persist_deleted_team_records( @@ -3373,7 +3539,7 @@ async def _add_team_member_budget_table( team_info_response_object: TeamInfoResponseObjectTeamTable, ) -> TeamInfoResponseObjectTeamTable: try: - team_budget = await prisma_client.db.litellm_budgettable.find_unique( + team_budget = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": team_member_budget_id} ) team_info_response_object.team_member_budget_table = team_budget @@ -3442,11 +3608,11 @@ async def team_info( ) try: - team_info: Optional[BaseModel] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - include={"object_permission": True}, - ) + team_info: Optional[BaseModel] = await TeamRepository( + prisma_client + ).table.find_unique( + where={"team_id": team_id}, + include={"object_permission": True}, ) if team_info is None: raise Exception @@ -3702,7 +3868,7 @@ async def block_team( if prisma_client is None: raise Exception("No DB Connected.") - existing_team = await prisma_client.db.litellm_teamtable.find_unique( + existing_team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) if existing_team is None: @@ -3717,7 +3883,7 @@ async def block_team( user_api_key_dict=user_api_key_dict, ) - record = await prisma_client.db.litellm_teamtable.update( + record = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"blocked": True} # type: ignore ) @@ -3754,7 +3920,7 @@ async def unblock_team( if prisma_client is None: raise Exception("No DB Connected.") - existing_team = await prisma_client.db.litellm_teamtable.find_unique( + existing_team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) if existing_team is None: @@ -3769,7 +3935,7 @@ async def unblock_team( user_api_key_dict=user_api_key_dict, ) - record = await prisma_client.db.litellm_teamtable.update( + record = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"blocked": False} # type: ignore ) @@ -3802,7 +3968,7 @@ async def list_available_teams( return [] # filter out teams that the user is already a member of - user_info = await prisma_client.db.litellm_usertable.find_unique( + user_info = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) if user_info is None: @@ -3816,7 +3982,7 @@ async def list_available_teams( team for team in available_teams if team not in user_info_correct_type.teams ] - available_teams_db = await prisma_client.db.litellm_teamtable.find_many( + available_teams_db = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": available_teams}} ) @@ -3962,7 +4128,7 @@ async def _batch_resolve_access_group_resources( return {} unique_ids = list(set(all_access_group_ids)) - rows = await _prisma_client.db.litellm_accessgrouptable.find_many( + rows = await AccessGroupRepository(_prisma_client).table.find_many( where={"access_group_id": {"in": unique_ids}}, ) @@ -4027,7 +4193,7 @@ async def _get_keys_count_by_team( if not page_team_ids: return {} - grouped = await prisma_client.db.litellm_verificationtoken.group_by( + grouped = await VerificationTokenRepository(prisma_client).table.group_by( by=["team_id"], where={"team_id": {"in": page_team_ids}}, count={"team_id": True}, @@ -4241,25 +4407,25 @@ async def list_team_v2( # Get teams with pagination if use_deleted_table: - teams = await prisma_client.db.litellm_deletedteamtable.find_many( + teams = await DeletedTeamRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, order=order_by if order_by else {"created_at": "desc"}, # Default sort ) # Get total count for pagination - total_count = await prisma_client.db.litellm_deletedteamtable.count( + total_count = await DeletedTeamRepository(prisma_client).table.count( where=where_conditions ) else: - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, order=order_by if order_by else {"created_at": "desc"}, # Default sort ) # Get total count for pagination - total_count = await prisma_client.db.litellm_teamtable.count( + total_count = await TeamRepository(prisma_client).table.count( where=where_conditions ) @@ -4365,7 +4531,7 @@ async def _authorize_and_filter_teams( if allowed_org_ids is not None: # Org admin: query DB for teams in their orgs - org_teams = await prisma_client.db.litellm_teamtable.find_many( + org_teams = await TeamRepository(prisma_client).table.find_many( where={"organization_id": {"in": allowed_org_ids}}, include={"litellm_model_table": True}, ) @@ -4380,7 +4546,7 @@ async def _authorize_and_filter_teams( ] elif user_id: # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) - response = await prisma_client.db.litellm_teamtable.find_many( + response = await TeamRepository(prisma_client).table.find_many( include={"litellm_model_table": True} ) return [ @@ -4392,7 +4558,7 @@ async def _authorize_and_filter_teams( else: # Proxy admin: all teams return list( - await prisma_client.db.litellm_teamtable.find_many( + await TeamRepository(prisma_client).table.find_many( include={"litellm_model_table": True} ) ) @@ -4453,7 +4619,7 @@ async def list_team( _team_memberships.append(tm) # add all keys that belong to the team - keys = await prisma_client.db.litellm_verificationtoken.find_many( + keys = await VerificationTokenRepository(prisma_client).table.find_many( where={"team_id": team.team_id} ) @@ -4509,10 +4675,10 @@ async def get_paginated_teams( # Calculate skip for pagination skip = (page - 1) * page_size # Get total count - total_count = await prisma_client.db.litellm_teamtable.count() + total_count = await TeamRepository(prisma_client).table.count() # Get paginated teams - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( skip=skip, take=page_size, order={"team_alias": "asc"} # Sort by team_alias ) return teams, total_count @@ -4585,7 +4751,7 @@ async def ui_view_teams( } # Query users with pagination and filters - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, @@ -4657,7 +4823,7 @@ async def team_model_add( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Get existing team - team_row = await prisma_client.db.litellm_teamtable.find_unique( + team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) @@ -4690,7 +4856,7 @@ async def team_model_add( # null them out — see object_permission_utils.validate_key_search_tools_against_team # and the MCP/agent authz paths, which treat a missing object_permission # as "no team-level restriction". - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"models": updated_models}, include={"object_permission": True}, # type: ignore @@ -4744,7 +4910,7 @@ async def team_model_delete( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Get existing team - team_row = await prisma_client.db.litellm_teamtable.find_unique( + team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) @@ -4778,7 +4944,7 @@ async def team_model_delete( updated_models = [m for m in current_models if m not in data.models] # Update team. See team_model_add for the rationale on `include`. - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"models": updated_models}, include={"object_permission": True}, # type: ignore @@ -4925,7 +5091,7 @@ async def update_team_member_permissions( }, ) # Update the team member permissions - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"team_member_permissions": data.team_member_permissions}, ) @@ -5029,7 +5195,7 @@ async def _append_permissions_to_specific_teams( prisma_client, team_ids: List[str], permissions_to_add: set ) -> int: """Fetch specific teams by ID and append permissions.""" - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": team_ids}}, ) @@ -5061,7 +5227,7 @@ async def _append_permissions_to_all_teams( find_args["cursor"] = {"team_id": cursor} find_args["skip"] = 1 - teams = await prisma_client.db.litellm_teamtable.find_many(**find_args) + teams = await TeamRepository(prisma_client).table.find_many(**find_args) if not teams: break @@ -5167,7 +5333,7 @@ async def get_team_daily_activity( where_condition = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases = await prisma_client.db.litellm_teamtable.find_many( + team_aliases = await TeamRepository(prisma_client).table.find_many( where=where_condition ) team_alias_metadata = { @@ -5204,9 +5370,9 @@ async def get_team_daily_activity( # If user does not have full team view, filter by their API keys if not has_full_team_view: # Get all API keys for this user - user_keys = await prisma_client.db.litellm_verificationtoken.find_many( - where={"user_id": user_api_key_dict.user_id} - ) + user_keys = await VerificationTokenRepository( + prisma_client + ).table.find_many(where={"user_id": user_api_key_dict.user_id}) user_api_keys = [key.token for key in user_keys if key.token] # If user has no API keys, return empty result if not user_api_keys: diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 19ca2c9f6be..a9b57db8a6f 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -21,6 +21,15 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.table_repositories import ( + SpendLogsRepository, + SpendLogToolIndexRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.tool_management import ( LiteLLM_ToolTableRow, ToolDetailResponse, @@ -256,8 +265,10 @@ async def get_tool_usage_logs( if end_time_filter is not None: where["start_time"]["lte"] = end_time_filter - total = await prisma_client.db.litellm_spendlogtoolindex.count(where=where) - index_rows = await prisma_client.db.litellm_spendlogtoolindex.find_many( + total = await SpendLogToolIndexRepository(prisma_client).table.count( + where=where + ) + index_rows = await SpendLogToolIndexRepository(prisma_client).table.find_many( where=where, order={"start_time": "desc"}, skip=(page - 1) * page_size, @@ -269,7 +280,7 @@ async def get_tool_usage_logs( logs=[], total=total, page=page, page_size=page_size ) - spend_logs = await prisma_client.db.litellm_spendlogs.find_many( + spend_logs = await SpendLogsRepository(prisma_client).table.find_many( where={"request_id": {"in": request_ids}} ) log_by_id = {s.request_id: s for s in spend_logs} @@ -348,7 +359,7 @@ async def _resolve_key_hash_to_object_permission_id( hashed = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash) if not hashed: return None - row = await prisma_client.db.litellm_verificationtoken.find_unique( + row = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed} ) if row is None: @@ -357,18 +368,18 @@ async def _resolve_key_hash_to_object_permission_id( if op_id: return op_id new_id = str(uuid.uuid4()) - await prisma_client.db.litellm_objectpermissiontable.create( + await ObjectPermissionRepository(prisma_client).table.create( data={"object_permission_id": new_id, "blocked_tools": []} ) - updated_count = await prisma_client.db.litellm_verificationtoken.update_many( + updated_count = await VerificationTokenRepository(prisma_client).table.update_many( where={"token": hashed, "object_permission_id": None}, data={"object_permission_id": new_id}, ) if updated_count == 0: - await prisma_client.db.litellm_objectpermissiontable.delete( + await ObjectPermissionRepository(prisma_client).table.delete( where={"object_permission_id": new_id} ) - row = await prisma_client.db.litellm_verificationtoken.find_unique( + row = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed} ) return getattr(row, "object_permission_id", None) if row else None @@ -383,7 +394,7 @@ async def _resolve_team_id_to_object_permission_id( if not team_id or not team_id.strip(): return None team_id_clean = team_id.strip() - row = await prisma_client.db.litellm_teamtable.find_unique( + row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id_clean}, select={"object_permission_id": True}, ) @@ -393,18 +404,18 @@ async def _resolve_team_id_to_object_permission_id( if op_id: return op_id new_id = str(uuid.uuid4()) - await prisma_client.db.litellm_objectpermissiontable.create( + await ObjectPermissionRepository(prisma_client).table.create( data={"object_permission_id": new_id, "blocked_tools": []} ) - updated_count = await prisma_client.db.litellm_teamtable.update_many( + updated_count = await TeamRepository(prisma_client).table.update_many( where={"team_id": team_id_clean, "object_permission_id": None}, data={"object_permission_id": new_id}, ) if updated_count == 0: - await prisma_client.db.litellm_objectpermissiontable.delete( + await ObjectPermissionRepository(prisma_client).table.delete( where={"object_permission_id": new_id} ) - row = await prisma_client.db.litellm_teamtable.find_unique( + row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id_clean}, select={"object_permission_id": True}, ) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d6082899c02..4812bed2f21 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -15,8 +15,8 @@ import os import re import secrets -from html import escape from copy import deepcopy +from html import escape from typing import ( TYPE_CHECKING, Any, @@ -39,9 +39,9 @@ from fastapi.responses import RedirectResponse import litellm -from litellm.caching.dual_cache import DualCache from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.caching.dual_cache import DualCache from litellm.constants import ( CLI_SSO_CLAIM_MAP, CLI_SSO_CLAIM_MAX_SCALAR_LENGTH, @@ -77,7 +77,6 @@ UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.auth.auth_utils import ( _get_request_ip_address, _has_user_setup_sso, @@ -92,6 +91,7 @@ jwt_display_template, ) from litellm.proxy.common_utils.html_forms.ui_login import html_form +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.sso_helper_utils import ( @@ -110,6 +110,9 @@ get_custom_url, get_server_root_path, ) +from litellm.repositories.table_repositories import SSOConfigRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository from litellm.secret_managers.main import get_secret_bool, str_to_bool from litellm.types.proxy.management_endpoints.ui_sso import * # noqa: F403, F401 from litellm.types.proxy.management_endpoints.ui_sso import ( @@ -438,7 +441,7 @@ async def _persist_cli_sso_user_metadata( return try: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) existing_metadata: Dict[str, Any] = {} @@ -451,7 +454,7 @@ async def _persist_cli_sso_user_metadata( existing_metadata=existing_metadata, attribution_metadata=attribution_metadata, ) - await prisma_client.db.litellm_usertable.update_many( + await UserRepository(prisma_client).table.update_many( where={"user_id": user_id}, data={"metadata": merged_metadata}, ) @@ -859,7 +862,7 @@ async def google_login( if premium_user is not True: # Check if under 'free SSO user' limit if prisma_client is not None: - total_users = await prisma_client.db.litellm_usertable.count() + total_users = await UserRepository(prisma_client).table.count() if total_users and total_users > 5: raise ProxyException( message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this", @@ -1150,7 +1153,7 @@ async def _setup_team_mappings() -> Optional["TeamMappings"]: "Prisma client is None, connect a database to your proxy" ) - sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique( where={"id": "sso_config"} ) @@ -1188,7 +1191,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: "Prisma client is None, connect a database to your proxy" ) - sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique( where={"id": "sso_config"} ) @@ -1217,7 +1220,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: generic_role_mappings_group_claim = os.getenv( "GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", None ) - generic_role_mappoings_default_role = os.getenv( + generic_role_mappings_default_role = os.getenv( "GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", None ) if generic_role_mappings is not None: @@ -1236,7 +1239,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: role_mappings_data = { "provider": "generic", "group_claim": generic_role_mappings_group_claim, - "default_role": generic_role_mappoings_default_role, + "default_role": generic_role_mappings_default_role, "roles": generic_user_role_mappings_data, } @@ -1755,7 +1758,7 @@ async def _sync_user_role_from_jwt_role_map( # Update existing DB record if role differs if user_info is not None and user_info.user_role != mapped_role.value: - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": user_info.user_id}, data={"user_role": mapped_role.value}, ) @@ -1819,7 +1822,7 @@ async def check_and_update_if_proxy_admin_id( return user_role if prisma_client: - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": user_id}, data={"user_role": LitellmUserRoles.PROXY_ADMIN.value}, ) @@ -1976,7 +1979,7 @@ async def _fetch_cli_sso_team_details( team_details: List[Dict[str, Any]] = [] try: if teams: - prisma_teams = await prisma_client.db.litellm_teamtable.find_many( + prisma_teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": teams}} ) for team_row in prisma_teams: @@ -2884,7 +2887,7 @@ async def upsert_sso_user( user_id=user_id, ) - await prisma_client.db.litellm_usertable.update_many( + await UserRepository(prisma_client).table.update_many( where={"user_id": user_id}, data=update_data ) else: @@ -2986,7 +2989,7 @@ async def create_litellm_team_from_sso_group( code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) try: - team_obj = await prisma_client.db.litellm_teamtable.find_first( + team_obj = await TeamRepository(prisma_client).table.find_first( where={"team_id": litellm_team_id} ) verbose_proxy_logger.debug(f"Team object: {team_obj}") diff --git a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py index ebd276fbee5..661487577c3 100644 --- a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py +++ b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py @@ -19,6 +19,11 @@ from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import DailyTagSpendRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) # Constants for analytics periods MAX_DAYS = 7 # Number of days to show in DAU analytics @@ -676,7 +681,7 @@ async def get_per_user_analytics( where_clause["tag"] = {"contains": tag_filter} # Get all tag records in the date range with optional tag filtering - tag_records = await prisma_client.db.litellm_dailytagspend.find_many( + tag_records = await DailyTagSpendRepository(prisma_client).table.find_many( where=where_clause ) @@ -693,9 +698,9 @@ async def get_per_user_analytics( ) # Lookup user_id for each api_key - api_key_records = await prisma_client.db.litellm_verificationtoken.find_many( - where={"token": {"in": list(api_keys)}} - ) + api_key_records = await VerificationTokenRepository( + prisma_client + ).table.find_many(where={"token": {"in": list(api_keys)}}) # Create mapping from api_key to user_id api_key_to_user_id = { @@ -704,7 +709,7 @@ async def get_per_user_analytics( # Get user emails for the user_ids user_ids = list(set(api_key_to_user_id.values())) - user_records = await prisma_client.db.litellm_usertable.find_many( + user_records = await UserRepository(prisma_client).table.find_many( where={"user_id": {"in": user_ids}} ) diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py index a19af4dd484..57cc0dc6745 100644 --- a/litellm/proxy/management_endpoints/workflow_management_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -27,6 +27,11 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import ( + WorkflowEventRepository, + WorkflowMessageRepository, + WorkflowRunRepository, +) router = APIRouter() @@ -96,13 +101,13 @@ class WorkflowMessageCreateRequest(BaseModel): async def _get_next_sequence_number(prisma_client: Any, run_id: str, table: str) -> int: """Return MAX(sequence_number) + 1 for the given run, for either events or messages.""" if table == "events": - rows = await prisma_client.db.litellm_workflowevent.find_many( + rows = await WorkflowEventRepository(prisma_client).table.find_many( where={"run_id": run_id}, order={"sequence_number": "desc"}, take=1, ) else: - rows = await prisma_client.db.litellm_workflowmessage.find_many( + rows = await WorkflowMessageRepository(prisma_client).table.find_many( where={"run_id": run_id}, order={"sequence_number": "desc"}, take=1, @@ -116,7 +121,7 @@ async def _require_run( user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> Any: """Return the run or raise 404. For non-admin callers, also enforce key ownership.""" - run = await prisma_client.db.litellm_workflowrun.find_unique( + run = await WorkflowRunRepository(prisma_client).table.find_unique( where={"run_id": run_id} ) if run is None: @@ -163,7 +168,7 @@ async def create_workflow_run( create_data["input"] = _json(data.input) if data.metadata is not None: create_data["metadata"] = _json(data.metadata) - run = await prisma_client.db.litellm_workflowrun.create(data=create_data) + run = await WorkflowRunRepository(prisma_client).table.create(data=create_data) return run except Exception as e: verbose_proxy_logger.exception("Error creating workflow run: %s", e) @@ -206,7 +211,7 @@ async def list_workflow_runs( where["created_by"] = caller try: - runs = await prisma_client.db.litellm_workflowrun.find_many( + runs = await WorkflowRunRepository(prisma_client).table.find_many( where=where, order={"created_at": "desc"}, take=limit, @@ -235,7 +240,7 @@ async def get_workflow_run( ) try: - run = await prisma_client.db.litellm_workflowrun.find_unique( + run = await WorkflowRunRepository(prisma_client).table.find_unique( where={"run_id": run_id}, include={"events": {"order_by": {"sequence_number": "desc"}, "take": 1}}, ) @@ -286,7 +291,7 @@ async def update_workflow_run( await _require_run(prisma_client, run_id, user_api_key_dict) try: - run = await prisma_client.db.litellm_workflowrun.update( + run = await WorkflowRunRepository(prisma_client).table.update( where={"run_id": run_id}, data=update, ) @@ -391,7 +396,7 @@ async def list_workflow_events( await _require_run(prisma_client, run_id, user_api_key_dict) try: - events = await prisma_client.db.litellm_workflowevent.find_many( + events = await WorkflowEventRepository(prisma_client).table.find_many( where={"run_id": run_id}, order={"sequence_number": "asc"}, take=limit, @@ -436,7 +441,9 @@ async def append_workflow_message( } if data.session_id is not None: msg_data["session_id"] = data.session_id - msg = await prisma_client.db.litellm_workflowmessage.create(data=msg_data) + msg = await WorkflowMessageRepository(prisma_client).table.create( + data=msg_data + ) return msg except Exception as e: @@ -481,7 +488,7 @@ async def list_workflow_messages( await _require_run(prisma_client, run_id, user_api_key_dict) try: - messages = await prisma_client.db.litellm_workflowmessage.find_many( + messages = await WorkflowMessageRepository(prisma_client).table.find_many( where={"run_id": run_id}, order={"sequence_number": "asc"}, take=limit, diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index 439c3b2118d..33599c3c622 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -18,6 +18,7 @@ Optional, UserAPIKeyAuth, ) +from litellm.repositories.table_repositories import AuditLogRepository from litellm.types.utils import StandardAuditLogPayload _audit_log_callback_cache: Dict[str, CustomLogger] = {} @@ -244,7 +245,7 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): _request_data = request_data.model_dump(exclude_none=True) try: - await prisma_client.db.litellm_auditlog.create( + await AuditLogRepository(prisma_client).table.create( data={ **_request_data, # type: ignore } diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 4c966b25413..f2ddae40d8c 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -12,6 +12,8 @@ from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.utils import PrismaClient +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.table_repositories import MCPServerRepository if TYPE_CHECKING: from litellm.proxy._types import ( @@ -48,10 +50,10 @@ async def attach_object_permission_to_dict( object_permission_id = data_dict.get("object_permission_id") if object_permission_id: - object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": object_permission_id}, - ) + object_permission = await ObjectPermissionRepository( + prisma_client + ).table.find_unique( + where={"object_permission_id": object_permission_id}, ) if object_permission: # Convert to dict if needed @@ -106,10 +108,10 @@ async def handle_update_object_permission_common( ) existing_object_permissions_dict: Dict = {} - existing_object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": object_permission_id_to_use}, - ) + existing_object_permission = await ObjectPermissionRepository( + prisma_client + ).table.find_unique( + where={"object_permission_id": object_permission_id_to_use}, ) # Update the object permission @@ -137,14 +139,14 @@ async def handle_update_object_permission_common( ######################################################### # Commit the update to the LiteLLM_ObjectPermissionTable ######################################################### - created_object_permission_row = ( - await prisma_client.db.litellm_objectpermissiontable.upsert( - where={"object_permission_id": object_permission_id_to_use}, - data={ - "create": existing_object_permissions_dict, - "update": existing_object_permissions_dict, - }, - ) + created_object_permission_row = await ObjectPermissionRepository( + prisma_client + ).table.upsert( + where={"object_permission_id": object_permission_id_to_use}, + data={ + "create": existing_object_permissions_dict, + "update": existing_object_permissions_dict, + }, ) verbose_proxy_logger.debug( @@ -183,7 +185,7 @@ async def _set_object_permission( clean_data["mcp_tool_permissions"] ) - created_permission = await prisma_client.db.litellm_objectpermissiontable.create( + created_permission = await ObjectPermissionRepository(prisma_client).table.create( data=clean_data ) @@ -220,7 +222,7 @@ async def _get_db_mcp_servers_by_identifiers( return [] identifier_list = list(identifiers) - return await prisma_client.db.litellm_mcpservertable.find_many( + return await MCPServerRepository(prisma_client).table.find_many( where={ "OR": [ {"server_id": {"in": identifier_list}}, diff --git a/litellm/proxy/management_helpers/user_invitation.py b/litellm/proxy/management_helpers/user_invitation.py index d2d800aa77f..babc920189a 100644 --- a/litellm/proxy/management_helpers/user_invitation.py +++ b/litellm/proxy/management_helpers/user_invitation.py @@ -4,6 +4,7 @@ import litellm from litellm.proxy._types import CommonProxyErrors, InvitationNew, UserAPIKeyAuth +from litellm.repositories.table_repositories import InvitationLinkRepository async def create_invitation_for_user( @@ -25,7 +26,7 @@ async def create_invitation_for_user( expires_at = current_time + timedelta(days=7) try: - response = await prisma_client.db.litellm_invitationlink.create( + response = await InvitationLinkRepository(prisma_client).table.create( data={ "user_id": data.user_id, "created_at": current_time, diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 0b175db3c87..830d6f84b85 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -9,9 +9,8 @@ import litellm from litellm._logging import verbose_logger -from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm._uuid import uuid -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.proxy._types import ( # key request types; user request types; team request types; customer request types BudgetNewRequest, DeleteCustomerRequest, @@ -32,7 +31,11 @@ VirtualKeyEvent, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.utils import PrismaClient +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.table_repositories import TeamMembershipRepository +from litellm.repositories.user_repository import UserRepository def get_new_internal_user_defaults( @@ -111,7 +114,7 @@ async def handle_budget_for_entity( budget_row.model_dump(exclude_none=True) ) - _budget = await prisma_client.db.litellm_budgettable.create( + _budget = await BudgetRepository(prisma_client).table.create( data={ **new_budget_data, # type: ignore "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -174,7 +177,7 @@ async def _clone_team_default_budget_for_member( so the member starts with the team default's values but gets their own private budget row (which can be edited independently). """ - default_budget = await prisma_client.db.litellm_budgettable.find_unique( + default_budget = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": default_team_budget_id} ) if default_budget is None: @@ -202,7 +205,7 @@ async def _clone_team_default_budget_for_member( cloned_data["budget_duration"] ) - new_budget = await prisma_client.db.litellm_budgettable.create(data=cloned_data) + new_budget = await BudgetRepository(prisma_client).table.create(data=cloned_data) return new_budget.budget_id @@ -229,7 +232,7 @@ async def add_new_member( ## ADD TEAM ID, to USER TABLE IF NEW ## if new_member.user_id is not None: new_user_defaults = get_new_internal_user_defaults(user_id=new_member.user_id) - _returned_user = await prisma_client.db.litellm_usertable.upsert( + _returned_user = await UserRepository(prisma_client).table.upsert( where={"user_id": new_member.user_id}, data={ "update": {"teams": {"push": [team_id]}}, @@ -259,7 +262,7 @@ async def add_new_member( returned_user = LiteLLM_UserTable(**_returned_user.model_dump()) elif len(existing_user_row) == 1: user_info = existing_user_row[0] - _returned_user = await prisma_client.db.litellm_usertable.update( + _returned_user = await UserRepository(prisma_client).table.update( where={"user_id": user_info.user_id}, # type: ignore data={"teams": {"push": [team_id]}}, ) @@ -284,7 +287,7 @@ async def add_new_member( budget_data["max_budget"] = max_budget_in_team if allowed_models is not None: budget_data["allowed_models"] = allowed_models - response = await prisma_client.db.litellm_budgettable.create(data=budget_data) + response = await BudgetRepository(prisma_client).table.create(data=budget_data) _budget_id = response.budget_id elif default_team_budget_id is not None: @@ -303,15 +306,15 @@ async def add_new_member( _budget_id = None if _budget_id and returned_user is not None and returned_user.user_id is not None: - _returned_team_membership = ( - await prisma_client.db.litellm_teammembership.create( - data={ - "team_id": team_id, - "user_id": returned_user.user_id, - "budget_id": _budget_id, - }, - include={"litellm_budget_table": True}, - ) + _returned_team_membership = await TeamMembershipRepository( + prisma_client + ).table.create( + data={ + "team_id": team_id, + "user_id": returned_user.user_id, + "budget_id": _budget_id, + }, + include={"litellm_budget_table": True}, ) returned_team_membership = LiteLLM_TeamMembership( diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 4d161be4263..6f1ca3196fe 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -29,6 +29,8 @@ UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import MemoryRepository +from litellm.repositories.team_repository import TeamRepository from litellm.types.memory_management import ( LiteLLM_MemoryRow, MemoryCreateRequest, @@ -173,7 +175,7 @@ async def _is_team_admin_for( ) try: - team_obj = await prisma_client.db.litellm_teamtable.find_unique( + team_obj = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) except Exception as e: @@ -304,7 +306,7 @@ async def create_memory( create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata) try: - row = await prisma_client.db.litellm_memorytable.create(data=create_data) + row = await MemoryRepository(prisma_client).table.create(data=create_data) except Exception as e: # Key is globally unique. Any duplicate → 409. if _is_unique_violation(e): @@ -364,8 +366,8 @@ async def list_memory( where = {"AND": [key_filter, vis]} try: - total = await prisma_client.db.litellm_memorytable.count(where=where) - rows = await prisma_client.db.litellm_memorytable.find_many( + total = await MemoryRepository(prisma_client).table.count(where=where) + rows = await MemoryRepository(prisma_client).table.find_many( where=where, order={"updated_at": "desc"}, skip=(page - 1) * page_size, @@ -386,7 +388,7 @@ async def _find_memory_for_caller( key_filter: dict = {"key": key} vis = _visibility_filter(user_api_key_dict) where: dict = key_filter if vis is None else {"AND": [key_filter, vis]} - rows = await prisma_client.db.litellm_memorytable.find_many( + rows = await MemoryRepository(prisma_client).table.find_many( where=where, take=1, order={"updated_at": "desc"} ) if not rows: @@ -475,7 +477,7 @@ async def _find_existing() -> Any: # their team) — otherwise a teammate could overwrite a personal # entry through the OR-based visibility filter. await _assert_write_access(prisma_client, existing, user_api_key_dict) - row = await prisma_client.db.litellm_memorytable.update( + row = await MemoryRepository(prisma_client).table.update( where={"memory_id": existing.memory_id}, data=data, ) @@ -503,7 +505,7 @@ async def _find_existing() -> Any: if body.metadata is not None: create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata) try: - row = await prisma_client.db.litellm_memorytable.create( + row = await MemoryRepository(prisma_client).table.create( data=create_data ) except Exception as e: @@ -524,7 +526,7 @@ async def _find_existing() -> Any: await _assert_write_access( prisma_client, existing_after_race, user_api_key_dict ) - row = await prisma_client.db.litellm_memorytable.update( + row = await MemoryRepository(prisma_client).table.update( where={"memory_id": existing_after_race.memory_id}, data=data, ) @@ -554,7 +556,7 @@ async def delete_memory( # Visibility != write authority — see the upsert handler for the rationale. await _assert_write_access(prisma_client, row, user_api_key_dict) try: - await prisma_client.db.litellm_memorytable.delete( + await MemoryRepository(prisma_client).table.delete( where={"memory_id": row.memory_id} ) except Exception as e: diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index cc0d06e4f40..b2834e52306 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -5,6 +5,10 @@ from types import MappingProxyType from typing import TYPE_CHECKING, List, Literal, Optional, Union +from litellm.repositories.table_repositories import ( + ManagedFileRepository, + ManagedObjectRepository, +) from litellm.types.utils import SpecialEnums if TYPE_CHECKING: @@ -697,7 +701,7 @@ async def resolve_input_file_id_to_unified(response, prisma_client) -> None: and prisma_client ): try: - managed_file = await prisma_client.db.litellm_managedfiletable.find_first( + managed_file = await ManagedFileRepository(prisma_client).table.find_first( where={"flat_model_file_ids": {"has": response.input_file_id}} ) if managed_file: @@ -719,7 +723,7 @@ async def resolve_output_file_ids_to_unified(response, prisma_client) -> None: if not raw_id or _is_base64_encoded_unified_file_id(raw_id): continue try: - managed_file = await prisma_client.db.litellm_managedfiletable.find_first( + managed_file = await ManagedFileRepository(prisma_client).table.find_first( where={"flat_model_file_ids": {"has": raw_id}} ) if managed_file: @@ -821,6 +825,7 @@ async def get_batch_from_database( - response_batch: Parsed LiteLLMBatch object (or None) """ import json + from litellm.types.utils import LiteLLMBatch if managed_files_obj is None or not unified_batch_id: @@ -830,7 +835,7 @@ async def get_batch_from_database( if not prisma_client: return None, None - db_batch_object = await prisma_client.db.litellm_managedobjecttable.find_first( + db_batch_object = await ManagedObjectRepository(prisma_client).table.find_first( where={"unified_object_id": batch_id} ) @@ -942,7 +947,7 @@ async def update_batch_in_database( update_data["batch_processed"] = True try: - await prisma_client.db.litellm_managedobjecttable.update( + await ManagedObjectRepository(prisma_client).table.update( where={"unified_object_id": batch_id}, data=update_data, ) @@ -958,7 +963,7 @@ async def update_batch_in_database( f"batch_processed column not found, retrying update without it: {col_err}" ) update_data.pop("batch_processed", None) - await prisma_client.db.litellm_managedobjecttable.update( + await ManagedObjectRepository(prisma_client).table.update( where={"unified_object_id": batch_id}, data=update_data, ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 378cbbda89c..9eef7cd7e8b 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -21,6 +21,7 @@ UploadFile, status, ) + import litellm from litellm import CreateFileRequest, get_secret_str from litellm._logging import verbose_proxy_logger @@ -37,15 +38,6 @@ get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) -from litellm.proxy.utils import ProxyLogging, is_known_model -from litellm.router import Router -from litellm.types.llms.openai import ( - CREATE_FILE_REQUESTS_PURPOSE, - FileExpiresAfter, - OpenAIFileObject, - OpenAIFilesPurpose, -) - from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, encode_file_id_with_model, @@ -54,6 +46,15 @@ handle_model_based_routing, prepare_data_with_credentials, ) +from litellm.proxy.utils import ProxyLogging, is_known_model +from litellm.repositories.table_repositories import ManagedFileRepository +from litellm.router import Router +from litellm.types.llms.openai import ( + CREATE_FILE_REQUESTS_PURPOSE, + FileExpiresAfter, + OpenAIFileObject, + OpenAIFilesPurpose, +) router = APIRouter() @@ -666,7 +667,7 @@ async def get_file_content( # noqa: PLR0915 managed_files_obj, "prisma_client", None ): prisma_client = getattr(managed_files_obj, "prisma_client") - db_file = await prisma_client.db.litellm_managedfiletable.find_first( + db_file = await ManagedFileRepository(prisma_client).table.find_first( where={"unified_file_id": file_id} ) if db_file and db_file.storage_backend and db_file.storage_url: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 7c3a6f19013..c7db818a07e 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -426,12 +426,7 @@ async def mistral_proxy_route( ) ## check for streaming - is_streaming_request = False - # anthropic is streaming when 'stream' = True is in the body - if request.method == "POST": - _request_body = await request.json() - if _request_body.get("stream"): - is_streaming_request = True + is_streaming_request = await is_streaming_request_fn(request) ## CREATE PASS-THROUGH endpoint_func = create_pass_through_route( 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 29bbb37501f..9f353226dd0 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 @@ -5,7 +5,7 @@ """ from datetime import datetime -from typing import List, Optional, Union +from typing import List, Optional, Tuple, Union from urllib.parse import urlparse import httpx @@ -18,6 +18,7 @@ ) from litellm.llms.openai.openai import OpenAIConfig from litellm.llms.openai.openai import OpenAIConfig as OpenAIConfigType +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( BasePassthroughLoggingHandler, @@ -29,9 +30,75 @@ EndpointType, PassthroughStandardLoggingPayload, ) +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ImageResponse, LlmProviders, PassthroughCallTypes from litellm.utils import ModelResponse, TextCompletionResponse +# Hostnames that route to OpenAI-compatible APIs. +# +# `api.openai.com` is OpenAI proper. The two Azure domains below are *shared by +# every Azure Cognitive Service* (Speech, Vision, Language, ...), not just Azure +# OpenAI: `openai.azure.com` is the classic Azure OpenAI domain, while +# `cognitiveservices.azure.com` is used by newer "Azure AI Foundry" / +# Cognitive Services-hosted Azure OpenAI deployments. Because the hostname alone +# cannot tell Azure OpenAI apart from the other Cognitive Services on those +# domains, requests there must additionally carry an OpenAI-style path segment. +_OPENAI_HOSTNAMES = ("api.openai.com",) +_AZURE_OPENAI_HOSTNAMES = ("openai.azure.com", "cognitiveservices.azure.com") +# Path markers that identify an Azure request as Azure OpenAI rather than Speech +# / Vision / Language / ... `/openai/` is the native Azure OpenAI path prefix; +# `/v1/` is the OpenAI-v1 surface used by LiteLLM's pass-through routing. Other +# Cognitive Services use service-named prefixes and versions like `/v3.1/`, +# `/v1.0/`, so they do not collide with these markers. +_AZURE_OPENAI_PATH_MARKERS = ("/openai/", "/v1/") + + +def _hostname_matches(hostname: str, suffixes: tuple) -> bool: + """True if hostname equals one of `suffixes` or is a subdomain of it. + + Uses suffix matching (not a bare substring test) so look-alikes such as + `cognitiveservices.azure.com.attacker.example` are not accepted. + """ + return any( + hostname == suffix or hostname.endswith("." + suffix) for suffix in suffixes + ) + + +def _is_openai_compatible_host(hostname: Optional[str]) -> bool: + """True if the hostname is OpenAI proper or one of the Azure OpenAI domains. + + Hostname-only check, kept for the route-level helpers that additionally + require a specific OpenAI path (e.g. `/v1/chat/completions`). When only the + hostname would otherwise gate dispatch, use `_is_openai_compatible_url` so + non-OpenAI Azure Cognitive Services on the shared domains are excluded. + """ + if not hostname: + return False + return _hostname_matches(hostname, _OPENAI_HOSTNAMES) or _hostname_matches( + hostname, _AZURE_OPENAI_HOSTNAMES + ) + + +def _is_openai_compatible_url(url_route: Optional[str]) -> bool: + """True if the URL targets an OpenAI-compatible API surface. + + For the shared Azure Cognitive Services domains we additionally require an + OpenAI-style path segment (`/openai/` or `/v1/`) so non-OpenAI Azure services + (Speech, Vision, Language, ...) on the same domain are not misclassified as + OpenAI routes. + """ + if not url_route: + return False + parsed_url = urlparse(url_route) + hostname = parsed_url.hostname + if not hostname: + return False + if _hostname_matches(hostname, _OPENAI_HOSTNAMES): + return True + if _hostname_matches(hostname, _AZURE_OPENAI_HOSTNAMES): + return any(marker in parsed_url.path for marker in _AZURE_OPENAI_PATH_MARKERS) + return False + class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): """ @@ -52,12 +119,8 @@ def is_openai_chat_completions_route(url_route: str) -> bool: if not url_route: return False parsed_url = urlparse(url_route) - return bool( - parsed_url.hostname - and ( - "api.openai.com" in parsed_url.hostname - or "openai.azure.com" in parsed_url.hostname - ) + return ( + _is_openai_compatible_host(parsed_url.hostname) and "/v1/chat/completions" in parsed_url.path ) @@ -67,12 +130,8 @@ def is_openai_image_generation_route(url_route: str) -> bool: if not url_route: return False parsed_url = urlparse(url_route) - return bool( - parsed_url.hostname - and ( - "api.openai.com" in parsed_url.hostname - or "openai.azure.com" in parsed_url.hostname - ) + return ( + _is_openai_compatible_host(parsed_url.hostname) and "/v1/images/generations" in parsed_url.path ) @@ -82,12 +141,8 @@ def is_openai_image_editing_route(url_route: str) -> bool: if not url_route: return False parsed_url = urlparse(url_route) - return bool( - parsed_url.hostname - and ( - "api.openai.com" in parsed_url.hostname - or "openai.azure.com" in parsed_url.hostname - ) + return ( + _is_openai_compatible_host(parsed_url.hostname) and "/v1/images/edits" in parsed_url.path ) @@ -97,13 +152,8 @@ def is_openai_responses_route(url_route: str) -> bool: if not url_route: return False parsed_url = urlparse(url_route) - return bool( - parsed_url.hostname - and ( - "api.openai.com" in parsed_url.hostname - or "openai.azure.com" in parsed_url.hostname - ) - and ("/v1/responses" in parsed_url.path or "/responses" in parsed_url.path) + return _is_openai_compatible_host(parsed_url.hostname) and ( + "/v1/responses" in parsed_url.path or "/responses" in parsed_url.path ) def _get_user_from_metadata( @@ -188,6 +238,42 @@ def _calculate_image_editing_cost( ) return 0.0 + @staticmethod + def _build_responses_api_response_and_cost( + model: str, + httpx_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str, + ) -> Tuple[ResponsesAPIResponse, float]: + """Transform a Responses API raw response into a ResponsesAPIResponse + and compute its cost. + + The Responses API has a different on-the-wire shape from chat + completions (`output: [...]` instead of `choices: [...]`), so the + chat-completions `transform_response` raises KeyError 'choices' on + a Responses payload. Use the dedicated Responses-API transformer + (`OpenAIResponsesAPIConfig.transform_response_api_response`) here. + + Returns (litellm_model_response, response_cost) — symmetric with the + chat-completions branch which produces the same two values inline, + and analogous to the image branches' `_calculate_image_*_cost` helpers + (which return cost only because the image-response object is trivial + to build inline; the Responses payload needs a real transformer). + """ + responses_config = OpenAIResponsesAPIConfig() + litellm_model_response = responses_config.transform_response_api_response( + model=model, + raw_response=httpx_response, + logging_obj=logging_obj, + ) + response_cost = litellm.completion_cost( + completion_response=litellm_model_response, + model=model, + custom_llm_provider=custom_llm_provider, + call_type="responses", + ) + return litellm_model_response, response_cost + @staticmethod def openai_passthrough_handler( # noqa: PLR0915 httpx_response: httpx.Response, @@ -253,7 +339,12 @@ def openai_passthrough_handler( # noqa: PLR0915 try: response_cost = 0.0 litellm_model_response: Optional[ - Union[ModelResponse, TextCompletionResponse, ImageResponse] + Union[ + ModelResponse, + TextCompletionResponse, + ImageResponse, + ResponsesAPIResponse, + ] ] = None handler_instance = OpenAIPassthroughLoggingHandler() @@ -336,29 +427,18 @@ def openai_passthrough_handler( # noqa: PLR0915 litellm_model_response._hidden_params = {} litellm_model_response._hidden_params["response_cost"] = response_cost elif is_responses: - # Handle responses API cost calculation - provider_config = handler_instance.get_provider_config(model=model) - existing_litellm_params = kwargs.get("litellm_params", {}) or {} - litellm_model_response = provider_config.transform_response( - raw_response=httpx_response, - model_response=litellm.ModelResponse(), + # Responses-API cost tracking — see + # `_build_responses_api_response_and_cost` for why this needs + # a dedicated transformer (the chat-completions transform + # crashes on the Responses payload shape). + ( + litellm_model_response, + response_cost, + ) = OpenAIPassthroughLoggingHandler._build_responses_api_response_and_cost( model=model, - messages=request_body.get("messages", []), + httpx_response=httpx_response, logging_obj=logging_obj, - optional_params=request_body.get("optional_params", {}), - api_key="", - request_data=request_body, - encoding=litellm.encoding, - json_mode=False, - litellm_params=existing_litellm_params, - ) - - # Calculate cost using LiteLLM's cost calculator with responses call type - response_cost = litellm.completion_cost( - completion_response=litellm_model_response, - model=model, custom_llm_provider=custom_llm_provider, - call_type="responses", ) # Update kwargs with cost information diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index a267c97c0e8..9c0fbe30fc3 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -43,6 +43,10 @@ can_access_resource, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.repositories.table_repositories import ( + ManagedFileRepository, + ManagedObjectRepository, +) from litellm.types.llms.openai import OpenAIFileObject from .managed_id_codec import ManagedIdPayload, decode, is_managed, new_managed_id @@ -323,7 +327,7 @@ async def _resolve_one( ) if not found and prisma_client is not None: try: - db_row = await prisma_client.db.litellm_managedfiletable.find_first( + db_row = await ManagedFileRepository(prisma_client).table.find_first( where={"unified_file_id": managed_id} ) if db_row is not None: @@ -339,7 +343,7 @@ async def _resolve_one( # Object table (batches, responses) if prisma_client is not None: try: - obj_row = await prisma_client.db.litellm_managedobjecttable.find_first( + obj_row = await ManagedObjectRepository(prisma_client).table.find_first( where={"unified_object_id": managed_id} ) if obj_row is not None: @@ -399,7 +403,7 @@ async def _guard_raw_provider_id( # id and scope to the current provider in the application layer (same as # _mint_or_reuse_file's dedup). try: - candidates = await prisma_client.db.litellm_managedfiletable.find_many( + candidates = await ManagedFileRepository(prisma_client).table.find_many( where={"flat_model_file_ids": {"has": raw_id}}, ) except Exception: @@ -425,7 +429,7 @@ async def _guard_raw_provider_id( # Object rows store model_object_id as "passthrough:{provider}:{raw}", so # the lookup is exact and already provider-scoped. try: - existing = await prisma_client.db.litellm_managedobjecttable.find_first( + existing = await ManagedObjectRepository(prisma_client).table.find_first( where={"model_object_id": f"passthrough:{provider}:{raw_id}"} ) except Exception: @@ -492,7 +496,7 @@ async def _mint_or_reuse_file( # reuse a stable row instead of minting duplicate rows on every call. if prisma_client is not None: try: - candidates = await prisma_client.db.litellm_managedfiletable.find_many( + candidates = await ManagedFileRepository(prisma_client).table.find_many( where={"flat_model_file_ids": {"has": raw_id}}, order={"created_at": "asc"}, ) @@ -627,7 +631,7 @@ async def _reuse_existing(existing: Any, refresh_snapshot: bool) -> str: # the batch's latest state (e.g. output_file_id / error_file_id that # were null at creation but populated once the batch completed). try: - await prisma_client.db.litellm_managedobjecttable.update( + await ManagedObjectRepository(prisma_client).table.update( where={"unified_object_id": existing.unified_object_id}, data={ "file_object": json.dumps(body_snapshot), @@ -647,7 +651,7 @@ async def _reuse_existing(existing: Any, refresh_snapshot: bool) -> str: # Dedup: look up by the namespaced key — guaranteed unique per provider. try: - existing = await prisma_client.db.litellm_managedobjecttable.find_first( + existing = await ManagedObjectRepository(prisma_client).table.find_first( where={"model_object_id": namespaced_model_object_id} ) except Exception: @@ -666,7 +670,7 @@ async def _reuse_existing(existing: Any, refresh_snapshot: bool) -> str: raw_id.split("_", 1)[0], ) try: - await prisma_client.db.litellm_managedobjecttable.upsert( + await ManagedObjectRepository(prisma_client).table.upsert( where={"unified_object_id": managed_id}, data={ "create": { @@ -690,7 +694,7 @@ async def _reuse_existing(existing: Any, refresh_snapshot: bool) -> str: # the winner's managed ID so both callers converge on one ID instead of # the loser silently keeping the raw id. try: - raced = await prisma_client.db.litellm_managedobjecttable.find_first( + raced = await ManagedObjectRepository(prisma_client).table.find_first( where={"model_object_id": namespaced_model_object_id} ) except Exception: @@ -883,9 +887,9 @@ async def _build_list_where_with_cursor( return where, fetch_order cursor_table = ( - prisma_client.db.litellm_managedfiletable + ManagedFileRepository(prisma_client).table if resource_kind == "files" - else prisma_client.db.litellm_managedobjecttable + else ManagedObjectRepository(prisma_client).table ) cursor_field = ( "unified_file_id" if resource_kind == "files" else "unified_object_id" @@ -932,12 +936,12 @@ async def _fetch_list_rows( # across rows that share a created_at timestamp. try: if resource_kind == "files": - return await prisma_client.db.litellm_managedfiletable.find_many( + return await ManagedFileRepository(prisma_client).table.find_many( where=where, order=[{"created_at": fetch_order}, {"unified_file_id": fetch_order}], take=fetch_limit, ) - return await prisma_client.db.litellm_managedobjecttable.find_many( + return await ManagedObjectRepository(prisma_client).table.find_many( where={**where, "file_purpose": "batch"}, order=[{"created_at": fetch_order}, {"unified_object_id": fetch_order}], take=fetch_limit, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 6667010447b..d2b848e3c33 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -60,12 +60,13 @@ ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.utils import normalize_route_for_root_path +from litellm.repositories.team_repository import TeamRepository from litellm.secret_managers.main import get_secret_str from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - EndpointType, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, + EndpointType, PassthroughStandardLoggingPayload, ) @@ -1002,9 +1003,7 @@ async def pass_through_request( # noqa: PLR0915 is_passthrough_list_route, list_passthrough_ids_from_db, ) - from litellm.proxy.proxy_server import ( - prisma_client as _list_prisma, - ) + from litellm.proxy.proxy_server import prisma_client as _list_prisma if ( is_passthrough_list_route( @@ -1100,6 +1099,20 @@ async def pass_through_request( # noqa: PLR0915 status_code=e.response.status_code, detail=await e.response.aread() ) + # Call response headers hook for streaming pass-through + _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( + headers=response.headers, + litellm_call_id=litellm_call_id, + ) + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=_parsed_body or {}, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + _response_headers.update(callback_headers) + return StreamingResponse( PassThroughStreamingHandler.chunk_processor( response=response, @@ -1110,10 +1123,7 @@ async def pass_through_request( # noqa: PLR0915 passthrough_success_handler_obj=pass_through_endpoint_logging, url_route=str(url), ), - headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=response.headers, - litellm_call_id=litellm_call_id, - ), + headers=_response_headers, status_code=response.status_code, ) @@ -1152,6 +1162,20 @@ async def pass_through_request( # noqa: PLR0915 status_code=e.response.status_code, detail=await e.response.aread() ) + # Call response headers hook for detected streaming pass-through + _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( + headers=response.headers, + litellm_call_id=litellm_call_id, + ) + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=_parsed_body or {}, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + _response_headers.update(callback_headers) + return StreamingResponse( PassThroughStreamingHandler.chunk_processor( response=response, @@ -1162,10 +1186,7 @@ async def pass_through_request( # noqa: PLR0915 passthrough_success_handler_obj=pass_through_endpoint_logging, url_route=str(url), ), - headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=response.headers, - litellm_call_id=litellm_call_id, - ), + headers=_response_headers, status_code=response.status_code, ) @@ -1304,6 +1325,16 @@ async def pass_through_request( # noqa: PLR0915 api_base=str(url._uri_reference), ) + # Call response headers hook + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=_parsed_body or {}, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + custom_headers.update(callback_headers) + response_headers = HttpPassThroughEndpointHelpers.get_response_headers( headers=response.headers, custom_headers=custom_headers, @@ -2875,7 +2906,7 @@ async def _filter_endpoints_by_team_allowed_routes( HTTPException: If team is not found """ # retrieve team from db - team = await prisma_client.db.litellm_teamtable.find_unique( + team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id}, ) if team is None: diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 292871bae67..46043d10a06 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -434,15 +434,20 @@ def is_cursor_route( return False def is_openai_route(self, url_route: str): - """Check if the URL route is an OpenAI API route.""" + """Check if the URL route is an OpenAI API route. + + Uses the URL-aware helper so that non-OpenAI Azure Cognitive Services + (Speech, Vision, Language, ...) sharing the `*.cognitiveservices.azure.com` + / `*.openai.azure.com` domains are not misclassified as OpenAI routes. + """ if not url_route: return False - parsed_url = urlparse(url_route) - return parsed_url.hostname and ( - "api.openai.com" in parsed_url.hostname - or "openai.azure.com" in parsed_url.hostname + from .llm_provider_handlers.openai_passthrough_logging_handler import ( + _is_openai_compatible_url, ) + return _is_openai_compatible_url(url_route) + def is_gemini_route( self, url_route: str, custom_llm_provider: Optional[str] = None ): @@ -453,7 +458,16 @@ def is_gemini_route( return False def _is_supported_openai_endpoint(self, url_route: str) -> bool: - """Check if the OpenAI endpoint is supported by the passthrough logging handler.""" + """Check if the OpenAI endpoint is supported by the passthrough logging handler. + + The Responses API route is included because + `openai_passthrough_handler` has a dedicated `elif is_responses:` + branch that knows how to extract usage + cost from the + Responses-API on-the-wire shape. Without including it here, the + outer dispatch filters Responses calls out before reaching the + handler — the inner branch is then unreachable and Responses + calls land in `LiteLLM_SpendLogs` with zero tokens / zero spend. + """ from .llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, ) @@ -464,6 +478,7 @@ def _is_supported_openai_endpoint(self, url_route: str) -> bool: url_route ) or OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) + or OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) ) def _set_cost_per_request( diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 8d5d8116919..fb1e2652e8a 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional from litellm._logging import verbose_proxy_logger +from litellm.repositories.table_repositories import PolicyAttachmentRepository from litellm.types.proxy.policy_engine import ( PolicyAttachment, PolicyAttachmentCreateRequest, @@ -278,21 +279,21 @@ async def add_attachment_to_db( PolicyAttachmentDBResponse with the created attachment """ try: - created_attachment = ( - await prisma_client.db.litellm_policyattachmenttable.create( - data={ - "policy_name": attachment_request.policy_name, - "scope": attachment_request.scope, - "teams": attachment_request.teams or [], - "keys": attachment_request.keys or [], - "models": attachment_request.models or [], - "tags": attachment_request.tags or [], - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), - "created_by": created_by, - "updated_by": created_by, - } - ) + created_attachment = await PolicyAttachmentRepository( + prisma_client + ).table.create( + data={ + "policy_name": attachment_request.policy_name, + "scope": attachment_request.scope, + "teams": attachment_request.teams or [], + "keys": attachment_request.keys or [], + "models": attachment_request.models or [], + "tags": attachment_request.tags or [], + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": created_by, + "updated_by": created_by, + } ) # Also add to in-memory registry @@ -340,17 +341,15 @@ async def delete_attachment_from_db( """ try: # Get attachment before deleting - attachment = ( - await prisma_client.db.litellm_policyattachmenttable.find_unique( - where={"attachment_id": attachment_id} - ) - ) + attachment = await PolicyAttachmentRepository( + prisma_client + ).table.find_unique(where={"attachment_id": attachment_id}) if attachment is None: raise Exception(f"Attachment with ID {attachment_id} not found") # Delete from DB - await prisma_client.db.litellm_policyattachmenttable.delete( + await PolicyAttachmentRepository(prisma_client).table.delete( where={"attachment_id": attachment_id} ) @@ -379,11 +378,9 @@ async def get_attachment_by_id_from_db( PolicyAttachmentDBResponse if found, None otherwise """ try: - attachment = ( - await prisma_client.db.litellm_policyattachmenttable.find_unique( - where={"attachment_id": attachment_id} - ) - ) + attachment = await PolicyAttachmentRepository( + prisma_client + ).table.find_unique(where={"attachment_id": attachment_id}) if attachment is None: return None @@ -419,10 +416,10 @@ async def get_all_attachments_from_db( List of PolicyAttachmentDBResponse objects """ try: - attachments = ( - await prisma_client.db.litellm_policyattachmenttable.find_many( - order={"created_at": "desc"}, - ) + attachments = await PolicyAttachmentRepository( + prisma_client + ).table.find_many( + order={"created_at": "desc"}, ) return [ diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index 75017c46603..d6265516269 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from litellm._logging import verbose_proxy_logger +from litellm.repositories.table_repositories import PolicyRepository from litellm.types.proxy.policy_engine import ( GuardrailPipeline, PipelineStep, @@ -295,7 +296,7 @@ async def add_policy_to_db( validated_pipeline = GuardrailPipeline(**policy_request.pipeline) data["pipeline"] = json.dumps(validated_pipeline.model_dump()) - created_policy = await prisma_client.db.litellm_policytable.create( + created_policy = await PolicyRepository(prisma_client).table.create( data=data ) @@ -347,7 +348,7 @@ async def update_policy_in_db( Exception: If policy is not in draft status (only drafts are editable). """ try: - existing = await prisma_client.db.litellm_policytable.find_unique( + existing = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id} ) if existing is None: @@ -382,7 +383,7 @@ async def update_policy_in_db( validated_pipeline = GuardrailPipeline(**policy_request.pipeline) update_data["pipeline"] = json.dumps(validated_pipeline.model_dump()) - updated_policy = await prisma_client.db.litellm_policytable.update( + updated_policy = await PolicyRepository(prisma_client).table.update( where={"policy_id": policy_id}, data=update_data, ) @@ -413,7 +414,7 @@ async def delete_policy_from_db( Dict with "message" and optional "warning" if production was deleted. """ try: - policy = await prisma_client.db.litellm_policytable.find_unique( + policy = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id} ) @@ -424,7 +425,7 @@ async def delete_policy_from_db( policy_name = policy.policy_name # Delete from DB - await prisma_client.db.litellm_policytable.delete( + await PolicyRepository(prisma_client).table.delete( where={"policy_id": policy_id} ) @@ -461,7 +462,7 @@ async def get_policy_by_id_from_db( PolicyDBResponse if found, None otherwise """ try: - policy = await prisma_client.db.litellm_policytable.find_unique( + policy = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id} ) @@ -512,7 +513,7 @@ async def get_all_policies_from_db( if version_status is not None: where["version_status"] = version_status - policies = await prisma_client.db.litellm_policytable.find_many( + policies = await PolicyRepository(prisma_client).table.find_many( where=where if where else None, order={"created_at": "desc"}, ) @@ -554,7 +555,7 @@ async def sync_policies_from_db( self.add_policy(policy_response.policy_name, policy) self._policies_by_id = {} - non_production = await prisma_client.db.litellm_policytable.find_many( + non_production = await PolicyRepository(prisma_client).table.find_many( where={"version_status": {"in": ["draft", "published"]}}, order={"created_at": "desc"}, ) @@ -654,7 +655,7 @@ async def get_versions_by_policy_name( PolicyVersionListResponse with policy_name and list of versions """ try: - rows = await prisma_client.db.litellm_policytable.find_many( + rows = await PolicyRepository(prisma_client).table.find_many( where={"policy_name": policy_name}, order={"version_number": "desc"}, ) @@ -690,7 +691,7 @@ async def create_new_version( """ try: if source_policy_id is not None: - source = await prisma_client.db.litellm_policytable.find_unique( + source = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": source_policy_id} ) if source is None: @@ -701,7 +702,7 @@ async def create_new_version( ) else: # Find current production version for this policy_name - prod = await prisma_client.db.litellm_policytable.find_first( + prod = await PolicyRepository(prisma_client).table.find_first( where={ "policy_name": policy_name, "version_status": "production", @@ -714,7 +715,7 @@ async def create_new_version( source = prod # Next version number - latest = await prisma_client.db.litellm_policytable.find_first( + latest = await PolicyRepository(prisma_client).table.find_first( where={"policy_name": policy_name}, order={"version_number": "desc"}, ) @@ -722,7 +723,7 @@ async def create_new_version( now = datetime.now(timezone.utc) # Set is_latest=False on all existing versions for this policy_name - await prisma_client.db.litellm_policytable.update_many( + await PolicyRepository(prisma_client).table.update_many( where={"policy_name": policy_name}, data={"is_latest": False}, ) @@ -758,7 +759,7 @@ async def create_new_version( else source.pipeline ) - created = await prisma_client.db.litellm_policytable.create(data=data) + created = await PolicyRepository(prisma_client).table.create(data=data) return _row_to_policy_db_response(created) except Exception as e: verbose_proxy_logger.exception(f"Error creating new version: {e}") @@ -794,7 +795,7 @@ async def update_version_status( f"Invalid status '{new_status}'. Use 'published' or 'production'." ) - row = await prisma_client.db.litellm_policytable.find_unique( + row = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id} ) if row is None: @@ -809,7 +810,7 @@ async def update_version_status( raise Exception( f"Only draft versions can be published. Current status: '{current}'." ) - updated = await prisma_client.db.litellm_policytable.update( + updated = await PolicyRepository(prisma_client).table.update( where={"policy_id": policy_id}, data={ "version_status": "published", @@ -832,7 +833,7 @@ async def update_version_status( ) # Demote current production to published - await prisma_client.db.litellm_policytable.update_many( + await PolicyRepository(prisma_client).table.update_many( where={ "policy_name": policy_name, "version_status": "production", @@ -845,7 +846,7 @@ async def update_version_status( ) # Promote this version to production - updated = await prisma_client.db.litellm_policytable.update( + updated = await PolicyRepository(prisma_client).table.update( where={"policy_id": policy_id}, data={ "version_status": "production", @@ -895,10 +896,10 @@ async def compare_versions( PolicyVersionCompareResponse with both versions and field_diffs """ try: - a = await prisma_client.db.litellm_policytable.find_unique( + a = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id_a} ) - b = await prisma_client.db.litellm_policytable.find_unique( + b = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id_b} ) if a is None: @@ -950,7 +951,7 @@ async def delete_all_versions( Dict with success message """ try: - await prisma_client.db.litellm_policytable.delete_many( + await PolicyRepository(prisma_client).table.delete_many( where={"policy_name": policy_name} ) self.remove_policy(policy_name) diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index 54374d90a16..84dcbcfd746 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -16,6 +16,10 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.policy_engine import ( AttachmentImpactResponse, PolicyAttachmentCreateRequest, @@ -76,7 +80,7 @@ def _get_tags_from_metadata(metadata: object, json_metadata: object = None) -> l async def _fetch_all_teams(prisma_client: object) -> list: """Fetch teams from DB once. Reuse the result across tag and alias lookups.""" - return await prisma_client.db.litellm_teamtable.find_many( # type: ignore + return await TeamRepository(prisma_client).table.find_many( # type: ignore where={}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, @@ -159,7 +163,7 @@ async def _find_affected_by_team_patterns( new_keys: list = [] unnamed_keys_count = 0 if matched_team_ids: - keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore + keys = await VerificationTokenRepository(prisma_client).table.find_many( # type: ignore where={"team_id": {"in": matched_team_ids}}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, @@ -182,7 +186,7 @@ async def _find_affected_keys_by_alias( affected: list = [] - keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore + keys = await VerificationTokenRepository(prisma_client).table.find_many( # type: ignore where=_build_alias_where("key_alias", key_patterns), order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, @@ -367,7 +371,7 @@ async def estimate_attachment_impact( # Tag-based impact if tag_patterns: - keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore + keys = await VerificationTokenRepository(prisma_client).table.find_many( # type: ignore where={}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, diff --git a/litellm/proxy/policy_engine/policy_validator.py b/litellm/proxy/policy_engine/policy_validator.py index b587e3432bb..46796fbae28 100644 --- a/litellm/proxy/policy_engine/policy_validator.py +++ b/litellm/proxy/policy_engine/policy_validator.py @@ -12,6 +12,10 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set from litellm._logging import verbose_proxy_logger +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.policy_engine import ( Policy, PolicyValidationError, @@ -95,7 +99,7 @@ async def check_team_alias_exists(self, team_alias: str) -> bool: return True # Can't validate without DB, assume valid try: - team = await self.prisma_client.db.litellm_teamtable.find_first( + team = await TeamRepository(self.prisma_client).table.find_first( where={"team_alias": team_alias}, ) return team is not None @@ -119,7 +123,9 @@ async def check_key_alias_exists(self, key_alias: str) -> bool: return True # Can't validate without DB, assume valid try: - key = await self.prisma_client.db.litellm_verificationtoken.find_first( + key = await VerificationTokenRepository( + self.prisma_client + ).table.find_first( where={"key_alias": key_alias}, ) return key is not None diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index 399a0ff3af7..c0d6794108a 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -22,6 +22,7 @@ from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.path_utils import safe_filename +from litellm.repositories.table_repositories import PromptRepository from litellm.types.prompts.init_prompts import ( ListPromptsResponse, PromptInfo, @@ -208,7 +209,7 @@ async def get_next_version_for_prompt( Returns: Next version number (1 if no versions exist, max_version + 1 otherwise) """ - existing_prompts = await prisma_client.db.litellm_prompttable.find_many( + existing_prompts = await PromptRepository(prisma_client).table.find_many( where={"prompt_id": prompt_id, "environment": environment} ) @@ -441,7 +442,7 @@ async def get_prompt_versions( where_clause: Dict[str, Any] = {"prompt_id": base_prompt_id} if environment: where_clause["environment"] = environment - db_prompts = await prisma_client.db.litellm_prompttable.find_many( + db_prompts = await PromptRepository(prisma_client).table.find_many( where=where_clause, order={"version": "desc"}, ) @@ -612,7 +613,7 @@ async def get_prompt_info( # Query all environments this prompt exists in (lightweight: distinct on environment) all_environments: List[str] = [] if prisma_client is not None: - all_prompt_rows = await prisma_client.db.litellm_prompttable.find_many( + all_prompt_rows = await PromptRepository(prisma_client).table.find_many( where={"prompt_id": base_prompt_id}, distinct=["environment"], ) @@ -634,7 +635,7 @@ async def get_prompt_info( } if requested_version is not None: where_clause["version"] = requested_version - env_prompts = await prisma_client.db.litellm_prompttable.find_many( + env_prompts = await PromptRepository(prisma_client).table.find_many( where=where_clause, order={"version": "desc"}, take=1, @@ -752,7 +753,7 @@ async def create_prompt( ) # Store prompt in db with version - prompt_db_entry = await prisma_client.db.litellm_prompttable.create( + prompt_db_entry = await PromptRepository(prisma_client).table.create( data={ "prompt_id": request.prompt_id, "version": new_version, @@ -848,7 +849,7 @@ async def update_prompt( ) # Check if any version of this prompt exists (in any environment) - existing_prompts = await prisma_client.db.litellm_prompttable.find_many( + existing_prompts = await PromptRepository(prisma_client).table.find_many( where={"prompt_id": base_prompt_id} ) @@ -877,7 +878,7 @@ async def update_prompt( ) # Store new version in db - prompt_db_entry = await prisma_client.db.litellm_prompttable.create( + prompt_db_entry = await PromptRepository(prisma_client).table.create( data={ "prompt_id": base_prompt_id, "version": new_version, @@ -993,7 +994,7 @@ async def delete_prompt( delete_where["environment"] = environment # Delete versions from the database (scoped to environment if provided) - await prisma_client.db.litellm_prompttable.delete_many(where=delete_where) + await PromptRepository(prisma_client).table.delete_many(where=delete_where) # Remove matching prompts from memory — scope to environment if provided if environment: @@ -1105,7 +1106,7 @@ async def patch_prompt( if requested_version is not None: find_where["version"] = requested_version - db_rows = await prisma_client.db.litellm_prompttable.find_many( + db_rows = await PromptRepository(prisma_client).table.find_many( where=find_where, order={"version": "desc"}, take=1, @@ -1163,7 +1164,7 @@ async def patch_prompt( update_data["created_by"] = user_api_key_dict.user_id # Update by primary key (id) to target exactly one row - updated_prompt_db_entry = await prisma_client.db.litellm_prompttable.update( + updated_prompt_db_entry = await PromptRepository(prisma_client).table.update( where={"id": target_row.id}, data=update_data, ) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e4567b9f494..8c3fa952903 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -44,15 +44,19 @@ def _build_db_connection_url_params( pool_timeout: Optional[Union[int, float]], connect_timeout: Optional[Union[int, float]] = None, socket_timeout: Optional[Union[int, float]] = None, + disable_prepared_statements: bool = False, extra_params: Optional[dict] = None, ) -> dict: """Build the Prisma DATABASE_URL query params controlling connection pool behavior. `connect_timeout` / `socket_timeout` map to the Prisma URL params of the same name (https://www.prisma.io/docs/orm/overview/databases/postgresql) and are - omitted when None so Prisma's defaults apply. `extra_params` is an - untyped passthrough — keys it provides win over the named arguments above, - so it can be used to override any default we set here. + omitted when None so Prisma's defaults apply. `disable_prepared_statements` + sets `pgbouncer=true`, which makes Prisma stop using server-side prepared + statements (pgbouncer transaction-pool compatible; also sidesteps the + "cached plan must not change result type" error during rolling migrations). + `extra_params` is an untyped passthrough — keys it provides win over the + named arguments above, so it can be used to override any default we set here. """ params: dict = { "connection_limit": connection_limit, @@ -63,6 +67,8 @@ def _build_db_connection_url_params( params["connect_timeout"] = connect_timeout if socket_timeout is not None: params["socket_timeout"] = socket_timeout + if disable_prepared_statements: + params["pgbouncer"] = "true" if extra_params: params.update(extra_params) return params @@ -555,6 +561,7 @@ def _maybe_setup_prometheus_multiproc_dir( @click.command() +@click.argument("cli_args", nargs=-1) @click.option( "--host", default="0.0.0.0", help="Host for the server to listen on.", envvar="HOST" ) @@ -808,6 +815,7 @@ def _maybe_setup_prometheus_multiproc_dir( help="Enable uvicorn hot reload (dev only). Also reloads when the --config YAML file changes. Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.", ) def run_server( # noqa: PLR0915 + cli_args, host, port, api_base, @@ -854,6 +862,20 @@ def run_server( # noqa: PLR0915 use_v2_migration_resolver: bool, reload: bool, ): + if cli_args: + if cli_args == ("xai-oauth", "login"): + from litellm.llms.xai.oauth import XAIOAuthAuthenticator + + authenticator = XAIOAuthAuthenticator() + auth_data = authenticator.login() + click.echo( + f"xAI OAuth login successful. Credentials saved to {authenticator.auth_file}." + ) + if auth_data.get("expires_at"): + click.echo(f"Access token expires at {auth_data['expires_at']}.") + return + raise click.UsageError(f"Unknown command: {' '.join(cli_args)}") + if setup: from litellm.setup_wizard import run_setup_wizard @@ -947,6 +969,7 @@ def run_server( # noqa: PLR0915 db_connection_timeout: Optional[Union[int, float]] = 60 db_connect_timeout: Optional[Union[int, float]] = None db_socket_timeout: Optional[Union[int, float]] = None + db_disable_prepared_statements: bool = False db_extra_connection_params: Optional[dict] = None general_settings = {} ### GET DB TOKEN FOR IAM AUTH ### @@ -1067,6 +1090,17 @@ def run_server( # noqa: PLR0915 ) db_connect_timeout = general_settings.get("database_connect_timeout") db_socket_timeout = general_settings.get("database_socket_timeout") + _disable_prepared_statements = general_settings.get( + "database_disable_prepared_statements", False + ) + if isinstance(_disable_prepared_statements, str): + from litellm.secret_managers.main import str_to_bool + + db_disable_prepared_statements = ( + str_to_bool(_disable_prepared_statements) is True + ) + else: + db_disable_prepared_statements = bool(_disable_prepared_statements) db_extra_connection_params = general_settings.get( "database_extra_connection_params" ) @@ -1114,6 +1148,7 @@ def run_server( # noqa: PLR0915 pool_timeout=db_connection_timeout, connect_timeout=db_connect_timeout, socket_timeout=db_socket_timeout, + disable_prepared_statements=db_disable_prepared_statements, extra_params=db_extra_connection_params, ) if os.getenv("DATABASE_URL", None) is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 72423b2a796..267e388112d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -48,6 +48,7 @@ AIOHTTP_TTL_DNS_CACHE, AUDIO_SPEECH_CHUNK_SIZE, BASE_MCP_ROUTE, + DAILY_TAG_SPEND_BATCH_MULTIPLIER, DEFAULT_MAX_RECURSE_DEPTH, DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, DEFAULT_SHARED_HEALTH_CHECK_TTL, @@ -56,13 +57,13 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES, LITELLM_UI_ALLOW_HEADERS, LITELLM_UI_SESSION_DURATION, - DAILY_TAG_SPEND_BATCH_MULTIPLIER, ) from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( + UI_TEAM_ID, CallbackDelete, CallInfo, CommonProxyErrors, @@ -79,8 +80,8 @@ InvitationModel, InvitationNew, InvitationUpdate, - Litellm_EntityType, LiteLLM_EndUserTable, + Litellm_EntityType, LiteLLM_JWTAuth, LiteLLM_TagTable, LiteLLM_TeamTable, @@ -96,7 +97,6 @@ TeamDefaultSettings, TokenCountRequest, TransformRequestBody, - UI_TEAM_ID, UserAPIKeyAuth, ) from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec @@ -212,8 +212,6 @@ def generate_feedback_box(): from litellm._logging import verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( _REALTIME_BODY_CACHE_SIZE, APSCHEDULER_COALESCE, @@ -247,8 +245,8 @@ def generate_feedback_box(): ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.proxy._types import * from litellm.proxy._lazy_features import attach_lazy_features +from litellm.proxy._types import * from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, ) @@ -308,10 +306,15 @@ def generate_feedback_box(): from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.container_endpoints.endpoints import router as container_router from litellm.proxy.credential_endpoints.endpoints import router as credential_router from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import ( + PrismaDBExceptionHandler, + call_with_db_reconnect_retry, +) from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router @@ -361,7 +364,9 @@ def generate_feedback_box(): from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) -from litellm.proxy.management_endpoints.internal_user_endpoints import user_update +from litellm.proxy.management_endpoints.internal_user_endpoints import ( + user_update, +) from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, @@ -398,10 +403,6 @@ def generate_feedback_box(): update_team, validate_membership, ) -from litellm.proxy.management_endpoints.workflow_management_endpoints import ( - router as workflow_management_router, -) -from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, ) @@ -409,7 +410,11 @@ def generate_feedback_box(): from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( router as user_agent_analytics_router, ) +from litellm.proxy.management_endpoints.workflow_management_endpoints import ( + router as workflow_management_router, +) from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update +from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.middleware.in_flight_requests_middleware import ( InFlightRequestsMiddleware, ) @@ -417,12 +422,13 @@ def generate_feedback_box(): from litellm.proxy.middleware.request_size_limit_middleware import ( RequestSizeLimitMiddleware, ) -from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config +from litellm.proxy.openai_files_endpoints.files_endpoints import ( + set_files_config, +) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -444,6 +450,7 @@ def generate_feedback_box(): from litellm.proxy.response_api_endpoints.endpoints import router as response_router from litellm.proxy.route_llm_request import route_request from litellm.proxy.search_endpoints.endpoints import router as search_router +from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, ) @@ -478,6 +485,7 @@ def generate_feedback_box(): update_spend, ) from litellm.proxy.video_endpoints.endpoints import router as video_router +from litellm.repositories.credentials_repository import CredentialsRepository from litellm.router import ( AssistantsTypedDict, Deployment, @@ -511,7 +519,9 @@ def generate_feedback_box(): LiteLLM_UpperboundKeyGenerateParams, ) from litellm.types.realtime import RealtimeQueryParams -from litellm.types.router import DeploymentTypedDict +from litellm.types.router import ( + DeploymentTypedDict, +) from litellm.types.router import ModelInfo as RouterModelInfo from litellm.types.router import ( RouterGeneralSettings, @@ -2256,7 +2266,7 @@ async def _reconcile_budget_reservation_for_counter_update( ) except Exception: verbose_proxy_logger.warning( - "Failed to reconcile budget reservation after persisted spend; invalidating reserved counters and continuing", + "Failed to reconcile budget reservation after persisted spend; invalidating reserved counters and falling back to direct increment", exc_info=True, ) try: @@ -2267,6 +2277,7 @@ async def _reconcile_budget_reservation_for_counter_update( verbose_proxy_logger.exception( "Failed to invalidate reserved counters after reservation reconciliation failed" ) + return set() return reserved_counter_keys @@ -4068,6 +4079,7 @@ async def load_config( # noqa: PLR0915 premium_user=premium_user, config_file_path=config_file_path, litellm_settings=litellm_settings, + callback_specific_params=callback_settings, ) elif key == "model_group_settings": @@ -4082,6 +4094,8 @@ async def load_config( # noqa: PLR0915 verbose_proxy_logger.debug( f"litellm.post_call_rules: {litellm.post_call_rules}" ) + elif key == "max_budget": + litellm.max_budget = float(value) elif key == "max_internal_user_budget": litellm.max_internal_user_budget = float(value) # type: ignore elif key == "default_max_internal_user_budget": @@ -4248,13 +4262,13 @@ async def load_config( # noqa: PLR0915 ) setattr(litellm, key, value) if key in {"s3_audit_callback_params", "s3_callback_params"}: - from litellm.proxy.management_helpers.audit_logs import ( - reset_audit_log_callback_cache, - ) + from litellm.integrations.s3_v2 import S3Logger as S3V2Logger from litellm.litellm_core_utils.litellm_logging import ( _in_memory_loggers, ) - from litellm.integrations.s3_v2 import S3Logger as S3V2Logger + from litellm.proxy.management_helpers.audit_logs import ( + reset_audit_log_callback_cache, + ) reset_audit_log_callback_cache() _in_memory_loggers[:] = [ @@ -5335,7 +5349,7 @@ async def _add_router_settings_from_db_config( 4. Update router settings """ if llm_router is not None and prisma_client is not None: - db_router_settings = await prisma_client.db.litellm_config.find_first( + db_router_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "router_settings"} ) @@ -5761,7 +5775,7 @@ def _should_load_db_object( async def _get_models_from_db(self, prisma_client: PrismaClient) -> list: try: - new_models = await prisma_client.db.litellm_proxymodeltable.find_many() + new_models = await ModelRepository(prisma_client).table.find_many() except Exception as e: verbose_proxy_logger.exception( "litellm.proxy_server.py::add_deployment() - Error getting new models from DB - {}".format( @@ -5975,8 +5989,12 @@ async def _init_sso_settings_in_db(self, prisma_client: PrismaClient): """ try: - sso_settings = await prisma_client.db.litellm_ssoconfig.find_unique( - where={"id": "sso_config"} + sso_settings = await call_with_db_reconnect_retry( + prisma_client, + lambda: SSOConfigRepository(prisma_client).table.find_unique( + where={"id": "sso_config"} + ), + reason="init_sso_settings_in_db_lookup_failure", ) if sso_settings is not None: sso_settings.sso_settings.pop("role_mappings", None) @@ -6011,8 +6029,12 @@ async def _init_hashicorp_vault_config_override(self, prisma_client: PrismaClien ) try: - db_record = await prisma_client.db.litellm_configoverrides.find_unique( - where={"config_type": "hashicorp_vault"} + db_record = await call_with_db_reconnect_retry( + prisma_client, + lambda: ConfigOverridesRepository(prisma_client).table.find_unique( + where={"config_type": "hashicorp_vault"} + ), + reason="init_hashicorp_vault_config_override_lookup_failure", ) if db_record is None or db_record.config_value is None: @@ -6130,7 +6152,7 @@ async def _check_and_reload_model_cost_map(self, prisma_client: PrismaClient): last_model_cost_map_reload = current_time.isoformat() # Clear force reload flag in database - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "model_cost_map_reload_config"}, data={ "create": { @@ -6239,7 +6261,7 @@ async def _check_and_reload_anthropic_beta_headers( last_anthropic_beta_headers_reload = current_time.isoformat() # Clear force reload flag in database - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "anthropic_beta_headers_reload_config"}, data={ "create": { @@ -6299,7 +6321,7 @@ async def _init_prompts_in_db(self, prisma_client: PrismaClient): from litellm.types.prompts.init_prompts import PromptSpec try: - prompts_in_db = await prisma_client.db.litellm_prompttable.find_many() + prompts_in_db = await PromptRepository(prisma_client).table.find_many() for prompt in prompts_in_db: # Convert DB object to dict and create versioned prompt_id prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) @@ -6589,7 +6611,7 @@ async def delete_credentials(self, db_credentials: List[CredentialItem]): async def get_credentials(self, prisma_client: PrismaClient): try: - credentials = await prisma_client.db.litellm_credentialstable.find_many() + credentials = await CredentialsRepository(prisma_client).find_all() credentials = [self.decrypt_credentials(cred) for cred in credentials] await self.delete_credentials( credentials @@ -7352,7 +7374,7 @@ async def _upsert_proxy_budget_with_reset_at_backfill( # spend cap blocks forever once it's hit. if prisma_client is not None and litellm.budget_duration is not None: try: - await prisma_client.db.litellm_usertable.update_many( + await UserRepository(prisma_client).table.update_many( where={ "user_id": litellm_proxy_budget_name, "budget_reset_at": None, @@ -7420,7 +7442,7 @@ async def _sync_ui_settings_to_general_settings(cls): if prisma_client is None: return - db_record = await prisma_client.db.litellm_uisettings.find_unique( + db_record = await UISettingsRepository(prisma_client).table.find_unique( where={"id": "ui_settings"} ) if db_record and db_record.ui_settings: @@ -7569,7 +7591,7 @@ async def initialize_scheduled_background_jobs( # noqa: PLR0915 # but YAML config has False. if store_model_in_db is not True and prisma_client is not None: try: - _db_gs_record = await prisma_client.db.litellm_config.find_first( + _db_gs_record = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) if _db_gs_record is not None and isinstance( @@ -9184,6 +9206,16 @@ async def audio_speech( hidden_params=hidden_params, ) + # Call response headers hook (matches audio_transcription behavior) + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + custom_headers.update(callback_headers) + # Determine media type based on model type media_type = "audio/mpeg" # Default for OpenAI TTS request_model = data.get("model", "") @@ -10361,6 +10393,18 @@ async def run_thread( # ) # async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)): from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ( + AccessGroupRepository, + ConfigOverridesRepository, + InvitationLinkRepository, + PromptRepository, + SSOConfigRepository, + UISettingsRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository def _get_provider_token_counter( @@ -10666,7 +10710,7 @@ async def _check_if_model_is_user_added( id = model.get("model_info", {}).get("id", None) if id is None: continue - db_model = await prisma_client.db.litellm_proxymodeltable.find_unique( + db_model = await ModelRepository(prisma_client).table.find_unique( where={"model_id": id} ) if db_model is not None: @@ -10723,7 +10767,7 @@ async def non_admin_all_models( if user_api_key_dict.user_id: try: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) except Exception: @@ -10823,7 +10867,7 @@ async def _add_access_group_models_to_team_models( return team_models # Single batch fetch for all access groups - access_group_rows = await prisma_client.db.litellm_accessgrouptable.find_many( + access_group_rows = await AccessGroupRepository(prisma_client).table.find_many( where={"access_group_id": {"in": list(all_access_group_ids)}} ) ag_model_map: Dict[str, List[str]] = { @@ -10865,13 +10909,13 @@ async def get_all_team_models( team_db_objects_typed: List[LiteLLM_TeamTable] = [] if user_teams == "*": - team_db_objects = await prisma_client.db.litellm_teamtable.find_many() + team_db_objects = await TeamRepository(prisma_client).table.find_many() team_db_objects_typed = [ LiteLLM_TeamTable(**team_db_object.model_dump()) for team_db_object in team_db_objects ] else: - team_db_objects = await prisma_client.db.litellm_teamtable.find_many( + team_db_objects = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": user_teams}} ) @@ -10938,7 +10982,7 @@ async def get_all_team_and_direct_access_models( exclude_team_models=True ) # has access to all models elif user_api_key_dict.user_id is not None: - user_db_object = await prisma_client.db.litellm_usertable.find_unique( + user_db_object = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) if user_db_object is not None: @@ -11082,7 +11126,7 @@ async def _get_caller_byok_team_scope( if user_id is None: return set() try: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) except Exception: @@ -11096,6 +11140,22 @@ async def _get_caller_byok_team_scope( return set(user_row.teams or []) +def _byok_row_outside_caller_teams( + model_info_dict: Dict[str, Any], allowed_team_ids: Optional[Set[str]] +) -> bool: + """Whether a team BYOK row belongs to a team the caller is not a member of. + + `team_id` is only set on team BYOK rows; non-team rows fall through + unaffected. `allowed_team_ids is None` means no scoping (e.g. admins). + """ + if allowed_team_ids is None: + return False + team_id = model_info_dict.get("team_id") + if team_id is None: + return False + return team_id not in allowed_team_ids + + # Hard cap on rows the DB-side BYOK search may pull when results need to be # sorted across the full match set. Without this, an authenticated caller # can hit `/v2/model/info?search=&sortBy=` and force the @@ -11143,13 +11203,13 @@ async def _fetch_db_models_for_search( else: take_limit = max(0, page * size - router_models_count) - db_models_total_count = await prisma_client.db.litellm_proxymodeltable.count( + db_models_total_count = await ModelRepository(prisma_client).table.count( where=db_where_condition ) db_models_raw: list = [] if take_limit > 0: - db_models_raw = await prisma_client.db.litellm_proxymodeltable.find_many( + db_models_raw = await ModelRepository(prisma_client).table.find_many( where=db_where_condition, take=take_limit, ) @@ -11217,15 +11277,7 @@ async def _apply_search_filter_to_models( ) def _is_byok_outside_caller_teams(model_info_dict: Dict[str, Any]) -> bool: - # `team_id` is only set on team BYOK rows. Non-team rows fall - # through unaffected — they are gated by other paths (router - # membership, direct_access, include_team_models). - if allowed_team_ids is None: - return False - team_id = model_info_dict.get("team_id") - if team_id is None: - return False - return team_id not in allowed_team_ids + return _byok_row_outside_caller_teams(model_info_dict, allowed_team_ids) def _model_matches_search(m: Dict[str, Any]) -> bool: # Team BYOK models persist an internal `model_name` @@ -11484,7 +11536,7 @@ async def _load_team_object_for_model_filter( ) -> Optional[LiteLLM_TeamTable]: """Load team row from DB; returns None if missing or on error.""" try: - team_db_object = await prisma_client.db.litellm_teamtable.find_unique( + team_db_object = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) if team_db_object is None: @@ -11547,7 +11599,7 @@ async def _gather_team_accessible_model_ids( _resolved_names = _team_models_resolve_to_names( team_object.models, access_groups ) - db_models = await prisma_client.db.litellm_proxymodeltable.find_many( + db_models = await ModelRepository(prisma_client).table.find_many( where={"model_name": {"in": _resolved_names}} ) for db_model in db_models: @@ -11586,7 +11638,7 @@ async def _authorize_team_id_query( detail={"error": "Not authorized to view this team's models"}, ) try: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) except Exception: @@ -11693,7 +11745,7 @@ async def _find_model_by_id( # If not found in config, search in database if found_model is None: try: - db_model = await prisma_client.db.litellm_proxymodeltable.find_unique( + db_model = await ModelRepository(prisma_client).table.find_unique( where={"model_id": model_id} ) if db_model: @@ -11723,10 +11775,8 @@ async def _find_model_by_id( @router.get( "/v2/model/info", - description="v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true", tags=["model management"], dependencies=[Depends(user_api_key_auth)], - include_in_schema=False, ) async def model_info_v2( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -11762,7 +11812,49 @@ async def model_info_v2( ), ): """ - BETA ENDPOINT. Might change unexpectedly. Use `/v1/model/info` for now. + Paginated model metadata for proxy deployments (pricing, provider, team access). + + Returns configured router deployments with enriched `model_info` (costs, provider, + context window, etc.). Sensitive fields such as API keys and api_base are omitted. + + Query parameters: + model: Filter to a single public `model_name`. + user_models_only: When true, only return models created by the calling user. + include_team_models: When true, populate `access_via_team_ids` and `direct_access` + on each model and filter to deployments the caller can use. + page / size: Pagination controls (defaults: page=1, size=50). + search: Case-insensitive partial match on model name or team public name. + modelId: Return a single deployment by LiteLLM model id. + teamId: Filter to models with direct access or team membership for this team id. + sortBy / sortOrder: Sort by model_name, created_at, updated_at, costs, or status. + + Example request: + ``` + curl -X GET 'http://localhost:4000/v2/model/info?include_team_models=true&page=1&size=50' \\ + --header 'Authorization: Bearer sk-1234' + ``` + + Example response: + ```json + { + "data": [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4.1"}, + "model_info": { + "id": "abc123", + "litellm_provider": "openai", + "access_via_team_ids": ["team-1"], + "direct_access": true + } + } + ], + "total_count": 1, + "current_page": 1, + "total_pages": 1, + "size": 50 + } + ``` """ global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router @@ -12325,6 +12417,72 @@ async def model_metrics_exceptions( return {"data": response, "exception_types": list(exception_types)} +def _deployment_matches_allowed_model_names( + model: Dict[str, Any], allowed_model_names: Set[str] +) -> bool: + """Match a router deployment against allowed public model names. + + Team-scoped rows store an internal routing key in ``model_name``; callers + with key/team restrictions still refer to the public name in + ``model_info.team_public_model_name``. + """ + if model.get("model_name") in allowed_model_names: + return True + model_info = model.get("model_info") + if not isinstance(model_info, dict): + return False + team_public_model_name = model_info.get("team_public_model_name") + return ( + isinstance(team_public_model_name, str) + and team_public_model_name in allowed_model_names + ) + + +def _get_v1_model_info_allowed_model_names( + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router, +) -> Optional[Set[str]]: + """Return key/team allowlisted public model names, or None if unrestricted.""" + model_access_groups = llm_router.get_model_access_groups() + proxy_model_list = llm_router.get_model_names() + key_models = get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + ) + team_models = get_team_models( + team_models=user_api_key_dict.team_models, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + ) + if not key_models and not team_models: + return None + return set( + get_complete_model_list( + key_models=key_models, + team_models=team_models, + proxy_model_list=proxy_model_list, + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + llm_router=llm_router, + return_wildcard_routes=False, + ) + ) + + +def _filter_v1_model_info_deployments( + all_models: List[dict], + allowed_model_names: Optional[Set[str]], +) -> List[dict]: + if allowed_model_names is None: + return all_models + return [ + model + for model in all_models + if _deployment_matches_allowed_model_names(model, allowed_model_names) + ] + + def _translate_model_name_for_response(model: dict) -> dict: """For team-scoped DB rows, replace `model_name` with the public name in `model_info.team_public_model_name` before returning. The DB column @@ -12494,49 +12652,42 @@ async def model_info_v1( # noqa: PLR0915 ) return {"data": [_deployment_info_dict]} - all_models: List[dict] = [] - model_access_groups: Dict[str, List[str]] = defaultdict(list) - ## CHECK IF MODEL RESTRICTIONS ARE SET AT KEY/TEAM LEVEL ## - if llm_router is None: - proxy_model_list = [] - else: - proxy_model_list = llm_router.get_model_names() - model_access_groups = llm_router.get_model_access_groups() - key_models = get_key_models( + # Return router deployments (same source as /v2/model/info), not wildcard- + # expanded model names from get_complete_model_list(). Team-scoped rows + # use internal routing keys (model_name_{team_id}_{uuid}) and were omitted + # when v1 resolved models only via public model_name strings. + all_models: List[dict] = copy.deepcopy(llm_router.model_list) + allowed_model_names = _get_v1_model_info_allowed_model_names( user_api_key_dict=user_api_key_dict, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, - ) - team_models = get_team_models( - team_models=user_api_key_dict.team_models, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, - ) - all_models_str = get_complete_model_list( - key_models=key_models, - team_models=team_models, - proxy_model_list=proxy_model_list, - user_model=user_model, - infer_model_from_keys=general_settings.get("infer_model_from_keys", False), llm_router=llm_router, ) - if len(all_models_str) > 0: - _relevant_models = [] - for model in all_models_str: - router_models = llm_router.get_model_list(model_name=model) - if router_models is not None: - _relevant_models.extend(router_models) - if llm_model_list is not None: - all_models = copy.deepcopy(_relevant_models) # type: ignore - else: - all_models = [] + all_models = _filter_v1_model_info_deployments( + all_models=all_models, + allowed_model_names=allowed_model_names, + ) - # Reassign each entry: _get_proxy_model_info returns a (possibly new) - # dict via _translate_model_name_for_response, which does NOT mutate in - # place. Binding only the loop variable would drop the public-name swap - # for team-scoped rows and leak the internal routing key (#28382). - all_models = [_get_proxy_model_info(model=model) for model in all_models] + # Team BYOK deployments carry an internal routing key and other teams' + # public name/team_id/api_base; drop the ones the caller cannot access so + # listing the full router model_list does not leak cross-team metadata. + allowed_team_ids = await _get_caller_byok_team_scope( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + all_models = [ + model + for model in all_models + if not _byok_row_outside_caller_teams( + model.get("model_info") or {}, allowed_team_ids + ) + ] + + all_models = [ + _translate_model_name_for_response( + _enrich_model_info_with_litellm_data(model=model, llm_router=llm_router) + ) + for model in all_models + ] verbose_proxy_logger.debug("all_models: %s", all_models) return {"data": all_models} @@ -12845,7 +12996,7 @@ async def alerting_settings( ) ## get general settings from db - db_general_settings = await prisma_client.db.litellm_config.find_first( + db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) @@ -13384,7 +13535,7 @@ async def onboarding(invite_link: str, request: Request): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - invite_obj = await prisma_client.db.litellm_invitationlink.find_unique( + invite_obj = await InvitationLinkRepository(prisma_client).table.find_unique( where={"id": invite_link} ) if invite_obj is None: @@ -13408,7 +13559,7 @@ async def onboarding(invite_link: str, request: Request): ) ### GET USER OBJECT ### - user_obj = await prisma_client.db.litellm_usertable.find_unique( + user_obj = await UserRepository(prisma_client).table.find_unique( where={"user_id": invite_obj.user_id} ) @@ -13513,7 +13664,7 @@ async def _rollback_onboarding_invite_claim( return try: - await prisma_client.db.litellm_invitationlink.update_many( + await InvitationLinkRepository(prisma_client).table.update_many( where={"id": invitation_link, "is_accepted": True}, data={ "accepted_at": None, @@ -13547,10 +13698,10 @@ async def _generate_onboarding_ui_session_token(user_obj: Any) -> str: ) key = response["token"] # type: ignore - from litellm.types.proxy.ui_sso import ReturnedUITokenObject - import jwt + from litellm.types.proxy.ui_sso import ReturnedUITokenObject + disabled_non_admin_personal_key_creation = ( get_disabled_non_admin_personal_key_creation() ) @@ -13596,7 +13747,7 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - invite_obj = await prisma_client.db.litellm_invitationlink.find_unique( + invite_obj = await InvitationLinkRepository(prisma_client).table.find_unique( where={"id": data.invitation_link} ) if invite_obj is None: @@ -13956,7 +14107,7 @@ async def invitation_info( }, ) - response = await prisma_client.db.litellm_invitationlink.find_unique( + response = await InvitationLinkRepository(prisma_client).table.find_unique( where={"id": invitation_id} ) @@ -14010,7 +14161,7 @@ async def invitation_update( ) current_time = litellm.utils.get_utc_datetime() - response = await prisma_client.db.litellm_invitationlink.update( + response = await InvitationLinkRepository(prisma_client).table.update( where={"id": data.invitation_id}, data={ "id": data.invitation_id, @@ -14081,7 +14232,7 @@ async def invitation_delete( # Org admins can only delete invitations they created if is_other_admin and not is_proxy_admin: - invitation = await prisma_client.db.litellm_invitationlink.find_unique( + invitation = await InvitationLinkRepository(prisma_client).table.find_unique( where={"id": data.invitation_id} ) if invitation is None: @@ -14097,7 +14248,7 @@ async def invitation_delete( }, ) - response = await prisma_client.db.litellm_invitationlink.delete( + response = await InvitationLinkRepository(prisma_client).table.delete( where={"id": data.invitation_id} ) @@ -14139,7 +14290,7 @@ async def update_config( # noqa: PLR0915 raise Exception("No DB Connected") async def _read_section(param_name: str) -> dict: - row = await prisma_client.db.litellm_config.find_first( + row = await ConfigRepository(prisma_client).table.find_first( where={"param_name": param_name} ) if row is None or row.param_value is None: @@ -14148,7 +14299,7 @@ async def _read_section(param_name: str) -> dict: async def _upsert_section(param_name: str, value: dict) -> None: serialized = json.dumps(value) - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": param_name}, data={ "create": {"param_name": param_name, "param_value": serialized}, @@ -14323,7 +14474,7 @@ async def update_config_general_settings( ) ## get general settings from db - db_general_settings = await prisma_client.db.litellm_config.find_first( + db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) ### update value @@ -14337,7 +14488,7 @@ async def update_config_general_settings( general_settings[data.field_name] = data.field_value - response = await prisma_client.db.litellm_config.upsert( + response = await ConfigRepository(prisma_client).table.upsert( where={"param_name": "general_settings"}, data={ "create": {"param_name": "general_settings", "param_value": json.dumps(general_settings)}, # type: ignore @@ -14387,7 +14538,7 @@ async def get_config_general_settings( ) ## get general settings from db - db_general_settings = await prisma_client.db.litellm_config.find_first( + db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) ### pop the value @@ -14450,7 +14601,7 @@ async def get_config_list( ) ## get general settings from db - db_general_settings = await prisma_client.db.litellm_config.find_first( + db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) @@ -14604,7 +14755,7 @@ async def delete_config_general_settings( ) ## get general settings from db - db_general_settings = await prisma_client.db.litellm_config.find_first( + db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) ### pop the value @@ -14621,7 +14772,7 @@ async def delete_config_general_settings( general_settings.pop(data.field_name, None) - response = await prisma_client.db.litellm_config.upsert( + response = await ConfigRepository(prisma_client).table.upsert( where={"param_name": "general_settings"}, data={ "create": {"param_name": "general_settings", "param_value": json.dumps(general_settings)}, # type: ignore @@ -14975,14 +15126,14 @@ async def reload_model_cost_map( last_model_cost_map_reload = current_time.isoformat() # Set force reload flag in database for other pods, preserving existing interval_hours - existing_config = await prisma_client.db.litellm_config.find_unique( + existing_config = await ConfigRepository(prisma_client).table.find_unique( where={"param_name": "model_cost_map_reload_config"} ) existing_interval = None if existing_config and existing_config.param_value: existing_interval = existing_config.param_value.get("interval_hours") - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "model_cost_map_reload_config"}, data={ "create": { @@ -15052,7 +15203,7 @@ async def schedule_model_cost_map_reload( ) # Update database with new reload configuration - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "model_cost_map_reload_config"}, data={ "create": { @@ -15119,7 +15270,7 @@ async def cancel_model_cost_map_reload( ) # Remove reload configuration from database - await prisma_client.db.litellm_config.delete( + await ConfigRepository(prisma_client).table.delete( where={"param_name": "model_cost_map_reload_config"} ) await invalidate_config_param("model_cost_map_reload_config") @@ -15178,7 +15329,7 @@ async def get_model_cost_map_reload_status( } # Get reload configuration from database - config_record = await prisma_client.db.litellm_config.find_unique( + config_record = await ConfigRepository(prisma_client).table.find_unique( where={"param_name": "model_cost_map_reload_config"} ) @@ -15329,7 +15480,7 @@ async def reload_anthropic_beta_headers( last_anthropic_beta_headers_reload = current_time.isoformat() # Set force reload flag in database for other pods, preserving existing interval_hours - existing_beta_config = await prisma_client.db.litellm_config.find_unique( + existing_beta_config = await ConfigRepository(prisma_client).table.find_unique( where={"param_name": "anthropic_beta_headers_reload_config"} ) existing_beta_interval = None @@ -15338,7 +15489,7 @@ async def reload_anthropic_beta_headers( "interval_hours" ) - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "anthropic_beta_headers_reload_config"}, data={ "create": { @@ -15412,7 +15563,7 @@ async def schedule_anthropic_beta_headers_reload( ) # Update database with new reload configuration - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "anthropic_beta_headers_reload_config"}, data={ "create": { @@ -15479,7 +15630,7 @@ async def cancel_anthropic_beta_headers_reload( ) # Remove reload configuration from database - await prisma_client.db.litellm_config.delete( + await ConfigRepository(prisma_client).table.delete( where={"param_name": "anthropic_beta_headers_reload_config"} ) await invalidate_config_param("anthropic_beta_headers_reload_config") @@ -15539,7 +15690,7 @@ async def get_anthropic_beta_headers_reload_status( } # Get reload configuration from database - config_record = await prisma_client.db.litellm_config.find_unique( + config_record = await ConfigRepository(prisma_client).table.find_unique( where={"param_name": "anthropic_beta_headers_reload_config"} ) diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index d12e7a35fbf..78467c4b2e7 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -4,9 +4,9 @@ from importlib.resources import files from typing import Any, Dict, List, Optional -import litellm from fastapi import APIRouter, HTTPException, Request +import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_blog_posts import ( BlogPost, @@ -17,6 +17,7 @@ from litellm.proxy._types import ( CommonProxyErrors, ) +from litellm.repositories.table_repositories import ClaudeCodePluginRepository from litellm.types.agents import AgentCard from litellm.types.mcp import MCPPublicServer from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -159,14 +160,14 @@ def _load_endpoints() -> List[Dict[str, Any]]: ) async def public_model_hub(): import litellm + from litellm.proxy.health_endpoints._health_endpoints import ( + _convert_health_check_to_dict, + ) from litellm.proxy.proxy_server import ( _get_model_group_info, llm_router, prisma_client, ) - from litellm.proxy.health_endpoints._health_endpoints import ( - _convert_health_check_to_dict, - ) if llm_router is None: raise HTTPException( @@ -266,7 +267,7 @@ async def public_skill_hub(): try: prisma_client = await _get_prisma_client() - plugins = await prisma_client.db.litellm_claudecodeplugintable.find_many( + plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many( where={"enabled": True} ) items = [] diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index a44e4781491..7ff54ac4c5a 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -17,16 +17,17 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.proxy._types import * +from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, get_form_data, ) -from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, ) +from litellm.repositories.table_repositories import ManagedVectorStoresRepository router = APIRouter() @@ -230,11 +231,9 @@ async def _save_vector_store_to_db_from_rag_ingest( try: # Check if vector store already exists in database - existing_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": vector_store_id} - ) - ) + existing_vector_store = await ManagedVectorStoresRepository( + prisma_client + ).table.find_unique(where={"vector_store_id": vector_store_id}) # Only create if it doesn't exist if existing_vector_store is None: @@ -289,7 +288,7 @@ async def _save_vector_store_to_db_from_rag_ingest( # Update the vector store from litellm.proxy.utils import safe_dumps - await prisma_client.db.litellm_managedvectorstorestable.update( + await ManagedVectorStoresRepository(prisma_client).table.update( where={"vector_store_id": vector_store_id}, data={"vector_store_metadata": safe_dumps(existing_metadata)}, ) diff --git a/litellm/proxy/search_endpoints/search_tool_registry.py b/litellm/proxy/search_endpoints/search_tool_registry.py index d4adc2573ea..2ec2533211b 100644 --- a/litellm/proxy/search_endpoints/search_tool_registry.py +++ b/litellm/proxy/search_endpoints/search_tool_registry.py @@ -7,7 +7,9 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import SearchToolsRepository from litellm.types.search import SearchTool @@ -63,16 +65,16 @@ async def add_search_tool_to_db( search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {})) # Create search tool in DB - created_search_tool = ( - await prisma_client.db.litellm_searchtoolstable.create( - data={ - "search_tool_name": search_tool_name, - "litellm_params": litellm_params, - "search_tool_info": search_tool_info, - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), - } - ) + created_search_tool = await SearchToolsRepository( + prisma_client + ).table.create( + data={ + "search_tool_name": search_tool_name, + "litellm_params": litellm_params, + "search_tool_info": search_tool_info, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } ) # Add search_tool_id to the returned search tool object @@ -101,15 +103,15 @@ async def delete_search_tool_from_db( """ try: # Get search tool before deletion for response - existing_tool = await prisma_client.db.litellm_searchtoolstable.find_unique( - where={"search_tool_id": search_tool_id} - ) + existing_tool = await SearchToolsRepository( + prisma_client + ).table.find_unique(where={"search_tool_id": search_tool_id}) if not existing_tool: raise Exception(f"Search tool with ID {search_tool_id} not found") # Delete from DB - await prisma_client.db.litellm_searchtoolstable.delete( + await SearchToolsRepository(prisma_client).table.delete( where={"search_tool_id": search_tool_id} ) @@ -145,16 +147,16 @@ async def update_search_tool_in_db( search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {})) # Update in DB - updated_search_tool = ( - await prisma_client.db.litellm_searchtoolstable.update( - where={"search_tool_id": search_tool_id}, - data={ - "search_tool_name": search_tool_name, - "litellm_params": litellm_params, - "search_tool_info": search_tool_info, - "updated_at": datetime.now(timezone.utc), - }, - ) + updated_search_tool = await SearchToolsRepository( + prisma_client + ).table.update( + where={"search_tool_id": search_tool_id}, + data={ + "search_tool_name": search_tool_name, + "litellm_params": litellm_params, + "search_tool_info": search_tool_info, + "updated_at": datetime.now(timezone.utc), + }, ) # Convert to dict with ISO formatted datetimes @@ -179,10 +181,12 @@ async def get_all_search_tools_from_db( List of search tool configurations """ try: - search_tools_from_db = ( - await prisma_client.db.litellm_searchtoolstable.find_many( + search_tools_from_db = await call_with_db_reconnect_retry( + prisma_client, + lambda: SearchToolsRepository(prisma_client).table.find_many( order={"created_at": "desc"}, - ) + ), + reason="get_all_search_tools_from_db_lookup_failure", ) search_tools: List[SearchTool] = [] @@ -214,7 +218,7 @@ async def get_search_tool_by_id_from_db( Search tool configuration or None if not found """ try: - search_tool = await prisma_client.db.litellm_searchtoolstable.find_unique( + search_tool = await SearchToolsRepository(prisma_client).table.find_unique( where={"search_tool_id": search_tool_id} ) @@ -244,7 +248,7 @@ async def get_search_tool_by_name_from_db( Search tool configuration or None if not found """ try: - search_tool = await prisma_client.db.litellm_searchtoolstable.find_unique( + search_tool = await SearchToolsRepository(prisma_client).table.find_unique( where={"search_tool_name": search_tool_name} ) diff --git a/litellm/proxy/spend_tracking/cloudzero_endpoints.py b/litellm/proxy/spend_tracking/cloudzero_endpoints.py index 1f551d5ffea..71f4a8af111 100644 --- a/litellm/proxy/spend_tracking/cloudzero_endpoints.py +++ b/litellm/proxy/spend_tracking/cloudzero_endpoints.py @@ -6,11 +6,12 @@ from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.repositories.config_repository import ConfigRepository from litellm.types.proxy.cloudzero_endpoints import ( CloudZeroExportRequest, CloudZeroExportResponse, @@ -53,7 +54,7 @@ async def _set_cloudzero_settings(api_key: str, connection_id: str, timezone: st "timezone": timezone, } - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "cloudzero_settings"}, data={ "create": { @@ -80,7 +81,7 @@ async def _get_cloudzero_settings(): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - cloudzero_config = await prisma_client.db.litellm_config.find_first( + cloudzero_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "cloudzero_settings"} ) if cloudzero_config is None or cloudzero_config.param_value is None: @@ -282,7 +283,7 @@ async def is_cloudzero_setup_in_db() -> bool: return False # Check for CloudZero settings in database - cloudzero_config = await prisma_client.db.litellm_config.find_first( + cloudzero_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "cloudzero_settings"} ) @@ -548,7 +549,7 @@ async def delete_cloudzero_settings( ) # Check if CloudZero settings exist - cloudzero_config = await prisma_client.db.litellm_config.find_first( + cloudzero_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "cloudzero_settings"} ) @@ -560,7 +561,7 @@ async def delete_cloudzero_settings( # Delete only the CloudZero settings entry # This uses a specific where clause to target only the cloudzero_settings row - await prisma_client.db.litellm_config.delete( + await ConfigRepository(prisma_client).table.delete( where={"param_name": "cloudzero_settings"} ) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index ca5e0473659..aa85be6671a 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -21,6 +21,11 @@ get_spend_by_team_and_customer, ) from litellm.proxy.utils import handle_exception_on_proxy +from litellm.repositories.table_repositories import SpendLogsRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) if TYPE_CHECKING: from litellm.proxy.proxy_server import PrismaClient @@ -2010,7 +2015,7 @@ def parse_date(date_str: str) -> datetime: order_direction = (sort_order or "desc").lower() # Get total count of records - total_records = await prisma_client.db.litellm_spendlogs.count( + total_records = await SpendLogsRepository(prisma_client).table.count( where=where_conditions, ) @@ -2374,7 +2379,7 @@ async def view_spend_logs( # noqa: PLR0915 # Check if user wants unsummarized data if not summarize: # Return filtered individual log entries (similar to UI endpoint) - data = await prisma_client.db.litellm_spendlogs.find_many( + data = await SpendLogsRepository(prisma_client).table.find_many( where=filter_query, # type: ignore order={ "startTime": "desc", @@ -2384,7 +2389,7 @@ async def view_spend_logs( # noqa: PLR0915 # Legacy behavior: return summarized data (when summarize=true) # SQL query - response = await prisma_client.db.litellm_spendlogs.group_by( + response = await SpendLogsRepository(prisma_client).table.group_by( by=["api_key", "user", "model", "startTime"], where=filter_query, # type: ignore sum={ @@ -2462,7 +2467,7 @@ async def view_spend_logs( # noqa: PLR0915 ) return spend_logs - data = await prisma_client.db.litellm_spendlogs.find_many( + data = await SpendLogsRepository(prisma_client).table.find_many( where=scoped_filter, # type: ignore order={"startTime": "desc"}, ) @@ -2514,10 +2519,10 @@ async def global_spend_reset(): code=status.HTTP_401_UNAUTHORIZED, ) - await prisma_client.db.litellm_verificationtoken.update_many( + await VerificationTokenRepository(prisma_client).table.update_many( data={"spend": 0.0}, where={} ) - await prisma_client.db.litellm_teamtable.update_many(data={"spend": 0.0}, where={}) + await TeamRepository(prisma_client).table.update_many(data={"spend": 0.0}, where={}) return { "message": "Spend for all API Keys and Teams reset successfully", @@ -3384,7 +3389,7 @@ async def ui_view_session_spend_logs( skip = (page - 1) * page_size # Get total count for pagination metadata - total_records = await prisma_client.db.litellm_spendlogs.count( + total_records = await SpendLogsRepository(prisma_client).table.count( where=where_conditions ) @@ -3400,7 +3405,7 @@ async def ui_view_session_spend_logs( session_id, status, mcp_namespaced_tool_name, agent_id FROM "LiteLLM_SpendLogs" WHERE session_id = $1 - ORDER BY "startTime" ASC + ORDER BY "startTime" DESC LIMIT $2 OFFSET $3 """ result = await prisma_client.db.query_raw( @@ -3485,7 +3490,7 @@ async def _build_ui_spend_logs_response( # is bounded by page_size (typically 25-50 distinct session IDs). # If performance degrades at scale, consider short-lived caching or # folding the count into the main query via a window function. - counts = await prisma_client.db.litellm_spendlogs.group_by( + counts = await SpendLogsRepository(prisma_client).table.group_by( by=["session_id"], where={"session_id": {"in": session_ids}}, count={"session_id": True}, @@ -3572,7 +3577,7 @@ async def _can_team_member_view_log( if team_id is None: return False - team_row = await prisma_client.db.litellm_teamtable.find_unique( + team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) if team_row is None: @@ -3614,7 +3619,7 @@ async def _assert_user_can_view_request_id( permitted teams (admin or ``/spend/logs`` permission). Raises HTTP 403 if not. """ - row = await prisma_client.db.litellm_spendlogs.find_unique( + row = await SpendLogsRepository(prisma_client).table.find_unique( where={"request_id": request_id}, include=None, ) @@ -3669,7 +3674,7 @@ async def _get_permitted_team_ids_for_spend_logs( if user_obj is None or not user_obj.teams: return [] - team_rows = await prisma_client.db.litellm_teamtable.find_many( + team_rows = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": user_obj.teams}} ) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index 60e54d005b3..1dde31b54cb 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -1,17 +1,18 @@ import json -import litellm from fastapi import APIRouter, Depends, HTTPException +import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.repositories.config_repository import ConfigRepository from litellm.types.proxy.vantage_endpoints import ( VantageDryRunRequest, VantageExportRequest, @@ -60,7 +61,7 @@ async def _set_vantage_settings(api_key: str, integration_token: str, base_url: "base_url": base_url, } - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": VANTAGE_SETTINGS_PARAM_NAME}, data={ "create": { @@ -82,7 +83,7 @@ async def _get_vantage_settings(): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - vantage_config = await prisma_client.db.litellm_config.find_first( + vantage_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} ) if vantage_config is None or vantage_config.param_value is None: @@ -265,7 +266,7 @@ async def is_vantage_setup_in_db() -> bool: if prisma_client is None: return False - vantage_config = await prisma_client.db.litellm_config.find_first( + vantage_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} ) @@ -553,7 +554,7 @@ async def delete_vantage_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - vantage_config = await prisma_client.db.litellm_config.find_first( + vantage_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} ) @@ -563,7 +564,7 @@ async def delete_vantage_settings( detail={"error": "Vantage settings not found"}, ) - await prisma_client.db.litellm_config.delete( + await ConfigRepository(prisma_client).table.delete( where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} ) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 07e2ca71950..3a609eec127 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -12,6 +12,12 @@ from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.table_repositories import ( + DailyTagSpendRepository, + SSOConfigRepository, + UISettingsRepository, +) from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, InProductNudgeResponse, @@ -172,6 +178,11 @@ class UISettings(BaseModel): description="If true, org admins cannot generate API keys via /key/generate.", ) + disable_ui_nudges: bool = Field( + default=False, + description="If true, suppresses in-product UI nudges (survey and Claude Code feedback popups) for all users.", + ) + class UISettingsResponse(SettingsResponse): """Response model for UI settings""" @@ -195,6 +206,7 @@ class UISettingsResponse(SettingsResponse): "scope_user_search_to_org", "disable_custom_api_keys", "disable_key_generate_for_org_admin", + "disable_ui_nudges", } # Flags that must be synced from the persisted UISettings into @@ -665,7 +677,7 @@ async def get_sso_settings(): ) # Get SSO config from dedicated table - sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique( where={"id": "sso_config"} ) @@ -836,7 +848,7 @@ async def update_sso_settings(sso_config: SSOConfig): ) # Save to dedicated SSO table - await prisma_client.db.litellm_ssoconfig.upsert( + await SSOConfigRepository(prisma_client).table.upsert( where={"id": "sso_config"}, data={ "create": { @@ -851,7 +863,7 @@ async def update_sso_settings(sso_config: SSOConfig): # Remove SSO-related env vars from config.environment_variables try: - env_var_entry = await prisma_client.db.litellm_config.find_unique( + env_var_entry = await ConfigRepository(prisma_client).table.find_unique( where={"param_name": "environment_variables"} ) @@ -872,7 +884,7 @@ async def update_sso_settings(sso_config: SSOConfig): if key not in env_vars_to_remove } - await prisma_client.db.litellm_config.update( + await ConfigRepository(prisma_client).table.update( where={"param_name": "environment_variables"}, data={ "param_value": json.dumps(filtered_env_vars, default=str), @@ -1123,7 +1135,7 @@ async def get_in_product_nudges(): detail={"error": "Database not connected. Please connect a database."}, ) - db_record = await prisma_client.db.litellm_dailytagspend.find_first( + db_record = await DailyTagSpendRepository(prisma_client).table.find_first( where={"tag": "User-Agent: claude-cli"} ) @@ -1155,7 +1167,7 @@ async def get_ui_settings_cached() -> Dict[str, Any]: if prisma_client is None: return {} - db_record = await prisma_client.db.litellm_uisettings.find_unique( + db_record = await UISettingsRepository(prisma_client).table.find_unique( where={"id": "ui_settings"} ) ui_settings: Dict[str, Any] = {} @@ -1196,7 +1208,7 @@ async def get_ui_settings(): ui_settings: Dict[str, Any] = {} - db_record = await prisma_client.db.litellm_uisettings.find_unique( + db_record = await UISettingsRepository(prisma_client).table.find_unique( where={"id": "ui_settings"} ) @@ -1309,7 +1321,7 @@ async def update_ui_settings( # Merge with existing persisted settings so a partial PATCH doesn't # overwrite fields the caller didn't send. existing: dict = {} - db_existing = await prisma_client.db.litellm_uisettings.find_unique( + db_existing = await UISettingsRepository(prisma_client).table.find_unique( where={"id": "ui_settings"} ) if db_existing and db_existing.ui_settings: @@ -1318,7 +1330,7 @@ async def update_ui_settings( ui_settings = {**existing, **incoming} - await prisma_client.db.litellm_uisettings.upsert( + await UISettingsRepository(prisma_client).table.upsert( where={"id": "ui_settings"}, data={ "create": { diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e77e24c9e71..ebd5b5d90cd 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -9,10 +9,10 @@ import threading import time import traceback +from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText -from dataclasses import dataclass, field from typing import ( TYPE_CHECKING, Any, @@ -139,6 +139,19 @@ ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.table_repositories import ( + EndUserRepository, + HealthCheckRepository, + SpendLogsRepository, + UserNotificationsRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES from litellm.types.mcp import ( @@ -2483,7 +2496,8 @@ def _build_litellm_call_info(data: dict, response: Any) -> Dict[str, Any]: ) return { - "custom_llm_provider": hidden_params.get("custom_llm_provider"), + "custom_llm_provider": hidden_params.get("custom_llm_provider") + or getattr(response, "custom_llm_provider", None), "model_info": model_info, "api_base": hidden_params.get("api_base"), "model_id": hidden_params.get("model_id"), @@ -2831,7 +2845,7 @@ async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> if not param_names: return try: - rows = await prisma_client.db.litellm_config.find_many( + rows = await ConfigRepository(prisma_client).table.find_many( where={"param_name": {"in": param_names}} # type: ignore ) except Exception as e: @@ -3194,15 +3208,15 @@ async def get_generic_data( async def _do_query(): if table_name == "users": - return await self.db.litellm_usertable.find_first( + return await UserRepository(self).table.find_first( where={key: value} # type: ignore ) elif table_name == "keys": - return await self.db.litellm_verificationtoken.find_first( # type: ignore + return await VerificationTokenRepository(self).table.find_first( # type: ignore where={key: value} # type: ignore ) elif table_name == "config": - return await self.db.litellm_config.find_first( # type: ignore + return await ConfigRepository(self).table.find_first( # type: ignore where={key: value} # type: ignore ) elif table_name == "spend": @@ -3239,40 +3253,49 @@ async def _query_first_with_cached_plan_fallback( self, sql_query: str, *args ) -> Optional[dict]: """ - Execute a query with automatic fallback for PostgreSQL cached plan errors. - - This handles the "cached plan must not change result type" error that occurs - during rolling deployments when schema changes are applied while old pods - still have cached query plans expecting the old schema. - - Args: - sql_query: SQL query string to execute - - Returns: - Query result or None - - Raises: - Original exception if not a cached plan error + Execute a query, recovering once from PostgreSQL's "cached plan must not + change result type" error. + + That error surfaces during rolling deployments when a schema change + invalidates the prepared-statement plans that pooled connections still + hold. Clearing only the server-side plans with DEALLOCATE ALL makes + things worse: Prisma's query engine keeps a per-connection client-side + cache of prepared-statement names, so once the server drops a plan the + engine re-sends a name PostgreSQL no longer recognizes and the + connection breaks with `prepared statement "sN" does not exist`. With a + small pool that connection stays poisoned and every auth lookup fails. + + Recreating the Prisma client kills the engine subprocess and drops the + server-side plans and the engine's client-side name cache together, so + the retried query is prepared fresh. We reconnect through + `attempt_db_reconnect`, which is singleflight: when a schema change + poisons every pooled connection at once, the first cached-plan error + recreates the client and the concurrent waiters reuse that single + recreate instead of racing to kill each other's fresh engine. We then + retry the identical query exactly once. + + The retry reuses the original query byte-for-byte. Mutating the SQL + (e.g. injecting a unique comment) would defeat PostgreSQL's plan cache, + forcing a fresh plan on every request and pegging the database CPU. + + If the reconnect is skipped because a recent reconnect is still within + its cooldown, the retry runs against the same connection and may fail + again; the get_data backoff decorator re-runs the lookup and a later + attempt reconnects once the cooldown elapses. """ try: return await self.db.query_first(sql_query, *args) except Exception as e: - error_str = str(e) - if "cached plan must not change result type" in error_str: - # Force PostgreSQL to re-plan by invalidating the cache - # Add a unique comment to make the query different - sql_query_retry = sql_query.replace( - "SELECT", - f"SELECT /* cache_invalidated_{int(time.time() * 1000)} */", - ) - verbose_proxy_logger.warning( - "PostgreSQL cached plan error detected for token lookup, " - "retrying with fresh plan. This may occur during rolling deployments " - "when schema changes are applied." - ) - return await self.db.query_first(sql_query_retry, *args) - else: + if "cached plan must not change result type" not in str(e): raise + verbose_proxy_logger.warning( + "PostgreSQL cached plan error detected for token lookup; " + "recreating the database connection and retrying with the same " + "query. This may occur during rolling deployments when schema " + "changes are applied." + ) + await self.attempt_db_reconnect(reason="postgres_cached_plan_error") + return await self.db.query_first(sql_query, *args) @backoff.on_exception( backoff.expo, @@ -3336,7 +3359,9 @@ async def get_data( # noqa: PLR0915 status_code=400, detail={"error": f"No token passed in. Token={token}"}, ) - response = await self.db.litellm_verificationtoken.find_unique( + response = await VerificationTokenRepository( + self + ).table.find_unique( where={"token": hashed_token}, # type: ignore include={"litellm_budget_table": True}, ) @@ -3353,7 +3378,7 @@ async def get_data( # noqa: PLR0915 detail=f"Authentication Error: invalid user key - user key does not exist in db. User Key={token}", ) elif query_type == "find_all" and user_id is not None: - response = await self.db.litellm_verificationtoken.find_many( + response = await VerificationTokenRepository(self).table.find_many( where={"user_id": user_id}, include={"litellm_budget_table": True}, ) @@ -3362,7 +3387,7 @@ async def get_data( # noqa: PLR0915 if isinstance(r.expires, datetime): r.expires = r.expires.isoformat() elif query_type == "find_all" and team_id is not None: - response = await self.db.litellm_verificationtoken.find_many( + response = await VerificationTokenRepository(self).table.find_many( where={"team_id": team_id}, include={"litellm_budget_table": True}, ) @@ -3375,7 +3400,7 @@ async def get_data( # noqa: PLR0915 and expires is not None and reset_at is not None ): - response = await self.db.litellm_verificationtoken.find_many( + response = await VerificationTokenRepository(self).table.find_many( where={ # type: ignore "OR": [ {"expires": None}, @@ -3405,7 +3430,7 @@ async def get_data( # noqa: PLR0915 else: hashed_tokens.append(t) where_filter["token"]["in"] = hashed_tokens - response = await self.db.litellm_verificationtoken.find_many( + response = await VerificationTokenRepository(self).table.find_many( order={"spend": "desc"}, where=where_filter, # type: ignore include={"litellm_budget_table": True}, @@ -3425,28 +3450,28 @@ async def get_data( # noqa: PLR0915 if key_val is None: key_val = {"user_id": user_id} - response = await self.db.litellm_usertable.find_unique( # type: ignore + response = await UserRepository(self).table.find_unique( # type: ignore where=key_val, # type: ignore include={"organization_memberships": True}, ) elif query_type == "find_all" and key_val is not None: - response = await self.db.litellm_usertable.find_many( + response = await UserRepository(self).table.find_many( where=key_val # type: ignore ) # type: ignore elif query_type == "find_all" and reset_at is not None: - response = await self.db.litellm_usertable.find_many( + response = await UserRepository(self).table.find_many( where={ # type: ignore "budget_reset_at": {"lt": reset_at}, } ) elif query_type == "find_all" and user_id_list is not None: - response = await self.db.litellm_usertable.find_many( + response = await UserRepository(self).table.find_many( where={"user_id": {"in": user_id_list}} ) elif query_type == "find_all": if expires is not None: - response = await self.db.litellm_usertable.find_many( # type: ignore + response = await UserRepository(self).table.find_many( # type: ignore order={"spend": "desc"}, where={ # type: ignore "OR": [ @@ -3478,26 +3503,26 @@ async def get_data( # noqa: PLR0915 ) if key_val is not None: if query_type == "find_unique": - response = await self.db.litellm_spendlogs.find_unique( # type: ignore + response = await SpendLogsRepository(self).table.find_unique( # type: ignore where={ # type: ignore key_val["key"]: key_val["value"], # type: ignore } ) elif query_type == "find_all": - response = await self.db.litellm_spendlogs.find_many( # type: ignore + response = await SpendLogsRepository(self).table.find_many( # type: ignore where={ key_val["key"]: key_val["value"], # type: ignore } ) return response else: - response = await self.db.litellm_spendlogs.find_many( # type: ignore + response = await SpendLogsRepository(self).table.find_many( # type: ignore order={"startTime": "desc"}, ) return response elif table_name == "budget" and reset_at is not None: if query_type == "find_all": - response = await self.db.litellm_budgettable.find_many( + response = await BudgetRepository(self).table.find_many( where={ # type: ignore "OR": [ { @@ -3514,45 +3539,45 @@ async def get_data( # noqa: PLR0915 elif table_name == "enduser" and budget_id_list is not None: if query_type == "find_all": - response = await self.db.litellm_endusertable.find_many( + response = await EndUserRepository(self).table.find_many( where={"budget_id": {"in": budget_id_list}} ) return response elif table_name == "team": if query_type == "find_unique": - response = await self.db.litellm_teamtable.find_unique( + response = await TeamRepository(self).table.find_unique( where={"team_id": team_id}, # type: ignore include={"litellm_model_table": True}, # type: ignore ) elif query_type == "find_all" and reset_at is not None: - response = await self.db.litellm_teamtable.find_many( + response = await TeamRepository(self).table.find_many( where={ # type: ignore "budget_reset_at": {"lt": reset_at}, } ) elif query_type == "find_all" and user_id is not None: - response = await self.db.litellm_teamtable.find_many( + response = await TeamRepository(self).table.find_many( where={ "members": {"has": user_id}, }, include={"litellm_budget_table": True}, ) elif query_type == "find_all" and team_id_list is not None: - response = await self.db.litellm_teamtable.find_many( + response = await TeamRepository(self).table.find_many( where={"team_id": {"in": team_id_list}} ) elif query_type == "find_all" and team_id_list is None: - response = await self.db.litellm_teamtable.find_many( + response = await TeamRepository(self).table.find_many( take=MAX_TEAM_LIST_LIMIT ) return response elif table_name == "user_notification": if query_type == "find_unique": - response = await self.db.litellm_usernotifications.find_unique( # type: ignore + response = await UserNotificationsRepository(self).table.find_unique( # type: ignore where={"user_id": user_id} # type: ignore ) elif query_type == "find_all": - response = await self.db.litellm_usernotifications.find_many() # type: ignore + response = await UserNotificationsRepository(self).table.find_many() # type: ignore return response elif table_name == "combined_view": # check if plain text or hash @@ -3744,7 +3769,7 @@ async def insert_data( # noqa: PLR0915 print_verbose( "PrismaClient: Before upsert into litellm_verificationtoken" ) - new_verification_token = await self.db.litellm_verificationtoken.upsert( # type: ignore + new_verification_token = await VerificationTokenRepository(self).table.upsert( # type: ignore where={ "token": hashed_token, }, @@ -3759,7 +3784,7 @@ async def insert_data( # noqa: PLR0915 elif table_name == "user": db_data = self.jsonify_object(data=data) try: - new_user_row = await self.db.litellm_usertable.upsert( + new_user_row = await UserRepository(self).table.upsert( where={"user_id": data["user_id"]}, data={ "create": {**db_data}, # type: ignore @@ -3782,7 +3807,7 @@ async def insert_data( # noqa: PLR0915 return new_user_row elif table_name == "team": db_data = self.jsonify_team_object(db_data=data) - new_team_row = await self.db.litellm_teamtable.upsert( + new_team_row = await TeamRepository(self).table.upsert( where={"team_id": data["team_id"]}, data={ "create": {**db_data}, # type: ignore @@ -3804,7 +3829,7 @@ async def insert_data( # noqa: PLR0915 for k, v in data.items(): updated_data = v updated_data = json.dumps(updated_data) - updated_table_row = self.db.litellm_config.upsert( + updated_table_row = ConfigRepository(self).table.upsert( where={"param_name": k}, # type: ignore data={ "create": {"param_name": k, "param_value": updated_data}, # type: ignore @@ -3820,7 +3845,7 @@ async def insert_data( # noqa: PLR0915 verbose_proxy_logger.info("Data Inserted into Config Table") elif table_name == "spend": db_data = self.jsonify_object(data=data) - new_spend_row = await self.db.litellm_spendlogs.upsert( + new_spend_row = await SpendLogsRepository(self).table.upsert( where={"request_id": data["request_id"]}, data={ "create": {**db_data}, # type: ignore @@ -3831,14 +3856,14 @@ async def insert_data( # noqa: PLR0915 return new_spend_row elif table_name == "user_notification": db_data = self.jsonify_object(data=data) - new_user_notification_row = ( - await self.db.litellm_usernotifications.upsert( # type: ignore - where={"request_id": data["request_id"]}, - data={ - "create": {**db_data}, # type: ignore - "update": {}, # don't do anything if it already exists - }, - ) + new_user_notification_row = await UserNotificationsRepository( + self + ).table.upsert( # type: ignore + where={"request_id": data["request_id"]}, + data={ + "create": {**db_data}, # type: ignore + "update": {}, # don't do anything if it already exists + }, ) verbose_proxy_logger.info("Data Inserted into Model Request Table") return new_user_notification_row @@ -3899,7 +3924,7 @@ async def update_data( # noqa: PLR0915 # check if plain text or hash token = _hash_token_if_needed(token=token) db_data["token"] = token - response = await self.db.litellm_verificationtoken.update( + response = await VerificationTokenRepository(self).table.update( where={"token": token}, # type: ignore data={**db_data}, # type: ignore ) @@ -3930,7 +3955,7 @@ async def update_data( # noqa: PLR0915 update_key_values = update_key_values_custom_query else: update_key_values = db_data - update_user_row = await self.db.litellm_usertable.upsert( + update_user_row = await UserRepository(self).table.upsert( where={"user_id": user_id}, # type: ignore data={ "create": {**db_data}, # type: ignore @@ -3971,7 +3996,7 @@ async def update_data( # noqa: PLR0915 update_key_values["members_with_roles"] = json.dumps( update_key_values["members_with_roles"] ) - update_team_row = await self.db.litellm_teamtable.upsert( + update_team_row = await TeamRepository(self).table.upsert( where={"team_id": team_id}, # type: ignore data={ "create": {**db_data}, # type: ignore @@ -4196,7 +4221,9 @@ async def delete_data( else: filter_query = {"token": {"in": hashed_tokens}} - deleted_tokens = await self.db.litellm_verificationtoken.delete_many( + deleted_tokens = await VerificationTokenRepository( + self + ).table.delete_many( where=filter_query # type: ignore ) verbose_proxy_logger.debug("deleted_tokens: %s", deleted_tokens) @@ -4207,7 +4234,7 @@ async def delete_data( and isinstance(team_id_list, List) ): # admin only endpoint -> `/team/delete` - await self.db.litellm_teamtable.delete_many( + await TeamRepository(self).table.delete_many( where={"team_id": {"in": team_id_list}} ) return {"deleted_teams": team_id_list} @@ -4217,7 +4244,7 @@ async def delete_data( and isinstance(team_id_list, List) ): # admin only endpoint -> `/team/delete` - await self.db.litellm_verificationtoken.delete_many( + await VerificationTokenRepository(self).table.delete_many( where={"team_id": {"in": team_id_list}} ) except Exception as e: @@ -5024,7 +5051,9 @@ async def save_health_check_result( ) verbose_proxy_logger.debug(f"Saving health check data: {health_check_data}") - return await self.db.litellm_healthchecktable.create(data=health_check_data) + return await HealthCheckRepository(self).table.create( + data=health_check_data + ) except Exception as e: verbose_proxy_logger.error( @@ -5049,7 +5078,7 @@ async def get_health_check_history( if status_filter: where_clause["status"] = status_filter - results = await self.db.litellm_healthchecktable.find_many( + results = await HealthCheckRepository(self).table.find_many( where=where_clause, order={"checked_at": "desc"}, take=limit, @@ -5068,7 +5097,7 @@ async def get_all_latest_health_checks(self): (via Prisma ``distinct`` + ``order``) so we never load the full history into memory. """ try: - return await self.db.litellm_healthchecktable.find_many( + return await HealthCheckRepository(self).table.find_many( distinct=["model_id", "model_name"], order=[ {"model_id": "asc"}, @@ -5228,7 +5257,7 @@ async def migrate_passwords_to_scrypt_async(prisma_client) -> str: are left alone (they migrate on next login via the SHA256 fallback). Skips quickly if no plaintext passwords exist. """ - all_with_pw = await prisma_client.db.litellm_usertable.find_many( + all_with_pw = await UserRepository(prisma_client).table.find_many( where={"password": {"not": None}}, ) @@ -5246,7 +5275,7 @@ def _is_sha256_hex(s: str) -> bool: return "No plaintext passwords found" for user in plaintext_users: - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": user.user_id}, data={"password": hash_password(user.password)}, ) @@ -5370,7 +5399,7 @@ async def update_spend_logs( prisma_client.jsonify_object({**entry}) for entry in batch ] - await prisma_client.db.litellm_spendlogs.create_many( + await SpendLogsRepository(prisma_client).table.create_many( data=batch_with_dates, skip_duplicates=True ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index b3bdfecbe55..9c2d3050346 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -1,6 +1,7 @@ from typing import Any, Dict, Optional from fastapi import APIRouter, Depends, HTTPException, Request, Response + from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) @@ -16,6 +17,7 @@ assert_user_can_access_vector_store, get_litellm_managed_vector_store, ) +from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository from litellm.types.vector_stores import IndexCreateRequest router = APIRouter() @@ -587,11 +589,9 @@ async def index_create( detail=CommonProxyErrors.db_not_connected_error.value, ) ## 1. check if index already exists - existing_index = ( - await prisma_client.db.litellm_managedvectorstoreindextable.find_unique( - where={"index_name": index_create_request.index_name} - ) - ) + existing_index = await ManagedVectorStoreIndexRepository( + prisma_client + ).table.find_unique(where={"index_name": index_create_request.index_name}) ## 2. set created_by and updated_by @@ -605,7 +605,7 @@ async def index_create( index_data = index_create_request.model_dump(exclude_none=True) index_data["created_by"] = user_api_key_dict.user_id index_data["updated_by"] = user_api_key_dict.user_id - new_index = await prisma_client.db.litellm_managedvectorstoreindextable.create( + new_index = await ManagedVectorStoreIndexRepository(prisma_client).table.create( data=jsonify_object(index_data) ) diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index cbb3d927184..032a3302fdc 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -29,6 +29,8 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ManagedVectorStoresRepository from litellm.secret_managers.main import get_secret from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, @@ -122,7 +124,7 @@ async def _fetch_and_authorize_vector_store( Raises HTTPException(404) on miss and HTTPException(403) on access denial. """ - row = await prisma_client.db.litellm_managedvectorstorestable.find_unique( + row = await ManagedVectorStoresRepository(prisma_client).table.find_unique( where={"vector_store_id": vector_store_id} ) if row is None: @@ -252,7 +254,7 @@ async def _resolve_embedding_config_from_db( # Try to find model in database for model_name in model_name_candidates: try: - db_model = await prisma_client.db.litellm_proxymodeltable.find_first( + db_model = await ModelRepository(prisma_client).table.find_first( where={"model_name": model_name} ) @@ -437,11 +439,9 @@ async def create_vector_store_in_db( raise HTTPException(status_code=500, detail="Database not connected") # Check if vector store already exists - existing_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": vector_store_id} - ) - ) + existing_vector_store = await ManagedVectorStoresRepository( + prisma_client + ).table.find_unique(where={"vector_store_id": vector_store_id}) if existing_vector_store is not None: raise HTTPException( status_code=400, @@ -487,7 +487,7 @@ async def create_vector_store_in_db( data_to_create["litellm_params"] = safe_dumps({}) # Create in database - _new_vector_store = await prisma_client.db.litellm_managedvectorstorestable.create( + _new_vector_store = await ManagedVectorStoresRepository(prisma_client).table.create( data=data_to_create ) @@ -725,11 +725,9 @@ async def delete_vector_store( memory_vector_store_exists = False vector_store_to_check = None - existing_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": data.vector_store_id} - ) - ) + existing_vector_store = await ManagedVectorStoresRepository( + prisma_client + ).table.find_unique(where={"vector_store_id": data.vector_store_id}) if existing_vector_store is not None: db_vector_store_exists = True vector_store_to_check = LiteLLM_ManagedVectorStore( @@ -764,7 +762,7 @@ async def delete_vector_store( # Delete from database if exists if db_vector_store_exists: - await prisma_client.db.litellm_managedvectorstorestable.delete( + await ManagedVectorStoresRepository(prisma_client).table.delete( where={"vector_store_id": data.vector_store_id} ) @@ -921,7 +919,7 @@ async def update_vector_store( update_data["litellm_params"] = safe_dumps(litellm_params_dict) # Update in database - updated = await prisma_client.db.litellm_managedvectorstorestable.update( + updated = await ManagedVectorStoresRepository(prisma_client).table.update( where={"vector_store_id": vector_store_id}, data=update_data, ) diff --git a/litellm/repositories/__init__.py b/litellm/repositories/__init__.py new file mode 100644 index 00000000000..4451f0865da --- /dev/null +++ b/litellm/repositories/__init__.py @@ -0,0 +1,127 @@ +""" +Repository classes for database operations. +""" + +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.credentials_repository import CredentialsRepository +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.object_permission_repository import ( + ObjectPermissionRepository, +) +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.project_repository import ProjectRepository +from litellm.repositories.table_repositories import ( + AccessGroupRepository, + AdaptiveRouterSessionRepository, + AdaptiveRouterStateRepository, + AgentsRepository, + AuditLogRepository, + CacheConfigRepository, + ClaudeCodePluginRepository, + ConfigOverridesRepository, + DailyGuardrailMetricsRepository, + DailyPolicyMetricsRepository, + DailyTagSpendRepository, + DeletedTeamRepository, + DeletedVerificationTokenRepository, + DeprecatedVerificationTokenRepository, + EndUserRepository, + GuardrailsRepository, + HealthCheckRepository, + InvitationLinkRepository, + JWTKeyMappingRepository, + ManagedFileRepository, + ManagedObjectRepository, + ManagedVectorStoreIndexRepository, + ManagedVectorStoresRepository, + MCPServerRepository, + MCPToolsetRepository, + MCPUserCredentialsRepository, + MemoryRepository, + ModelTableRepository, + OrganizationMembershipRepository, + PolicyAttachmentRepository, + PolicyRepository, + PrismaTableRepository, + PromptRepository, + SearchToolsRepository, + SkillsRepository, + SpendLogGuardrailIndexRepository, + SpendLogsRepository, + SpendLogToolIndexRepository, + SSOConfigRepository, + TagRepository, + TeamMembershipRepository, + ToolRepository, + UISettingsRepository, + UserNotificationsRepository, + WorkflowEventRepository, + WorkflowMessageRepository, + WorkflowRunRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) + +__all__ = [ + "PrismaTableRepository", + "PolicyRepository", + "AgentsRepository", + "GuardrailsRepository", + "MCPServerRepository", + "ManagedObjectRepository", + "OrganizationMembershipRepository", + "SpendLogsRepository", + "ClaudeCodePluginRepository", + "TeamMembershipRepository", + "EndUserRepository", + "ManagedVectorStoresRepository", + "MCPUserCredentialsRepository", + "PromptRepository", + "TagRepository", + "InvitationLinkRepository", + "JWTKeyMappingRepository", + "ManagedFileRepository", + "MemoryRepository", + "SearchToolsRepository", + "ConfigOverridesRepository", + "MCPToolsetRepository", + "ToolRepository", + "DeletedVerificationTokenRepository", + "WorkflowRunRepository", + "ModelTableRepository", + "AccessGroupRepository", + "SSOConfigRepository", + "UISettingsRepository", + "DailyGuardrailMetricsRepository", + "PolicyAttachmentRepository", + "DeletedTeamRepository", + "SkillsRepository", + "CacheConfigRepository", + "ManagedVectorStoreIndexRepository", + "WorkflowMessageRepository", + "DailyTagSpendRepository", + "SpendLogToolIndexRepository", + "SpendLogGuardrailIndexRepository", + "UserNotificationsRepository", + "HealthCheckRepository", + "DeprecatedVerificationTokenRepository", + "WorkflowEventRepository", + "DailyPolicyMetricsRepository", + "AdaptiveRouterStateRepository", + "AuditLogRepository", + "AdaptiveRouterSessionRepository", + "BudgetRepository", + "ConfigRepository", + "CredentialsRepository", + "ModelRepository", + "ObjectPermissionRepository", + "OrganizationRepository", + "ProjectRepository", + "TeamRepository", + "UserRepository", + "VerificationTokenRepository", +] diff --git a/litellm/repositories/base_repository.py b/litellm/repositories/base_repository.py new file mode 100644 index 00000000000..a25620c7b4d --- /dev/null +++ b/litellm/repositories/base_repository.py @@ -0,0 +1,117 @@ +""" +Base repository class with common functionality. +""" + +from abc import ABC, abstractmethod +from typing import Any, Dict, Generic, List, Optional, Type, TypeVar + +from pydantic import BaseModel + +T = TypeVar("T", bound=BaseModel) + + +def _record_to_dict(record: Any) -> Dict[str, Any]: + if isinstance(record, dict): + return record + if hasattr(record, "model_dump") and callable(record.model_dump): + return record.model_dump() + if hasattr(record, "dict") and callable(record.dict): + return record.dict() + return dict(record) + + +class BaseRepository(ABC, Generic[T]): + """Abstract base class for all repositories.""" + + def __init__(self, prisma_client: Any): + self._prisma_client = prisma_client + + @property + def prisma_client(self) -> Any: + if self._prisma_client is None: + raise RuntimeError( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + return self._prisma_client + + @property + @abstractmethod + def table(self) -> Any: + """Return the Prisma table for this repository.""" + ... + + @property + @abstractmethod + def model_class(self) -> Type[T]: + """Return the domain model class for this repository.""" + ... + + def _to_model(self, record: Any) -> Optional[T]: + """Convert a database record to a domain model.""" + if record is None: + return None + return self.model_class(**_record_to_dict(record)) + + def _to_model_list(self, records: List[Any]) -> List[T]: + """Convert a list of database records to domain models.""" + result: List[T] = [] + for r in records: + if r is not None: + model = self._to_model(r) + if model is not None: + result.append(model) + return result + + async def find_by_id(self, id_value: str, id_field: str = "id") -> Optional[T]: + """Find a record by its primary key.""" + record = await self.table.find_unique(where={id_field: id_value}) + return self._to_model(record) + + async def find_many( + self, + where: Optional[Dict[str, Any]] = None, + skip: Optional[int] = None, + take: Optional[int] = None, + order: Optional[Dict[str, str]] = None, + ) -> List[T]: + """Find multiple records matching the criteria.""" + kwargs: Dict[str, Any] = {} + if where: + kwargs["where"] = where + if skip is not None: + kwargs["skip"] = skip + if take is not None: + kwargs["take"] = take + if order: + kwargs["order"] = order + + records = await self.table.find_many(**kwargs) + return self._to_model_list(records) + + async def create(self, data: Dict[str, Any]) -> T: + """Create a new record.""" + record = await self.table.create(data=data) + model = self._to_model(record) + assert model is not None + return model + + async def update( + self, id_value: str, data: Dict[str, Any], id_field: str = "id" + ) -> Optional[T]: + """Update an existing record.""" + record = await self.table.update(where={id_field: id_value}, data=data) + return self._to_model(record) + + async def delete(self, id_value: str, id_field: str = "id") -> Optional[T]: + """Delete a record by its primary key.""" + record = await self.table.delete(where={id_field: id_value}) + return self._to_model(record) + + async def count(self, where: Optional[Dict[str, Any]] = None) -> int: + """Count records matching the criteria.""" + return await self.table.count(where=where) + + async def exists(self, id_value: str, id_field: str = "id") -> bool: + """Check if a record exists.""" + record = await self.table.find_unique(where={id_field: id_value}) + return record is not None diff --git a/litellm/repositories/budget_repository.py b/litellm/repositories/budget_repository.py new file mode 100644 index 00000000000..5947701fb4e --- /dev/null +++ b/litellm/repositories/budget_repository.py @@ -0,0 +1,99 @@ +""" +Budget repository for database operations on LiteLLM_BudgetTable. +""" + +from typing import Any, Dict, List, Optional, Type + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.repositories.base_repository import BaseRepository + + +class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]): + """Repository for budget database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_budgettable + + @property + def model_class(self) -> Type[LiteLLM_BudgetTable]: + return LiteLLM_BudgetTable + + async def find_by_id( + self, budget_id: str, id_field: str = "budget_id" + ) -> Optional[LiteLLM_BudgetTable]: + return await super().find_by_id(budget_id, id_field) + + async def create_budget( + self, + created_by: str, + max_budget: Optional[float] = None, + soft_budget: Optional[float] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + model_max_budget: Optional[Dict[str, Any]] = None, + budget_duration: Optional[str] = None, + allowed_models: Optional[List[str]] = None, + ) -> LiteLLM_BudgetTable: + """Create a new budget record.""" + data: Dict[str, Any] = { + "created_by": created_by, + "updated_by": created_by, + } + if max_budget is not None: + data["max_budget"] = max_budget + if soft_budget is not None: + data["soft_budget"] = soft_budget + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if model_max_budget is not None: + data["model_max_budget"] = model_max_budget + if budget_duration is not None: + data["budget_duration"] = budget_duration + if allowed_models is not None: + data["allowed_models"] = allowed_models + + return await self.create(data) + + async def update_budget( + self, + budget_id: str, + updated_by: str, + max_budget: Optional[float] = None, + soft_budget: Optional[float] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + model_max_budget: Optional[Dict[str, Any]] = None, + budget_duration: Optional[str] = None, + allowed_models: Optional[List[str]] = None, + ) -> Optional[LiteLLM_BudgetTable]: + """Update an existing budget record.""" + data: Dict[str, Any] = {"updated_by": updated_by} + if max_budget is not None: + data["max_budget"] = max_budget + if soft_budget is not None: + data["soft_budget"] = soft_budget + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if model_max_budget is not None: + data["model_max_budget"] = model_max_budget + if budget_duration is not None: + data["budget_duration"] = budget_duration + if allowed_models is not None: + data["allowed_models"] = allowed_models + + return await self.update(budget_id, data, id_field="budget_id") + + async def delete_budget(self, budget_id: str) -> Optional[LiteLLM_BudgetTable]: + """Delete a budget record.""" + return await self.delete(budget_id, id_field="budget_id") diff --git a/litellm/repositories/config_repository.py b/litellm/repositories/config_repository.py new file mode 100644 index 00000000000..eba7ebe26ca --- /dev/null +++ b/litellm/repositories/config_repository.py @@ -0,0 +1,241 @@ +""" +Config repository for database operations on LiteLLM_Config. + +This repository handles config reconciliation between database values and +YAML configmap values. DB values override configmap values except for +None values and empty lists. +""" + +import asyncio +import copy +import json +import os +from typing import Any, Dict, List, Literal, Optional, cast + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + + +class ConfigParam: + """Simple wrapper for config parameter from DB.""" + + def __init__(self, param_name: str, param_value: Any): + self.param_name = param_name + self.param_value = param_value + + +class ConfigRepository: + """Repository for config database operations with reconciliation support.""" + + CONFIG_PARAMS = [ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + ] + + def __init__(self, prisma_client: Any): + self._prisma_client = prisma_client + + @property + def prisma_client(self) -> Any: + if self._prisma_client is None: + raise RuntimeError( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + return self._prisma_client + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_config + + async def get_param(self, param_name: str) -> Optional[ConfigParam]: + """Get a config parameter from the database.""" + record = await self.table.find_unique(where={"param_name": param_name}) + if record is None: + return None + param_value = record.param_value + if isinstance(param_value, str): + param_value = json.loads(param_value) + return ConfigParam(param_name=param_name, param_value=param_value) + + async def set_param(self, param_name: str, param_value: Any) -> ConfigParam: + """Set a config parameter in the database.""" + value_json = ( + json.dumps(param_value) if not isinstance(param_value, str) else param_value + ) + await self.table.upsert( + where={"param_name": param_name}, + data={ + "create": {"param_name": param_name, "param_value": value_json}, + "update": {"param_value": value_json}, + }, + ) + return ConfigParam(param_name=param_name, param_value=param_value) + + async def delete_param(self, param_name: str) -> bool: + """Delete a config parameter from the database.""" + try: + await self.table.delete(where={"param_name": param_name}) + return True + except Exception: + return False + + async def get_all_params(self) -> Dict[str, Any]: + """Get all config parameters from the database.""" + records = await self.table.find_many() + result = {} + for record in records: + param_value = record.param_value + if isinstance(param_value, str): + param_value = json.loads(param_value) + result[record.param_name] = param_value + return result + + def _deep_merge_dicts(self, dst: dict, src: dict) -> None: + """Deep-merge src into dst, skipping None values and empty lists from src. + + On conflicts, src (DB) wins, but empty lists are treated as "no value" + and don't overwrite the destination. + """ + stack = [(dst, src)] + while stack: + d, s = stack.pop() + for k, v in s.items(): + if v is None: + continue + if isinstance(v, list) and len(v) == 0: + continue + if isinstance(v, dict) and isinstance(d.get(k), dict): + stack.append((d[k], v)) + else: + d[k] = v + + def _decrypt_env_variables( + self, env_vars: Dict[str, Any], return_original_value: bool = True + ) -> Dict[str, str]: + """Decrypt environment variables from database.""" + decrypted: Dict[str, str] = {} + for key, value in env_vars.items(): + if isinstance(value, str): + decrypted_value = decrypt_value_helper( + value=value, + key=key, + exception_type="debug", + return_original_value=return_original_value, + ) + if decrypted_value is not None: + decrypted[key] = decrypted_value + else: + decrypted[key] = str(value) + return decrypted + + def _normalize_env_variable_keys(self, env_vars: Dict[str, str]) -> Dict[str, str]: + """Normalize env variable keys to include both original and uppercase versions.""" + normalized: Dict[str, str] = {} + for key, value in env_vars.items(): + normalized[key] = value + upper_key = key.upper() + normalized[upper_key] = value + return normalized + + def _update_config_fields( + self, + current_config: dict, + param_name: Literal[ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + ], + db_param_value: Any, + ) -> dict: + """Update config fields with DB values, handling the merge strategy.""" + if param_name == "environment_variables": + decrypted_env_vars = self._decrypt_env_variables( + db_param_value, return_original_value=True + ) + merged_env_vars = self._normalize_env_variable_keys(decrypted_env_vars) + for env_key, value in merged_env_vars.items(): + os.environ[env_key] = value + + current_config.setdefault("environment_variables", {}).update( + merged_env_vars + ) + return current_config + + if param_name not in current_config: + current_config[param_name] = db_param_value + return current_config + + if isinstance(current_config[param_name], dict) and isinstance( + db_param_value, dict + ): + self._deep_merge_dicts(current_config[param_name], db_param_value) + else: + current_config[param_name] = db_param_value + + return current_config + + async def reconcile_config( + self, + yaml_config: dict, + store_model_in_db: Optional[bool] = None, + ) -> dict: + """Reconcile config from YAML with database overrides. + + This is the main config reconciliation method that loads config params + from the database and merges them with the YAML config. DB values + override YAML values except for None values and empty lists. + + Args: + yaml_config: The configuration loaded from YAML file + store_model_in_db: Whether to load config from DB + + Returns: + The merged configuration with DB overrides applied + """ + if store_model_in_db is not True: + verbose_proxy_logger.info( + "'store_model_in_db' is not True, skipping db config reconciliation" + ) + return yaml_config + + tasks = [self.get_param(k) for k in self.CONFIG_PARAMS] + responses = await asyncio.gather(*tasks) + + config = copy.deepcopy(yaml_config) + for response in responses: + if response is None: + continue + + param_name = response.param_name + param_value = response.param_value + verbose_proxy_logger.debug( + f"param_name={param_name}, param_value={param_value}" + ) + + if param_name is not None and param_value is not None: + config = self._update_config_fields( + current_config=config, + param_name=cast( + Literal[ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + ], + param_name, + ), + db_param_value=param_value, + ) + + return config + + async def prefetch_params(self, param_names: List[str]) -> None: + """Prefetch config params to warm the cache. + + This can be called before reconcile_config to ensure all needed + params are loaded in a single batch. + """ + await asyncio.gather(*[self.get_param(k) for k in param_names]) diff --git a/litellm/repositories/credentials_repository.py b/litellm/repositories/credentials_repository.py new file mode 100644 index 00000000000..dd53c753307 --- /dev/null +++ b/litellm/repositories/credentials_repository.py @@ -0,0 +1,61 @@ +""" +Credentials repository for database operations on LiteLLM_CredentialsTable. + +This is the only place that talks to ``litellm_credentialstable``. Encryption of +credential values is the caller's responsibility (see ``CredentialHelperUtils``), +so reads return the stored values verbatim. +""" + +from typing import Any, Dict, Optional + +from litellm.models.credentials import CredentialItem + + +class CredentialsRepository: + """Repository for credentials database operations, keyed by credential name.""" + + def __init__(self, prisma_client: Any): + self._prisma_client = prisma_client + + @property + def prisma_client(self) -> Any: + if self._prisma_client is None: + raise RuntimeError( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + return self._prisma_client + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_credentialstable + + @staticmethod + def _to_model(record: Any) -> Optional[CredentialItem]: + if record is None: + return None + data = record.dict() if hasattr(record, "dict") else dict(record) + return CredentialItem( + credential_name=data["credential_name"], + credential_values=data.get("credential_values") or {}, + credential_info=data.get("credential_info") or {}, + ) + + async def find_all(self) -> Any: + return await self.table.find_many() + + async def create(self, data: Dict[str, Any]) -> Any: + return await self.table.create(data=data) + + async def find_by_name(self, credential_name: str) -> Optional[CredentialItem]: + record = await self.table.find_unique( + where={"credential_name": credential_name} + ) + return self._to_model(record) + + async def update_by_name(self, credential_name: str, data: Dict[str, Any]) -> Any: + return await self.table.update( + where={"credential_name": credential_name}, data=data + ) + + async def delete_by_name(self, credential_name: str) -> Any: + return await self.table.delete(where={"credential_name": credential_name}) diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py new file mode 100644 index 00000000000..893cf342d71 --- /dev/null +++ b/litellm/repositories/model_repository.py @@ -0,0 +1,171 @@ +""" +Model repository for database operations on LiteLLM_ProxyModelTable. +""" + +import json +from typing import Any, Dict, List, Optional, Type + +from litellm.models.model import LiteLLM_ProxyModelTable +from litellm.repositories.base_repository import BaseRepository +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + encrypt_value_helper, +) + + +class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): + """Repository for proxy model database operations with encryption support.""" + + def __init__(self, prisma_client: Any, encryption_key: Optional[str] = None): + super().__init__(prisma_client) + self._encryption_key = encryption_key + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_proxymodeltable + + @property + def model_class(self) -> Type[LiteLLM_ProxyModelTable]: + return LiteLLM_ProxyModelTable + + def _encrypt_litellm_params(self, litellm_params: Dict[str, Any]) -> Dict[str, Any]: + """Encrypt sensitive values in litellm_params.""" + encrypted = {} + for key, value in litellm_params.items(): + if isinstance(value, str): + encrypted[key] = encrypt_value_helper( + value, new_encryption_key=self._encryption_key + ) + else: + encrypted[key] = value + return encrypted + + def _decrypt_litellm_params(self, litellm_params: Dict[str, Any]) -> Dict[str, Any]: + """Decrypt sensitive values in litellm_params.""" + decrypted = {} + for key, value in litellm_params.items(): + if isinstance(value, str): + decrypted[key] = decrypt_value_helper( + value, key=key, exception_type="debug", return_original_value=True + ) + else: + decrypted[key] = value + return decrypted + + def _to_model(self, record: Any) -> Optional[LiteLLM_ProxyModelTable]: + """Convert a database record to a Model with decryption.""" + if record is None: + return None + + data = record.dict() if hasattr(record, "dict") else dict(record) + + if isinstance(data.get("litellm_params"), str): + data["litellm_params"] = json.loads(data["litellm_params"]) + if isinstance(data.get("model_info"), str): + data["model_info"] = json.loads(data["model_info"]) + + if data.get("litellm_params"): + data["litellm_params"] = self._decrypt_litellm_params( + data["litellm_params"] + ) + + return LiteLLM_ProxyModelTable(**data) + + async def find_by_id( + self, model_id: str, id_field: str = "model_id" + ) -> Optional[LiteLLM_ProxyModelTable]: + return await super().find_by_id(model_id, id_field) + + async def find_by_name(self, model_name: str) -> List[LiteLLM_ProxyModelTable]: + """Find models by name.""" + records = await self.table.find_many(where={"model_name": model_name}) + return self._to_model_list(records) + + async def find_all(self) -> List[LiteLLM_ProxyModelTable]: + """Find all models.""" + records = await self.table.find_many() + return self._to_model_list(records) + + async def find_unblocked(self) -> List[LiteLLM_ProxyModelTable]: + """Find all models that are not blocked.""" + records = await self.table.find_many(where={"blocked": False}) + return self._to_model_list(records) + + async def find_by_team_id(self, team_id: str) -> List[LiteLLM_ProxyModelTable]: + """Find models associated with a specific team. + + Note: This filters in-memory since team_id is stored within litellm_params + JSON. For large deployments with many models, consider adding a dedicated + team_id column with a database index. + """ + all_models = await self.find_all() + return [m for m in all_models if m.team_id == team_id] + + async def create_model( + self, + model_name: str, + litellm_params: Dict[str, Any], + created_by: str, + model_id: Optional[str] = None, + model_info: Optional[Dict[str, Any]] = None, + blocked: bool = False, + ) -> LiteLLM_ProxyModelTable: + """Create a new model with encryption.""" + encrypted_params = self._encrypt_litellm_params(litellm_params) + + data: Dict[str, Any] = { + "model_name": model_name, + "litellm_params": json.dumps(encrypted_params), + "created_by": created_by, + "updated_by": created_by, + "blocked": blocked, + } + if model_id is not None: + data["model_id"] = model_id + if model_info is not None: + data["model_info"] = json.dumps(model_info) + + record = await self.table.create(data=data) + model = self._to_model(record) + assert model is not None + return model + + async def update_model( + self, + model_id: str, + updated_by: str, + model_name: Optional[str] = None, + litellm_params: Optional[Dict[str, Any]] = None, + model_info: Optional[Dict[str, Any]] = None, + blocked: Optional[bool] = None, + ) -> Optional[LiteLLM_ProxyModelTable]: + """Update a model with encryption.""" + data: Dict[str, Any] = {"updated_by": updated_by} + if model_name is not None: + data["model_name"] = model_name + if litellm_params is not None: + encrypted_params = self._encrypt_litellm_params(litellm_params) + data["litellm_params"] = json.dumps(encrypted_params) + if model_info is not None: + data["model_info"] = json.dumps(model_info) + if blocked is not None: + data["blocked"] = blocked + + record = await self.table.update(where={"model_id": model_id}, data=data) + return self._to_model(record) + + async def delete_model(self, model_id: str) -> Optional[LiteLLM_ProxyModelTable]: + """Delete a model.""" + return await self.delete(model_id, id_field="model_id") + + async def block_model( + self, model_id: str, updated_by: str + ) -> Optional[LiteLLM_ProxyModelTable]: + """Block a model.""" + return await self.update_model(model_id, updated_by, blocked=True) + + async def unblock_model( + self, model_id: str, updated_by: str + ) -> Optional[LiteLLM_ProxyModelTable]: + """Unblock a model.""" + return await self.update_model(model_id, updated_by, blocked=False) diff --git a/litellm/repositories/object_permission_repository.py b/litellm/repositories/object_permission_repository.py new file mode 100644 index 00000000000..f4d9a8bb90a --- /dev/null +++ b/litellm/repositories/object_permission_repository.py @@ -0,0 +1,110 @@ +""" +ObjectPermission repository for database operations on LiteLLM_ObjectPermissionTable. +""" + +from typing import Any, Dict, List, Optional, Type + +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.repositories.base_repository import BaseRepository + + +class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): + """Repository for object permission database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_objectpermissiontable + + @property + def model_class(self) -> Type[LiteLLM_ObjectPermissionTable]: + return LiteLLM_ObjectPermissionTable + + async def find_by_id( + self, object_permission_id: str, id_field: str = "object_permission_id" + ) -> Optional[LiteLLM_ObjectPermissionTable]: + return await super().find_by_id(object_permission_id, id_field) + + async def create_permission( + self, + mcp_servers: Optional[List[str]] = None, + mcp_access_groups: Optional[List[str]] = None, + mcp_tool_permissions: Optional[Dict[str, List[str]]] = None, + vector_stores: Optional[List[str]] = None, + agents: Optional[List[str]] = None, + agent_access_groups: Optional[List[str]] = None, + models: Optional[List[str]] = None, + blocked_tools: Optional[List[str]] = None, + mcp_toolsets: Optional[List[str]] = None, + search_tools: Optional[List[str]] = None, + ) -> LiteLLM_ObjectPermissionTable: + """Create a new object permission record.""" + data: Dict[str, Any] = {} + if mcp_servers is not None: + data["mcp_servers"] = mcp_servers + if mcp_access_groups is not None: + data["mcp_access_groups"] = mcp_access_groups + if mcp_tool_permissions is not None: + data["mcp_tool_permissions"] = mcp_tool_permissions + if vector_stores is not None: + data["vector_stores"] = vector_stores + if agents is not None: + data["agents"] = agents + if agent_access_groups is not None: + data["agent_access_groups"] = agent_access_groups + if models is not None: + data["models"] = models + if blocked_tools is not None: + data["blocked_tools"] = blocked_tools + if mcp_toolsets is not None: + data["mcp_toolsets"] = mcp_toolsets + if search_tools is not None: + data["search_tools"] = search_tools + + return await self.create(data) + + async def update_permission( + self, + object_permission_id: str, + mcp_servers: Optional[List[str]] = None, + mcp_access_groups: Optional[List[str]] = None, + mcp_tool_permissions: Optional[Dict[str, List[str]]] = None, + vector_stores: Optional[List[str]] = None, + agents: Optional[List[str]] = None, + agent_access_groups: Optional[List[str]] = None, + models: Optional[List[str]] = None, + blocked_tools: Optional[List[str]] = None, + mcp_toolsets: Optional[List[str]] = None, + search_tools: Optional[List[str]] = None, + ) -> Optional[LiteLLM_ObjectPermissionTable]: + """Update an object permission record.""" + data: Dict[str, Any] = {} + if mcp_servers is not None: + data["mcp_servers"] = mcp_servers + if mcp_access_groups is not None: + data["mcp_access_groups"] = mcp_access_groups + if mcp_tool_permissions is not None: + data["mcp_tool_permissions"] = mcp_tool_permissions + if vector_stores is not None: + data["vector_stores"] = vector_stores + if agents is not None: + data["agents"] = agents + if agent_access_groups is not None: + data["agent_access_groups"] = agent_access_groups + if models is not None: + data["models"] = models + if blocked_tools is not None: + data["blocked_tools"] = blocked_tools + if mcp_toolsets is not None: + data["mcp_toolsets"] = mcp_toolsets + if search_tools is not None: + data["search_tools"] = search_tools + + return await self.update( + object_permission_id, data, id_field="object_permission_id" + ) + + async def delete_permission( + self, object_permission_id: str + ) -> Optional[LiteLLM_ObjectPermissionTable]: + """Delete an object permission record.""" + return await self.delete(object_permission_id, id_field="object_permission_id") diff --git a/litellm/repositories/organization_repository.py b/litellm/repositories/organization_repository.py new file mode 100644 index 00000000000..2d25a43e836 --- /dev/null +++ b/litellm/repositories/organization_repository.py @@ -0,0 +1,103 @@ +""" +Organization repository for database operations on LiteLLM_OrganizationTable. +""" + +from typing import Any, Dict, List, Optional, Type + +from litellm.models.organization import LiteLLM_OrganizationTable +from litellm.repositories.base_repository import BaseRepository + + +class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]): + """Repository for organization database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_organizationtable + + @property + def model_class(self) -> Type[LiteLLM_OrganizationTable]: + return LiteLLM_OrganizationTable + + async def find_by_id( + self, organization_id: str, id_field: str = "organization_id" + ) -> Optional[LiteLLM_OrganizationTable]: + return await super().find_by_id(organization_id, id_field) + + async def find_by_alias( + self, organization_alias: str + ) -> Optional[LiteLLM_OrganizationTable]: + """Find an organization by alias.""" + records = await self.table.find_many( + where={"organization_alias": organization_alias} + ) + if records: + return self._to_model(records[0]) + return None + + async def create_organization( + self, + organization_alias: str, + budget_id: str, + created_by: str, + organization_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + models: Optional[List[str]] = None, + object_permission_id: Optional[str] = None, + ) -> LiteLLM_OrganizationTable: + """Create a new organization.""" + data: Dict[str, Any] = { + "organization_alias": organization_alias, + "budget_id": budget_id, + "created_by": created_by, + "updated_by": created_by, + } + if organization_id is not None: + data["organization_id"] = organization_id + if metadata is not None: + data["metadata"] = metadata + if models is not None: + data["models"] = models + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.create(data) + + async def update_organization( + self, + organization_id: str, + updated_by: str, + organization_alias: Optional[str] = None, + budget_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + models: Optional[List[str]] = None, + object_permission_id: Optional[str] = None, + ) -> Optional[LiteLLM_OrganizationTable]: + """Update an organization.""" + data: Dict[str, Any] = {"updated_by": updated_by} + if organization_alias is not None: + data["organization_alias"] = organization_alias + if budget_id is not None: + data["budget_id"] = budget_id + if metadata is not None: + data["metadata"] = metadata + if models is not None: + data["models"] = models + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.update(organization_id, data, id_field="organization_id") + + async def delete_organization( + self, organization_id: str + ) -> Optional[LiteLLM_OrganizationTable]: + """Delete an organization.""" + return await self.delete(organization_id, id_field="organization_id") + + async def update_spend( + self, organization_id: str, spend: float + ) -> Optional[LiteLLM_OrganizationTable]: + """Update organization spend.""" + return await self.update( + organization_id, {"spend": spend}, id_field="organization_id" + ) diff --git a/litellm/repositories/project_repository.py b/litellm/repositories/project_repository.py new file mode 100644 index 00000000000..86567dd05fb --- /dev/null +++ b/litellm/repositories/project_repository.py @@ -0,0 +1,129 @@ +""" +Project repository for database operations on LiteLLM_ProjectTable. +""" + +from typing import Any, Dict, List, Optional, Type + +from litellm.models.project import LiteLLM_ProjectTable +from litellm.repositories.base_repository import BaseRepository + + +class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): + """Repository for project database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_projecttable + + @property + def model_class(self) -> Type[LiteLLM_ProjectTable]: + return LiteLLM_ProjectTable + + async def find_by_id( + self, project_id: str, id_field: str = "project_id" + ) -> Optional[LiteLLM_ProjectTable]: + return await super().find_by_id(project_id, id_field) + + async def find_by_alias(self, project_alias: str) -> Optional[LiteLLM_ProjectTable]: + """Find a project by alias.""" + records = await self.table.find_many(where={"project_alias": project_alias}) + if records: + return self._to_model(records[0]) + return None + + async def find_by_team_id(self, team_id: str) -> List[LiteLLM_ProjectTable]: + """Find all projects belonging to a team.""" + records = await self.table.find_many(where={"team_id": team_id}) + return self._to_model_list(records) + + async def create_project( + self, + created_by: str, + project_id: Optional[str] = None, + project_alias: Optional[str] = None, + description: Optional[str] = None, + team_id: Optional[str] = None, + budget_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + models: Optional[List[str]] = None, + model_rpm_limit: Optional[Dict[str, int]] = None, + model_tpm_limit: Optional[Dict[str, int]] = None, + object_permission_id: Optional[str] = None, + ) -> LiteLLM_ProjectTable: + """Create a new project.""" + data: Dict[str, Any] = { + "created_by": created_by, + "updated_by": created_by, + } + if project_id is not None: + data["project_id"] = project_id + if project_alias is not None: + data["project_alias"] = project_alias + if description is not None: + data["description"] = description + if team_id is not None: + data["team_id"] = team_id + if budget_id is not None: + data["budget_id"] = budget_id + if metadata is not None: + data["metadata"] = metadata + if models is not None: + data["models"] = models + if model_rpm_limit is not None: + data["model_rpm_limit"] = model_rpm_limit + if model_tpm_limit is not None: + data["model_tpm_limit"] = model_tpm_limit + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.create(data) + + async def update_project( + self, + project_id: str, + updated_by: str, + project_alias: Optional[str] = None, + description: Optional[str] = None, + team_id: Optional[str] = None, + budget_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + models: Optional[List[str]] = None, + model_rpm_limit: Optional[Dict[str, int]] = None, + model_tpm_limit: Optional[Dict[str, int]] = None, + blocked: Optional[bool] = None, + object_permission_id: Optional[str] = None, + ) -> Optional[LiteLLM_ProjectTable]: + """Update a project.""" + data: Dict[str, Any] = {"updated_by": updated_by} + if project_alias is not None: + data["project_alias"] = project_alias + if description is not None: + data["description"] = description + if team_id is not None: + data["team_id"] = team_id + if budget_id is not None: + data["budget_id"] = budget_id + if metadata is not None: + data["metadata"] = metadata + if models is not None: + data["models"] = models + if model_rpm_limit is not None: + data["model_rpm_limit"] = model_rpm_limit + if model_tpm_limit is not None: + data["model_tpm_limit"] = model_tpm_limit + if blocked is not None: + data["blocked"] = blocked + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.update(project_id, data, id_field="project_id") + + async def delete_project(self, project_id: str) -> Optional[LiteLLM_ProjectTable]: + """Delete a project.""" + return await self.delete(project_id, id_field="project_id") + + async def update_spend( + self, project_id: str, spend: float + ) -> Optional[LiteLLM_ProjectTable]: + """Update project spend.""" + return await self.update(project_id, {"spend": spend}, id_field="project_id") diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py new file mode 100644 index 00000000000..47ea11c0592 --- /dev/null +++ b/litellm/repositories/table_repositories.py @@ -0,0 +1,215 @@ +""" +Passthrough table repositories. + +Each repository centralizes access to a single Prisma table behind a ``table`` +property, making the repository the one place that names the underlying table. +These are thin wrappers for tables that do not (yet) need domain-specific query +methods; richer repositories live in their own modules. +""" + +from typing import Any + + +class PrismaTableRepository: + """Base for repositories that expose a single Prisma table.""" + + table_name: str + + def __init__(self, prisma_client: Any): + self._prisma_client = prisma_client + + @property + def prisma_client(self) -> Any: + if self._prisma_client is None: + raise RuntimeError( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + return self._prisma_client + + @property + def table(self) -> Any: + return getattr(self.prisma_client.db, self.table_name) + + +class PolicyRepository(PrismaTableRepository): + table_name = "litellm_policytable" + + +class AgentsRepository(PrismaTableRepository): + table_name = "litellm_agentstable" + + +class GuardrailsRepository(PrismaTableRepository): + table_name = "litellm_guardrailstable" + + +class MCPServerRepository(PrismaTableRepository): + table_name = "litellm_mcpservertable" + + +class ManagedObjectRepository(PrismaTableRepository): + table_name = "litellm_managedobjecttable" + + +class OrganizationMembershipRepository(PrismaTableRepository): + table_name = "litellm_organizationmembership" + + +class SpendLogsRepository(PrismaTableRepository): + table_name = "litellm_spendlogs" + + +class ClaudeCodePluginRepository(PrismaTableRepository): + table_name = "litellm_claudecodeplugintable" + + +class TeamMembershipRepository(PrismaTableRepository): + table_name = "litellm_teammembership" + + +class EndUserRepository(PrismaTableRepository): + table_name = "litellm_endusertable" + + +class ManagedVectorStoresRepository(PrismaTableRepository): + table_name = "litellm_managedvectorstorestable" + + +class MCPUserCredentialsRepository(PrismaTableRepository): + table_name = "litellm_mcpusercredentials" + + +class PromptRepository(PrismaTableRepository): + table_name = "litellm_prompttable" + + +class TagRepository(PrismaTableRepository): + table_name = "litellm_tagtable" + + +class InvitationLinkRepository(PrismaTableRepository): + table_name = "litellm_invitationlink" + + +class JWTKeyMappingRepository(PrismaTableRepository): + table_name = "litellm_jwtkeymapping" + + +class ManagedFileRepository(PrismaTableRepository): + table_name = "litellm_managedfiletable" + + +class MemoryRepository(PrismaTableRepository): + table_name = "litellm_memorytable" + + +class SearchToolsRepository(PrismaTableRepository): + table_name = "litellm_searchtoolstable" + + +class ConfigOverridesRepository(PrismaTableRepository): + table_name = "litellm_configoverrides" + + +class MCPToolsetRepository(PrismaTableRepository): + table_name = "litellm_mcptoolsettable" + + +class ToolRepository(PrismaTableRepository): + table_name = "litellm_tooltable" + + +class DeletedVerificationTokenRepository(PrismaTableRepository): + table_name = "litellm_deletedverificationtoken" + + +class WorkflowRunRepository(PrismaTableRepository): + table_name = "litellm_workflowrun" + + +class ModelTableRepository(PrismaTableRepository): + table_name = "litellm_modeltable" + + +class AccessGroupRepository(PrismaTableRepository): + table_name = "litellm_accessgrouptable" + + +class SSOConfigRepository(PrismaTableRepository): + table_name = "litellm_ssoconfig" + + +class UISettingsRepository(PrismaTableRepository): + table_name = "litellm_uisettings" + + +class DailyGuardrailMetricsRepository(PrismaTableRepository): + table_name = "litellm_dailyguardrailmetrics" + + +class PolicyAttachmentRepository(PrismaTableRepository): + table_name = "litellm_policyattachmenttable" + + +class DeletedTeamRepository(PrismaTableRepository): + table_name = "litellm_deletedteamtable" + + +class SkillsRepository(PrismaTableRepository): + table_name = "litellm_skillstable" + + +class CacheConfigRepository(PrismaTableRepository): + table_name = "litellm_cacheconfig" + + +class ManagedVectorStoreIndexRepository(PrismaTableRepository): + table_name = "litellm_managedvectorstoreindextable" + + +class WorkflowMessageRepository(PrismaTableRepository): + table_name = "litellm_workflowmessage" + + +class DailyTagSpendRepository(PrismaTableRepository): + table_name = "litellm_dailytagspend" + + +class SpendLogToolIndexRepository(PrismaTableRepository): + table_name = "litellm_spendlogtoolindex" + + +class SpendLogGuardrailIndexRepository(PrismaTableRepository): + table_name = "litellm_spendlogguardrailindex" + + +class UserNotificationsRepository(PrismaTableRepository): + table_name = "litellm_usernotifications" + + +class HealthCheckRepository(PrismaTableRepository): + table_name = "litellm_healthchecktable" + + +class DeprecatedVerificationTokenRepository(PrismaTableRepository): + table_name = "litellm_deprecatedverificationtoken" + + +class WorkflowEventRepository(PrismaTableRepository): + table_name = "litellm_workflowevent" + + +class DailyPolicyMetricsRepository(PrismaTableRepository): + table_name = "litellm_dailypolicymetrics" + + +class AdaptiveRouterStateRepository(PrismaTableRepository): + table_name = "litellm_adaptiverouterstate" + + +class AuditLogRepository(PrismaTableRepository): + table_name = "litellm_auditlog" + + +class AdaptiveRouterSessionRepository(PrismaTableRepository): + table_name = "litellm_adaptiveroutersession" diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py new file mode 100644 index 00000000000..2ae6647060c --- /dev/null +++ b/litellm/repositories/team_repository.py @@ -0,0 +1,351 @@ +""" +Team repository for database operations on LiteLLM_TeamTable. +""" + +import json +from datetime import datetime +from typing import Any, Dict, List, Optional, Type + +from litellm.models.team import LiteLLM_TeamTable +from litellm.repositories.base_repository import BaseRepository + + +class TeamRepository(BaseRepository[LiteLLM_TeamTable]): + """Repository for team database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_teamtable + + @property + def deleted_table(self) -> Any: + return self.prisma_client.db.litellm_deletedteamtable + + @property + def model_class(self) -> Type[LiteLLM_TeamTable]: + return LiteLLM_TeamTable + + def _to_model(self, record: Any) -> Optional[LiteLLM_TeamTable]: + """Convert a database record to a Team model.""" + if record is None: + return None + + data = record.dict() if hasattr(record, "dict") else dict(record) + + json_fields = [ + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + "budget_limits", + "members_with_roles", + ] + for field in json_fields: + if isinstance(data.get(field), str): + data[field] = json.loads(data[field]) + + return LiteLLM_TeamTable(**data) + + async def find_by_id( + self, team_id: str, id_field: str = "team_id" + ) -> Optional[LiteLLM_TeamTable]: + return await super().find_by_id(team_id, id_field) + + async def find_by_alias(self, team_alias: str) -> Optional[LiteLLM_TeamTable]: + """Find a team by alias.""" + records = await self.table.find_many(where={"team_alias": team_alias}) + if records: + return self._to_model(records[0]) + return None + + async def find_by_organization_id( + self, organization_id: str + ) -> List[LiteLLM_TeamTable]: + """Find all teams belonging to an organization.""" + records = await self.table.find_many(where={"organization_id": organization_id}) + return self._to_model_list(records) + + async def find_by_member(self, user_id: str) -> List[LiteLLM_TeamTable]: + """Find all teams where user is a member.""" + records = await self.table.find_many(where={"members": {"has": user_id}}) + return self._to_model_list(records) + + async def find_by_admin(self, user_id: str) -> List[LiteLLM_TeamTable]: + """Find all teams where user is an admin.""" + records = await self.table.find_many(where={"admins": {"has": user_id}}) + return self._to_model_list(records) + + async def create_team( + self, + team_id: str, + team_alias: Optional[str] = None, + organization_id: Optional[str] = None, + admins: Optional[List[str]] = None, + members: Optional[List[str]] = None, + members_with_roles: Optional[Dict[str, Any]] = None, + metadata: Optional[Dict[str, Any]] = None, + max_budget: Optional[float] = None, + soft_budget: Optional[float] = None, + models: Optional[List[str]] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + object_permission_id: Optional[str] = None, + ) -> LiteLLM_TeamTable: + """Create a new team.""" + data: Dict[str, Any] = {"team_id": team_id} + if team_alias is not None: + data["team_alias"] = team_alias + if organization_id is not None: + data["organization_id"] = organization_id + if admins is not None: + data["admins"] = admins + if members is not None: + data["members"] = members + if members_with_roles is not None: + data["members_with_roles"] = json.dumps(members_with_roles) + if metadata is not None: + data["metadata"] = json.dumps(metadata) + if max_budget is not None: + data["max_budget"] = max_budget + if soft_budget is not None: + data["soft_budget"] = soft_budget + if models is not None: + data["models"] = models + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if budget_duration is not None: + data["budget_duration"] = budget_duration + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.create(data) + + async def update_team( + self, + team_id: str, + team_alias: Optional[str] = None, + organization_id: Optional[str] = None, + admins: Optional[List[str]] = None, + members: Optional[List[str]] = None, + members_with_roles: Optional[Dict[str, Any]] = None, + metadata: Optional[Dict[str, Any]] = None, + max_budget: Optional[float] = None, + soft_budget: Optional[float] = None, + models: Optional[List[str]] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + blocked: Optional[bool] = None, + object_permission_id: Optional[str] = None, + ) -> Optional[LiteLLM_TeamTable]: + """Update a team.""" + data: Dict[str, Any] = {} + if team_alias is not None: + data["team_alias"] = team_alias + if organization_id is not None: + data["organization_id"] = organization_id + if admins is not None: + data["admins"] = admins + if members is not None: + data["members"] = members + if members_with_roles is not None: + data["members_with_roles"] = json.dumps(members_with_roles) + if metadata is not None: + data["metadata"] = json.dumps(metadata) + if max_budget is not None: + data["max_budget"] = max_budget + if soft_budget is not None: + data["soft_budget"] = soft_budget + if models is not None: + data["models"] = models + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if budget_duration is not None: + data["budget_duration"] = budget_duration + if blocked is not None: + data["blocked"] = blocked + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.update(team_id, data, id_field="team_id") + + async def delete_team( + self, + team_id: str, + deleted_by: Optional[str] = None, + deleted_by_api_key: Optional[str] = None, + litellm_changed_by: Optional[str] = None, + ) -> Optional[LiteLLM_TeamTable]: + """Delete a team and archive it to the deleted teams table. + + Uses a transaction to ensure atomicity of the archive-then-delete operation. + """ + team = await self.find_by_id(team_id) + if team is None: + return None + + archive_data = self._build_archive_data(team) + archive_data["deleted_by"] = deleted_by + archive_data["deleted_by_api_key"] = deleted_by_api_key + archive_data["litellm_changed_by"] = litellm_changed_by + archive_data["deleted_at"] = datetime.utcnow() + + async with self.prisma_client.db.tx() as tx: + await tx.litellm_deletedteamtable.create(data=archive_data) + await tx.litellm_teamtable.delete(where={"team_id": team_id}) + + return team + + def _build_archive_data(self, team: LiteLLM_TeamTable) -> Dict[str, Any]: + """Build archive data dict with only columns that exist in LiteLLM_DeletedTeamTable.""" + data: Dict[str, Any] = {"team_id": team.team_id} + if team.team_alias is not None: + data["team_alias"] = team.team_alias + if team.organization_id is not None: + data["organization_id"] = team.organization_id + if team.object_permission_id is not None: + data["object_permission_id"] = team.object_permission_id + data["admins"] = team.admins + data["members"] = team.members + if team.members_with_roles: + data["members_with_roles"] = json.dumps( + [m.model_dump() for m in team.members_with_roles] + ) + if team.metadata: + data["metadata"] = json.dumps(team.metadata) + if team.max_budget is not None: + data["max_budget"] = team.max_budget + if team.soft_budget is not None: + data["soft_budget"] = team.soft_budget + data["spend"] = team.spend if team.spend is not None else 0.0 + data["models"] = team.models + if team.max_parallel_requests is not None: + data["max_parallel_requests"] = team.max_parallel_requests + if team.tpm_limit is not None: + data["tpm_limit"] = team.tpm_limit + if team.rpm_limit is not None: + data["rpm_limit"] = team.rpm_limit + if team.budget_duration is not None: + data["budget_duration"] = team.budget_duration + if team.budget_reset_at is not None: + data["budget_reset_at"] = team.budget_reset_at + data["blocked"] = team.blocked + if team.model_spend: + data["model_spend"] = json.dumps(team.model_spend) + if team.model_max_budget: + data["model_max_budget"] = json.dumps(team.model_max_budget) + if team.router_settings is not None: + data["router_settings"] = json.dumps(team.router_settings) + data["team_member_permissions"] = team.team_member_permissions or [] + data["access_group_ids"] = team.access_group_ids or [] + data["policies"] = team.policies or [] + if team.model_id is not None: + data["model_id"] = team.model_id + data["allow_team_guardrail_config"] = team.allow_team_guardrail_config + return data + + async def update_spend( + self, team_id: str, spend: float + ) -> Optional[LiteLLM_TeamTable]: + """Update team spend.""" + return await self.update(team_id, {"spend": spend}, id_field="team_id") + + async def add_member( + self, team_id: str, user_id: str + ) -> Optional[LiteLLM_TeamTable]: + """Add a member to a team using atomic array push operation.""" + if not await self.exists(team_id, id_field="team_id"): + return None + + record = await self.table.update( + where={"team_id": team_id}, + data={"members": {"push": user_id}}, + ) + return self._to_model(record) + + async def remove_member( + self, team_id: str, user_id: str + ) -> Optional[LiteLLM_TeamTable]: + """Remove a member from a team. + + Note: Prisma doesn't support atomic array removal, so we use a + read-modify-write pattern here. For high-concurrency scenarios, + consider using raw SQL with array_remove(). + """ + team = await self.find_by_id(team_id) + if team is None: + return None + + members = [m for m in team.members if m != user_id] + return await self.update(team_id, {"members": members}, id_field="team_id") + + async def add_admin( + self, team_id: str, user_id: str + ) -> Optional[LiteLLM_TeamTable]: + """Add an admin to a team using atomic array push operation.""" + if not await self.exists(team_id, id_field="team_id"): + return None + + record = await self.table.update( + where={"team_id": team_id}, + data={"admins": {"push": user_id}}, + ) + return self._to_model(record) + + async def remove_admin( + self, team_id: str, user_id: str + ) -> Optional[LiteLLM_TeamTable]: + """Remove an admin from a team. + + Note: Prisma doesn't support atomic array removal, so we use a + read-modify-write pattern here. For high-concurrency scenarios, + consider using raw SQL with array_remove(). + """ + team = await self.find_by_id(team_id) + if team is None: + return None + + admins = [a for a in team.admins if a != user_id] + return await self.update(team_id, {"admins": admins}, id_field="team_id") + + async def add_models( + self, team_id: str, models: List[str] + ) -> Optional[LiteLLM_TeamTable]: + """Add models to a team's allowed models list using atomic array push.""" + if not await self.exists(team_id, id_field="team_id"): + return None + + record = await self.table.update( + where={"team_id": team_id}, + data={"models": {"push": models}}, + ) + return self._to_model(record) + + async def remove_models( + self, team_id: str, models: List[str] + ) -> Optional[LiteLLM_TeamTable]: + """Remove models from a team's allowed models list. + + Note: Prisma doesn't support atomic array removal, so we use a + read-modify-write pattern here. For high-concurrency scenarios, + consider using raw SQL with array_remove(). + """ + team = await self.find_by_id(team_id) + if team is None: + return None + + current_models = [m for m in team.models if m not in models] + return await self.update( + team_id, {"models": current_models}, id_field="team_id" + ) diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py new file mode 100644 index 00000000000..4d28b58f0ab --- /dev/null +++ b/litellm/repositories/user_repository.py @@ -0,0 +1,229 @@ +""" +User repository for database operations on LiteLLM_UserTable. +""" + +import json +from typing import Any, Dict, List, Optional, Type + +from litellm.models.user import LiteLLM_UserTable +from litellm.repositories.base_repository import BaseRepository + + +class UserRepository(BaseRepository[LiteLLM_UserTable]): + """Repository for user database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_usertable + + @property + def model_class(self) -> Type[LiteLLM_UserTable]: + return LiteLLM_UserTable + + def _to_model(self, record: Any) -> Optional[LiteLLM_UserTable]: + """Convert a database record to a User model.""" + if record is None: + return None + + data = record.dict() if hasattr(record, "dict") else dict(record) + + json_fields = ["metadata", "model_spend", "model_max_budget"] + for field in json_fields: + if isinstance(data.get(field), str): + data[field] = json.loads(data[field]) + + return LiteLLM_UserTable(**data) + + async def find_by_id( + self, user_id: str, id_field: str = "user_id" + ) -> Optional[LiteLLM_UserTable]: + return await super().find_by_id(user_id, id_field) + + async def find_by_email(self, user_email: str) -> Optional[LiteLLM_UserTable]: + """Find a user by email.""" + records = await self.table.find_many(where={"user_email": user_email}) + if records: + return self._to_model(records[0]) + return None + + async def find_by_sso_id(self, sso_user_id: str) -> Optional[LiteLLM_UserTable]: + """Find a user by SSO ID.""" + record = await self.table.find_unique(where={"sso_user_id": sso_user_id}) + return self._to_model(record) + + async def find_by_organization_id( + self, organization_id: str + ) -> List[LiteLLM_UserTable]: + """Find all users in an organization.""" + records = await self.table.find_many(where={"organization_id": organization_id}) + return self._to_model_list(records) + + async def find_by_team_id(self, team_id: str) -> List[LiteLLM_UserTable]: + """Find all users in a team.""" + records = await self.table.find_many(where={"teams": {"has": team_id}}) + return self._to_model_list(records) + + async def create_user( + self, + user_id: str, + user_alias: Optional[str] = None, + team_id: Optional[str] = None, + sso_user_id: Optional[str] = None, + organization_id: Optional[str] = None, + password: Optional[str] = None, + teams: Optional[List[str]] = None, + user_role: Optional[str] = None, + max_budget: Optional[float] = None, + user_email: Optional[str] = None, + models: Optional[List[str]] = None, + metadata: Optional[Dict[str, Any]] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + allowed_cache_controls: Optional[List[str]] = None, + policies: Optional[List[str]] = None, + object_permission_id: Optional[str] = None, + ) -> LiteLLM_UserTable: + """Create a new user.""" + data: Dict[str, Any] = {"user_id": user_id} + if user_alias is not None: + data["user_alias"] = user_alias + if team_id is not None: + data["team_id"] = team_id + if sso_user_id is not None: + data["sso_user_id"] = sso_user_id + if organization_id is not None: + data["organization_id"] = organization_id + if password is not None: + data["password"] = password + if teams is not None: + data["teams"] = teams + if user_role is not None: + data["user_role"] = user_role + if max_budget is not None: + data["max_budget"] = max_budget + if user_email is not None: + data["user_email"] = user_email + if models is not None: + data["models"] = models + if metadata is not None: + data["metadata"] = json.dumps(metadata) + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if budget_duration is not None: + data["budget_duration"] = budget_duration + if allowed_cache_controls is not None: + data["allowed_cache_controls"] = allowed_cache_controls + if policies is not None: + data["policies"] = policies + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.create(data) + + async def update_user( + self, + user_id: str, + user_alias: Optional[str] = None, + team_id: Optional[str] = None, + sso_user_id: Optional[str] = None, + organization_id: Optional[str] = None, + password: Optional[str] = None, + teams: Optional[List[str]] = None, + user_role: Optional[str] = None, + max_budget: Optional[float] = None, + user_email: Optional[str] = None, + models: Optional[List[str]] = None, + metadata: Optional[Dict[str, Any]] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + allowed_cache_controls: Optional[List[str]] = None, + policies: Optional[List[str]] = None, + object_permission_id: Optional[str] = None, + ) -> Optional[LiteLLM_UserTable]: + """Update a user.""" + data: Dict[str, Any] = {} + if user_alias is not None: + data["user_alias"] = user_alias + if team_id is not None: + data["team_id"] = team_id + if sso_user_id is not None: + data["sso_user_id"] = sso_user_id + if organization_id is not None: + data["organization_id"] = organization_id + if password is not None: + data["password"] = password + if teams is not None: + data["teams"] = teams + if user_role is not None: + data["user_role"] = user_role + if max_budget is not None: + data["max_budget"] = max_budget + if user_email is not None: + data["user_email"] = user_email + if models is not None: + data["models"] = models + if metadata is not None: + data["metadata"] = json.dumps(metadata) + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if budget_duration is not None: + data["budget_duration"] = budget_duration + if allowed_cache_controls is not None: + data["allowed_cache_controls"] = allowed_cache_controls + if policies is not None: + data["policies"] = policies + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.update(user_id, data, id_field="user_id") + + async def delete_user(self, user_id: str) -> Optional[LiteLLM_UserTable]: + """Delete a user.""" + return await self.delete(user_id, id_field="user_id") + + async def update_spend( + self, user_id: str, spend: float + ) -> Optional[LiteLLM_UserTable]: + """Update user spend.""" + return await self.update(user_id, {"spend": spend}, id_field="user_id") + + async def add_to_team( + self, user_id: str, team_id: str + ) -> Optional[LiteLLM_UserTable]: + """Add a user to a team using atomic array push operation.""" + if not await self.exists(user_id, id_field="user_id"): + return None + + record = await self.table.update( + where={"user_id": user_id}, + data={"teams": {"push": team_id}}, + ) + return self._to_model(record) + + async def remove_from_team( + self, user_id: str, team_id: str + ) -> Optional[LiteLLM_UserTable]: + """Remove a user from a team. + + Note: Prisma doesn't support atomic array removal, so we use a + read-modify-write pattern here. For high-concurrency scenarios, + consider using raw SQL with array_remove(). + """ + user = await self.find_by_id(user_id) + if user is None: + return None + + teams = [t for t in user.teams if t != team_id] + return await self.update(user_id, {"teams": teams}, id_field="user_id") diff --git a/litellm/repositories/verification_token_repository.py b/litellm/repositories/verification_token_repository.py new file mode 100644 index 00000000000..56c3e0714aa --- /dev/null +++ b/litellm/repositories/verification_token_repository.py @@ -0,0 +1,375 @@ +""" +VerificationToken repository for database operations on LiteLLM_VerificationToken. +""" + +import json +from datetime import datetime +from typing import Any, Dict, List, Optional, Type + +from litellm.models.verification_token import ( + LiteLLM_VerificationToken, +) +from litellm.repositories.base_repository import BaseRepository + + +class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): + """Repository for verification token (API key) database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_verificationtoken + + @property + def deleted_table(self) -> Any: + return self.prisma_client.db.litellm_deletedverificationtoken + + @property + def model_class(self) -> Type[LiteLLM_VerificationToken]: + return LiteLLM_VerificationToken + + def _to_model(self, record: Any) -> Optional[LiteLLM_VerificationToken]: + """Convert a database record to a VerificationToken model.""" + if record is None: + return None + + data = record.dict() if hasattr(record, "dict") else dict(record) + + json_fields = [ + "aliases", + "config", + "permissions", + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + "budget_limits", + "litellm_budget_table", + ] + for field in json_fields: + if isinstance(data.get(field), str): + data[field] = json.loads(data[field]) + + if data.get("org_id") is None and data.get("organization_id") is not None: + data["org_id"] = data["organization_id"] + + return LiteLLM_VerificationToken(**data) + + async def find_by_id( + self, token: str, id_field: str = "token" + ) -> Optional[LiteLLM_VerificationToken]: + return await super().find_by_id(token, id_field) + + async def find_by_alias( + self, key_alias: str + ) -> Optional[LiteLLM_VerificationToken]: + """Find a token by key alias.""" + records = await self.table.find_many(where={"key_alias": key_alias}) + if records: + return self._to_model(records[0]) + return None + + async def find_by_user_id(self, user_id: str) -> List[LiteLLM_VerificationToken]: + """Find all tokens belonging to a user.""" + records = await self.table.find_many(where={"user_id": user_id}) + return self._to_model_list(records) + + async def find_by_team_id(self, team_id: str) -> List[LiteLLM_VerificationToken]: + """Find all tokens belonging to a team.""" + records = await self.table.find_many(where={"team_id": team_id}) + return self._to_model_list(records) + + async def find_by_project_id( + self, project_id: str + ) -> List[LiteLLM_VerificationToken]: + """Find all tokens belonging to a project.""" + records = await self.table.find_many(where={"project_id": project_id}) + return self._to_model_list(records) + + async def find_active_tokens(self) -> List[LiteLLM_VerificationToken]: + """Find all active (non-expired, non-blocked) tokens.""" + records = await self.table.find_many( + where={ + "blocked": {"not": True}, + "OR": [{"expires": None}, {"expires": {"gt": datetime.utcnow()}}], + } + ) + return self._to_model_list(records) + + def _build_token_data( + self, + token: str, + key_name: Optional[str] = None, + key_alias: Optional[str] = None, + max_budget: Optional[float] = None, + expires: Optional[datetime] = None, + models: Optional[List[str]] = None, + aliases: Optional[Dict[str, str]] = None, + config: Optional[Dict[str, Any]] = None, + user_id: Optional[str] = None, + team_id: Optional[str] = None, + agent_id: Optional[str] = None, + project_id: Optional[str] = None, + max_parallel_requests: Optional[int] = None, + metadata: Optional[Dict[str, Any]] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + allowed_cache_controls: Optional[List[str]] = None, + allowed_routes: Optional[List[str]] = None, + permissions: Optional[Dict[str, Any]] = None, + org_id: Optional[str] = None, + created_by: Optional[str] = None, + object_permission_id: Optional[str] = None, + access_group_ids: Optional[List[str]] = None, + budget_id: Optional[str] = None, + ) -> Dict[str, Any]: + """Build data dictionary for token creation.""" + json_fields = { + "aliases": aliases, + "config": config, + "metadata": metadata, + "permissions": permissions, + } + simple_fields = { + "token": token, + "key_name": key_name, + "key_alias": key_alias, + "max_budget": max_budget, + "expires": expires, + "models": models, + "user_id": user_id, + "team_id": team_id, + "agent_id": agent_id, + "project_id": project_id, + "max_parallel_requests": max_parallel_requests, + "tpm_limit": tpm_limit, + "rpm_limit": rpm_limit, + "budget_duration": budget_duration, + "allowed_cache_controls": allowed_cache_controls, + "allowed_routes": allowed_routes, + "object_permission_id": object_permission_id, + "access_group_ids": access_group_ids, + "budget_id": budget_id, + } + data: Dict[str, Any] = {k: v for k, v in simple_fields.items() if v is not None} + for key, val in json_fields.items(): + if val is not None: + data[key] = json.dumps(val) + if org_id is not None: + data["organization_id"] = org_id + if created_by is not None: + data["created_by"] = created_by + data["updated_by"] = created_by + return data + + async def create_token( + self, + token: str, + key_name: Optional[str] = None, + key_alias: Optional[str] = None, + max_budget: Optional[float] = None, + expires: Optional[datetime] = None, + models: Optional[List[str]] = None, + aliases: Optional[Dict[str, str]] = None, + config: Optional[Dict[str, Any]] = None, + user_id: Optional[str] = None, + team_id: Optional[str] = None, + agent_id: Optional[str] = None, + project_id: Optional[str] = None, + max_parallel_requests: Optional[int] = None, + metadata: Optional[Dict[str, Any]] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + allowed_cache_controls: Optional[List[str]] = None, + allowed_routes: Optional[List[str]] = None, + permissions: Optional[Dict[str, Any]] = None, + org_id: Optional[str] = None, + created_by: Optional[str] = None, + object_permission_id: Optional[str] = None, + access_group_ids: Optional[List[str]] = None, + budget_id: Optional[str] = None, + ) -> LiteLLM_VerificationToken: + """Create a new verification token.""" + data = self._build_token_data( + token=token, + key_name=key_name, + key_alias=key_alias, + max_budget=max_budget, + expires=expires, + models=models, + aliases=aliases, + config=config, + user_id=user_id, + team_id=team_id, + agent_id=agent_id, + project_id=project_id, + max_parallel_requests=max_parallel_requests, + metadata=metadata, + tpm_limit=tpm_limit, + rpm_limit=rpm_limit, + budget_duration=budget_duration, + allowed_cache_controls=allowed_cache_controls, + allowed_routes=allowed_routes, + permissions=permissions, + org_id=org_id, + created_by=created_by, + object_permission_id=object_permission_id, + access_group_ids=access_group_ids, + budget_id=budget_id, + ) + return await self.create(data) + + async def update_token( + self, + token: str, + updated_by: Optional[str] = None, + key_name: Optional[str] = None, + key_alias: Optional[str] = None, + max_budget: Optional[float] = None, + expires: Optional[datetime] = None, + models: Optional[List[str]] = None, + aliases: Optional[Dict[str, str]] = None, + config: Optional[Dict[str, Any]] = None, + max_parallel_requests: Optional[int] = None, + metadata: Optional[Dict[str, Any]] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + allowed_cache_controls: Optional[List[str]] = None, + allowed_routes: Optional[List[str]] = None, + permissions: Optional[Dict[str, Any]] = None, + blocked: Optional[bool] = None, + object_permission_id: Optional[str] = None, + access_group_ids: Optional[List[str]] = None, + ) -> Optional[LiteLLM_VerificationToken]: + """Update a verification token.""" + data: Dict[str, Any] = {} + if updated_by is not None: + data["updated_by"] = updated_by + if key_name is not None: + data["key_name"] = key_name + if key_alias is not None: + data["key_alias"] = key_alias + if max_budget is not None: + data["max_budget"] = max_budget + if expires is not None: + data["expires"] = expires + if models is not None: + data["models"] = models + if aliases is not None: + data["aliases"] = json.dumps(aliases) + if config is not None: + data["config"] = json.dumps(config) + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if metadata is not None: + data["metadata"] = json.dumps(metadata) + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if budget_duration is not None: + data["budget_duration"] = budget_duration + if allowed_cache_controls is not None: + data["allowed_cache_controls"] = allowed_cache_controls + if allowed_routes is not None: + data["allowed_routes"] = allowed_routes + if permissions is not None: + data["permissions"] = json.dumps(permissions) + if blocked is not None: + data["blocked"] = blocked + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + if access_group_ids is not None: + data["access_group_ids"] = access_group_ids + + return await self.update(token, data, id_field="token") + + async def delete_token( + self, + token: str, + deleted_by: Optional[str] = None, + deleted_by_api_key: Optional[str] = None, + litellm_changed_by: Optional[str] = None, + ) -> Optional[LiteLLM_VerificationToken]: + """Delete a token and archive it to the deleted tokens table. + + Uses a transaction to ensure atomicity of the archive-then-delete operation. + """ + token_record = await self.find_by_id(token) + if token_record is None: + return None + + archive_data = self._build_archive_data(token_record) + archive_data["deleted_by"] = deleted_by + archive_data["deleted_by_api_key"] = deleted_by_api_key + archive_data["litellm_changed_by"] = litellm_changed_by + archive_data["deleted_at"] = datetime.utcnow() + + async with self.prisma_client.db.tx() as tx: + await tx.litellm_deletedverificationtoken.create(data=archive_data) + await tx.litellm_verificationtoken.delete(where={"token": token}) + + return token_record + + def _build_archive_data(self, token: LiteLLM_VerificationToken) -> Dict[str, Any]: + """Build archive data with only columns present in LiteLLM_DeletedVerificationToken. + + Serializes JSON columns to strings (the archive table stores them as JSON + columns the same way the live table does) and maps ``org_id`` onto the + ``organization_id`` column so the foreign key is preserved. + """ + data = token.model_dump(exclude_none=True) + for field in ("object_permission", "litellm_budget_table", "budget_limits"): + data.pop(field, None) + + org_id = data.pop("org_id", None) + if org_id is not None: + data["organization_id"] = org_id + + json_fields = [ + "aliases", + "config", + "permissions", + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + ] + for field in json_fields: + if field in data: + data[field] = json.dumps(data[field]) + return data + + async def update_spend( + self, token: str, spend: float + ) -> Optional[LiteLLM_VerificationToken]: + """Update token spend.""" + return await self.update(token, {"spend": spend}, id_field="token") + + async def update_last_active( + self, token: str + ) -> Optional[LiteLLM_VerificationToken]: + """Update the last_active timestamp.""" + return await self.update( + token, {"last_active": datetime.utcnow()}, id_field="token" + ) + + async def block_token( + self, token: str, updated_by: Optional[str] = None + ) -> Optional[LiteLLM_VerificationToken]: + """Block a token.""" + data: Dict[str, Any] = {"blocked": True} + if updated_by is not None: + data["updated_by"] = updated_by + return await self.update(token, data, id_field="token") + + async def unblock_token( + self, token: str, updated_by: Optional[str] = None + ) -> Optional[LiteLLM_VerificationToken]: + """Unblock a token.""" + data: Dict[str, Any] = {"blocked": False} + if updated_by is not None: + data["updated_by"] = updated_by + return await self.update(token, data, id_field="token") diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index e2ba8353591..d3d30642216 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -148,7 +148,9 @@ def _transform_tool_choice( # which is equivalent to "required" in OpenAI format return "required" elif tool_choice_type == "function": - # function type without name - fall back to required + function_name = tool_choice.get("name") + if function_name: + return {"type": "function", "function": {"name": function_name}} return "required" # Return as-is for unknown formats diff --git a/litellm/router.py b/litellm/router.py index d0f4e5ff44d..d1c8e227bea 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1658,6 +1658,67 @@ def validate_fallbacks(self, fallback_param: Optional[List]): f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys." ) + def _add_encrypted_content_affinity_check( + self, enable_global_affinity: bool + ) -> None: + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + def _move_before_deployment_affinity( + callback_list: List[Any], + callback_to_move: EncryptedContentAffinityCheck, + ) -> None: + if callback_to_move not in callback_list: + return + callback_list.remove(callback_to_move) + insert_index = next( + ( + idx + for idx, callback in enumerate(callback_list) + if isinstance(callback, DeploymentAffinityCheck) + ), + len(callback_list), + ) + callback_list.insert(insert_index, callback_to_move) + + if ( + enable_global_affinity + or EncryptedContentAffinityCheck.has_model_group_affinity_enabled( + self.model_group_affinity_config + ) + ): + if self.optional_callbacks is None: + self.optional_callbacks = [] + + existing_ec_callback: Optional[EncryptedContentAffinityCheck] = None + for cb in self.optional_callbacks: + if isinstance(cb, EncryptedContentAffinityCheck): + existing_ec_callback = cb + break + + if existing_ec_callback is not None: + existing_ec_callback.router = self + existing_ec_callback.enable_global_affinity = ( + existing_ec_callback.enable_global_affinity + or enable_global_affinity + ) + existing_ec_callback.model_group_affinity_config = ( + self.model_group_affinity_config or {} + ) + ec_callback = existing_ec_callback + else: + ec_callback = EncryptedContentAffinityCheck( + router=self, + enable_global_affinity=enable_global_affinity, + model_group_affinity_config=self.model_group_affinity_config, + ) + self.optional_callbacks.append(ec_callback) + litellm.logging_callback_manager.add_litellm_callback(ec_callback) + + _move_before_deployment_affinity(self.optional_callbacks, ec_callback) + _move_before_deployment_affinity(litellm.callbacks, ec_callback) + def add_optional_pre_call_checks( self, optional_pre_call_checks: Optional[OptionalPreCallChecks] ): @@ -1721,22 +1782,11 @@ def add_optional_pre_call_checks( # --------------------------------------------------------------------- # Encrypted content affinity # --------------------------------------------------------------------- - if "encrypted_content_affinity" in optional_pre_call_checks: - from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, - ) - - if self.optional_callbacks is None: - self.optional_callbacks = [] - - already_registered = any( - isinstance(cb, EncryptedContentAffinityCheck) - for cb in self.optional_callbacks + self._add_encrypted_content_affinity_check( + enable_global_affinity=( + "encrypted_content_affinity" in optional_pre_call_checks ) - if not already_registered: - ec_callback = EncryptedContentAffinityCheck(router=self) - self.optional_callbacks.append(ec_callback) - litellm.logging_callback_manager.add_litellm_callback(ec_callback) + ) # --------------------------------------------------------------------- # Remaining optional pre-call checks @@ -7738,6 +7788,39 @@ def _generate_model_id(self, model_group: str, litellm_params: dict): return hash_object.hexdigest() + @staticmethod + def _inherit_builtin_cache_pricing( + model_info: dict, backend_model: str, custom_llm_provider: Optional[str] + ) -> None: + """Fill missing cache pricing on a custom-priced deployment entry from + the backend model's built-in cost map entry, so a deployment that + only spells out ``input_cost_per_token``/``output_cost_per_token`` + does not silently bill cache_read/cache_creation at 0. + + User-specified cache fields always win; only ``None``/missing entries + are inherited. No-op when the backend model has no canonical entry. + """ + cache_fields = ( + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost", + "cache_read_input_token_cost_above_200k_tokens", + ) + if all(model_info.get(f) is not None for f in cache_fields): + return + try: + backend_info = litellm.get_model_info( + model=backend_model, custom_llm_provider=custom_llm_provider + ) + except Exception: + return + for field in cache_fields: + if model_info.get(field) is None: + backend_value = backend_info.get(field) + if backend_value is not None: + model_info[field] = backend_value + def _create_deployment( self, deployment_info: dict, @@ -7766,6 +7849,13 @@ def _create_deployment( if deployment.litellm_params.get(field) is not None: _model_info[field] = deployment.litellm_params[field] + if _model_info.get("input_cost_per_token") is not None: + Router._inherit_builtin_cache_pricing( + model_info=_model_info, + backend_model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) + ## REGISTER MODEL INFO IN LITELLM MODEL COST MAP model_id = deployment.model_info.id if model_id is not None: @@ -8471,6 +8561,13 @@ def _initialize_deployment_for_pass_through( credential_values.get("api_key") or deployment.litellm_params.api_key ) + if api_key is None: + verbose_router_logger.debug( + "Skipping pass-through credential setup for deployment model=%s, custom_llm_provider=%s; no api_key set. Providers like bedrock resolve credentials at request time.", + model, + custom_llm_provider, + ) + return passthrough_endpoint_router.set_pass_through_credentials( custom_llm_provider=custom_llm_provider, api_base=api_base, @@ -8505,6 +8602,13 @@ def add_deployment(self, deployment: Deployment) -> Optional[Deployment]: if field_value is not None: _model_info_dict[field] = field_value + if _model_info_dict.get("input_cost_per_token") is not None: + Router._inherit_builtin_cache_pricing( + model_info=_model_info_dict, + backend_model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) + # Register custom pricing in litellm.model_cost. # Mirrors _create_deployment() logic to ensure dynamically-added deployments # (e.g., loaded from DB) also have their custom pricing registered. diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 3bccef36e68..4856d7ff4cd 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -55,6 +55,7 @@ _SESSION_STATE_SWEEP_THRESHOLD: int = 1024 # Same pattern for the owner cache. _OWNER_CACHE_SWEEP_THRESHOLD: int = 1024 +from litellm.repositories.table_repositories import AdaptiveRouterStateRepository from litellm.types.llms.openai import AllMessageValues from litellm.types.router import ( AdaptiveRouterConfig, @@ -113,7 +114,7 @@ async def load_state_from_db(self, prisma_client: Any) -> None: if prisma_client is None: return try: - rows = await prisma_client.db.litellm_adaptiverouterstate.find_many( + rows = await AdaptiveRouterStateRepository(prisma_client).table.find_many( where={"router_name": self.router_name} ) loaded = 0 diff --git a/litellm/router_strategy/adaptive_router/update_queue.py b/litellm/router_strategy/adaptive_router/update_queue.py index b667f3a53a7..1d87feddd84 100644 --- a/litellm/router_strategy/adaptive_router/update_queue.py +++ b/litellm/router_strategy/adaptive_router/update_queue.py @@ -22,6 +22,10 @@ from typing import Any, Dict, Tuple from litellm._logging import verbose_router_logger +from litellm.repositories.table_repositories import ( + AdaptiveRouterSessionRepository, + AdaptiveRouterStateRepository, +) StateKey = Tuple[str, str, str] # (router_name, request_type, model_name) SessionKey = Tuple[str, str, str] # (session_id, router_name, model_name) @@ -112,7 +116,7 @@ async def flush_state_to_db(self, prisma_client: Any) -> int: # other. The upsert creates the row with the delta as the # initial value on first write, then increments on subsequent # writes — no read-modify-write race. - await prisma_client.db.litellm_adaptiverouterstate.upsert( + await AdaptiveRouterStateRepository(prisma_client).table.upsert( where={ "router_name_request_type_model_name": { "router_name": router, @@ -174,7 +178,7 @@ async def flush_session_to_db(self, prisma_client: Any) -> int: for k, v in payload.items() if k not in ("session_id", "router_name", "model_name") } - await prisma_client.db.litellm_adaptiveroutersession.upsert( + await AdaptiveRouterSessionRepository(prisma_client).table.upsert( where={ "session_id_router_name_model_name": { "session_id": session_id, diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 148b7fce0ee..d3e7e2ffa34 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -39,7 +39,12 @@ class DeploymentAffinityCheck(CustomLogger): CACHE_KEY_PREFIX = "deployment_affinity:v1" VALID_FLAGS = frozenset( - {"deployment_affinity", "responses_api_deployment_check", "session_affinity"} + { + "deployment_affinity", + "responses_api_deployment_check", + "session_affinity", + "encrypted_content_affinity", + } ) def __init__( diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 4ed19c5cd26..5fd2be9c6dd 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -37,7 +37,7 @@ """ import time -from typing import TYPE_CHECKING, Any, List, Optional, cast +from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast import httpx @@ -64,17 +64,45 @@ class EncryptedContentAffinityCheck(CustomLogger): The ``model_id`` is decoded directly from the litellm-encoded item IDs – no caching or TTL management needed. - Wired via ``Router(optional_pre_call_checks=["encrypted_content_affinity"])``. + Wired via ``Router(optional_pre_call_checks=["encrypted_content_affinity"])`` or + per-model group ``model_group_affinity_config``. """ - def __init__(self, router: Optional["Router"] = None) -> None: + def __init__( + self, + router: Optional["Router"] = None, + enable_global_affinity: bool = True, + model_group_affinity_config: Optional[Dict[str, List[str]]] = None, + ) -> None: super().__init__() self.router = router + self.enable_global_affinity = enable_global_affinity + self.model_group_affinity_config: Dict[str, List[str]] = ( + model_group_affinity_config or {} + ) # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ + @staticmethod + def has_model_group_affinity_enabled( + model_group_affinity_config: Optional[Dict[str, List[str]]], + ) -> bool: + if not model_group_affinity_config: + return False + + return any( + "encrypted_content_affinity" in checks + for checks in model_group_affinity_config.values() + ) + + def _is_enabled_for_model_group(self, model_group: str) -> bool: + group_checks = self.model_group_affinity_config.get(model_group) + return self.enable_global_affinity or ( + group_checks is not None and "encrypted_content_affinity" in group_checks + ) + @staticmethod def _extract_model_id_from_input(request_input: Any) -> Optional[str]: """ @@ -213,6 +241,8 @@ async def async_filter_deployments( """ request_kwargs = request_kwargs or {} typed_healthy_deployments = cast(List[dict], healthy_deployments) + if not self._is_enabled_for_model_group(model): + return typed_healthy_deployments # Signal to the response post-processor that encrypted item IDs should be # encoded in the output of this request. Only set the flag when diff --git a/litellm/setup_wizard.py b/litellm/setup_wizard.py index 862ca13e7ba..2f0cb1233ae 100644 --- a/litellm/setup_wizard.py +++ b/litellm/setup_wizard.py @@ -52,11 +52,12 @@ { "id": "anthropic", "name": "Anthropic", - "description": "Claude Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5", + "description": "Claude Fable 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5", "env_key": "ANTHROPIC_API_KEY", "key_hint": "sk-ant-...", "test_model": "claude-haiku-4-5-20251001", "models": [ + "claude-fable-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6", diff --git a/litellm/types/caching.py b/litellm/types/caching.py index f8050b292c7..10453c74a15 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -118,4 +118,5 @@ class CachedEmbedding(TypedDict): index: Optional[int] object: Optional[str] model: Optional[str] + prompt_tokens: Optional[int] prompt_tokens_details: Optional[dict] diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 0d81e25592d..52fcb4c934f 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -357,6 +357,15 @@ class PresidioConfigModel(PresidioPresidioConfigModelUserInterface): default=None, description="Path to a JSON file containing ad-hoc recognizers for Presidio", ) + unmask_streamed_tool_calls: Optional[bool] = Field( + default=None, + description=( + "Only used with output_parse_pii on Anthropic native streaming. When True, " + "masked PII placeholders inside streamed tool-call arguments " + "(input_json_delta) are also restored. Off by default so streamed tool " + "executions never receive restored PII unless explicitly allowed." + ), + ) mock_redacted_text: Optional[dict] = Field( default=None, description="Mock redacted text for testing" ) diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 55f4fc96504..5b1d32cd93c 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -115,6 +115,8 @@ def all_error_messages(self) -> List[str]: REQUESTED_MODEL = "requested_model" EXCEPTION_STATUS = "exception_status" EXCEPTION_CLASS = "exception_class" +RATE_LIMIT_CATEGORY = "rate_limit_category" +RATE_LIMIT_TYPE = "rate_limit_type" STATUS_CODE = "status_code" EXCEPTION_LABELS = [EXCEPTION_STATUS, EXCEPTION_CLASS] LATENCY_BUCKETS = ( @@ -174,6 +176,8 @@ class UserAPIKeyLabelNames(Enum): API_PROVIDER = "api_provider" EXCEPTION_STATUS = EXCEPTION_STATUS EXCEPTION_CLASS = EXCEPTION_CLASS + RATE_LIMIT_CATEGORY = RATE_LIMIT_CATEGORY + RATE_LIMIT_TYPE = RATE_LIMIT_TYPE STATUS_CODE = "status_code" FALLBACK_MODEL = "fallback_model" ROUTE = "route" @@ -343,6 +347,10 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER_EMAIL.value, UserAPIKeyLabelNames.EXCEPTION_STATUS.value, UserAPIKeyLabelNames.EXCEPTION_CLASS.value, + # ``rate_limit_category`` / ``rate_limit_type`` are appended in + # ``get_labels()`` when ``litellm.prometheus_emit_rate_limit_labels`` + # is True. Kept opt-in so existing dashboards keyed on this metric's + # historical label set keep matching after upgrade. UserAPIKeyLabelNames.ROUTE.value, UserAPIKeyLabelNames.CLIENT_IP.value, UserAPIKeyLabelNames.USER_AGENT.value, @@ -745,6 +753,25 @@ def get_labels(label_name: DEFINED_PROMETHEUS_METRICS) -> List[str]: ): custom_labels.append(UserAPIKeyLabelNames.STREAM.value) + # Conditionally add unified rate-limit labels to + # litellm_proxy_failed_requests_metric. Off by default so the metric's + # historical label set is preserved across upgrade; enable via + # ``litellm.prometheus_emit_rate_limit_labels`` once downstream + # dashboards include the new labels in their matchers / aggregations. + if ( + label_name == "litellm_proxy_failed_requests_metric" + and litellm.prometheus_emit_rate_limit_labels is True + ): + for _rate_limit_label in ( + UserAPIKeyLabelNames.RATE_LIMIT_CATEGORY.value, + UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value, + ): + if ( + _rate_limit_label not in default_labels + and _rate_limit_label not in custom_labels + ): + custom_labels.append(_rate_limit_label) + _user_budget_metrics = { "litellm_remaining_user_budget_metric", "litellm_user_max_budget_metric", @@ -807,6 +834,8 @@ class UserAPIKeyLabelValues: api_provider: Optional[str] = None exception_status: Optional[str] = None exception_class: Optional[str] = None + rate_limit_category: Optional[str] = None + rate_limit_type: Optional[str] = None status_code: Optional[str] = None fallback_model: Optional[str] = None route: Optional[str] = None diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 60e640636aa..fa8c3a93ef3 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -250,6 +250,7 @@ class ToolJsonSchemaBlock(TypedDict, total=False): type: Literal["object"] properties: dict required: List[str] + additionalProperties: bool class ToolInputSchemaBlock(TypedDict): @@ -260,6 +261,7 @@ class ToolSpecBlock(TypedDict, total=False): inputSchema: Required[ToolInputSchemaBlock] name: Required[str] description: str + strict: bool class SystemToolBlock(TypedDict, total=False): @@ -283,6 +285,36 @@ class ToolBlock(TypedDict, total=False): cachePoint: Optional[CachePointBlock] +class BedrockToolSpec(dict): + def __init__( + self, + *, + name: str, + description: str, + parameters: dict, + strict: Optional[bool], + supports_strict_tools: bool, + ) -> None: + json_schema: ToolJsonSchemaBlock = { + "type": parameters["type"], + "properties": parameters.get("properties", {}), + "required": parameters.get("required", []), + } + additional_properties = parameters.get("additionalProperties") + if supports_strict_tools and additional_properties is not None: + json_schema["additionalProperties"] = additional_properties + + tool_spec: ToolSpecBlock = { + "inputSchema": {"json": json_schema}, + "name": name, + "description": description, + } + if supports_strict_tools and strict is not None: + tool_spec["strict"] = strict + + super().__init__(toolSpec=tool_spec) + + class SpecificToolChoiceBlock(TypedDict): name: str diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 51429d0769e..b28fee51284 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -232,6 +232,7 @@ class VoiceConfig(TypedDict): class SpeechConfig(TypedDict, total=False): voiceConfig: VoiceConfig + languageCode: str class GenerationConfig(TypedDict, total=False): @@ -757,3 +758,12 @@ class VertexPartnerProvider(str, Enum): llama = "llama" ai21 = "ai21" claude = "claude" + + +VERTEX_AI_PROVIDER_METADATA_FIELDS = ( + "vertex_ai_grounding_metadata", + "vertex_ai_url_context_metadata", + "vertex_ai_safety_ratings", + "vertex_ai_safety_results", + "vertex_ai_citation_metadata", +) diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 92ca027c5bf..809da6418d7 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -3,8 +3,7 @@ from pydantic import BaseModel, ConfigDict -from litellm.proxy._types import MCPAuthType, MCPTransportType -from litellm.types.mcp import MCPAuth +from litellm.types.mcp import MCPAuth, MCPAuthType, MCPTransportType # MCPInfo now allows arbitrary additional fields for custom metadata MCPInfo = Dict[str, Any] diff --git a/litellm/types/router.py b/litellm/types/router.py index ef7eb05d087..5047cee424b 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -178,6 +178,7 @@ class CredentialLiteLLMParams(BaseModel): aws_secret_access_key: Optional[str] = None aws_region_name: Optional[str] = None aws_bedrock_runtime_endpoint: Optional[str] = None + aws_bedrock_project_id: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None @@ -220,6 +221,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): use_in_pass_through: Optional[bool] = False use_litellm_proxy: Optional[bool] = False use_chat_completions_api: Optional[bool] = None + use_xai_oauth: Optional[bool] = Field( + default=False, + description="Use stored xAI OAuth credentials when no xAI API key is configured.", + ) model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) merge_reasoning_content_in_choices: Optional[bool] = False model_info: Optional[Dict] = None @@ -360,6 +365,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): aws_access_key_id: Optional[str] aws_secret_access_key: Optional[str] aws_region_name: Optional[str] + aws_bedrock_project_id: Optional[str] ## AWS S3 VECTORS ## vector_bucket_name: Optional[str] index_name: Optional[str] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a7a0b0f6238..21eb0c9a173 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -38,7 +38,6 @@ Field, PrivateAttr, field_validator, - model_validator, ) from typing_extensions import Required, TypedDict @@ -198,6 +197,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): ] # OpenAI priority service tier pricing cache_read_input_token_cost_above_200k_tokens: Optional[float] cache_read_input_token_cost_above_272k_tokens: Optional[float] + cache_read_input_token_cost_above_512k_tokens: Optional[float] input_cost_per_character: Optional[float] # only for vertex ai models input_cost_per_audio_token: Optional[float] input_cost_per_token_above_128k_tokens: Optional[float] # only for vertex ai models @@ -207,6 +207,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_token_above_272k_tokens: Optional[ float ] # GPT-5.4/5.4-pro: prompts >272K priced at 2x input + input_cost_per_token_above_512k_tokens: Optional[ + float + ] # MiniMax-M3: prompts >512K priced at 2x input input_cost_per_character_above_128k_tokens: Optional[ float ] # only for vertex ai models @@ -240,6 +243,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token_above_272k_tokens: Optional[ float ] # GPT-5.4/5.4-pro: prompts >272K priced at 1.5x output + output_cost_per_token_above_512k_tokens: Optional[ + float + ] # MiniMax-M3: prompts >512K priced at 2x output output_cost_per_character_above_128k_tokens: Optional[ float ] # only for vertex ai models @@ -1541,6 +1547,11 @@ class ServerToolUse(BaseModel): web_search_requests: Optional[int] = None tool_search_requests: Optional[int] = None + def __getitem__(self, key: str) -> Optional[int]: + if key not in self.__class__.model_fields: + raise KeyError(key) + return getattr(self, key) + class Usage(SafeAttributeModel, CompletionUsage): _cache_creation_input_tokens: int = PrivateAttr( @@ -1571,7 +1582,7 @@ def __init__( # noqa: PLR0915 completion_tokens_details: Optional[ Union[CompletionTokensDetailsWrapper, dict] ] = None, - server_tool_use: Optional[ServerToolUse] = None, + server_tool_use: Optional[Union[ServerToolUse, dict]] = None, cost: Optional[float] = None, **params, ): @@ -1672,6 +1683,9 @@ def __init__( # noqa: PLR0915 prompt_tokens_details=_prompt_tokens_details or None, ) + if isinstance(server_tool_use, dict): + server_tool_use = ServerToolUse(**server_tool_use) + if server_tool_use is not None: self.server_tool_use = server_tool_use else: # maintain openai compatibility in usage object if possible @@ -2720,6 +2734,23 @@ class StandardLoggingPayloadErrorInformation(TypedDict, total=False): llm_provider: Optional[str] traceback: Optional[str] error_message: Optional[str] + # error_rate_limit_category: + # For 429 / rate-limit errors, the source of the rate limit. One of the + # string values defined by `litellm.exceptions.RateLimitErrorCategory` + # (vendor_rate_limit, vendor_batch_rate_limit, litellm_rate_limit, + # litellm_batch_rate_limit). None for non-rate-limit exceptions. + # Surfaced here so custom callbacks / metrics consumers can switch on + # the rate-limit source without reaching for the raw exception. + error_rate_limit_category: Optional[str] + # error_rate_limit_type: + # For 429 / rate-limit errors, the dimension that was exceeded. One of + # the string values defined by `litellm.exceptions.RateLimitType` + # (requests, tokens, concurrent_requests, budget, max_iterations). + # None for non-rate-limit exceptions and for rate-limit exceptions that + # did not classify the failure (e.g. legacy vendor 429 with no header + # hints). Lets dashboards split rate-limit failures by cause without + # parsing free-text error messages. + error_rate_limit_type: Optional[str] class GuardrailMode(TypedDict, total=False): @@ -3193,6 +3224,7 @@ class CustomPricingLiteLLMParams(BaseModel): "search_tool_name", "order", "enable_json_schema_validation", + "use_xai_oauth", ] + list(StandardCallbackDynamicParams.__annotations__.keys()) + list(CustomPricingLiteLLMParams.model_fields.keys()) @@ -3376,6 +3408,7 @@ class LlmProviders(str, Enum): POE = "poe" CHUTES = "chutes" NEOSANTARA = "neosantara" + PARASAIL = "parasail" XIAOMI_MIMO = "xiaomi_mimo" TENSORMESH = "tensormesh" LITELLM_AGENT = "litellm_agent" @@ -3560,25 +3593,11 @@ class RawRequestTypedDict(TypedDict, total=False): error: Optional[str] -class CredentialBase(BaseModel): - credential_name: str - credential_info: dict - - -class CredentialItem(CredentialBase): - credential_values: dict - - -class CreateCredentialItem(CredentialBase): - credential_values: Optional[dict] = None - model_id: Optional[str] = None - - @model_validator(mode="before") - @classmethod - def check_credential_params(cls, values): - if not values.get("credential_values") and not values.get("model_id"): - raise ValueError("Either credential_values or model_id must be set") - return values +from litellm.models.credentials import CredentialBase as CredentialBase # noqa: E402 +from litellm.models.credentials import CredentialItem as CredentialItem # noqa: E402 +from litellm.models.credentials import ( # noqa: E402 + CreateCredentialItem as CreateCredentialItem, +) class ExtractedFileData(TypedDict): diff --git a/litellm/utils.py b/litellm/utils.py index 7312e71bbd1..a0b66234a70 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2887,6 +2887,61 @@ def _convert_stringified_numbers(value): return value +_BEDROCK_REGION_PREFIXES = ( + "us.", + "eu.", + "apac.", + "jp.", + "au.", + "us-gov.", + "global.", + "ap-northeast-1.", +) + +_CACHE_PRICING_FIELDS = ( + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost", + "cache_read_input_token_cost_above_200k_tokens", +) + + +def _resolve_builtin_model_cost_entry( + key: str, provider: str +) -> Optional[Dict[str, Any]]: + """Best-effort lookup of a built-in ``model_cost`` entry for a custom key + whose shape ``get_model_info`` cannot resolve (double provider prefixes + like ``bedrock/bedrock/us.anthropic.claude-sonnet-4-6`` or region aliases). + + Returns a copy of the matching entry so the caller can inherit its defaults + (most importantly cache pricing) without mutating the shared built-in. + Returns ``None`` when no safe match exists. + """ + candidates: List[str] = [] + segments = key.split("/") + idx = 0 + while idx < len(segments) - 1 and segments[idx] in LlmProvidersSet: + idx += 1 + candidates.append("/".join(segments[idx:])) + + base = candidates[-1] if candidates else key + for region_prefix in _BEDROCK_REGION_PREFIXES: + if base.startswith(region_prefix): + candidates.append(base[len(region_prefix) :]) + + if provider: + stripped = _strip_model_name(model=base, custom_llm_provider=provider) + if stripped != base: + candidates.append(stripped) + + for candidate in candidates: + entry = litellm.model_cost.get(candidate) + if entry is not None and entry.get("litellm_provider") is not None: + return dict(entry) + return None + + def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 """ Register new / Override existing models (and their pricing) to specific providers. @@ -2933,6 +2988,26 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 except Exception: existing_model = {} model_cost_key = key + builtin_entry = _resolve_builtin_model_cost_entry( + key=_key_str, provider=provider + ) + if builtin_entry is not None: + for field in _CACHE_PRICING_FIELDS: + if ( + value.get(field) is None + and builtin_entry.get(field) is not None + ): + existing_model[field] = builtin_entry[field] + elif ( + value.get("cache_creation_input_token_cost") is None + and value.get("cache_read_input_token_cost") is None + ): + verbose_logger.warning( + f"register_model: model={key} not in built-in cost map and no " + "prefix/region variant matched; cache cost fields will default " + "to 0. To track cache cost, add cache_creation_input_token_cost " + "and cache_read_input_token_cost to model_info" + ) # ``get_model_info`` returns ``litellm_provider: None`` when the # provider is unknown (e.g. custom deployments registered via # ``Router.add_deployment``). Persisting that None into @@ -3762,6 +3837,10 @@ def base_pre_process_non_default_params( additional_endpoint_specific_params: List[str], ) -> dict: for k, v in special_params.items(): + if k == "aws_bedrock_project_id": + # sent as a request header (read from litellm_params by the + # bedrock-mantle configs), never as a request body field + continue if k.startswith("aws_") and ( custom_llm_provider != "bedrock" and not custom_llm_provider.startswith("sagemaker") @@ -5775,6 +5854,7 @@ def _get_model_info_helper( # noqa: PLR0915 ] split_model = potential_model_names["split_model"] custom_llm_provider = potential_model_names["custom_llm_provider"] + model_cost_custom_llm_provider = custom_llm_provider ######################### provider_config: Optional[BaseLLMModelInfo] = None if custom_llm_provider and custom_llm_provider in LlmProvidersSet: @@ -5840,7 +5920,8 @@ def _get_model_info_helper( # noqa: PLR0915 key = _matched_key _model_info = _get_model_info_from_model_cost(key=cast(str, key)) if not _check_provider_match( - model_info=_model_info, custom_llm_provider=custom_llm_provider + model_info=_model_info, + custom_llm_provider=model_cost_custom_llm_provider, ): _model_info = None if _model_info is None: @@ -5849,7 +5930,8 @@ def _get_model_info_helper( # noqa: PLR0915 key = _matched_key _model_info = _get_model_info_from_model_cost(key=cast(str, key)) if not _check_provider_match( - model_info=_model_info, custom_llm_provider=custom_llm_provider + model_info=_model_info, + custom_llm_provider=model_cost_custom_llm_provider, ): _model_info = None if _model_info is None: @@ -5858,7 +5940,8 @@ def _get_model_info_helper( # noqa: PLR0915 key = _matched_key _model_info = _get_model_info_from_model_cost(key=cast(str, key)) if not _check_provider_match( - model_info=_model_info, custom_llm_provider=custom_llm_provider + model_info=_model_info, + custom_llm_provider=model_cost_custom_llm_provider, ): _model_info = None if _model_info is None: @@ -5867,7 +5950,8 @@ def _get_model_info_helper( # noqa: PLR0915 key = _matched_key _model_info = _get_model_info_from_model_cost(key=cast(str, key)) if not _check_provider_match( - model_info=_model_info, custom_llm_provider=custom_llm_provider + model_info=_model_info, + custom_llm_provider=model_cost_custom_llm_provider, ): _model_info = None if _model_info is None: @@ -5876,7 +5960,8 @@ def _get_model_info_helper( # noqa: PLR0915 key = _matched_key _model_info = _get_model_info_from_model_cost(key=cast(str, key)) if not _check_provider_match( - model_info=_model_info, custom_llm_provider=custom_llm_provider + model_info=_model_info, + custom_llm_provider=model_cost_custom_llm_provider, ): _model_info = None @@ -5884,7 +5969,6 @@ def _get_model_info_helper( # noqa: PLR0915 raise ValueError( "This model isn't mapped yet. Add it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json" ) - _input_cost_per_token: Optional[float] = _model_info.get( "input_cost_per_token" ) @@ -5936,6 +6020,9 @@ def _get_model_info_helper( # noqa: PLR0915 cache_read_input_token_cost_above_272k_tokens=_model_info.get( "cache_read_input_token_cost_above_272k_tokens", None ), + cache_read_input_token_cost_above_512k_tokens=_model_info.get( + "cache_read_input_token_cost_above_512k_tokens", None + ), cache_read_input_token_cost_flex=_model_info.get( "cache_read_input_token_cost_flex", None ), @@ -5957,6 +6044,9 @@ def _get_model_info_helper( # noqa: PLR0915 input_cost_per_token_above_272k_tokens=_model_info.get( "input_cost_per_token_above_272k_tokens", None ), + input_cost_per_token_above_512k_tokens=_model_info.get( + "input_cost_per_token_above_512k_tokens", None + ), input_cost_per_query=_model_info.get("input_cost_per_query", None), input_cost_per_second=_model_info.get("input_cost_per_second", None), input_cost_per_audio_token=_model_info.get( @@ -6012,6 +6102,9 @@ def _get_model_info_helper( # noqa: PLR0915 output_cost_per_token_above_272k_tokens=_model_info.get( "output_cost_per_token_above_272k_tokens", None ), + output_cost_per_token_above_512k_tokens=_model_info.get( + "output_cost_per_token_above_512k_tokens", None + ), output_cost_per_second=_model_info.get("output_cost_per_second", None), output_cost_per_second_1080p=_model_info.get( "output_cost_per_second_1080p", None @@ -8895,7 +8988,13 @@ def _get_python_responses_api_config( elif litellm.LlmProviders.XAI == provider: return litellm.XAIResponsesAPIConfig() elif litellm.LlmProviders.GITHUB_COPILOT == provider: - return litellm.GithubCopilotResponsesAPIConfig() + from litellm.llms.github_copilot.responses.transformation import ( + github_copilot_supports_responses_api, + ) + + if model is None or github_copilot_supports_responses_api(model=model): + return litellm.GithubCopilotResponsesAPIConfig() + return None elif litellm.LlmProviders.CHATGPT == provider: return litellm.ChatGPTResponsesAPIConfig() elif litellm.LlmProviders.LITELLM_PROXY == provider: @@ -8916,14 +9015,33 @@ def _get_python_responses_api_config( elif litellm.LlmProviders.HOSTED_VLLM == provider: return litellm.HostedVLLMResponsesAPIConfig() elif litellm.LlmProviders.BEDROCK_MANTLE == provider: - # Only OpenAI gpt frontier models (gpt-5.x, and future gpt-6 etc.) are - # served on the /openai/v1/responses path. gpt-oss and every non-OpenAI - # model on Mantle (nvidia, mistral, google, zai, ...) are chat-completions - # only and 400 on that path, so they fall through to None to keep the - # chat-completions emulation (see litellm/responses/main.py "config is None"). - model_lower = model.lower() if model else "" - if "openai.gpt-" in model_lower and "gpt-oss" not in model_lower: - return litellm.BedrockMantleResponsesAPIConfig() + # Mantle serves Responses on two upstream paths. A model takes the + # /openai/v1/responses path when its price-map entry declares + # use_openai_responses_path (data-driven, so a non-gpt-named frontier + # model can be onboarded by JSON alone), or, as a fallback needing no + # price-map entry, when its name matches the openai.gpt- frontier + # convention (minus gpt-oss) -- this keeps a future gpt-6 routing + # correctly before its entry loads. Any other model declared + # mode=responses takes the standard /v1/responses path. Everything + # else returns None and keeps the chat-completions emulation (see + # responses/main.py "config is None"). + if not model: + return None + model_lower = model.lower() + entry = litellm.model_cost.get(f"bedrock_mantle/{model}", {}) + on_openai_path = entry.get("use_openai_responses_path") is True + name_is_frontier = ( + "openai.gpt-" in model_lower and "gpt-oss" not in model_lower + ) + if on_openai_path or name_is_frontier: + return litellm.BedrockMantleResponsesAPIConfig(use_openai_path=True) + try: + if get_model_info(model, "bedrock_mantle").get("mode") == "responses": + return litellm.BedrockMantleResponsesAPIConfig( + use_openai_path=False + ) + except Exception: + pass return None return None diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 1fd95b16309..94f0483e1cc 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -5,6 +5,10 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import remove_items_at_indices +from litellm.repositories.table_repositories import ( + ManagedVectorStoreIndexRepository, + ManagedVectorStoresRepository, +) from litellm.types.vector_stores import ( VECTOR_STORE_OPENAI_PARAMS, LiteLLM_ManagedVectorStore, @@ -91,10 +95,10 @@ async def _get_vector_store_indexes_from_db( """ vector_stores_from_db: List[LiteLLM_ManagedVectorStoreIndex] = [] if prisma_client is not None: - _vector_stores_from_db = ( - await prisma_client.db.litellm_managedvectorstoreindextable.find_many( - order={"created_at": "desc"}, - ) + _vector_stores_from_db = await ManagedVectorStoreIndexRepository( + prisma_client + ).table.find_many( + order={"created_at": "desc"}, ) for vector_store in _vector_stores_from_db: _dict_vector_store = dict(vector_store) @@ -374,9 +378,9 @@ async def pop_vector_stores_to_run_with_db_fallback( if vector_store is not None and prisma_client is not None: try: # Check if it still exists in database - db_vector_store = await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": vector_store_id} - ) + db_vector_store = await ManagedVectorStoresRepository( + prisma_client + ).table.find_unique(where={"vector_store_id": vector_store_id}) if db_vector_store is None: # Vector store was deleted from database, remove from cache verbose_logger.debug( @@ -541,10 +545,10 @@ async def _get_vector_stores_from_db( """ vector_stores_from_db: List[LiteLLM_ManagedVectorStore] = [] if prisma_client is not None: - _vector_stores_from_db = ( - await prisma_client.db.litellm_managedvectorstorestable.find_many( - order={"created_at": "desc"}, - ) + _vector_stores_from_db = await ManagedVectorStoresRepository( + prisma_client + ).table.find_many( + order={"created_at": "desc"}, ) for vector_store in _vector_stores_from_db: _dict_vector_store = dict(vector_store) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b2836a096b7..f0b2432ddc8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1156,6 +1156,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1202,6 +1203,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1233,6 +1235,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1264,6 +1267,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1295,6 +1299,139 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "global.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "us.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "eu.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1327,6 +1464,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1359,6 +1497,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1391,6 +1530,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1423,6 +1563,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1455,6 +1596,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1485,6 +1627,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2208,6 +2351,37 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "azure_ai/claude-fable-5": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2237,6 +2411,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -6889,6 +7064,43 @@ "/v1/images/generations" ] }, + "azure_ai/MAI-Image-2.5": { + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.05, + "output_cost_per_image_token": 4.7e-05, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/new-mai-models-in-microsoft-foundry-across-text-image-voice-and-speech/4524632", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "azure_ai/MAI-Image-2.5-Flash": { + "input_cost_per_image_token": 1.75e-06, + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.0338, + "output_cost_per_image_token": 3.3e-05, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/new-mai-models-in-microsoft-foundry-across-text-image-voice-and-speech/4524632", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "azure_ai/MAI-Image-2e": { + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.02, + "output_cost_per_image_token": 1.95e-05, + "source": "https://aka.ms/mai-image-2e-foundryblog", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "azure_ai/Llama-3.2-11B-Vision-Instruct": { "input_cost_per_token": 3.7e-07, "litellm_provider": "azure_ai", @@ -10133,6 +10345,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10167,6 +10380,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10177,6 +10391,40 @@ }, "supports_output_config": true }, + "claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true + }, "claude-opus-4-8": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -10201,6 +10449,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -14286,10 +14535,10 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://fireworks.ai/models/fireworks/glm-5p1", - "supports_function_calling": false, + "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": false, - "supports_tool_choice": false + "supports_response_schema": true, + "supports_tool_choice": true }, "fireworks_ai/accounts/fireworks/models/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, @@ -14567,10 +14816,10 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://fireworks.ai/models/fireworks/glm-5p1", - "supports_function_calling": false, + "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": false, - "supports_tool_choice": false + "supports_response_schema": true, + "supports_tool_choice": true }, "fireworks_ai/kimi-k2p5": { "cache_read_input_token_cost": 1e-07, @@ -24143,9 +24392,12 @@ "max_output_tokens": 8192 }, "minimax/MiniMax-M3": { - "input_cost_per_token": 6e-07, - "output_cost_per_token": 2.4e-06, - "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 3e-07, + "input_cost_per_token_above_512k_tokens": 6e-07, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_512k_tokens": 2.4e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_512k_tokens": 1.2e-07, "litellm_provider": "minimax", "mode": "chat", "supports_function_calling": true, @@ -24154,7 +24406,7 @@ "supports_reasoning": true, "supports_system_messages": true, "supports_vision": true, - "max_input_tokens": 512000, + "max_input_tokens": 1000000, "max_output_tokens": 128000 }, "mistral.devstral-2-123b": { @@ -34007,6 +34259,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34035,6 +34288,67 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5@default": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34064,6 +34378,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34093,6 +34408,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -41373,6 +41689,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "use_openai_responses_path": true, "supported_endpoints": ["/v1/responses"], "supported_modalities": ["text", "image"], "supported_output_modalities": ["text"], @@ -41392,6 +41709,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "use_openai_responses_path": true, "supported_endpoints": ["/v1/responses"], "supported_modalities": ["text", "image"], "supported_output_modalities": ["text"], @@ -41865,5 +42183,164 @@ "source": "https://soniox.com/pricing", "supported_endpoints": ["/v1/audio/transcriptions"], "supports_audio_input": true + }, + "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 6e-07, + "output_cost_per_token": 3.6e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 4.5e-07, + "output_cost_per_token": 1.8e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3.6-27B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3.2e-07, + "output_cost_per_token": 3.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 202752, + "max_output_tokens": 202752, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/deepseek-ai/DeepSeek-V4-Flash": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/moonshotai/Kimi-K2.6": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 9.6e-07, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/MiniMaxAI/MiniMax-M2.5": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 196608, + "max_output_tokens": 196608, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/google/gemma-4-31B-it": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 5.6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-120b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-20b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" } } diff --git a/packaging/homebrew/README.md b/packaging/homebrew/README.md new file mode 100644 index 00000000000..ef441ded304 --- /dev/null +++ b/packaging/homebrew/README.md @@ -0,0 +1,27 @@ +# Homebrew formula for the `lite` CLI + +[`lite.rb`](./lite.rb) is the canonical source for the Homebrew formula that installs the thin LiteLLM CLI (`litellm[cli]`). It lives here so it is versioned with the code, but Homebrew serves formulae from a tap, so it has to be published to the `BerriAI/homebrew-litellm` tap to be installable. + +Once published, end users install with + +```shell +brew install BerriAI/litellm/lite +``` + +which gives them the `lite` command (`lite login`, `lite claude`, `lite models list`, ...) without the proxy server runtime. For the full proxy server, they keep using pip/uv with `litellm[proxy]` or the Docker image. + +## Why a tap and not homebrew-core + +The formula builds the published `litellm` sdist with the `cli` extra and resolves that extra's dependencies from PyPI at build time. homebrew-core forbids network access during `install` and would require every transitive dependency declared as a pinned `resource`, regenerated on each release. For a fast-moving CLI that tradeoff is not worth it, so this stays a tap formula. + +## Release runbook + +The formula can only point at a published artifact, so it activates with the first `litellm` release that ships the `cli` extra (added in [pyproject.toml](../../pyproject.toml)). + +1. Cut a `litellm` release whose `pyproject.toml` includes the `cli` extra and confirm it is on PyPI. +2. Fetch the sdist URL and checksum for that version: `curl -fsSL https://pypi.org/pypi/litellm//json | jq -r '.urls[] | select(.packagetype=="sdist") | "\(.url)\n\(.digests.sha256)"'` +3. Set `url` and `sha256` in `lite.rb` to those values; `version` is parsed from `url`. +4. Copy `lite.rb` into the tap repo under `Formula/lite.rb`, then run `brew install --build-from-source ./Formula/lite.rb` and `brew test lite` to verify a clean build and that `lite --help` works. +5. Commit and push to `BerriAI/homebrew-litellm`. + +Keep `lite.rb` here in sync with the tap copy so the in-repo formula stays the source of truth. diff --git a/packaging/homebrew/lite.rb b/packaging/homebrew/lite.rb new file mode 100644 index 00000000000..d0d61bb5b43 --- /dev/null +++ b/packaging/homebrew/lite.rb @@ -0,0 +1,33 @@ +# Homebrew formula for the thin LiteLLM `lite` CLI (litellm[cli]). +# +# Ships in the BerriAI/homebrew-litellm tap, not homebrew-core: it builds the +# published litellm sdist with the `cli` extra into a dedicated virtualenv and +# pulls the extra's deps from PyPI. That is the low-maintenance path for a +# fast-moving Python CLI; the resource-stanza alternative would need every +# transitive dep re-pinned with a fresh sha256 on each release. +# +# RELEASE STEP (see README.md in this directory): point `url` + `sha256` at the +# PyPI sdist of the first litellm version that ships the `cli` extra. `version` +# is parsed from `url`, and the build installs exactly that version, so the three +# stay in lockstep automatically. +class Lite < Formula + include Language::Python::Virtualenv + + desc "Thin client for the LiteLLM proxy: lite login, lite claude/codex/opencode" + homepage "https://docs.litellm.ai/docs/proxy/management_cli" + url "https://files.pythonhosted.org/packages/source/l/litellm/litellm-REPLACE_AT_RELEASE.tar.gz" + sha256 "REPLACE_AT_RELEASE" + license "MIT" + + depends_on "python@3.13" + + def install + virtualenv_create(libexec, "python3.13") + system libexec/"bin/pip", "install", "#{buildpath}[cli]" + bin.install_symlink libexec/"bin/lite" + end + + test do + assert_match "login", shell_output("#{bin}/lite --help") + end +end diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index a1ad20fffd1..6caab585ac9 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1834,6 +1834,23 @@ "search": true } }, + "parasail": { + "display_name": "Parasail (`parasail`)", + "url": "https://docs.litellm.ai/docs/providers/parasail", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "perplexity": { "display_name": "Perplexity AI (`perplexity`)", "url": "https://docs.litellm.ai/docs/providers/perplexity", diff --git a/pyproject.toml b/pyproject.toml index 577e800d79b..b9d76379faf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,7 @@ proxy = [ "orjson>=3.11.6,<4.0", "apscheduler>=3.11.2,<4.0", "fastapi-sso>=0.19.0,<1.0", - "PyJWT>=2.12.0,<3.0", + "PyJWT>=2.13.0,<3.0", "python-multipart>=0.0.27,<1.0", "cryptography>=46.0.7,<47.0", "pynacl>=1.6.2,<2.0", @@ -71,6 +71,14 @@ proxy = [ "pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'", "pydantic-settings>=2.14.1,<3.0", ] +# Thin client install for the `lite` CLI on developer laptops. The CLI's heavy +# imports (fastapi, cryptography, ...) are all guarded, so it runs on the base +# SDK plus just these three; none of the server runtime in `proxy` is pulled in. +cli = [ + "rich>=13.9.4,<14.0", + "pyyaml>=6.0.3,<7.0", + "requests>=2.32.0,<3.0", +] extra_proxy = [ "prisma>=0.11.0,<1.0", "azure-identity>=1.25.2,<2.0", @@ -132,6 +140,7 @@ proxy-runtime = [ [project.scripts] litellm = "litellm:run_server" +lite = "litellm.proxy.client.cli:cli" litellm-proxy = "litellm.proxy.client.cli:cli" [dependency-groups] diff --git a/scripts/install-cli.sh b/scripts/install-cli.sh new file mode 100755 index 00000000000..d147286fcac --- /dev/null +++ b/scripts/install-cli.sh @@ -0,0 +1,128 @@ +#!/usr/bin/env bash +# LiteLLM CLI Installer (the thin `lite` client) +# Usage: curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install-cli.sh | sh +# +# Installs only litellm[cli]: the `lite` command for authenticating to a LiteLLM +# proxy and running coding agents (lite claude / codex / opencode) through it. +# None of the proxy server runtime is pulled in. To run a proxy server instead, +# use scripts/install.sh, which installs litellm[proxy]. +# +# Needs only curl: uv is bootstrapped if missing, and uv provisions a compatible +# Python itself (honouring litellm's requires-python), downloading a managed one +# when the host has no suitable interpreter. +# +# NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian +# ignores the shebang when invoked as `sh` and does not support `pipefail`). +set -eu + +# NOTE: before merging, this must stay as "litellm[cli]" to install from PyPI. +LITELLM_PACKAGE="litellm[cli]" +UV_VERSION="0.10.9" + +# ── colours ──────────────────────────────────────────────────────────────── +if [ -t 1 ]; then + BOLD='\033[1m' + GREEN='\033[38;2;78;186;101m' + GREY='\033[38;2;153;153;153m' + RESET='\033[0m' +else + BOLD='' GREEN='' GREY='' RESET='' +fi + +info() { printf "${GREY} %s${RESET}\n" "$*"; } +success() { printf "${GREEN} ✔ %s${RESET}\n" "$*"; } +header() { printf "${BOLD} %s${RESET}\n" "$*"; } +die() { printf "\n Error: %s\n\n" "$*" >&2; exit 1; } + +# ── banner ───────────────────────────────────────────────────────────────── +echo "" +cat << 'EOF' + ██╗ ██╗████████╗███████╗ + ██║ ██║╚══██╔══╝██╔════╝ + ██║ ██║ ██║ █████╗ + ██║ ██║ ██║ ██╔══╝ + ███████╗██║ ██║ ███████╗ + ╚══════╝╚═╝ ╚═╝ ╚══════╝ +EOF +printf " ${BOLD}LiteLLM CLI Installer${RESET} ${GREY}the thin 'lite' client for your proxy${RESET}\n\n" + +# ── OS detection ─────────────────────────────────────────────────────────── +OS="$(uname -s)" +ARCH="$(uname -m)" + +case "$OS" in + Darwin) PLATFORM="macOS ($ARCH)" ;; + Linux) PLATFORM="Linux ($ARCH)" ;; + *) die "Unsupported OS: $OS. LiteLLM supports macOS and Linux." ;; +esac + +info "Platform: $PLATFORM" + +# ── uv detection / install ──────────────────────────────────────────────── +UV_BIN="" +CURRENT_UV_VERSION="" +for candidate in uv "$HOME/.local/bin/uv"; do + if command -v "$candidate" >/dev/null 2>&1; then + UV_BIN="$(command -v "$candidate")" + break + elif [ -x "$candidate" ]; then + UV_BIN="$candidate" + break + fi +done + +if [ -n "$UV_BIN" ]; then + CURRENT_UV_VERSION="$("$UV_BIN" --version 2>/dev/null | awk '{print $2}' | head -1 || true)" +fi + +if [ -z "$UV_BIN" ] || [ "${CURRENT_UV_VERSION:-}" != "$UV_VERSION" ]; then + header "Installing uv…" + if [ -n "${CURRENT_UV_VERSION:-}" ]; then + info "Upgrading uv from ${CURRENT_UV_VERSION} to ${UV_VERSION}" + fi + curl -LsSf "https://astral.sh/uv/${UV_VERSION}/install.sh" | env UV_NO_MODIFY_PATH=1 sh \ + || die "uv installation failed. Try manually: curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh" + UV_BIN="$HOME/.local/bin/uv" +fi + +# ── install ──────────────────────────────────────────────────────────────── +# --python-preference system: reuse a compatible system Python when present, +# otherwise download a managed one. Either way uv honours litellm's requires-python, +# so a too-old (3.9) or too-new (3.14+) system Python is skipped, not forced. +echo "" +header "Installing litellm[cli]…" +echo "" + +"$UV_BIN" tool install --python-preference system --force "${LITELLM_PACKAGE}" \ + || die "uv tool install failed. Try manually: $UV_BIN tool install '${LITELLM_PACKAGE}'" + +# ── find the lite binary installed by uv tool ────────────────────────────── +SCRIPTS_DIR="$("$UV_BIN" tool dir --bin)" +LITE_BIN="${SCRIPTS_DIR}/lite" + +if [ ! -x "$LITE_BIN" ]; then + die "lite binary not found after install. Try: $UV_BIN tool install '${LITELLM_PACKAGE}'" +fi + +# ── success banner ───────────────────────────────────────────────────────── +echo "" +success "LiteLLM CLI installed" + +installed_ver="$("$LITE_BIN" --version 2>&1 | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -1 || true)" +[ -n "$installed_ver" ] && info "Version: $installed_ver" + +# ── PATH hint ────────────────────────────────────────────────────────────── +if ! command -v lite >/dev/null 2>&1; then + info "Note: add lite to your PATH: export PATH=\"\$PATH:${SCRIPTS_DIR}\"" +fi + +# ── next steps ───────────────────────────────────────────────────────────── +echo "" +header "Next steps:" +echo "" +info " export LITELLM_PROXY_URL=https://your-proxy # point at your gateway" +info " lite login # authenticate via SSO" +info " lite claude # run Claude Code through the proxy" +echo "" +info "Docs: https://docs.litellm.ai/docs/proxy/management_cli" +echo "" diff --git a/scripts/install.sh b/scripts/install.sh index c28d7da872f..06e6249c9ba 100755 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -2,13 +2,13 @@ # LiteLLM Installer # Usage: curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh # +# Needs only curl: uv is bootstrapped if missing, and uv provisions a compatible +# Python itself (reusing a suitable system one, else downloading a managed build). +# # NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian # ignores the shebang when invoked as `sh` and does not support `pipefail`). set -eu -MIN_PYTHON_MAJOR=3 -MIN_PYTHON_MINOR=9 - # NOTE: before merging, this must stay as "litellm[proxy]" to install from PyPI. LITELLM_PACKAGE="litellm[proxy]" UV_VERSION="0.10.9" @@ -52,27 +52,6 @@ esac info "Platform: $PLATFORM" -# ── Python detection ─────────────────────────────────────────────────────── -PYTHON_BIN="" -for candidate in python3 python; do - if command -v "$candidate" >/dev/null 2>&1; then - major="$("$candidate" -c 'import sys; print(sys.version_info.major)' 2>/dev/null || true)" - minor="$("$candidate" -c 'import sys; print(sys.version_info.minor)' 2>/dev/null || true)" - if [ "${major:-0}" -ge "$MIN_PYTHON_MAJOR" ] && [ "${minor:-0}" -ge "$MIN_PYTHON_MINOR" ]; then - PYTHON_BIN="$(command -v "$candidate")" - info "Python: $("$candidate" --version 2>&1)" - break - fi - fi -done - -if [ -z "$PYTHON_BIN" ]; then - die "Python ${MIN_PYTHON_MAJOR}.${MIN_PYTHON_MINOR}+ is required but not found. - Install it from https://python.org/downloads or via your package manager: - macOS: brew install python@3 - Ubuntu: sudo apt install python3" -fi - # ── uv detection / install ──────────────────────────────────────────────── UV_BIN="" CURRENT_UV_VERSION="" @@ -105,15 +84,18 @@ echo "" header "Installing litellm[proxy]…" echo "" -"$UV_BIN" tool install --python "$PYTHON_BIN" --force "${LITELLM_PACKAGE}" \ - || die "uv tool install failed. Try manually: $UV_BIN tool install --python '$PYTHON_BIN' '${LITELLM_PACKAGE}'" +# --python-preference system: reuse a compatible system Python when present, +# otherwise download a managed one. Either way uv honours litellm's requires-python, +# so a too-old (3.9) or too-new (3.14+) system Python is skipped, not forced. +"$UV_BIN" tool install --python-preference system --force "${LITELLM_PACKAGE}" \ + || die "uv tool install failed. Try manually: $UV_BIN tool install '${LITELLM_PACKAGE}'" # ── find the litellm binary installed by uv tool ─────────────────────────── SCRIPTS_DIR="$("$UV_BIN" tool dir --bin)" LITELLM_BIN="${SCRIPTS_DIR}/litellm" if [ ! -x "$LITELLM_BIN" ]; then - die "litellm binary not found after install. Try: $UV_BIN tool install --python '$PYTHON_BIN' '${LITELLM_PACKAGE}'" + die "litellm binary not found after install. Try: $UV_BIN tool install '${LITELLM_PACKAGE}'" fi # ── success banner ───────────────────────────────────────────────────────── diff --git a/security.md b/security.md index c6cd64ddaac..cb5eda7ee22 100644 --- a/security.md +++ b/security.md @@ -3,12 +3,20 @@ ## Security Vulnerability Reporting Guidelines +> [!WARNING] +> Reports that do not include a video demonstrating the exploit will be closed without review. See [Reproduction Video Requirement](#reproduction-video-requirement) below. + We value the security community's role in protecting our systems and users. To report a security vulnerability: - File a private vulnerability report on GitHub: [Report a vulnerability](https://github.com/BerriAI/litellm/security/advisories/new) - Include steps to reproduce the issue +- Include a video or screen recording demonstrating the full exploit against a live LiteLLM instance, from initial access through to impact. A terminal recording (for example asciinema) is fine for CLI-only exploits. - Provide any relevant additional information +### Reproduction Video Requirement + +A video demonstrating the exploit is required for every report. AI tools have made it easy to produce plausible-sounding vulnerability reports that do not reproduce in practice, and triaging them takes time away from real issues. Reports submitted without a working reproduction video will be closed without review. If you add a video to a closed report, we will reopen and triage it. + ### Vulnerability Categories We classify vulnerabilities into the following categories: @@ -38,7 +46,7 @@ We offer bounties for responsibly disclosed vulnerabilities based on severity: | **Medium** | N/A | P2 authenticated privilege escalation | | **Low** | N/A | Minor information disclosure, low-impact misconfigurations | -To qualify for a bounty, reports must include clear reproduction steps and must not involve systems or accounts you do not own. We review all submissions promptly and will follow up within 5 business days. +To qualify for a bounty, reports must include clear reproduction steps, a reproduction video as described above, and must not involve systems or accounts you do not own. We review all submissions promptly and will follow up within 5 business days. ### Known Non-Issues diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index f8bac820582..d0ad1cc8f82 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -783,6 +783,16 @@ async def test_async_post_call_failure_hook(prometheus_logger): it should increment the litellm_proxy_failed_requests_metric and litellm_proxy_total_requests_metric """ + # Opt into the unified rate-limit labels so this test exercises the + # full label set surfaced when `prometheus_emit_rate_limit_labels` is on. + # The logger caches each metric's label set at construction time (so the + # labels passed to ``counter.labels(...)`` stay in lock step with the + # labels used to register the metric), so we must invalidate the cache + # after flipping the toggle for the cache to pick up the new label set. + original_emit = litellm.prometheus_emit_rate_limit_labels + litellm.prometheus_emit_rate_limit_labels = True + prometheus_logger._cached_metric_labels.clear() + # Mock the prometheus metrics prometheus_logger.litellm_proxy_failed_requests_metric = MagicMock() prometheus_logger.litellm_proxy_total_requests_metric = MagicMock() @@ -804,32 +814,38 @@ async def test_async_post_call_failure_hook(prometheus_logger): request_route="/chat/completions", ) - # Call the function - await prometheus_logger.async_post_call_failure_hook( - request_data=request_data, - original_exception=original_exception, - user_api_key_dict=user_api_key_dict, - ) + try: + # Call the function + await prometheus_logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=original_exception, + user_api_key_dict=user_api_key_dict, + ) - # Assert failed requests metric was incremented with correct labels - prometheus_logger.litellm_proxy_failed_requests_metric.labels.assert_called_once_with( - end_user=None, - user="test_user", - user_email=None, - hashed_api_key="test_key", - api_key_alias="test_alias", - team="test_team", - team_alias="test_team_alias", - org_id=None, - org_alias=None, - requested_model="gpt-5-mini", - exception_status="429", - exception_class="Openai.RateLimitError", - route=user_api_key_dict.request_route, - model_id=None, - client_ip=None, - user_agent=None, - ) + # Assert failed requests metric was incremented with correct labels + prometheus_logger.litellm_proxy_failed_requests_metric.labels.assert_called_once_with( + end_user=None, + user="test_user", + user_email=None, + hashed_api_key="test_key", + api_key_alias="test_alias", + team="test_team", + team_alias="test_team_alias", + org_id=None, + org_alias=None, + requested_model="gpt-5-mini", + exception_status="429", + exception_class="Openai.RateLimitError", + rate_limit_category="vendor_rate_limit", + rate_limit_type=None, + route=user_api_key_dict.request_route, + model_id=None, + client_ip=None, + user_agent=None, + ) + finally: + litellm.prometheus_emit_rate_limit_labels = original_emit + prometheus_logger._cached_metric_labels.clear() prometheus_logger.litellm_proxy_failed_requests_metric.labels().inc.assert_called_once() # Assert total requests metric was incremented with correct labels @@ -1962,6 +1978,10 @@ def test_set_team_budget_metrics_with_custom_labels(prometheus_logger, monkeypat # Set custom prometheus labels custom_labels = ["metadata.organization", "metadata.environment"] monkeypatch.setattr("litellm.custom_prometheus_metadata_labels", custom_labels) + # Logger caches each metric's label set at construction time (fixture + # runs before this monkeypatch), so invalidate so the cached label set + # picks up the freshly-configured custom metadata labels. + prometheus_logger._cached_metric_labels.clear() # Create test team with custom metadata team = MagicMock( diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py index a08013cd439..83a2c286d64 100644 --- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py +++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py @@ -1,7 +1,6 @@ from dataclasses import dataclass, field from typing import Dict, FrozenSet, List, Optional, Tuple - OMIT = object() @@ -136,6 +135,13 @@ def expected(model: ModelEntry, effort: str) -> CellExpectation: ANTHROPIC_DIRECT_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="claude-fable-5", + model="anthropic/claude-fable-5", + mode="adaptive", + required_env=_ANTHROPIC_REQ, + caps=_CAPS_XHIGH_MAX, + ), ModelEntry( alias="claude-opus-4-8", model="anthropic/claude-opus-4-8", @@ -168,6 +174,19 @@ def expected(model: ModelEntry, effort: str) -> CellExpectation: AZURE_AI_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="azure-claude-fable-5", + model="azure_ai/claude-fable-5", + mode="adaptive", + required_env=_AZURE_FOUNDRY_REQ, + caps=_CAPS_XHIGH_MAX, + fail_reason=( + "claude-fable-5 has no deployment on the CI Microsoft Foundry " + "resource yet; Foundry returns DeploymentNotFound until someone " + "creates the fable-5 deployment, so this cell stays loud in CI. " + "Remove this fail_reason once the deployment exists." + ), + ), ModelEntry( alias="azure-claude-opus-4-8", model="azure_ai/claude-opus-4-8", @@ -213,6 +232,20 @@ def expected(model: ModelEntry, effort: str) -> CellExpectation: VERTEX_AI_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="vertex-claude-fable-5", + model="vertex_ai/claude-fable-5", + mode="adaptive", + extra_params=(("vertex_location", "global"),), + required_env=_VERTEX_REQ, + caps=_CAPS_XHIGH_MAX, + fail_reason=( + "claude-fable-5 availability on the CI Vertex project is not yet " + "confirmed for this brand-new release, so this cell stays loud in " + "CI until verified. Remove this fail_reason once the model is " + "confirmed available on the global Vertex endpoint." + ), + ), ModelEntry( alias="vertex-claude-opus-4-8", model="vertex_ai/claude-opus-4-8", @@ -263,6 +296,23 @@ def expected(model: ModelEntry, effort: str) -> CellExpectation: BEDROCK_CONVERSE_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="bedrock-claude-fable-5", + model="bedrock/converse/us.anthropic.claude-fable-5", + mode="adaptive", + extra_params=(("aws_region_name", "us-east-1"),), + required_env=_BEDROCK_REQ, + caps=_CAPS_XHIGH_MAX, + bedrock_effort_ceiling="xhigh", + unavailable_error="is not available for this account", + fail_reason=( + "claude-fable-5 on Bedrock requires the account to opt in to " + "provider data sharing (data retention mode " + "'provider_data_sharing' via the Data Retention API); the CI " + "account has not opted in yet, so this cell stays loud in CI. " + "Remove this fail_reason once the opt-in is done." + ), + ), ModelEntry( alias="bedrock-claude-opus-4-8", model="bedrock/converse/us.anthropic.claude-opus-4-8", diff --git a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py index 551ab8459d1..a5f16f928e5 100644 --- a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py +++ b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py @@ -15,7 +15,6 @@ all_cells, ) - _PROMPT_MESSAGES: List[Dict[str, str]] = [ {"role": "user", "content": "Step by step, calculate 47 * 53. Show your work."} ] @@ -201,8 +200,8 @@ async def test_reasoning_effort_grid( def test_grid_cell_count() -> None: - assert len(_PARAMS) == 25 * 11, ( - f"expected 275 cells (25 provider x model combos x 11 efforts), " + assert len(_PARAMS) == 29 * 11, ( + f"expected 319 cells (29 provider x model combos x 11 efforts), " f"got {len(_PARAMS)}" ) diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index 9cf253c379d..fa22ff6b392 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -2712,6 +2712,10 @@ def test_bedrock_top_k_param(model, expected_params): data = json.loads(mock_post.call_args.kwargs["data"]) if "mistral" in model: assert data["top_k"] == 2 + elif expected_params == {}: + # Models that don't support top_k produce no additionalModelRequestFields; + # the empty block is now omitted entirely rather than sent as `{}`. + assert "additionalModelRequestFields" not in data else: assert data["additionalModelRequestFields"] == expected_params @@ -3059,8 +3063,6 @@ async def test_bedrock_max_completion_tokens(model: str): assert request_body == { "messages": [{"role": "user", "content": [{"text": "Hello!"}]}], - "additionalModelRequestFields": {}, - "system": [], "inferenceConfig": {"maxTokens": 10}, } diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index 01bcb1a247a..9a69f513069 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -2425,6 +2425,42 @@ def test_reasoning_content_extracted(self): assert result.choices[0].message.content == "The answer is 4." assert result.choices[0].message.reasoning_content == "2+2=4" + def test_reasoning_content_not_mirrored_into_provider_specific_fields(self): + """Mirroring reasoning_content into provider_specific_fields made + cache-replayed messages diverge from live Anthropic messages, which + only set it top-level, breaking cache key stability (issue #27337).""" + response_object = { + "id": "chatcmpl-5", + "model": "claude-sonnet-4-5", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "The answer is 4.", + "role": "assistant", + "reasoning_content": "2+2=4", + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "2+2=4", + "signature": "sig", + } + ], + }, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}, + } + + result = convert_to_model_response_object( + response_object=response_object, + model_response_object=ModelResponse(), + ) + message = result.choices[0].message + assert message.reasoning_content == "2+2=4" + assert "reasoning_content" not in (message.provider_specific_fields or {}) + def test_response_none_raises(self): with pytest.raises(Exception): convert_to_model_response_object( diff --git a/tests/local_testing/test_basic_python_version.py b/tests/local_testing/test_basic_python_version.py index 8308e0d6033..e31c3953714 100644 --- a/tests/local_testing/test_basic_python_version.py +++ b/tests/local_testing/test_basic_python_version.py @@ -92,6 +92,56 @@ def test_package_dependencies(): ) +def test_cli_extra_is_a_thin_client_install(): + """The `cli` extra must install a working `lite` client without dragging in the + proxy server runtime. It therefore has to declare the CLI's real third-party + deps (rich, pyyaml, requests) and must never contain a server-only dependency + from the `proxy` extra; a leak there silently re-bloats the laptop install. + """ + import pathlib + + import litellm + from packaging.requirements import Requirement + + try: + import tomllib as tomli + except ImportError: + try: + import tomli + except ImportError: + pytest.skip("tomli/tomllib not available - skipping dependency check") + + pyproject_path = pathlib.Path(litellm.__file__).parent.parent / "pyproject.toml" + with open(pyproject_path, "rb") as f: + optional_deps = tomli.load(f)["project"]["optional-dependencies"] + + assert "cli" in optional_deps, "Expected a `cli` extra for the thin lite install" + + cli_names = {Requirement(req).name.lower() for req in optional_deps["cli"]} + + missing = {"rich", "pyyaml", "requests"} - cli_names + assert not missing, f"`cli` extra is missing deps the lite CLI imports: {missing}" + + server_only = { + "fastapi", + "uvicorn", + "gunicorn", + "granian", + "starlette", + "boto3", + "polars", + "soundfile", + "mcp", + "cryptography", + "apscheduler", + "rq", + "litellm-enterprise", + "litellm-proxy-extras", + } + leaked = cli_names & server_only + assert not leaked, f"`cli` extra leaks proxy-server deps onto laptops: {leaked}" + + import os import subprocess import time diff --git a/tests/local_testing/test_router_debug_logs.py b/tests/local_testing/test_router_debug_logs.py index f1c7e9d722e..ad807539bf2 100644 --- a/tests/local_testing/test_router_debug_logs.py +++ b/tests/local_testing/test_router_debug_logs.py @@ -82,7 +82,9 @@ async def _make_request(): asyncio.run(_make_request()) captured_logs = [rec.message for rec in caplog.records] - # on circle ci the captured logs get some async task exception logs - filter them out "Task exception was never retrieved" + # on circle ci the captured logs get async cleanup noise from the gc (leaked + # task warnings, plus aiohttp "Unclosed client session"/"Unclosed connector" + # warnings from cached clients other router tests evicted) - filter it out captured_logs = [ log for log in captured_logs @@ -90,6 +92,8 @@ async def _make_request(): and "Task was destroyed but it is pending" not in log and "get_available_deployment" not in log and "in the Langfuse queue" not in log + and "Unclosed client session" not in log + and "Unclosed connector" not in log ] print("\n Captured caplog records - ", captured_logs) diff --git a/tests/proxy_behavior/management/test_team_budget_limits.py b/tests/proxy_behavior/management/test_team_budget_limits.py index 1534cee2b2e..dad775370ad 100644 --- a/tests/proxy_behavior/management/test_team_budget_limits.py +++ b/tests/proxy_behavior/management/test_team_budget_limits.py @@ -28,7 +28,7 @@ from litellm.proxy.utils import hash_token from .actors import Actor -from .conftest import create_scratch_org, create_scratch_team +from .conftest import MASTER_KEY, create_scratch_org, create_scratch_team pytestmark = pytest.mark.asyncio(loop_scope="session") @@ -288,34 +288,130 @@ async def test_check_user_team_limits( # --------------------------------------------------------------------------- -# /team/update path — _check_user_team_limits on existing team, no-org. -# Pin one over-budget rejection here so the update-side wiring is also -# covered (the update path is a second call site with its own data shape). +# /team/update path — budget authority. +# +# The caller's PERSONAL limits are never applied on update (that compared the +# wrong thing). But raising a team's spend ceiling is reserved for proxy admins: +# a team admin may keep or LOWER the budget, only a proxy admin may RAISE it. +# _check_user_team_limits() only runs on /team/new. # --------------------------------------------------------------------------- -async def test_team_update_user_limit_rejected(proxy_client, prisma, scratch): +async def test_team_admin_raise_budget_blocked(proxy_client, prisma, scratch): + """A team admin cannot raise the team's budget; the block is NOT based on + their personal budget (which here is higher than the requested value).""" caller_cleartext = await _seed_scratch_actor_with_caps( prisma, scratch.prefix, - max_budget=100.0, + max_budget=100000.0, # generous personal budget; must not matter ) creator_user_id = f"{scratch.prefix}-team-creator" - # Team must exist before /team/update; seed a standalone scratch team - # owned by the same actor so the update authz gate passes. team_id = await create_scratch_team( prisma, team_id=scratch.tag("team"), admin_user_ids=[creator_user_id], max_budget=50.0, ) + # Raise the team budget 50 -> 999 as a team admin. resp = await proxy_client.post( "/team/update", headers={"Authorization": f"Bearer {caller_cleartext}"}, json={"team_id": team_id, "max_budget": 999.0}, ) - assert resp.status_code == 400, resp.text + assert resp.status_code == 403, resp.text row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id}) assert row is not None - assert row.max_budget == 50.0, "row max_budget mutated despite rejection" + assert row.max_budget == 50.0, "team budget must not change on a blocked raise" + + +async def test_team_admin_lower_budget_allowed(proxy_client, prisma, scratch): + """A team admin may freely lower (or keep) the team's budget.""" + caller_cleartext = await _seed_scratch_actor_with_caps( + prisma, + scratch.prefix, + max_budget=10.0, # below both the old and new team budget; must not matter + ) + creator_user_id = f"{scratch.prefix}-team-creator" + team_id = await create_scratch_team( + prisma, + team_id=scratch.tag("team"), + admin_user_ids=[creator_user_id], + max_budget=500.0, + ) + # Lower the team budget 500 -> 300 as a team admin. + resp = await proxy_client.post( + "/team/update", + headers={"Authorization": f"Bearer {caller_cleartext}"}, + json={"team_id": team_id, "max_budget": 300.0}, + ) + assert resp.status_code == 200, resp.text + + row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + assert row is not None + assert row.max_budget == 300.0, "team admin should be able to lower the budget" + + +async def test_proxy_admin_raise_budget_allowed(proxy_client, prisma, scratch): + """A proxy admin may raise a team's budget.""" + team_id = await create_scratch_team( + prisma, + team_id=scratch.tag("team"), + admin_user_ids=[f"{scratch.prefix}-team-creator"], + max_budget=50.0, + ) + # MASTER_KEY acts as proxy admin. + resp = await proxy_client.post( + "/team/update", + headers={"Authorization": f"Bearer {MASTER_KEY}"}, + json={"team_id": team_id, "max_budget": 999.0}, + ) + assert resp.status_code == 200, resp.text + + row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + assert row is not None + assert row.max_budget == 999.0, "proxy admin should be able to raise the budget" + + +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 + ) + team_id = await create_scratch_team( + prisma, + team_id=scratch.tag("team"), + admin_user_ids=[f"{scratch.prefix}-team-creator"], + max_budget=50.0, + ) + resp = await proxy_client.post( + "/team/update", + headers={"Authorization": f"Bearer {caller_cleartext}"}, + json={"team_id": team_id, "max_budget": None}, + ) + assert resp.status_code == 403, resp.text + + row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + assert row is not None + assert row.max_budget == 50.0, "team budget cap must not be removed by a team admin" + + +async def test_proxy_admin_remove_budget_cap_allowed(proxy_client, prisma, scratch): + """A proxy admin may remove a team's cap (max_budget=null).""" + team_id = await create_scratch_team( + prisma, + team_id=scratch.tag("team"), + admin_user_ids=[f"{scratch.prefix}-team-creator"], + max_budget=50.0, + ) + resp = await proxy_client.post( + "/team/update", + headers={"Authorization": f"Bearer {MASTER_KEY}"}, + json={"team_id": team_id, "max_budget": None}, + ) + assert resp.status_code == 200, resp.text + + row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + assert row is not None + assert row.max_budget is None, "proxy admin should be able to remove the cap" diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index 3f0afe2a5a6..c170972d984 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -1236,3 +1236,60 @@ async def test_init_containers_api_endpoints_managed_id_without_model_id_applies assert call_kw["container_id"] == "cfile_upstream_abc" assert call_kw["file_id"] == "cfile_xyz" assert call_kw["custom_llm_provider"] == "azure" + + +def test_router_model_group_encrypted_content_affinity_callback_registration(): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + model_group_affinity_config = { + model_group: ["encrypted_content_affinity"], + } + router = Router( + model_list=[ + { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key", + }, + } + ], + model_group_affinity_config=model_group_affinity_config, + num_retries=0, + ) + + try: + callbacks = router.optional_callbacks or [] + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) + assert len(encrypted_content_callbacks) == 1 + assert encrypted_content_callbacks[0].enable_global_affinity is False + assert ( + encrypted_content_callbacks[0].model_group_affinity_config + == model_group_affinity_config + ) + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( + deployment_callback + ) + + router._add_encrypted_content_affinity_check(enable_global_affinity=True) + + callbacks = router.optional_callbacks or [] + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] + assert len(encrypted_content_callbacks) == 1 + assert encrypted_content_callbacks[0].enable_global_affinity is True + assert encrypted_content_callbacks[0].router is router + finally: + router.discard() diff --git a/tests/test_litellm/caching/test_caching.py b/tests/test_litellm/caching/test_caching.py new file mode 100644 index 00000000000..20614103ed2 --- /dev/null +++ b/tests/test_litellm/caching/test_caching.py @@ -0,0 +1,78 @@ +import logging +import re + +from litellm.caching.caching import Cache +from litellm.types.caching import LiteLLMCacheType +from litellm.types.utils import Embedding, EmbeddingResponse, Usage + + +def test_cache_key_debug_log_does_not_include_prompt_material(caplog): + cache = Cache(type=LiteLLMCacheType.LOCAL) + prompt_marker = "secret prompt material " + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + cache_key = cache.get_cache_key( + model="gpt-4.1-mini", + messages=[ + {"role": "system", "content": prompt_marker * 100}, + {"role": "user", "content": "hello"}, + ], + tools=[ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, + } + ], + response_format={ + "type": "json_schema", + "json_schema": { + "name": "lookup_response", + "schema": {"type": "object"}, + }, + }, + stream=True, + ) + + assert re.fullmatch(r"[0-9a-f]{64}", cache_key) + + created_cache_key_logs = [ + record.getMessage() + for record in caplog.records + if "Created cache key:" in record.getMessage() + ] + assert created_cache_key_logs + assert all(prompt_marker not in message for message in created_cache_key_logs) + assert any(cache_key in message for message in created_cache_key_logs) + + +def _embedding_response(prompt_tokens, num_items): + return EmbeddingResponse( + model="amazon.titan-embed-image-v1", + data=[ + Embedding(embedding=[0.0], index=i, object="embedding") + for i in range(num_items) + ], + usage=Usage( + prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens + ), + ) + + +def test_get_per_item_prompt_tokens_single_item_returns_full_value(): + cache = Cache(type=LiteLLMCacheType.LOCAL) + result = _embedding_response(prompt_tokens=0, num_items=1) + assert cache._get_per_item_prompt_tokens(result, 0) == 0 + + +def test_get_per_item_prompt_tokens_distributes_with_remainder(): + cache = Cache(type=LiteLLMCacheType.LOCAL) + result = _embedding_response(prompt_tokens=10, num_items=3) + per_item = [cache._get_per_item_prompt_tokens(result, i) for i in range(3)] + assert sum(per_item) == 10 # 4 + 3 + 3 + assert per_item == [4, 3, 3] diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py index 3eb949d7f29..01327529410 100644 --- a/tests/test_litellm/caching/test_caching_handler.py +++ b/tests/test_litellm/caching/test_caching_handler.py @@ -436,3 +436,123 @@ def test_convert_cached_responses_legacy_stream_path(): ) assert isinstance(result, CachedResponsesAPIStreamingIterator) + + +@pytest.mark.asyncio +async def test_embedding_cache_restores_stored_prompt_tokens_for_image_input(): + """Image-embedding cache hit restores prompt_tokens=0 from the stored value + instead of recomputing a bogus count by tokenizing the base64 input.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # base64-like blob — token_counter over this would return a large nonzero count + image_input = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk" * 50 + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens": 0, + "prompt_tokens_details": {"image_count": 1}, + } + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "amazon.titan-embed-image-v1", "input": image_input}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="amazon.titan-embed-image-v1", + ) + + assert cache_hit + assert response.usage is not None + assert response.usage.prompt_tokens == 0 + assert response.usage.total_tokens == 0 + assert response.usage.prompt_tokens_details.image_count == 1 + + +@pytest.mark.asyncio +async def test_embedding_cache_sums_stored_prompt_tokens_across_items(): + """A multi-item cache hit sums the stored per-item prompt_tokens back to the total.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.01], + "index": 0, + "object": "embedding", + "model": "text-embedding-3-small", + "prompt_tokens": 5, + }, + { + "embedding": [-0.02], + "index": 1, + "object": "embedding", + "model": "text-embedding-3-small", + "prompt_tokens": 4, + }, + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "text-embedding-3-small", "input": ["hello world", "foo bar"]}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="text-embedding-3-small", + ) + + assert cache_hit + assert response.usage.prompt_tokens == 9 + assert response.usage.total_tokens == 9 + + +@pytest.mark.asyncio +async def test_embedding_cache_falls_back_to_token_counter_for_legacy_entries(): + """Legacy cache entries with no stored prompt_tokens still recompute via token_counter + for str inputs (backward compatibility).""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # No prompt_tokens key — pre-fix entry + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "text-embedding-ada-002", + }, + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "text-embedding-ada-002", "input": "hello world"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="text-embedding-ada-002", + ) + + assert cache_hit + # token_counter over "hello world" yields a nonzero count — fallback path still runs + assert response.usage.prompt_tokens > 0 diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index b50a35ef50e..13f9d00136d 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -523,3 +523,468 @@ async def test_redis_semantic_cache_async_set_cache_stores_cache_key_filter( filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, ttl=60, ) + + +def test_redis_semantic_cache_set_cache_uses_responses_string_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache._get_cache_filters = MagicMock( + return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} + ) + redis_semantic_cache._get_ttl = MagicMock(return_value=None) + + redis_semantic_cache.set_cache( + key="test_key", + value={"content": "Paris"}, + input="What is the capital of France?", + ) + + redis_semantic_cache.llmcache.store.assert_called_once_with( + "What is the capital of France?", + "{'content': 'Paris'}", + filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, + ) + + +def test_redis_semantic_cache_get_cache_uses_responses_string_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.similarity_threshold = 0.8 + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache.llmcache.check = MagicMock( + return_value=[ + { + "prompt": "What is the capital of France?", + "response": '{"content": "Paris"}', + "vector_distance": 0.1, + RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key", + } + ] + ) + + with patch.object( + redis_semantic_cache, + "_get_cache_key_filter_expression", + return_value="cache-key-filter", + ): + metadata = {} + result = redis_semantic_cache.get_cache( + key="test_key", + input="What is the capital of France?", + metadata=metadata, + ) + + assert result == {"content": "Paris"} + assert metadata["semantic-similarity"] == pytest.approx(0.9) + redis_semantic_cache.llmcache.check.assert_called_once_with( + prompt="What is the capital of France?", + filter_expression="cache-key-filter", + ) + + +def test_redis_semantic_cache_set_cache_flattens_structured_responses_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache._get_cache_filters = MagicMock( + return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} + ) + redis_semantic_cache._get_ttl = MagicMock(return_value=None) + + redis_semantic_cache.set_cache( + key="test_key", + value={"content": "Paris"}, + input=[ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "What is the capital of France?"}, + {"type": "input_text", "text": "Answer briefly."}, + { + "type": "input_image", + "image_url": "https://example.com/paris.png", + }, + ], + } + ], + ) + + redis_semantic_cache.llmcache.store.assert_called_once_with( + "What is the capital of France?\nAnswer briefly.", + "{'content': 'Paris'}", + filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, + ) + + +def test_redis_semantic_cache_prompt_extraction_prefers_messages(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + prompt = RedisSemanticCache._get_prompt_from_kwargs( + messages=[{"content": "message prompt"}], + input="responses prompt", + ) + + assert prompt == "message prompt" + + +def test_redis_semantic_cache_prompt_extraction_handles_model_objects(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + class ModelDumpInput: + def model_dump(self): + return {"content": [{"text": "model dump prompt"}]} + + class DictInput: + def dict(self): + return {"content": [{"output_text": "dict prompt"}]} + + prompt = RedisSemanticCache._get_prompt_from_kwargs( + input=[ + ModelDumpInput(), + DictInput(), + {"content": [{"input_text": "inline prompt"}]}, + {"content": [{"type": "input_image", "image_url": "https://example.com"}]}, + ] + ) + + assert prompt == "model dump prompt\ndict prompt\ninline prompt" + + +def test_redis_semantic_cache_prompt_extraction_returns_none_without_text(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + assert RedisSemanticCache._get_prompt_from_kwargs() is None + assert RedisSemanticCache._get_prompt_from_kwargs(input=None) is None + assert RedisSemanticCache._get_prompt_from_kwargs(input=" ") is None + assert ( + RedisSemanticCache._get_prompt_from_kwargs( + input=[{"type": "input_image", "image_url": "https://example.com"}] + ) + is None + ) + + +def test_redis_semantic_cache_prompt_extraction_skips_blank_dict_text_keys(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + prompt = RedisSemanticCache._get_prompt_from_kwargs( + input={"text": " ", "input_text": "fallback prompt"} + ) + + assert prompt == "fallback prompt" + + +def test_redis_semantic_cache_prompt_extraction_skips_blank_object_text_keys(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + class ResponseInput: + text = " " + input_text = "fallback prompt" + + prompt = RedisSemanticCache._get_prompt_from_kwargs(input=ResponseInput()) + + assert prompt == "fallback prompt" + + +def test_redis_semantic_cache_prompt_extraction_handles_object_content(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + class ResponseInput: + content = [{"text": "object content prompt"}] + + prompt = RedisSemanticCache._get_prompt_from_kwargs(input=ResponseInput()) + + assert prompt == "object content prompt" + + +def test_redis_semantic_cache_set_cache_skips_blank_responses_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + + redis_semantic_cache.set_cache( + key="test_key", + value={"content": "Paris"}, + input=" ", + ) + + redis_semantic_cache.llmcache.store.assert_not_called() + + +def test_redis_semantic_cache_get_cache_sets_similarity_on_blank_responses_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + metadata = {} + + result = redis_semantic_cache.get_cache( + key="test_key", + input=" ", + metadata=metadata, + ) + + assert result is None + assert metadata["semantic-similarity"] == 0.0 + redis_semantic_cache.llmcache.check.assert_not_called() + + +def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache.llmcache.check = MagicMock(return_value=[]) + + with patch.object( + redis_semantic_cache, + "_get_cache_key_filter_expression", + return_value="cache-key-filter", + ): + metadata = {} + result = redis_semantic_cache.get_cache( + key="test_key", + input="What is the capital of France?", + metadata=metadata, + ) + + assert result is None + assert metadata["semantic-similarity"] == 0.0 + redis_semantic_cache.llmcache.check.assert_called_once_with( + prompt="What is the capital of France?", + filter_expression="cache-key-filter", + ) + + +@pytest.mark.asyncio +async def test_redis_semantic_cache_async_paths_use_responses_string_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.similarity_threshold = 0.8 + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache.llmcache.astore = AsyncMock() + redis_semantic_cache.llmcache.acheck = AsyncMock( + return_value=[ + { + "prompt": "What is the capital of France?", + "response": '{"content": "Paris"}', + "vector_distance": 0.1, + RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key", + } + ] + ) + redis_semantic_cache._get_cache_filters = MagicMock( + return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} + ) + redis_semantic_cache._get_ttl = MagicMock(return_value=None) + redis_semantic_cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3]) + + await redis_semantic_cache.async_set_cache( + key="test_key", + value={"content": "Paris"}, + input="What is the capital of France?", + ) + + with patch.object( + redis_semantic_cache, + "_get_cache_key_filter_expression", + return_value="cache-key-filter", + ): + metadata = {} + result = await redis_semantic_cache.async_get_cache( + key="test_key", + input="What is the capital of France?", + metadata=metadata, + ) + + redis_semantic_cache.llmcache.astore.assert_called_once_with( + "What is the capital of France?", + "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], + filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, + ) + assert result == {"content": "Paris"} + assert metadata["semantic-similarity"] == pytest.approx(0.9) + redis_semantic_cache.llmcache.acheck.assert_called_once_with( + prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], + filter_expression="cache-key-filter", + ) + + +@pytest.mark.asyncio +async def test_redis_semantic_cache_async_paths_set_similarity_on_misses(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache.llmcache.astore = AsyncMock() + redis_semantic_cache.llmcache.acheck = AsyncMock(return_value=[]) + redis_semantic_cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3]) + + await redis_semantic_cache.async_set_cache( + key="test_key", + value={"content": "Paris"}, + input=" ", + ) + + redis_semantic_cache.llmcache.astore.assert_not_called() + redis_semantic_cache._get_async_embedding.assert_not_called() + + blank_metadata = {} + blank_result = await redis_semantic_cache.async_get_cache( + key="test_key", + input=" ", + metadata=blank_metadata, + ) + + assert blank_result is None + assert blank_metadata["semantic-similarity"] == 0.0 + redis_semantic_cache.llmcache.acheck.assert_not_called() + redis_semantic_cache._get_async_embedding.assert_not_called() + + with patch.object( + redis_semantic_cache, + "_get_cache_key_filter_expression", + return_value="cache-key-filter", + ): + miss_metadata = {} + miss_result = await redis_semantic_cache.async_get_cache( + key="test_key", + input="What is the capital of France?", + metadata=miss_metadata, + ) + + assert miss_result is None + assert miss_metadata["semantic-similarity"] == 0.0 + redis_semantic_cache.llmcache.acheck.assert_called_once_with( + prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], + filter_expression="cache-key-filter", + ) + + +def test_cache_get_cache_passes_responses_input_to_backend_cache(): + from litellm.caching.caching import Cache + + cache = Cache.__new__(Cache) + cache.cache = MagicMock() + cache.cache.get_cache = MagicMock(return_value=None) + cache.should_use_cache = MagicMock(return_value=True) + cache.get_cache_key = MagicMock(return_value="test_key") + + metadata = {} + cache.get_cache( + input="What is the capital of France?", + metadata=metadata, + cache={}, + ) + + cache.cache.get_cache.assert_called_once_with( + "test_key", + input="What is the capital of France?", + metadata=metadata, + ) + + +def test_cache_get_cache_filters_sensitive_kwargs_from_backend_cache(): + from litellm.caching.caching import Cache + + cache = Cache.__new__(Cache) + cache.cache = MagicMock() + cache.should_use_cache = MagicMock(return_value=True) + cache.get_cache_key = MagicMock(return_value="test_key") + cache._get_cache_logic = MagicMock(return_value={"content": "Paris"}) + + def _cache_hit(_cache_key, **cache_kwargs): + cache_kwargs["metadata"]["semantic-similarity"] = 0.7 + return {"content": "Paris"} + + cache.cache.get_cache = MagicMock(side_effect=_cache_hit) + + metadata = {"user_api_key": "sk-secret", "trace_id": "trace-id"} + result = cache.get_cache( + input="What is the capital of France?", + metadata=metadata, + cache={"s-maxage": 10}, + api_key="sk-secret", + headers={"authorization": "Bearer sk-secret"}, + ) + + assert result == {"content": "Paris"} + assert metadata == { + "user_api_key": "sk-secret", + "trace_id": "trace-id", + "semantic-similarity": 0.7, + } + + forwarded_kwargs = cache.cache.get_cache.call_args.kwargs + assert forwarded_kwargs == { + "input": "What is the capital of France?", + "metadata": {"semantic-similarity": 0.7}, + } + assert forwarded_kwargs["metadata"] is not metadata + cache._get_cache_logic.assert_called_once_with( + cached_result={"content": "Paris"}, + max_age=10, + ) + + +def test_cache_get_cache_filters_sensitive_kwargs_without_metadata(): + from litellm.caching.caching import Cache + + cache = Cache.__new__(Cache) + cache.cache = MagicMock() + cache.cache.get_cache = MagicMock(return_value={"content": "Paris"}) + cache.should_use_cache = MagicMock(return_value=True) + cache.get_cache_key = MagicMock(return_value="test_key") + cache._get_cache_logic = MagicMock(return_value={"content": "Paris"}) + + result = cache.get_cache( + input="What is the capital of France?", + cache={"s-maxage": 10}, + api_key="sk-secret", + headers={"authorization": "Bearer sk-secret"}, + ) + + assert result == {"content": "Paris"} + cache.cache.get_cache.assert_called_once_with( + "test_key", + input="What is the capital of France?", + ) + + +def test_cache_get_cache_passes_responses_input_to_dynamic_cache(): + from litellm.caching.caching import Cache + + cache = Cache.__new__(Cache) + cache.should_use_cache = MagicMock(return_value=True) + cache.get_cache_key = MagicMock(return_value="test_key") + cache._get_cache_logic = MagicMock(return_value={"content": "Paris"}) + dynamic_cache_object = MagicMock() + dynamic_cache_object.get_cache = MagicMock(return_value={"content": "Paris"}) + + metadata = {} + result = cache.get_cache( + dynamic_cache_object=dynamic_cache_object, + input="What is the capital of France?", + metadata=metadata, + cache={}, + ) + + assert result == {"content": "Paris"} + dynamic_cache_object.get_cache.assert_called_once_with( + "test_key", + input="What is the capital of France?", + metadata=metadata, + ) + cache._get_cache_logic.assert_called_once_with( + cached_result={"content": "Paris"}, + max_age=float("inf"), + ) diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index d335c359aa0..06457dfebff 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -13,6 +13,9 @@ 0, os.path.abspath("../../..") ) # Adds the parent directory to the system-path import litellm +from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, +) def test_convert_chat_completion_messages_to_responses_api_image_input(): @@ -860,6 +863,39 @@ def test_extract_extra_body_params_reasoning_effort_override(): assert "extra_body" not in result +def test_transform_request_system_only_message_maps_to_system_input_item(): + """System-only requests must not send input=[] to the Responses API. + + OpenAI rejects both input=[] and input="". When the only message is a + system message, carry it as a system-role input item (single copy, correct + role) rather than leaving input empty or duplicating it into instructions. + """ + handler = LiteLLMResponsesTransformationHandler() + logging_obj = Mock() + messages = [{"role": "system", "content": "You are a helpful assistant."}] + + result = handler.transform_request( + model="gpt-5.3-codex", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=logging_obj, + ) + + assert result["input"] == [ + { + "type": "message", + "role": "system", + "content": [ + {"type": "input_text", "text": "You are a helpful assistant."} + ], + } + ] + # System content lives in input only; not duplicated into instructions. + assert not result.get("instructions") + + def test_transform_request_single_char_keys_not_matched(): """Test that single-character keys are not incorrectly matched to 'metadata' or 'previous_response_id' diff --git a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py index 56e5a94cd49..ffa81abf86c 100644 --- a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py +++ b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py @@ -32,6 +32,40 @@ def test_initialize_from_proxy_config(): assert logger.compression_target == 789 +def test_initialize_from_proxy_config_ignores_non_dict_callback_specific_params(): + """Regression (#29590): a non-dict value under + callback_settings.compression_interception must not crash initialization. + + Forwarding callback_settings as callback_specific_params activates this + branch; without the isinstance(dict) guard a non-dict value reached + from_config_yaml(...).get(...) and raised AttributeError at proxy startup. + The value is ignored and the logger falls back to defaults. + """ + logger = CompressionInterceptionLogger.initialize_from_proxy_config( + litellm_settings={}, + callback_specific_params={"compression_interception": True}, + ) + + assert logger.enabled is True + assert logger.compression_trigger == 200_000 + + +def test_initialize_from_proxy_config_honors_dict_callback_specific_params(): + """A valid dict under callback_settings.compression_interception is applied.""" + logger = CompressionInterceptionLogger.initialize_from_proxy_config( + litellm_settings={}, + callback_specific_params={ + "compression_interception": { + "enabled": False, + "compression_trigger": 12345, + } + }, + ) + + assert logger.enabled is False + assert logger.compression_trigger == 12345 + + @pytest.mark.asyncio async def test_pre_call_hook_compresses_messages_and_injects_tool(monkeypatch): """Test pre-call hook compresses and stores per-call cache.""" diff --git a/tests/test_litellm/integrations/focus/test_focus_database.py b/tests/test_litellm/integrations/focus/test_focus_database.py index 5ee98cc9dd0..d77af2dd170 100644 --- a/tests/test_litellm/integrations/focus/test_focus_database.py +++ b/tests/test_litellm/integrations/focus/test_focus_database.py @@ -72,3 +72,18 @@ async def test_should_reject_invalid_limit(monkeypatch: pytest.MonkeyPatch): await db.get_usage_data(limit="invalid") assert query_mock.await_count == 0 + + +@pytest.mark.asyncio +async def test_should_join_organization_table(monkeypatch: pytest.MonkeyPatch): + db, query_mock = _setup_db(monkeypatch, []) + + await db.get_usage_data() + + query_text, *_ = query_mock.await_args.args + assert ( + "COALESCE(vt.organization_id, tt.organization_id) as organization_id" + in query_text + ) + assert "ot.organization_alias as organization_alias" in query_text + assert 'LEFT JOIN "LiteLLM_OrganizationTable" ot' in query_text diff --git a/tests/test_litellm/integrations/focus/test_focus_gcs_destination.py b/tests/test_litellm/integrations/focus/test_focus_gcs_destination.py new file mode 100644 index 00000000000..35cdb18326a --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_focus_gcs_destination.py @@ -0,0 +1,180 @@ +"""Tests for FocusGCSDestination.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.integrations.focus.destinations.base import FocusTimeWindow + + +def _make_window(frequency: str = "hourly") -> FocusTimeWindow: + return FocusTimeWindow( + start_time=datetime(2026, 1, 1, 10, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2026, 1, 1, 11, 0, 0, tzinfo=timezone.utc), + frequency=frequency, + ) + + +@pytest.mark.asyncio +async def test_deliver_posts_to_gcs_upload_endpoint(): + """deliver() must POST raw bytes to the GCS upload endpoint.""" + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination( + prefix="focus_exports", + config={"bucket_name": "my-bucket", "service_account_json": None}, + ) + + mock_response = MagicMock() + mock_response.status_code = 200 + + mock_client = MagicMock() + mock_client.post = AsyncMock(return_value=mock_response) + dest.async_httpx_client = mock_client + + with patch.object( + dest, + "construct_request_headers", + new=AsyncMock(return_value={"Authorization": "Bearer tok-123"}), + ): + await dest.deliver( + content=b"col1,col2\nval1,val2\n", + time_window=_make_window(), + filename="usage_20260101T100000Z_20260101T110000Z.csv", + ) + + mock_client.post.assert_called_once() + call_kwargs = mock_client.post.call_args + url = call_kwargs.kwargs.get("url") or call_kwargs.args[0] + assert "my-bucket" in url + assert "uploadType=media" in url + headers = call_kwargs.kwargs["headers"] + assert headers["Authorization"] == "Bearer tok-123" + + +@pytest.mark.asyncio +async def test_deliver_raises_on_gcs_error(): + """deliver() must raise RuntimeError when GCS returns non-200.""" + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination( + prefix="focus_exports", + config={"bucket_name": "my-bucket"}, + ) + + mock_response = MagicMock() + mock_response.status_code = 403 + mock_response.text = "Permission denied" + + mock_client = MagicMock() + mock_client.post = AsyncMock(return_value=mock_response) + dest.async_httpx_client = mock_client + + with patch.object( + dest, + "construct_request_headers", + new=AsyncMock(return_value={"Authorization": "Bearer tok-bad"}), + ): + with pytest.raises(RuntimeError, match="GCS upload failed"): + await dest.deliver( + content=b"data", + time_window=_make_window(), + filename="usage.csv", + ) + + +def test_build_object_key_hourly(): + """Hourly key must include date= and hour= components.""" + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination(prefix="focus_exports", config={"bucket_name": "b"}) + key = dest._build_object_key( + time_window=_make_window("hourly"), filename="usage.parquet" + ) + + assert key == "focus_exports/date=2026-01-01/hour=10/usage.parquet" + + +def test_build_object_key_daily(): + """Daily key must include date= but not hour=.""" + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination(prefix="focus_exports", config={"bucket_name": "b"}) + window = FocusTimeWindow( + start_time=datetime(2026, 1, 1, 0, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2026, 1, 2, 0, 0, 0, tzinfo=timezone.utc), + frequency="daily", + ) + key = dest._build_object_key(time_window=window, filename="usage.parquet") + + assert key == "focus_exports/date=2026-01-01/usage.parquet" + + +def test_missing_bucket_name_raises(): + """Constructing without bucket_name must raise ValueError.""" + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + with pytest.raises(ValueError, match="bucket_name"): + FocusGCSDestination(prefix="focus_exports", config={}) + + +def test_global_gcs_service_account_not_overwritten_when_absent(monkeypatch): + """service_account_json absent from config must not overwrite GCS_PATH_SERVICE_ACCOUNT. + + GCSBucketBase sets self.path_service_account_json from GCS_PATH_SERVICE_ACCOUNT. + If config has no service_account_json key, we must leave the parent value intact + so deployments using the global credential don't silently fall back to ADC. + """ + monkeypatch.setenv("GCS_PATH_SERVICE_ACCOUNT", "/global/sa.json") + + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination(prefix="focus_exports", config={"bucket_name": "b"}) + + assert dest.path_service_account_json == "/global/sa.json" + + +def test_explicit_service_account_overrides_global(monkeypatch): + """Explicit service_account_json in config must take precedence over GCS_PATH_SERVICE_ACCOUNT.""" + monkeypatch.setenv("GCS_PATH_SERVICE_ACCOUNT", "/global/sa.json") + + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination( + prefix="focus_exports", + config={"bucket_name": "b", "service_account_json": "/focus/sa.json"}, + ) + + assert dest.path_service_account_json == "/focus/sa.json" + + +def test_factory_creates_gcs_destination(monkeypatch): + """FocusDestinationFactory.create(provider='gcs') must return FocusGCSDestination.""" + monkeypatch.setenv("FOCUS_GCS_BUCKET_NAME", "env-bucket") + + from litellm.integrations.focus.destinations.factory import FocusDestinationFactory + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusDestinationFactory.create(provider="gcs", prefix="focus_exports") + + assert isinstance(dest, FocusGCSDestination) + assert dest.BUCKET_NAME == "env-bucket" diff --git a/tests/test_litellm/integrations/focus/test_transformer.py b/tests/test_litellm/integrations/focus/test_transformer.py new file mode 100644 index 00000000000..4461d19efde --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_transformer.py @@ -0,0 +1,62 @@ +"""Tests for FocusTransformer organization metadata in Tags.""" + +from __future__ import annotations + +import json +from datetime import date + +import polars as pl + +from litellm.integrations.focus.transformer import FocusTransformer + + +def test_should_include_organization_fields_in_tags(): + frame = pl.DataFrame( + { + "date": [date(2024, 1, 2)], + "spend": [1.25], + "api_requests": [1], + "api_key": ["hashed-key"], + "api_key_alias": ["prod-key"], + "model": ["gpt-4o"], + "model_group": ["gpt-4o"], + "custom_llm_provider": ["openai"], + "team_id": ["team-1"], + "team_alias": ["Platform"], + "organization_id": ["org-123"], + "organization_alias": ["Acme Corp"], + "user_id": ["user-1"], + "user_email": ["user@example.com"], + } + ) + + normalized = FocusTransformer().transform(frame) + + tags = json.loads(normalized["Tags"][0]) + assert tags["organization_id"] == "org-123" + assert tags["organization_alias"] == "Acme Corp" + assert tags["team_id"] == "team-1" + + +def test_should_omit_missing_organization_fields_from_tags(): + frame = pl.DataFrame( + { + "date": [date(2024, 1, 2)], + "spend": [0.5], + "api_requests": [1], + "api_key": ["hashed-key"], + "api_key_alias": ["prod-key"], + "model": ["gpt-4o-mini"], + "model_group": ["gpt-4o-mini"], + "custom_llm_provider": ["openai"], + "team_id": ["team-1"], + "team_alias": ["Platform"], + } + ) + + normalized = FocusTransformer().transform(frame) + + tags = json.loads(normalized["Tags"][0]) + assert "organization_id" not in tags + assert "organization_alias" not in tags + assert tags["team_id"] == "team-1" diff --git a/tests/test_litellm/integrations/test_galileo.py b/tests/test_litellm/integrations/test_galileo.py index 8ce5eb776f9..0533b7ca7d1 100644 --- a/tests/test_litellm/integrations/test_galileo.py +++ b/tests/test_litellm/integrations/test_galileo.py @@ -357,12 +357,18 @@ def test_galileo_record_to_v2_span_with_tags_and_offset(): def test_galileo_get_output_str_variants(galileo_v2_env): logger = GalileoObserve() - assert logger.get_output_str_from_response(None, {}) is None + assert logger.get_output_str_from_response(None, {}) == "" assert ( logger.get_output_str_from_response( EmbeddingResponse(), {"call_type": "embedding"} ) - is None + == "embedding-output" + ) + assert ( + logger.get_output_str_from_response( + EmbeddingResponse(), {"call_type": "aembedding"} + ) + == "embedding-output" ) text_resp = TextCompletionResponse() @@ -414,7 +420,7 @@ def test_galileo_get_output_str_variants(galileo_v2_env): {"call_type": "acompletion", "messages": [{"role": "user", "content": "hi"}]}, ) - assert logger.get_output_str_from_response("not-a-supported-type", {}) is None + assert logger.get_output_str_from_response("not-a-supported-type", {}) == "" def test_galileo_get_input_output_error_status_message(galileo_v2_env): @@ -445,6 +451,48 @@ def test_galileo_get_output_str_rerank_response(galileo_v2_env): assert '"relevance_score": 0.98' in output +@pytest.mark.asyncio +async def test_galileo_async_log_success_embedding(galileo_v2_env): + import datetime + + logger = GalileoObserve() + embedding_response = EmbeddingResponse( + data=[{"object": "embedding", "embedding": [0.1, 0.2, 0.3], "index": 0}] + ) + + mock_response = MagicMock() + mock_response.is_success = True + mock_response.status_code = 201 + + with patch.object(logger.async_httpx_handler, "post", return_value=mock_response): + await logger.async_log_success_event( + kwargs={ + "call_type": "aembedding", + "model": "text-embedding-3-small", + "input": "hello world", + "standard_logging_object": { + "call_type": "aembedding", + "model": "text-embedding-3-small", + "prompt_tokens": 2, + "completion_tokens": 0, + "total_tokens": 2, + "response_cost": 0.0, + "startTime": datetime.datetime( + 2026, 5, 25, 12, 0, 0, tzinfo=datetime.timezone.utc + ).timestamp(), + "endTime": datetime.datetime( + 2026, 5, 25, 12, 0, 1, tzinfo=datetime.timezone.utc + ).timestamp(), + }, + }, + response_obj=embedding_response, + start_time=datetime.datetime(2026, 5, 25, 12, 0, 0), + end_time=datetime.datetime(2026, 5, 25, 12, 0, 1), + ) + + assert logger.in_memory_records == [] + + @pytest.mark.asyncio async def test_galileo_async_log_success_rerank(galileo_v2_env): import datetime @@ -524,6 +572,158 @@ def test_galileo_get_ingest_request_legacy(monkeypatch): assert payload["traces"][0]["input"] == "hi" +@pytest.mark.asyncio +async def test_galileo_async_health_check_success(galileo_v2_env): + logger = GalileoObserve() + current_user_resp = MagicMock() + current_user_resp.status_code = 200 + + with patch.object( + logger.async_httpx_handler, "get", new_callable=AsyncMock + ) as mock_get: + mock_get.return_value = current_user_resp + result = await logger.async_health_check() + + assert result["status"] == "healthy" + mock_get.assert_awaited_once_with( + url="https://api.galileo.ai/current_user", + headers={ + "accept": "application/json", + "Content-Type": "application/json", + "Galileo-API-Key": "test-api-key", + }, + ) + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_api_error(galileo_v2_env): + logger = GalileoObserve() + current_user_resp = MagicMock() + current_user_resp.status_code = 401 + + with patch.object( + logger.async_httpx_handler, "get", new_callable=AsyncMock + ) as mock_get: + mock_get.return_value = current_user_resp + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert "HTTP 401" in result["error_message"] + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_missing_project_id(monkeypatch): + monkeypatch.setenv("GALILEO_API_KEY", "test-api-key") + monkeypatch.setenv("GALILEO_BASE_URL", "https://api.galileo.ai") + monkeypatch.delenv("GALILEO_PROJECT_ID", raising=False) + logger = GalileoObserve() + + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert "GALILEO_PROJECT_ID" in result["error_message"] + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_missing_base_url(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_BASE_URL", raising=False) + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + logger = GalileoObserve() + + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert "GALILEO_BASE_URL" in result["error_message"] + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_missing_credentials(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_USERNAME", raising=False) + monkeypatch.delenv("GALILEO_PASSWORD", raising=False) + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + logger = GalileoObserve() + + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert "GALILEO_USERNAME" in result["error_message"] + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_auth_failed(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + logger = GalileoObserve() + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.side_effect = Exception("login failed") + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert result["error_message"] == "Galileo authentication failed" + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_request_exception(galileo_v2_env): + logger = GalileoObserve() + + with patch.object( + logger.async_httpx_handler, "get", new_callable=AsyncMock + ) as mock_get: + mock_get.side_effect = Exception("connection refused") + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert "connection refused" in result["error_message"] + + +@pytest.mark.asyncio +async def test_galileo_async_log_success_empty_model_response(galileo_v2_env): + import datetime + + logger = GalileoObserve() + logger.batch_size = 2 + empty_response = ModelResponse(choices=[]) + + await logger.async_log_success_event( + kwargs={ + "call_type": "acompletion", + "model": "gpt-5.2", + "messages": [{"role": "user", "content": "hi"}], + "standard_logging_object": { + "call_type": "acompletion", + "model": "gpt-5.2", + "prompt_tokens": 1, + "completion_tokens": 0, + "total_tokens": 1, + "response_cost": 0.0, + "startTime": datetime.datetime( + 2026, 5, 25, 12, 0, 0, tzinfo=datetime.timezone.utc + ).timestamp(), + "endTime": datetime.datetime( + 2026, 5, 25, 12, 0, 1, tzinfo=datetime.timezone.utc + ).timestamp(), + }, + }, + response_obj=empty_response, + start_time=datetime.datetime(2026, 5, 25, 12, 0, 0), + end_time=datetime.datetime(2026, 5, 25, 12, 0, 1), + ) + + assert len(logger.in_memory_records) == 1 + assert logger.in_memory_records[0]["output_text"] == "" + + @pytest.mark.asyncio async def test_galileo_ensure_headers_v2_missing_key(monkeypatch): monkeypatch.delenv("GALILEO_API_KEY", raising=False) diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/test_litellm/integrations/test_prometheus_labels.py index 1ba332a341b..c83d89e87c4 100644 --- a/tests/test_litellm/integrations/test_prometheus_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_labels.py @@ -284,6 +284,12 @@ def test_prometheus_metrics_use_normalized_routes(): # Create a mock PrometheusLogger prometheus_logger = MagicMock() + # ``get_labels_for_metric`` reads ``_cached_metric_labels`` and + # ``label_filters`` off ``self``; default MagicMock attribute access + # returns Mocks that masquerade as a populated cache, so seed real + # containers before binding the real method. + prometheus_logger._cached_metric_labels = {} + prometheus_logger.label_filters = {} prometheus_logger.get_labels_for_metric = ( PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger) ) @@ -327,6 +333,8 @@ def test_prometheus_label_value_sanitization(): from unittest.mock import MagicMock prometheus_logger = MagicMock() + prometheus_logger._cached_metric_labels = {} + prometheus_logger.label_filters = {} prometheus_logger.get_labels_for_metric = ( PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger) ) diff --git a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py new file mode 100644 index 00000000000..bb035c4c3ee --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py @@ -0,0 +1,328 @@ +""" +Tests for the Prometheus rate-limit labels added on top of PR #27687. + +Covers two follow-up gaps to the unified rate-limit error work: + +1. ``litellm_proxy_failed_requests_metric`` now carries + ``rate_limit_category`` and ``rate_limit_type`` labels populated from + :class:`litellm.RateLimitError` (vendor + ``ProxyRateLimitError`` + subclass). Closes the Prometheus side of LIT-2718. +2. ``_get_exception_class_name`` keeps emitting the literal string + ``"HTTPException"`` for ``ProxyRateLimitError`` so existing dashboards + that key off ``exception_class="HTTPException"`` for litellm-internal + 429s don't silently break when the new class lands. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.exceptions import ( + RateLimitError, + RateLimitErrorCategory, + RateLimitType, +) +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.types.integrations.prometheus import ( + PrometheusMetricLabels, + UserAPIKeyLabelNames, + UserAPIKeyLabelValues, +) + + +# --------------------------------------------------------------------------- +# Label / enum wiring +# --------------------------------------------------------------------------- + + +def test_should_register_rate_limit_label_names_on_enum(): + assert UserAPIKeyLabelNames.RATE_LIMIT_CATEGORY.value == "rate_limit_category" + assert UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value == "rate_limit_type" + + +def test_should_include_rate_limit_labels_on_failed_requests_metric(): + import litellm + + original = litellm.prometheus_emit_rate_limit_labels + try: + litellm.prometheus_emit_rate_limit_labels = True + labels = PrometheusMetricLabels.get_labels( + "litellm_proxy_failed_requests_metric" + ) + assert "rate_limit_category" in labels + assert "rate_limit_type" in labels + # These must coexist with the legacy exception labels (back-compat). + assert "exception_class" in labels + assert "exception_status" in labels + finally: + litellm.prometheus_emit_rate_limit_labels = original + + +def test_should_omit_rate_limit_labels_by_default_for_back_compat(): + """Default-off preserves the metric's historical label set so existing + dashboards / recording rules keyed on `litellm_proxy_failed_requests_metric` + keep matching after upgrade.""" + import litellm + + assert litellm.prometheus_emit_rate_limit_labels is False + labels = PrometheusMetricLabels.get_labels("litellm_proxy_failed_requests_metric") + assert "rate_limit_category" not in labels + assert "rate_limit_type" not in labels + # Pre-PR labels must still be present. + assert "exception_class" in labels + assert "exception_status" in labels + + +def test_should_accept_rate_limit_fields_on_user_api_key_label_values(): + enum_values = UserAPIKeyLabelValues( + rate_limit_category="litellm_rate_limit", + rate_limit_type="requests", + ) + assert enum_values.rate_limit_category == "litellm_rate_limit" + assert enum_values.rate_limit_type == "requests" + + +# --------------------------------------------------------------------------- +# _extract_rate_limit_labels helper +# --------------------------------------------------------------------------- + + +def test_should_extract_vendor_category_for_vanilla_rate_limit_error(): + err = RateLimitError(message="vendor 429", llm_provider="openai", model="gpt-4o") + category, rate_limit_type = PrometheusLogger._extract_rate_limit_labels(err) + assert category == "vendor_rate_limit" + assert rate_limit_type is None + + +def test_should_extract_litellm_category_and_type_for_proxy_rate_limit_error(): + err = ProxyRateLimitError( + detail={"error": "tpm exceeded"}, + category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, + rate_limit_type=RateLimitType.TOKENS, + ) + category, rate_limit_type = PrometheusLogger._extract_rate_limit_labels(err) + assert category == "litellm_rate_limit" + assert rate_limit_type == "tokens" + + +def test_should_return_none_for_non_rate_limit_exception(): + assert PrometheusLogger._extract_rate_limit_labels(ValueError("nope")) == ( + None, + None, + ) + + +def test_should_return_none_for_none_exception(): + assert PrometheusLogger._extract_rate_limit_labels(None) == (None, None) + + +def test_should_extract_budget_dimension_for_budget_exceeded_error(): + # Virtual-key / team / org / end-user budget caps raise + # `litellm.BudgetExceededError` (a bare Exception subclass), which sets + # the same `.category` / `.rate_limit_type` attributes as the unified + # RateLimitError path so Prometheus can split budget 429s from other + # 429s without the customer parsing free-text error messages. + import litellm + + err = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + category, rate_limit_type = PrometheusLogger._extract_rate_limit_labels(err) + assert category == "litellm_rate_limit" + assert rate_limit_type == "budget" + + +@pytest.mark.parametrize( + "category_enum,rate_limit_enum,expected_category,expected_type", + [ + ( + RateLimitErrorCategory.LITELLM_RATE_LIMIT, + RateLimitType.REQUESTS, + "litellm_rate_limit", + "requests", + ), + ( + RateLimitErrorCategory.LITELLM_RATE_LIMIT, + RateLimitType.TOKENS, + "litellm_rate_limit", + "tokens", + ), + ( + RateLimitErrorCategory.LITELLM_RATE_LIMIT, + RateLimitType.CONCURRENT_REQUESTS, + "litellm_rate_limit", + "concurrent_requests", + ), + ( + RateLimitErrorCategory.LITELLM_RATE_LIMIT, + RateLimitType.BUDGET, + "litellm_rate_limit", + "budget", + ), + ( + RateLimitErrorCategory.LITELLM_RATE_LIMIT, + RateLimitType.MAX_ITERATIONS, + "litellm_rate_limit", + "max_iterations", + ), + ( + RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT, + RateLimitType.REQUESTS, + "litellm_batch_rate_limit", + "requests", + ), + ], +) +def test_should_serialize_rate_limit_enums_as_underlying_string_values( + category_enum, rate_limit_enum, expected_category, expected_type +): + err = ProxyRateLimitError( + detail="boom", category=category_enum, rate_limit_type=rate_limit_enum + ) + category, rate_limit_type = PrometheusLogger._extract_rate_limit_labels(err) + assert category == expected_category + assert rate_limit_type == expected_type + + +# --------------------------------------------------------------------------- +# _get_exception_class_name back-compat +# --------------------------------------------------------------------------- + + +def test_should_emit_legacy_http_exception_label_for_proxy_rate_limit_error(): + """ + ``ProxyRateLimitError`` multi-inherits from ``HTTPException`` + + ``RateLimitError``. The ``exception_class`` label MUST keep emitting + "HTTPException" for back-compat with existing dashboards (see Slack + thread + PR #27687 review). Distinguishing vendor vs. litellm 429s + is now the job of the new ``rate_limit_category`` label. + """ + err = ProxyRateLimitError(detail={"error": "boom"}) + assert PrometheusLogger._get_exception_class_name(err) == "HTTPException" + + +def test_should_keep_provider_prefixed_exception_class_for_vendor_rate_limit_errors(): + err = RateLimitError(message="vendor 429", llm_provider="openai", model="gpt-4o") + # Vendor-side errors keep the historical "Provider.ClassName" formatting. + assert PrometheusLogger._get_exception_class_name(err) == "Openai.RateLimitError" + + +def test_should_preserve_exception_class_name_for_unrelated_exceptions(): + assert PrometheusLogger._get_exception_class_name(ValueError("nope")) == ( + "ValueError" + ) + + +# --------------------------------------------------------------------------- +# End-to-end wiring through async_post_call_failure_hook +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_should_populate_rate_limit_labels_for_proxy_rate_limit_error_on_failure_hook(): + """ + When a proxy hook raises ``ProxyRateLimitError`` and the failure flows + through ``async_post_call_failure_hook``, the resulting + ``UserAPIKeyLabelValues`` must carry both new labels AND keep + ``exception_class="HTTPException"`` for back-compat. + """ + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_proxy_failed_requests_metric = MagicMock() + logger.litellm_proxy_total_requests_metric = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=PrometheusMetricLabels.get_labels( + "litellm_proxy_failed_requests_metric" + ) + ) + + err = ProxyRateLimitError( + detail={"error": "rpm exceeded"}, + category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, + rate_limit_type=RateLimitType.REQUESTS, + ) + + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + await logger.async_post_call_failure_hook( + request_data={"model": "gpt-4o-mini", "metadata": {}}, + original_exception=err, + user_api_key_dict=UserAPIKeyAuth(token="t"), + ) + + enum_values = mock_label_factory.call_args_list[0].kwargs["enum_values"] + assert isinstance(enum_values, UserAPIKeyLabelValues) + assert enum_values.rate_limit_category == "litellm_rate_limit" + assert enum_values.rate_limit_type == "requests" + # Back-compat: exception_class on a ProxyRateLimitError stays "HTTPException". + assert enum_values.exception_class == "HTTPException" + assert enum_values.exception_status == "429" + + +@pytest.mark.asyncio +async def test_should_populate_rate_limit_labels_for_vendor_rate_limit_error_on_failure_hook(): + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_proxy_failed_requests_metric = MagicMock() + logger.litellm_proxy_total_requests_metric = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=PrometheusMetricLabels.get_labels( + "litellm_proxy_failed_requests_metric" + ) + ) + + err = RateLimitError(message="upstream 429", llm_provider="openai", model="gpt-4o") + + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + await logger.async_post_call_failure_hook( + request_data={"model": "gpt-4o", "metadata": {}}, + original_exception=err, + user_api_key_dict=UserAPIKeyAuth(token="t"), + ) + + enum_values = mock_label_factory.call_args_list[0].kwargs["enum_values"] + assert isinstance(enum_values, UserAPIKeyLabelValues) + assert enum_values.rate_limit_category == "vendor_rate_limit" + assert enum_values.rate_limit_type is None + # Vendor errors keep the historical Provider.ClassName label. + assert enum_values.exception_class == "Openai.RateLimitError" + assert enum_values.exception_status == "429" + + +@pytest.mark.asyncio +async def test_should_leave_rate_limit_labels_blank_for_non_rate_limit_failure(): + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_proxy_failed_requests_metric = MagicMock() + logger.litellm_proxy_total_requests_metric = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=PrometheusMetricLabels.get_labels( + "litellm_proxy_failed_requests_metric" + ) + ) + + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + await logger.async_post_call_failure_hook( + request_data={"model": "gpt-4o", "metadata": {}}, + original_exception=RuntimeError("boom"), + user_api_key_dict=UserAPIKeyAuth(token="t"), + ) + + enum_values = mock_label_factory.call_args_list[0].kwargs["enum_values"] + assert isinstance(enum_values, UserAPIKeyLabelValues) + assert enum_values.rate_limit_category is None + assert enum_values.rate_limit_type is None diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py index 12f30ab6024..361ab7332f8 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -511,29 +511,34 @@ def test_set_user_budget_metrics_default_no_email_alias_labels( ) -def test_set_user_budget_metrics_includes_user_email_and_alias_labels_when_opted_in( - prometheus_logger, -): - """When prometheus_user_budget_label_include_email_alias=True, email+alias labels appear.""" +def test_set_user_budget_metrics_includes_user_email_and_alias_labels_when_opted_in(): + """When prometheus_user_budget_label_include_email_alias=True, email+alias labels appear. + + The flag is read once per metric at logger construction time and snapshotted, + so it must be enabled before the PrometheusLogger is built (mirroring how the + proxy applies config at startup before instantiating callbacks). + """ import litellm from litellm.proxy._types import LiteLLM_UserTable litellm.prometheus_user_budget_label_include_email_alias = True - user = LiteLLM_UserTable( - user_id="user-abc-123", - user_email="alice@example.com", - user_alias="Alice", - spend=25.0, - max_budget=100.0, - budget_reset_at=datetime(2026, 3, 1, tzinfo=timezone.utc), - ) + try: + prometheus_logger = PrometheusLogger() - prometheus_logger.litellm_remaining_user_budget_metric = MagicMock() - prometheus_logger.litellm_user_max_budget_metric = MagicMock() - prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock() + user = LiteLLM_UserTable( + user_id="user-abc-123", + user_email="alice@example.com", + user_alias="Alice", + spend=25.0, + max_budget=100.0, + budget_reset_at=datetime(2026, 3, 1, tzinfo=timezone.utc), + ) + + prometheus_logger.litellm_remaining_user_budget_metric = MagicMock() + prometheus_logger.litellm_user_max_budget_metric = MagicMock() + prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock() - try: prometheus_logger._set_user_budget_metrics(user) prometheus_logger.litellm_remaining_user_budget_metric.labels.assert_called_once_with( diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index 10951265115..c2a502b34eb 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -34,6 +34,35 @@ def test_initialize_from_proxy_config(): assert logger.search_tool_name == "my-search" +def test_initialize_from_proxy_config_ignores_non_dict_callback_specific_params(): + """Regression (#29590): a non-dict value under + callback_settings.websearch_interception must not crash initialization. + + Forwarding callback_settings as callback_specific_params activates this + branch; without the isinstance(dict) guard a non-dict value reached + from_config_yaml(...).get(...) and raised AttributeError at proxy startup. + The value is ignored and the logger falls back to defaults. + """ + logger = WebSearchInterceptionLogger.initialize_from_proxy_config( + litellm_settings={}, + callback_specific_params={"websearch_interception": True}, + ) + + assert logger.search_tool_name is None + + +def test_initialize_from_proxy_config_honors_dict_callback_specific_params(): + """A valid dict under callback_settings.websearch_interception is applied.""" + logger = WebSearchInterceptionLogger.initialize_from_proxy_config( + litellm_settings={}, + callback_specific_params={ + "websearch_interception": {"search_tool_name": "ws-tool"} + }, + ) + + assert logger.search_tool_name == "ws-tool" + + @pytest.mark.asyncio async def test_async_should_run_agentic_loop(): """Test that agentic loop is NOT triggered for wrong provider or missing WebSearch tool""" diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 2b47a232262..fe49b930c10 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -328,6 +328,41 @@ def test_generic_cost_per_token_gpt54_above_272k_tokens(): assert round(completion_cost, 10) == round(expected_completion, 10) +def test_generic_cost_per_token_minimax_m3_above_512k_tokens(): + """MiniMax-M3: prompts >512K input tokens priced at 2x input, output, and cache read.""" + model = "minimax/MiniMax-M3" + custom_llm_provider = "minimax" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model_cost_map = litellm.model_cost[model] + prompt_tokens = 600000 + cached_tokens = 100000 + completion_tokens = 1000 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + expected_prompt = ( + model_cost_map["input_cost_per_token_above_512k_tokens"] + * (prompt_tokens - cached_tokens) + + model_cost_map["cache_read_input_token_cost_above_512k_tokens"] + * cached_tokens + ) + expected_completion = ( + model_cost_map["output_cost_per_token_above_512k_tokens"] * completion_tokens + ) + assert round(prompt_cost, 10) == round(expected_prompt, 10) + assert round(completion_cost, 10) == round(expected_completion, 10) + + def test_generic_cost_per_token_gpt55(): """gpt-5.5: base pricing — $5/1M input, $30/1M output, $0.50/1M cached input.""" model = "gpt-5.5" diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index a04f6407e4b..c43291566b6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,17 +1,14 @@ -import json import os import sys -from unittest.mock import MagicMock import pytest -from fastapi.testclient import TestClient import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) from litellm.types.llms.openai import FileSearchTool, WebSearchOptions -from litellm.types.utils import ModelInfo, ModelResponse, StandardBuiltInToolsParams +from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams sys.path.insert( 0, os.path.abspath("../../..") @@ -139,6 +136,22 @@ def test_get_cost_for_anthropic_web_search(): assert cost > 0.0 +def test_get_cost_for_anthropic_web_search_with_server_tool_use_dict(): + """ + Anthropic-compatible passthrough responses can construct Usage from a raw + usage payload. Ensure dict server_tool_use values are normalized before + built-in tool cost tracking reads server_tool_use.web_search_requests. + """ + from litellm.types.utils import ServerToolUse, Usage + + usage = Usage(server_tool_use={"web_search_requests": 1}) + + assert isinstance(usage.server_tool_use, ServerToolUse) + assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=None, usage=usage + ) + + @pytest.mark.parametrize( "model", ["gemini/gemini-2.0-flash-001", "gemini-2.0-flash-001"] ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py new file mode 100644 index 00000000000..4eee6b59d34 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py @@ -0,0 +1,88 @@ +""" +Tests that the cost-tracking call sites tolerate ``server_tool_use`` being +either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( + StandardBuiltInToolCostTracking, + _get_web_search_requests, +) +from litellm.types.utils import ModelResponse, ServerToolUse, Usage + + +class _UsageWithDictServerToolUse: + """ + Tiny stand-in that mimics the broken streaming-rebuild shape: + ``server_tool_use`` is a plain dict. + """ + + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + self.prompt_tokens_details = None + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 5}) == 5 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + stu = ServerToolUse(web_search_requests=7) + assert _get_web_search_requests(stu) == 7 + + +def test_get_web_search_requests_handles_pydantic_with_none_value(): + stu = ServerToolUse() + assert _get_web_search_requests(stu) is None + + +def test_response_object_includes_web_search_call_with_dict_server_tool_use(): + """ + The exact bug: ``usage.server_tool_use`` is a dict and the check in + ``response_object_includes_web_search_call`` used to crash with + ``AttributeError``. + """ + response = ModelResponse() + usage = _UsageWithDictServerToolUse({"web_search_requests": 2}) + + # Must not raise — and must correctly detect the web search call. + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_pydantic_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(ServerToolUse(web_search_requests=2)) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_none_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(None) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is False diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 91fd07dcffc..ed2dfc9440e 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -11,6 +11,7 @@ BedrockImageProcessor, _bedrock_converse_messages_pt, _bedrock_tools_pt, + _rename_duplicate_bedrock_document_names, _convert_to_bedrock_tool_call_invoke, _convert_to_bedrock_tool_call_result, anthropic_messages_pt, @@ -898,6 +899,61 @@ def test_bedrock_tools_unpack_defs(): _bedrock_tools_pt(tools=tools) +def test_bedrock_tools_pt_strict_parameter(): + """Regression for strict tools on the Bedrock Converse path. + + Claude on Bedrock honours strict in toolSpec (with additionalProperties, which + Bedrock requires alongside strict); without forwarding it the model ignores the + enum constraint the caller asked for. Every other Bedrock family (Nova, Llama, + GPT-OSS) rejects the strict field, so it must only be forwarded for Claude. + """ + tools_with_strict = [ + { + "type": "function", + "function": { + "name": "generate_sql", + "strict": True, + "description": "Generate a SQL query", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + "additionalProperties": False, + }, + }, + } + ] + result = _bedrock_tools_pt( + tools_with_strict, model="anthropic.claude-sonnet-4-5-20250929-v1:0" + ) + assert result[0]["toolSpec"]["strict"] is True + assert result[0]["toolSpec"]["inputSchema"]["json"]["additionalProperties"] is False + + result = _bedrock_tools_pt(tools_with_strict, model="us.amazon.nova-micro-v1:0") + assert "strict" not in result[0]["toolSpec"] + assert "additionalProperties" not in result[0]["toolSpec"]["inputSchema"]["json"] + + tools_without_strict = [ + { + "type": "function", + "function": { + "name": "generate_sql", + "description": "Generate a SQL query", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, + } + ] + result = _bedrock_tools_pt( + tools_without_strict, model="anthropic.claude-sonnet-4-5-20250929-v1:0" + ) + assert "strict" not in result[0]["toolSpec"] + assert "additionalProperties" not in result[0]["toolSpec"]["inputSchema"]["json"] + + def test_bedrock_image_processor_content_type_fallback_url_extension(): """ Test that _post_call_image_processing falls back to URL extension @@ -2754,6 +2810,93 @@ def test_bedrock_converse_messages_pt_document_deterministic_name(): assert name1 == name2 +def test_bedrock_converse_messages_pt_renames_duplicate_document_names(): + """ + The same document in multiple turns must not produce duplicate names; + Bedrock rejects requests with "Messages can not contain duplicate + document names". The first occurrence keeps its hash-based name and + later occurrences get a deterministic positional suffix. + """ + document_block = { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "dGVzdA==", + }, + } + messages = [ + { + "role": "user", + "content": [document_block, {"type": "text", "text": "summarize this"}], + }, + {"role": "assistant", "content": "It says test."}, + { + "role": "user", + "content": [document_block, {"type": "text", "text": "summarize again"}], + }, + ] + + result1 = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + result2 = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + + names1 = [ + block["document"]["name"] + for message in result1 + for block in message["content"] + if "document" in block + ] + names2 = [ + block["document"]["name"] + for message in result2 + for block in message["content"] + if "document" in block + ] + + assert len(names1) == 2 + assert len(set(names1)) == 2 + assert names1[1] == f"{names1[0]}_2" + assert names1 == names2 + + single_turn = _bedrock_converse_messages_pt( + [messages[0]], "anthropic.claude-sonnet-4-6", "bedrock" + ) + assert names1[0] == single_turn[0]["content"][0]["document"]["name"] + + +def test_rename_duplicate_bedrock_document_names_skips_organic_suffixes(): + """ + A renamed duplicate must not collide with a document whose organic name + already carries the would-be suffix (e.g. an existing ``report_2``), + regardless of whether that document appears before or after the rename. + """ + + def _contents(names): + return [ + { + "role": "user", + "content": [{"document": {"name": name}} for name in names], + } + ] + + def _names(contents): + return [block["document"]["name"] for block in contents[0]["content"]] + + organic_first = _rename_duplicate_bedrock_document_names( + _contents(["report", "report_2", "report"]) + ) + assert _names(organic_first) == ["report", "report_2", "report_3"] + + organic_last = _rename_duplicate_bedrock_document_names( + _contents(["report", "report", "report_2"]) + ) + assert _names(organic_last) == ["report", "report_3", "report_2"] + + def test_bedrock_converse_messages_pt_document_rejects_url_source(): """Test that a URL-type document source raises a clear error instead of KeyError.""" messages = [ diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index d57d8dafdbd..34edd6eccf3 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2165,6 +2165,41 @@ def test_get_assembled_streaming_response_returns_result_for_streaming(): assert assembled is result +def test_streaming_success_handler_includes_vertex_ai_metadata_in_standard_logging(): + """Assembled streaming responses should include Vertex AI metadata in logging payload.""" + import datetime + + from litellm.types.utils import Choices, Message + + logging_obj = _make_logging_obj(stream=True) + grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] + url_context_metadata = [{"urlMetadata": [{"retrievedUrl": "https://example.com"}]}] + result = ModelResponse( + id="resp-1", + choices=[ + Choices( + index=0, + message=Message(role="assistant", content="hello"), + finish_reason="stop", + ) + ], + model="gemini-2.5-flash", + ) + setattr(result, "vertex_ai_grounding_metadata", grounding_metadata) + setattr(result, "vertex_ai_url_context_metadata", url_context_metadata) + result._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata + result._hidden_params["vertex_ai_url_context_metadata"] = url_context_metadata + + start = datetime.datetime.now() + end = datetime.datetime.now() + logging_obj.success_handler(result=result, start_time=start, end_time=end) + + payload = logging_obj.model_call_details.get("standard_logging_object") + assert payload is not None + assert payload["response"]["vertex_ai_grounding_metadata"] == grounding_metadata + assert payload["response"]["vertex_ai_url_context_metadata"] == url_context_metadata + + def test_get_assembled_streaming_response_returns_none_for_non_streaming_text_completion(): """Non-streaming TextCompletionResponse should also return None.""" import datetime diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 3424bfd801c..0f8d5cfd85a 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -773,17 +773,15 @@ async def apply_guardrail( guardrail_items = [ e for e in sent_to_backend if e.get("type") == "conversation.item.create" ] - assert len(guardrail_items) == 1, ( - f"Guardrail should inject a conversation.item.create with violation message, " - f"got: {guardrail_items}" - ) + assert ( + len(guardrail_items) == 1 + ), f"Guardrail should inject a conversation.item.create with violation message, got: {guardrail_items}" response_creates = [ e for e in sent_to_backend if e.get("type") == "response.create" ] - assert len(response_creates) == 1, ( - f"Guardrail should send exactly one response.create to voice the violation, " - f"got: {response_creates}" - ) + assert ( + len(response_creates) == 1 + ), f"Guardrail should send exactly one response.create to voice the violation, got: {response_creates}" # ASSERT 2: error event was sent directly to the client WebSocket sent_to_client = [ @@ -1050,10 +1048,9 @@ async def apply_guardrail( # every toolCall with a toolResponse (Gemini/Vertex Live) exit their # pending-tool-call state instead of stalling. The placeholder must NOT # contain any of the blocked content. - assert len(forwarded_tool_outputs) == 1, ( - f"Sanitized function_call_output should be forwarded, got: " - f"{forwarded_tool_outputs}" - ) + assert ( + len(forwarded_tool_outputs) == 1 + ), f"Sanitized function_call_output should be forwarded, got: {forwarded_tool_outputs}" sanitized_item = forwarded_tool_outputs[0]["item"] assert sanitized_item["call_id"] == "call_123" assert "test@example.com" not in sanitized_item["output"] @@ -2110,3 +2107,202 @@ async def test_deferred_setup_caps_non_audio_buffered_bytes(monkeypatch): assert ( streaming._pending_messages_byte_total <= RealTimeStreaming._MAX_BUFFERED_BYTES ) + + +def _beta_client_ws(): + ws = MagicMock() + ws.scope = {"headers": [(b"openai-beta", b"realtime=v1")]} + ws.send_text = AsyncMock() + return ws + + +def _ga_client_ws(): + ws = MagicMock() + ws.scope = {"headers": []} + ws.send_text = AsyncMock() + return ws + + +def _streaming_with(client_ws): + backend_ws = MagicMock() + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + return RealTimeStreaming(client_ws, backend_ws, logging_obj) + + +def test_parse_backend_event_returns_none_for_non_json(): + assert RealTimeStreaming._parse_backend_event("not json") is None + + +def test_parse_backend_event_returns_none_for_non_dict_json(): + assert RealTimeStreaming._parse_backend_event("[1, 2, 3]") is None + assert RealTimeStreaming._parse_backend_event('"a string"') is None + + +def test_parse_backend_event_returns_dict(): + parsed = RealTimeStreaming._parse_backend_event('{"type": "x", "v": 1}') + assert parsed == {"type": "x", "v": 1} + + +def test_translate_event_to_beta_returns_identity_when_no_translation(): + """An event with no renamed type and no item/response is returned unchanged + (same object), so the caller can forward the raw frame without re-serializing.""" + ev = {"type": "error", "error": {"message": "boom"}} + out = RealTimeStreaming._translate_event_to_beta(ev) + assert out is ev + + +def test_translate_event_to_beta_preserves_audio_delta_payload(): + payload = "QUJDREVG" * 200 + out = RealTimeStreaming._translate_event_to_beta( + {"type": "response.output_audio.delta", "delta": payload, "event_id": "e1"} + ) + assert out is not None + assert out["type"] == "response.audio.delta" + assert out["delta"] == payload + + +def test_translate_event_to_beta_remaps_response_done_output_content_types(): + out = RealTimeStreaming._translate_event_to_beta( + { + "type": "response.done", + "response": { + "output": [ + { + "type": "message", + "content": [{"type": "output_audio", "transcript": "hi"}], + } + ] + }, + } + ) + assert out is not None + assert out["response"]["output"][0]["content"][0]["type"] == "audio" + + +@pytest.mark.asyncio +async def test_beta_client_receives_translated_audio_delta(): + client_ws = _beta_client_ws() + frame = json.dumps( + {"type": "response.output_audio.delta", "delta": "QUJD", "event_id": "e1"} + ) + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[frame.encode(), ConnectionClosed(None, None)] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.backend_to_client_send_messages() + + assert client_ws.send_text.await_count == 1 + sent = json.loads(client_ws.send_text.await_args.args[0]) + assert sent["type"] == "response.audio.delta" + assert sent["delta"] == "QUJD" + + +@pytest.mark.asyncio +async def test_ga_client_receives_raw_passthrough(): + client_ws = _ga_client_ws() + frame = json.dumps( + {"type": "response.output_audio.delta", "delta": "QUJD", "event_id": "e1"} + ) + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[frame.encode(), ConnectionClosed(None, None)] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.backend_to_client_send_messages() + + assert client_ws.send_text.await_count == 1 + # GA client gets the byte-identical frame, no re-serialization. + assert client_ws.send_text.await_args.args[0] == frame + + +@pytest.mark.asyncio +async def test_beta_client_non_translated_event_forwarded_raw(): + """For a beta client, an event needing no translation is forwarded as the + original raw frame (identity return path), not a re-serialized copy.""" + client_ws = _beta_client_ws() + frame = json.dumps({"type": "error", "error": {"message": "boom"}}) + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[frame.encode(), ConnectionClosed(None, None)] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.backend_to_client_send_messages() + + assert client_ws.send_text.await_count == 1 + assert client_ws.send_text.await_args.args[0] == frame + + +@pytest.mark.asyncio +async def test_beta_client_drops_conversation_item_done(): + client_ws = _beta_client_ws() + frame = json.dumps({"type": "conversation.item.done", "item": {"id": "i1"}}) + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[frame.encode(), ConnectionClosed(None, None)] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.backend_to_client_send_messages() + + assert client_ws.send_text.await_count == 0 + + +def test_store_message_skips_pydantic_for_unlogged_audio_delta(): + """Audio deltas are not in DefaultLoggedRealTimeEventTypes; store_message must + skip the Pydantic build entirely (no append, no validation).""" + streaming = _streaming_with(_ga_client_ws()) + with patch( + "litellm.litellm_core_utils.realtime_streaming.OpenAIRealtimeStreamResponseBaseObject" + ) as base_obj: + streaming.store_message({"type": "response.output_audio.delta", "delta": "x"}) + base_obj.assert_not_called() + assert streaming.messages == [] + + +@pytest.mark.asyncio +async def test_audio_delta_frame_parsed_at_most_once(): + client_ws = _beta_client_ws() + frame = json.dumps( + {"type": "response.output_audio.delta", "delta": "QUJD", "event_id": "e1"} + ) + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[frame.encode(), ConnectionClosed(None, None)] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + real_loads = json.loads + calls = {"n": 0} + + def counting_loads(*args, **kwargs): + calls["n"] += 1 + return real_loads(*args, **kwargs) + + with patch( + "litellm.litellm_core_utils.realtime_streaming.json.loads", + side_effect=counting_loads, + ): + await streaming.backend_to_client_send_messages() + + assert calls["n"] == 1 diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 60cfff6e4a0..36f220f9a2c 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -349,3 +349,96 @@ def test_redacts_responses_api_response_object(self): assert redacted.output[0].content[0].text == "redacted-by-litellm" assert response.output[0].content[0].text == "sensitive output" + + def test_redacts_vertex_provider_metadata_in_standard_logging_response(self): + details = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "sensitive prompt"}], + "response": { + "choices": [ + { + "message": { + "content": "sensitive answer", + "role": "assistant", + } + } + ], + "vertex_ai_grounding_metadata": [ + {"webSearchQueries": ["sensitive search term"]} + ], + "vertex_ai_url_context_metadata": [ + {"urlMetadata": [{"retrievedUrl": "https://example.com"}]} + ], + }, + } + } + + perform_redaction(details, None) + + response = details["standard_logging_object"]["response"] + assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" + assert response["vertex_ai_grounding_metadata"] == [] + assert response["vertex_ai_url_context_metadata"] == [] + + def test_redacts_vertex_provider_metadata_on_streaming_model_response(self): + response = litellm.ModelResponse( + id="resp-1", + choices=[ + litellm.Choices( + message=litellm.Message( + content="sensitive answer", + role="assistant", + ) + ) + ], + model="gemini-2.5-flash", + ) + setattr( + response, + "vertex_ai_grounding_metadata", + [{"webSearchQueries": ["sensitive search term"]}], + ) + response._hidden_params["vertex_ai_grounding_metadata"] = [ + {"webSearchQueries": ["sensitive search term"]} + ] + + details = { + "stream": True, + "complete_streaming_response": response, + } + + perform_redaction(details, response) + + assert response.choices[0].message.content == "redacted-by-litellm" + assert getattr(response, "vertex_ai_grounding_metadata") == [] + assert "vertex_ai_grounding_metadata" not in response._hidden_params + + def test_redacts_vertex_provider_metadata_from_metadata_hidden_params(self): + """Streaming success_handler copies _hidden_params into metadata before redaction.""" + details = { + "stream": True, + "litellm_params": { + "metadata": { + "hidden_params": { + "response_cost": 0.01, + "vertex_ai_grounding_metadata": [ + {"webSearchQueries": ["sensitive search term"]} + ], + "vertex_ai_url_context_metadata": [ + {"urlMetadata": [{"retrievedUrl": "https://example.com"}]} + ], + "vertex_ai_safety_ratings": [{"category": "HARM"}], + "vertex_ai_citation_metadata": [{"citations": ["source"]}], + } + } + }, + } + + perform_redaction(details, None) + + hidden_params = details["litellm_params"]["metadata"]["hidden_params"] + assert hidden_params["response_cost"] == 0.01 + assert "vertex_ai_grounding_metadata" not in hidden_params + assert "vertex_ai_url_context_metadata" not in hidden_params + assert "vertex_ai_safety_ratings" not in hidden_params + assert "vertex_ai_citation_metadata" not in hidden_params diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py new file mode 100644 index 00000000000..4e28d5ba7d2 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py @@ -0,0 +1,130 @@ +""" +Regression tests for https://github.com/BerriAI/litellm/issues/26153 + +``stream_chunk_builder`` used to leave ``usage.server_tool_use`` as a plain +``dict`` when reconstructing a streaming response. Downstream cost-calculation +code (``StandardBuiltInToolCostTracking.response_object_includes_web_search_call`` +and ``get_cost_for_anthropic_web_search``) accesses +``usage.server_tool_use.web_search_requests`` as an attribute, which raised +``AttributeError: 'dict' object has no attribute 'web_search_requests'``. + +These tests reconstruct streaming chunks for an Anthropic-style web_search +response and assert: + +1. ``stream_chunk_builder`` returns ``ServerToolUse`` (not ``dict``) for + ``usage.server_tool_use``. +2. ``completion_cost`` runs end-to-end on the rebuilt response without + raising ``AttributeError``. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm import completion_cost, stream_chunk_builder +from litellm.types.utils import ( + Delta, + ModelResponseStream, + ServerToolUse, + StreamingChoices, + Usage, +) + + +def _make_text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(role="assistant", content=text), + ) + ], + ) + + +def _make_finish_chunk_with_usage_dict_server_tool_use() -> ModelResponseStream: + """Final chunk where server_tool_use is a *dict* — reproduces the bug shape.""" + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ], + usage=Usage( + prompt_tokens=42, + completion_tokens=11, + total_tokens=53, + # NOTE: passed as a dict on purpose — this is the shape that + # historically slipped through stream_chunk_builder unchanged. + server_tool_use={"web_search_requests": 3}, + ), + ) + + +def test_stream_chunk_builder_coerces_server_tool_use_to_pydantic(): + """ + Regression: stream_chunk_builder must produce ServerToolUse, not dict. + """ + chunks = [ + _make_text_chunk("Otters "), + _make_text_chunk("are great."), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + + assert rebuilt is not None + assert rebuilt.usage is not None # type: ignore[attr-defined] + server_tool_use = rebuilt.usage.server_tool_use # type: ignore[attr-defined] + + assert ( + server_tool_use is not None + ), "server_tool_use should be carried through from the final chunk" + assert isinstance(server_tool_use, ServerToolUse), ( + f"expected ServerToolUse, got {type(server_tool_use).__name__}: " + f"{server_tool_use!r}" + ) + # Attribute access must not raise (this is exactly what was broken). + assert server_tool_use.web_search_requests == 3 + + +def test_completion_cost_does_not_raise_on_streaming_web_search_response(): + """ + Regression: completion_cost(...) must not raise AttributeError when the + response was reconstructed by stream_chunk_builder from a streaming + Anthropic web_search call. + """ + chunks = [ + _make_text_chunk("hello"), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + assert rebuilt is not None + + # The exact dollar amount depends on the model-pricing table; what matters + # for this regression is that it does NOT raise AttributeError on + # `dict has no attribute 'web_search_requests'`. + try: + cost = completion_cost(completion_response=rebuilt) + except AttributeError as e: # pragma: no cover - regression guard + pytest.fail( + "completion_cost raised AttributeError after stream_chunk_builder " + f"(issue #26153 regression): {e}" + ) + + assert isinstance(cost, (int, float)) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index e40a0817fd9..c5794194528 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -520,7 +520,10 @@ def test_stream_chunk_builder_anthropic_web_search(): assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 assert usage.total_tokens == 77 - assert usage.server_tool_use["web_search_requests"] == 2 + # server_tool_use must be a ServerToolUse pydantic so downstream cost-calc + # (which uses attribute access) works. See issue #26153. + assert isinstance(usage.server_tool_use, ServerToolUse) + assert usage.server_tool_use.web_search_requests == 2 def test_sort_chunks_handles_dict_hidden_params_created_at(): @@ -613,3 +616,153 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): assert ( response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" ) + + +def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): + """Vertex AI metadata on streaming chunks must appear on assembled response.""" + grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] + url_context_metadata = [{"urlMetadata": [{"retrievedUrl": "https://example.com"}]}] + + chunk1 = ModelResponseStream( + id="chatcmpl-vertex-1", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="The weather", role="assistant"), + ) + ], + ) + setattr(chunk1, "vertex_ai_grounding_metadata", grounding_metadata) + chunk1._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata + + chunk2 = ModelResponseStream( + id="chatcmpl-vertex-1", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" is sunny.", role="assistant"), + ) + ], + ) + setattr(chunk2, "vertex_ai_url_context_metadata", url_context_metadata) + chunk2._hidden_params["vertex_ai_url_context_metadata"] = url_context_metadata + + response = stream_chunk_builder(chunks=[chunk1, chunk2]) + assert response is not None + assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata + assert getattr(response, "vertex_ai_url_context_metadata") == url_context_metadata + assert response._hidden_params["vertex_ai_grounding_metadata"] == grounding_metadata + assert ( + response._hidden_params["vertex_ai_url_context_metadata"] + == url_context_metadata + ) + + dumped = response.model_dump() + assert dumped["vertex_ai_grounding_metadata"] == grounding_metadata + assert dumped["vertex_ai_url_context_metadata"] == url_context_metadata + + +def test_stream_chunk_builder_uses_assembled_model_for_provider_metadata(): + grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] + + chunk1 = ModelResponseStream( + id="chatcmpl-vertex-router", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="The weather", role="assistant"), + ) + ], + ) + chunk2 = ModelResponseStream( + id="chatcmpl-vertex-router", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" is sunny.", role=None), + ) + ], + ) + setattr(chunk2, "vertex_ai_grounding_metadata", grounding_metadata) + chunk2._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata + + response = stream_chunk_builder(chunks=[chunk1, chunk2]) + assert response is not None + assert response.model == "gemini-2.5-flash" + assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata + + +def test_stream_chunk_builder_propagates_vertex_ai_safety_results(): + """Assembled response must expose safety data under the non-streaming field name.""" + safety_ratings = [ + [{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}] + ] + + chunk = ModelResponseStream( + id="chatcmpl-vertex-safety", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + setattr(chunk, "vertex_ai_safety_ratings", safety_ratings) + setattr(chunk, "vertex_ai_safety_results", safety_ratings) + chunk._hidden_params["vertex_ai_safety_ratings"] = safety_ratings + chunk._hidden_params["vertex_ai_safety_results"] = safety_ratings + + response = stream_chunk_builder(chunks=[chunk]) + assert response is not None + assert getattr(response, "vertex_ai_safety_results") == safety_ratings + assert response._hidden_params["vertex_ai_safety_results"] == safety_ratings + assert response.model_dump()["vertex_ai_safety_results"] == safety_ratings + + +def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks(): + """Dict snapshot chunks (model_dump) should also propagate Vertex AI metadata.""" + chunk_dict = ModelResponseStream( + id="chatcmpl-vertex-2", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ).model_dump() + chunk_dict["_hidden_params"] = { + "vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}] + } + + response = stream_chunk_builder(chunks=[chunk_dict]) + assert response is not None + assert getattr(response, "vertex_ai_grounding_metadata") == [ + {"webSearchQueries": ["test query"]} + ] + assert response.model_dump()["vertex_ai_grounding_metadata"] == [ + {"webSearchQueries": ["test query"]} + ] diff --git a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py new file mode 100644 index 00000000000..03790b220eb --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py @@ -0,0 +1,81 @@ +import os +import sys + +sys.path.insert(0, os.path.abspath("../../..")) + +import litellm +from litellm import LlmProviders +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.litellm_core_utils.get_llm_provider_logic import ( + _get_openai_compatible_provider_info, +) +from litellm.llms.xai.chat.transformation import XAIChatConfig +from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import ( + ProviderConfigManager, + get_optional_params, + validate_environment, +) + + +def test_xai_provider_config_routing(): + chat_config = ProviderConfigManager.get_provider_chat_config( + model="grok-3-mini", + provider=LlmProviders.XAI, + ) + responses_config = ProviderConfigManager.get_provider_responses_api_config( + model="grok-3-mini", + provider=LlmProviders.XAI, + ) + + assert isinstance(chat_config, XAIChatConfig) + assert isinstance(responses_config, XAIResponsesAPIConfig) + + +def test_xai_openai_compatible_provider_info(): + model, custom_llm_provider, dynamic_api_key, api_base = ( + _get_openai_compatible_provider_info( + model="xai/grok-3-mini", + api_base="https://api.x.ai/v1", + api_key="api-key", + dynamic_api_key=None, + ) + ) + + assert model == "grok-3-mini" + assert custom_llm_provider == "xai" + assert api_base == "https://api.x.ai/v1" + assert dynamic_api_key == "api-key" + + +def test_xai_get_model_info_uses_xai_pricing_metadata(): + model_info = litellm.get_model_info("xai/grok-3-mini") + + assert model_info["litellm_provider"] == "xai" + assert model_info["key"] == "xai/grok-3-mini" + assert model_info["mode"] == "chat" + + +def test_xai_validate_environment_reads_api_key(monkeypatch): + monkeypatch.setenv("XAI_API_KEY", "api-key") + + result = validate_environment(model="xai/grok-3-mini") + + assert result == {"keys_in_environment": True, "missing_keys": []} + + +def test_xai_oauth_flag_is_generic_litellm_param(): + litellm_params = GenericLiteLLMParams(use_xai_oauth=True) + runtime_params = get_litellm_params(use_xai_oauth=True) + result = get_optional_params( + model="grok-3-mini", + custom_llm_provider="xai", + temperature=0.2, + drop_params=True, + ) + + assert result["temperature"] == 0.2 + assert litellm_params.use_xai_oauth is True + assert runtime_params["use_xai_oauth"] is True + assert "use_xai_oauth" not in result diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 4c330312930..abb162e9ddb 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -5092,6 +5092,177 @@ def test_map_tool_helper_collision_prefers_definitions_over_components_schemas() assert transformed["input_schema"]["properties"]["from_components"] == expected +BILLING_HEADER_BLOCK = { + "type": "text", + "text": "x-anthropic-billing-header: cc_version=1.0.abc; cc_entrypoint=cli; cch=00000;", +} + + +def _system_with_billing_header(real_text: str) -> list: + return [ + { + "role": "system", + "content": [BILLING_HEADER_BLOCK, {"type": "text", "text": real_text}], + } + ] + + +def test_translate_system_message_keeps_billing_header_for_first_party_anthropic(): + config = AnthropicConfig() + assert config.should_strip_billing_metadata() is False + + result = config.translate_system_message( + messages=_system_with_billing_header( + "You are Claude Code, Anthropic's official CLI for Claude." + ) + ) + + texts = [block["text"] for block in result] + assert any(t.startswith("x-anthropic-billing-header:") for t in texts) + assert "You are Claude Code, Anthropic's official CLI for Claude." in texts + + +def test_translate_system_message_strips_billing_header_for_bedrock(): + from litellm.llms.bedrock.claude_platform.transformation import ( + BedrockClaudePlatformConfig, + ) + + config = BedrockClaudePlatformConfig() + assert config.should_strip_billing_metadata() is True + + result = config.translate_system_message( + messages=_system_with_billing_header("real system prompt") + ) + + texts = [block["text"] for block in result] + assert all(not t.startswith("x-anthropic-billing-header:") for t in texts) + assert "real system prompt" in texts + + +def test_anthropic_messages_request_keeps_billing_header_for_first_party(): + from litellm.types.router import GenericLiteLLMParams + + config = AnthropicMessagesConfig() + assert config.should_strip_billing_metadata() is False + + optional_params = { + "max_tokens": 16, + "system": [ + BILLING_HEADER_BLOCK, + {"type": "text", "text": "real system prompt"}, + ], + } + result = config.transform_anthropic_messages_request( + model="claude-3-5-sonnet-latest", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + texts = [block["text"] for block in result["system"]] + assert any(t.startswith("x-anthropic-billing-header:") for t in texts) + + +def test_anthropic_messages_request_strips_billing_header_for_minimax(): + from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig + from litellm.types.router import GenericLiteLLMParams + + config = MinimaxMessagesConfig() + assert config.should_strip_billing_metadata() is True + + optional_params = { + "max_tokens": 16, + "system": [ + BILLING_HEADER_BLOCK, + {"type": "text", "text": "real system prompt"}, + ], + } + result = config.transform_anthropic_messages_request( + model="MiniMax-M2", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + texts = [block["text"] for block in result.get("system", [])] + assert all(not t.startswith("x-anthropic-billing-header:") for t in texts) + + +def test_translate_system_message_strips_billing_header_for_bedrock_invoke(): + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeConfig, + ) + + config = AmazonAnthropicClaudeConfig() + assert config.should_strip_billing_metadata() is True + + result = config.translate_system_message( + messages=_system_with_billing_header("real system prompt") + ) + + texts = [block["text"] for block in result] + assert all(not t.startswith("x-anthropic-billing-header:") for t in texts) + assert "real system prompt" in texts + + +@pytest.mark.parametrize( + "module_path, class_name, expected_strip", + [ + ("litellm.llms.anthropic.chat.transformation", "AnthropicConfig", False), + ( + "litellm.llms.anthropic.experimental_pass_through.messages.transformation", + "AnthropicMessagesConfig", + False, + ), + ( + "litellm.llms.bedrock.claude_platform.transformation", + "BedrockClaudePlatformConfig", + True, + ), + ( + "litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation", + "AmazonAnthropicClaudeConfig", + True, + ), + ( + "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation", + "VertexAIAnthropicConfig", + True, + ), + ( + "litellm.llms.azure_ai.anthropic.transformation", + "AzureAnthropicConfig", + True, + ), + ("litellm.llms.minimax.messages.transformation", "MinimaxMessagesConfig", True), + ( + "litellm.llms.azure_ai.anthropic.messages_transformation", + "AzureAnthropicMessagesConfig", + True, + ), + ( + "litellm.llms.deepseek.messages.transformation", + "DeepSeekAnthropicMessagesConfig", + True, + ), + ( + "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation", + "VertexAIPartnerModelsAnthropicMessagesConfig", + True, + ), + ], +) +def test_should_strip_billing_metadata_by_provider( + module_path, class_name, expected_strip +): + import importlib + + config_cls = getattr(importlib.import_module(module_path), class_name) + assert config_cls().should_strip_billing_metadata() is expected_strip + + def test_namespace_tool_flat_nested_tools_are_extracted(): """Codex sends nested tools in flat format {type, name, description, parameters} with no 'function' wrapper. These must be normalized and mapped without raising KeyError: 'function'.""" @@ -5188,3 +5359,140 @@ def test_client_metadata_stripped_from_anthropic_request(): headers={}, ) assert "client_metadata" not in result + + +@pytest.mark.parametrize( + "model", + ["claude-fable-5", "claude-opus-4-7", "claude-opus-4-8-20260120"], +) +def test_sampling_params_dropped_for_models_that_removed_them(model): + """Fable 5 / Opus 4.7 / 4.8 reject temperature != 1 and any top_p with a + 400; with drop_params set they must be dropped, not forwarded (#30064).""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model=model, + drop_params=True, + ) + + assert "temperature" not in result + assert "top_p" not in result + + +@pytest.mark.parametrize("params", [{"temperature": 0.5}, {"top_p": 0.9}, {"top_p": 1}]) +def test_sampling_params_raise_clean_error_without_drop_params(params, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AnthropicConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.map_openai_params( + non_default_params=params, + optional_params={}, + model="claude-fable-5", + drop_params=False, + ) + + +def test_temperature_1_forwarded_on_models_that_removed_sampling_params(): + """temperature=1 (the API default) is still accepted and must pass through.""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 1}, + optional_params={}, + model="claude-fable-5", + drop_params=False, + ) + + assert result["temperature"] == 1 + + +@pytest.mark.parametrize("model", ["claude-opus-4-6", "claude-sonnet-4-6"]) +def test_sampling_params_forwarded_on_models_that_accept_them(model): + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model=model, + drop_params=True, + ) + + assert result["temperature"] == 0.5 + assert result["top_p"] == 0.9 + + +def test_sampling_param_gating_driven_by_model_map_flag(monkeypatch): + """The drop/raise decision must come from ``supports_sampling_params`` in + the model map, not just name matching: a flagged entry gates a model whose + name says nothing, and an explicit ``true`` overrides the name fallback.""" + monkeypatch.setitem( + litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False} + ) + monkeypatch.setitem( + litellm.model_cost, "claude-fable-5-test", {"supports_sampling_params": True} + ) + config = AnthropicConfig() + + flagged_off = config.map_openai_params( + non_default_params={"top_p": 0.9}, + optional_params={}, + model="claude-zeta-9", + drop_params=True, + ) + assert "top_p" not in flagged_off + + flagged_on = config.map_openai_params( + non_default_params={"top_p": 0.9}, + optional_params={}, + model="claude-fable-5-test", + drop_params=True, + ) + assert flagged_on["top_p"] == 0.9 + + +def test_top_k_dropped_at_transform_for_models_that_removed_it(): + """``top_k`` is a provider-specific kwarg that bypasses + ``map_openai_params``, so it must be stripped at the transform_request + boundary shared by the direct, invoke, Vertex, and Azure paths (#30064).""" + config = AnthropicConfig() + + result = config.transform_request( + model="claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"max_tokens": 10, "top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert "top_k" not in result + + +def test_top_k_raises_at_transform_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AnthropicConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.transform_request( + model="claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"max_tokens": 10, "top_k": 40}, + litellm_params={}, + headers={}, + ) + + +def test_top_k_forwarded_at_transform_on_models_that_accept_it(): + config = AnthropicConfig() + + result = config.transform_request( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "hello"}], + optional_params={"max_tokens": 10, "top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert result["top_k"] == 40 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py new file mode 100644 index 00000000000..f74c5b61300 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -0,0 +1,150 @@ +""" +Regression tests for fake-streamed providers routed through `/v1/messages`. + +A fake-streaming provider (e.g. Vertex AI Gemma `:predict`) collapses its whole +response into a single `MockResponseIterator` chunk that carries content text AND a +`finish_reason` together. `AnthropicStreamWrapper` previously dropped all content in +this case — `translate_streaming_openai_response_to_anthropic` sees the finish_reason +and emits only a `message_delta`. `_CombinedChunkSplitter` splits such chunks so the +content survives. +""" + +import asyncio +import json +from types import SimpleNamespace + +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, + _CombinedChunkSplitter, +) +from litellm.llms.base_llm.base_model_iterator import MockResponseIterator +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + Usage, +) + + +def _build_fake_stream( + content: str, finish_reason: str = "stop" +) -> MockResponseIterator: + """Mimic a Vertex Gemma `:predict` fake stream: one collapsed chunk.""" + model_response = ModelResponse() + model_response.choices = [ + Choices( + index=0, + message=Message(role="assistant", content=content), + finish_reason=finish_reason, + ) + ] + model_response.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + model_response.model = "gemma4" + return MockResponseIterator(model_response=model_response) + + +def _collect_async(wrapper: AnthropicStreamWrapper) -> str: + async def _run() -> str: + out = [] + async for raw in wrapper.async_anthropic_sse_wrapper(): + out.append(raw.decode() if isinstance(raw, bytes) else raw) + return "".join(out) + + return asyncio.run(_run()) + + +def test_fake_stream_content_reaches_anthropic_sse(): + """Content from a collapsed fake-stream chunk must be emitted as a delta.""" + wrapper = AnthropicStreamWrapper( + completion_stream=_build_fake_stream("Hello, the answer is 2."), + model="gemma4", + ) + sse = _collect_async(wrapper) + + assert "content_block_delta" in sse + assert "Hello, the answer is 2." in sse + assert "message_delta" in sse + assert "message_stop" in sse + + +def test_fake_stream_usage_preserved(): + """The finish chunk keeps usage so output_tokens is non-zero.""" + wrapper = AnthropicStreamWrapper( + completion_stream=_build_fake_stream("Two."), + model="gemma4", + ) + sse = _collect_async(wrapper) + + message_delta = next( + json.loads(line[len("data: ") :]) + for block in sse.split("\n\n") + for line in block.splitlines() + if line.startswith("data: ") and '"message_delta"' in line + ) + assert message_delta["usage"]["output_tokens"] == 5 + assert message_delta["usage"]["input_tokens"] == 10 + + +def test_splitter_passes_through_non_combined_chunks(): + """A chunk with content but no finish_reason is not split.""" + chunk = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, delta=Delta(content="partial"), finish_reason=None + ) + ] + ) + chunks = list(_CombinedChunkSplitter(iter([chunk]))) + assert len(chunks) == 1 + assert chunks[0].choices[0].delta.content == "partial" + + +def test_splitter_splits_combined_chunk_into_content_then_finish(): + """A chunk with both content and finish_reason becomes two chunks.""" + chunk = ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="done"), finish_reason="stop") + ] + ) + content_chunk, finish_chunk = list(_CombinedChunkSplitter(iter([chunk]))) + + assert content_chunk.choices[0].delta.content == "done" + assert content_chunk.choices[0].finish_reason is None + + assert finish_chunk.choices[0].finish_reason == "stop" + assert finish_chunk.choices[0].delta.content is None + + +def test_is_combined_false_when_choices_empty(): + """A metadata-only chunk with no choices is never treated as combined.""" + assert _CombinedChunkSplitter._is_combined(SimpleNamespace(choices=[])) is False + + +def test_is_combined_false_when_delta_missing(): + """A finish chunk whose choice has no delta is not combined.""" + chunk = SimpleNamespace(choices=[SimpleNamespace(finish_reason="stop", delta=None)]) + assert _CombinedChunkSplitter._is_combined(chunk) is False + + +def test_split_clears_reasoning_and_thinking_on_finish_chunk(): + """When the combined delta carries reasoning/thinking, only the content + chunk keeps them — the finish chunk is cleared.""" + delta = SimpleNamespace( + content="hi", + tool_calls=None, + reasoning_content="some reasoning", + thinking_blocks=[{"type": "thinking"}], + ) + chunk = SimpleNamespace( + choices=[SimpleNamespace(finish_reason="stop", delta=delta)] + ) + + content_chunk, finish_chunk = _CombinedChunkSplitter._split(chunk) + + assert content_chunk.choices[0].delta.reasoning_content == "some reasoning" + assert content_chunk.choices[0].delta.thinking_blocks == [{"type": "thinking"}] + assert finish_chunk.choices[0].delta.reasoning_content is None + assert finish_chunk.choices[0].delta.thinking_blocks is None diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py new file mode 100644 index 00000000000..450f69fb87c --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -0,0 +1,79 @@ +""" +Tests for AnthropicResponsesStreamWrapper +(litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py) +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../..")) +) + +from litellm.llms.anthropic.experimental_pass_through.responses_adapters.streaming_iterator import ( + AnthropicResponsesStreamWrapper, +) + + +def _process_all(events: list) -> list: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=None, model="m") + for event in events: + wrapper._process_event(event) + return list(wrapper._chunk_queue) + + +class TestProcessEventTextDeltaWithoutOutputItemAdded: + """Streams that skip response.output_item.added (e.g. LMStudio) must still + open a text block before any delta and never emit index -1.""" + + def test_process_event_synthesizes_content_block_start_before_delta(self): + chunks = _process_all( + [ + {"type": "response.output_text.delta", "item_id": "i1", "delta": "Hel"}, + {"type": "response.output_text.delta", "item_id": "i1", "delta": "lo"}, + ] + ) + assert [c["type"] for c in chunks] == [ + "content_block_start", + "content_block_delta", + "content_block_delta", + ] + assert chunks[0]["content_block"] == {"type": "text", "text": ""} + assert [c["index"] for c in chunks] == [0, 0, 0] + assert chunks[1]["delta"] == {"type": "text_delta", "text": "Hel"} + + def test_process_event_delta_without_item_id_never_yields_negative_index(self): + chunks = _process_all([{"type": "response.output_text.delta", "delta": "Hi"}]) + assert [(c["type"], c["index"]) for c in chunks] == [ + ("content_block_start", 0), + ("content_block_delta", 0), + ] + + def test_process_event_unregistered_item_id_opens_new_text_block(self): + chunks = _process_all( + [ + { + "type": "response.output_item.added", + "item": {"type": "reasoning", "id": "rs_1"}, + }, + {"type": "response.output_text.delta", "item_id": "m1", "delta": "Hi"}, + ] + ) + assert chunks[1]["type"] == "content_block_start" + assert chunks[1]["content_block"] == {"type": "text", "text": ""} + assert [c["index"] for c in chunks[1:]] == [1, 1] + + def test_process_event_registered_item_id_does_not_synthesize_start(self): + chunks = _process_all( + [ + { + "type": "response.output_item.added", + "item": {"type": "message", "id": "m1"}, + }, + {"type": "response.output_text.delta", "item_id": "m1", "delta": "Hi"}, + ] + ) + assert [(c["type"], c["index"]) for c in chunks] == [ + ("content_block_start", 0), + ("content_block_delta", 0), + ] diff --git a/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py new file mode 100644 index 00000000000..70fef0162e6 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py @@ -0,0 +1,94 @@ +""" +Tests that ``get_cost_for_anthropic_web_search`` tolerates ``server_tool_use`` +being either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.llms.anthropic.cost_calculation import ( + _get_web_search_requests, + get_cost_for_anthropic_web_search, +) +from litellm.types.utils import ModelInfo, ServerToolUse + + +class _UsageWithServerToolUse: + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + + +def _make_model_info(cost_per_query: float = 0.01) -> ModelInfo: + info: ModelInfo = { # type: ignore[typeddict-item] + "search_context_cost_per_query": { + "search_context_size_low": cost_per_query, + "search_context_size_medium": cost_per_query, + "search_context_size_high": cost_per_query, + } + } + return info + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 4}) == 4 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + assert _get_web_search_requests(ServerToolUse(web_search_requests=2)) == 2 + + +def test_get_cost_for_anthropic_web_search_with_dict_server_tool_use(): + """ + Regression: ``server_tool_use`` was a dict from ``stream_chunk_builder`` and + direct attribute access on it raised ``AttributeError``. + """ + usage = _UsageWithServerToolUse({"web_search_requests": 3}) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_pydantic_server_tool_use(): + usage = _UsageWithServerToolUse(ServerToolUse(web_search_requests=3)) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_none_server_tool_use(): + usage = _UsageWithServerToolUse(None) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == 0.0 + + +def test_get_cost_for_anthropic_web_search_with_no_usage(): + info = _make_model_info(cost_per_query=0.01) + cost = get_cost_for_anthropic_web_search(model_info=info, usage=None) + assert cost == 0.0 diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index f49b3f09d6d..a211a69b9c7 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -57,6 +57,22 @@ def test_azure_providers_image_generation_json_body_keeps_model(): assert out == data +def test_azure_image_generation_mai_base_model_uses_mai_url(): + azure_chat = AzureChatCompletion() + url = azure_chat.create_azure_base_url( + azure_client_params={ + "azure_endpoint": "https://my-resource.services.ai.azure.com", + "api_version": "preview", + }, + model="image-deployment-alias", + base_model="MAI-Image-2.5", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" + ) + + def test_azure_image_generation_flattens_extra_body(): """ Test that Azure image generation correctly flattens extra_body parameters. diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py new file mode 100644 index 00000000000..d5256be02d7 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py @@ -0,0 +1,171 @@ +import io +import os +import sys +from unittest.mock import MagicMock + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../../..")) + +from litellm.llms.azure_ai.image_edit import ( + AzureFoundryMAIImageEditConfig, + get_azure_ai_image_edit_config, +) +from litellm.llms.azure_ai.image_generation.mai_transformation import ( + AzureFoundryMAIImageGenerationConfig, +) + + +class TestAzureMAIImageEdit: + def test_get_mai_image_edit_url(self): + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url( + api_base="https://my-resource.services.ai.azure.com", + api_version="preview", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/edits?api-version=preview" + ) + + def test_get_mai_image_edit_url_rewrites_generation_url(self): + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url( + api_base=( + "https://my-resource.services.ai.azure.com/mai/v1/images/generations" + "?api-version=preview" + ), + api_version="preview", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/edits?api-version=preview" + ) + + def test_get_mai_image_edit_url_appends_edits_to_mai_root(self): + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url( + api_base="https://my-resource.services.ai.azure.com/mai/v1", + api_version="preview", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/edits?api-version=preview" + ) + + def test_get_azure_ai_image_edit_config_returns_mai(self): + config = get_azure_ai_image_edit_config("MAI-Image-2.5") + assert isinstance(config, AzureFoundryMAIImageEditConfig) + + def test_validate_environment_uses_api_key_header(self): + config = AzureFoundryMAIImageEditConfig() + headers: dict = {} + config.validate_environment(headers, "MAI-Image-2.5", api_key="test-key") + assert headers["api-key"] == "test-key" + assert "Api-Key" not in headers + + def test_get_complete_url(self): + config = AzureFoundryMAIImageEditConfig() + url = config.get_complete_url( + model="MAI-Image-2.5", + api_base="https://my-resource.services.ai.azure.com", + litellm_params={"api_version": "preview"}, + ) + assert "/mai/v1/images/edits" in url + assert "api-version=preview" in url + + def test_map_openai_params_keeps_size(self): + config = AzureFoundryMAIImageEditConfig() + optional_params = config.map_openai_params( + image_edit_optional_params={"size": "1792x1024", "n": 1}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["size"] == "1792x1024" + assert optional_params["n"] == 1 + assert "width" not in optional_params + assert "height" not in optional_params + + def test_map_openai_params_defaults_size(self): + config = AzureFoundryMAIImageEditConfig() + optional_params = config.map_openai_params( + image_edit_optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["size"] == "1024x1024" + + def test_map_openai_params_unsupported_size_raises(self): + config = AzureFoundryMAIImageEditConfig() + with pytest.raises(ValueError, match="Unsupported size value: 'auto'"): + config.map_openai_params( + image_edit_optional_params={"size": "auto"}, + model="MAI-Image-2.5", + drop_params=True, + ) + + def test_map_openai_params_invalid_size_format_raises(self): + config = AzureFoundryMAIImageEditConfig() + with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"): + config.map_openai_params( + image_edit_optional_params={"size": "1024xabc"}, + model="MAI-Image-2.5", + drop_params=True, + ) + + def test_transform_image_edit_request_uses_image_field(self): + config = AzureFoundryMAIImageEditConfig() + image_bytes = io.BytesIO(b"fake-image-bytes") + + data, files = config.transform_image_edit_request( + model="MAI-Image-2.5", + prompt="Turn this into a studio product shot", + image=image_bytes, + image_edit_optional_request_params={"size": "1024x1024", "n": 1}, + litellm_params={}, + headers={}, + ) + + assert data["model"] == "MAI-Image-2.5" + assert data["prompt"] == "Turn this into a studio product shot" + assert data["size"] == "1024x1024" + assert data["n"] == 1 + assert len(files) == 1 + assert files[0][0] == "image" + assert files[0][0] != "image[]" + + def test_normalize_mai_image_usage_maps_edit_response_fields(self): + usage = AzureFoundryMAIImageGenerationConfig.normalize_mai_image_usage( + { + "num_output_tokens": 1024, + "output_image_tokens": 1024, + } + ) + assert usage["output_tokens"] == 1024 + assert usage["input_tokens"] == 0 + assert usage["total_tokens"] == 1024 + assert usage["input_tokens_details"]["text_tokens"] == 0 + assert usage["input_tokens_details"]["image_tokens"] == 0 + + def test_transform_image_edit_response_parses_mai_usage(self): + config = AzureFoundryMAIImageEditConfig() + raw_response = MagicMock(spec=httpx.Response) + raw_response.status_code = 200 + raw_response.text = "" + raw_response.json.return_value = { + "created": 1780897477, + "data": [{"b64_json": "abc123"}], + "usage": { + "num_output_tokens": 1024, + "output_image_tokens": 1024, + }, + } + + logging_obj = MagicMock() + image_response = config.transform_image_edit_response( + model="MAI-Image-2.5", + raw_response=raw_response, + logging_obj=logging_obj, + ) + + assert image_response.data[0].b64_json == "abc123" + assert image_response.usage.output_tokens == 1024 + assert image_response.usage.total_tokens == 1024 diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py new file mode 100644 index 00000000000..f7ad333293c --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -0,0 +1,380 @@ +import os +import sys +from unittest.mock import MagicMock + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../../..")) + +import litellm +from litellm.llms.azure.azure import AzureChatCompletion +from litellm.llms.azure.image_generation import get_azure_image_generation_config +from litellm.llms.azure.image_generation.http_utils import ( + azure_deployment_image_generation_json_body, +) +from litellm.llms.azure_ai.image_generation import ( + AzureFoundryMAIImageGenerationConfig, + get_azure_ai_image_generation_config, +) +from litellm.llms.azure_ai.image_generation.cost_calculator import ( + cost_calculator as azure_ai_image_cost_calculator, +) +from litellm.types.utils import ( + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, +) +from litellm.utils import get_optional_params_image_gen + + +class TestAzureMAIImageGeneration: + def test_is_mai_model(self): + assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5") + assert AzureFoundryMAIImageGenerationConfig.is_mai_model( + "azure_ai/MAI-Image-2.5" + ) + assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5-Flash") + assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2e") + assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro") + assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-DS-R1") + + def test_mai_flash_and_2e_model_pricing_in_cost_map(self): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + flash_info = litellm.get_model_info( + model="azure_ai/MAI-Image-2.5-Flash", + custom_llm_provider="azure_ai", + ) + assert flash_info["input_cost_per_token"] == 1.75e-06 + assert flash_info["input_cost_per_image_token"] == 1.75e-06 + assert flash_info["output_cost_per_image_token"] == 3.3e-05 + + image_2e_info = litellm.get_model_info( + model="azure_ai/MAI-Image-2e", + custom_llm_provider="azure_ai", + ) + assert image_2e_info["input_cost_per_token"] == 5e-06 + assert image_2e_info["output_cost_per_image_token"] == 1.95e-05 + + def test_get_mai_image_generation_url(self): + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( + api_base="https://my-resource.services.ai.azure.com", + api_version="preview", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" + ) + + def test_get_mai_image_generation_url_preserves_full_path(self): + api = ( + "https://my-resource.services.ai.azure.com/mai/v1/images/generations" + "?api-version=preview" + ) + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( + api_base=api, + api_version="preview", + ) + assert url == api + + def test_get_mai_image_generation_url_appends_generations_to_mai_root(self): + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( + api_base="https://my-resource.services.ai.azure.com/mai/v1", + api_version="preview", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" + ) + + def test_get_azure_ai_image_generation_config_returns_mai(self): + config = get_azure_ai_image_generation_config("MAI-Image-2.5") + assert isinstance(config, AzureFoundryMAIImageGenerationConfig) + + def test_azure_image_generation_config_returns_mai(self): + config = get_azure_image_generation_config("MAI-Image-2.5") + assert isinstance(config, AzureFoundryMAIImageGenerationConfig) + + def test_map_openai_params_size_to_width_height(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"size": "1024x1024", "n": 1}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 1024 + assert optional_params["height"] == 1024 + assert optional_params["n"] == 1 + assert "size" not in optional_params + + def test_map_openai_params_defaults(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 1024 + assert optional_params["height"] == 1024 + + def test_get_optional_params_image_gen_mai(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = get_optional_params_image_gen( + model="MAI-Image-2.5", + size="1792x1024", + n=1, + custom_llm_provider="azure_ai", + provider_config=config, + drop_params=True, + ) + assert optional_params["width"] == 1792 + assert optional_params["height"] == 1024 + assert "size" not in optional_params + + def test_azure_create_azure_base_url_mai(self): + azure_chat = AzureChatCompletion() + url = azure_chat.create_azure_base_url( + azure_client_params={ + "azure_endpoint": "https://my-resource.services.ai.azure.com", + "api_version": "preview", + }, + model="MAI-Image-2.5", + ) + assert "/mai/v1/images/generations" in url + assert "api-version=preview" in url + + def test_mai_json_body_keeps_model(self): + api = ( + "https://my-resource.services.ai.azure.com/mai/v1/images/generations" + "?api-version=preview" + ) + data = { + "model": "MAI-Image-2.5", + "prompt": "A photograph of a red fox", + "width": 1024, + "height": 1024, + "n": 1, + } + out = azure_deployment_image_generation_json_body(api, data) + assert out == data + + def test_map_openai_params_custom_size(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"size": "768x768"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 768 + assert optional_params["height"] == 768 + + def test_map_openai_params_width_only_gets_height_default(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"width": 1792}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 1792 + assert optional_params["height"] == config.DEFAULT_HEIGHT + + def test_map_openai_params_height_only_gets_width_default(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"height": 1792}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == config.DEFAULT_WIDTH + assert optional_params["height"] == 1792 + + def test_map_openai_params_unsupported_size_raises(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="Unsupported size value: 'auto'"): + config.map_openai_params( + non_default_params={"size": "auto"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + def test_map_openai_params_invalid_custom_size_raises(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"): + config.map_openai_params( + non_default_params={"size": "1024xabc"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + def test_map_openai_params_unsupported_param_raises(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="Parameter quality is not supported"): + config.map_openai_params( + non_default_params={"quality": "hd"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + + def test_transform_image_generation_response_normalizes_mai_usage(self): + config = AzureFoundryMAIImageGenerationConfig() + raw_response = MagicMock(spec=httpx.Response) + raw_response.json.return_value = { + "created": 1780897477, + "data": [{"b64_json": "abc123"}], + "usage": { + "num_output_tokens": 1024, + "num_input_text_tokens": 22, + "output_image_tokens": 1024, + }, + } + + logging_obj = MagicMock() + image_response = config.transform_image_generation_response( + model="MAI-Image-2.5", + raw_response=raw_response, + model_response=ImageResponse(), + logging_obj=logging_obj, + request_data={"prompt": "A red fox"}, + optional_params={"width": 1024, "height": 1024}, + litellm_params={}, + encoding=None, + ) + + assert image_response.data[0].b64_json == "abc123" + assert image_response.usage.output_tokens == 1024 + assert image_response.usage.input_tokens == 22 + assert image_response.usage.total_tokens == 1046 + + def test_transform_image_generation_response_non_json_raises_openai_error(self): + from litellm.llms.openai.common_utils import OpenAIError + + config = AzureFoundryMAIImageGenerationConfig() + raw_response = MagicMock(spec=httpx.Response) + raw_response.json.side_effect = ValueError("not json") + raw_response.text = "upstream gateway error" + raw_response.status_code = 502 + + with pytest.raises(OpenAIError) as exc_info: + config.transform_image_generation_response( + model="MAI-Image-2.5", + raw_response=raw_response, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={"prompt": "A red fox"}, + optional_params={"width": 1024, "height": 1024}, + litellm_params={}, + encoding=None, + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.message == "upstream gateway error" + + def test_normalize_mai_usage_preserves_zero_output_tokens(self): + config = AzureFoundryMAIImageGenerationConfig() + normalized = config.normalize_mai_image_usage( + { + "num_output_tokens": 0, + "output_image_tokens": 1024, + "num_input_text_tokens": 22, + } + ) + assert normalized["output_tokens"] == 0 + assert normalized["input_tokens"] == 22 + assert normalized["total_tokens"] == 22 + + def test_azure_sync_image_generation_uses_mai_response_transform(self): + raw_response = MagicMock(spec=httpx.Response) + raw_response.json.return_value = { + "created": 1780897477, + "data": [{"b64_json": "abc123"}], + "usage": { + "num_output_tokens": 1024, + "num_input_text_tokens": 22, + }, + } + + class MAIImageGenerationAzureChatCompletion(AzureChatCompletion): + def make_sync_azure_httpx_request(self, **kwargs): + return raw_response + + logging_obj = MagicMock() + image_response = MAIImageGenerationAzureChatCompletion().image_generation( + prompt="A red fox", + timeout=60.0, + optional_params={"width": 1792, "height": 1024}, + logging_obj=logging_obj, + headers={}, + model="MAI-Image-2.5", + api_key="test-key", + api_base="https://my-resource.services.ai.azure.com", + api_version="preview", + litellm_params={}, + ) + + assert image_response.data[0].b64_json == "abc123" + assert image_response.usage.output_tokens == 1024 + assert image_response.usage.input_tokens == 22 + assert image_response.usage.total_tokens == 1046 + assert image_response.size == "1792x1024" + + def test_mai_image_cost_calculator_token_based(self): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "azure_ai/MAI-Image-2.5" + model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") + input_text_tokens = 100 + output_image_tokens = 1024 + + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=input_text_tokens, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=input_text_tokens, + image_tokens=0, + ), + output_tokens=output_image_tokens, + total_tokens=input_text_tokens + output_image_tokens, + ), + ) + + cost = azure_ai_image_cost_calculator( + model=model, + image_response=image_response, + ) + + expected_cost = ( + input_text_tokens * model_info["input_cost_per_token"] + + output_image_tokens * model_info["output_cost_per_image_token"] + ) + assert round(cost, 10) == round(expected_cost, 10) + + def test_mai_image_cost_calculator_falls_back_to_flat_image_pricing(self): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "azure_ai/MAI-Image-2.5" + model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") + image_response = ImageResponse( + data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")] + ) + + cost = azure_ai_image_cost_calculator( + model=model, + image_response=image_response, + ) + + assert ( + cost == len(image_response.data or []) * model_info["output_cost_per_image"] + ) + assert cost > 0 diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index a6aa35ee6d1..5c83f8b34f7 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -1467,11 +1467,10 @@ def test_transform_request_with_function_tool(): ) # Verify the structure - assert "additionalModelRequestFields" in request_data - additional_fields = request_data["additionalModelRequestFields"] + # Function tools are not computer use tools, so they don't get anthropic_beta — + # additionalModelRequestFields should be absent (not serialized as empty {}) + assert "additionalModelRequestFields" not in request_data - # Function tools are not computer use tools, so they don't get anthropic_beta - # They are processed through the regular tool config assert "toolConfig" in request_data assert "tools" in request_data["toolConfig"] assert len(request_data["toolConfig"]["tools"]) == 1 @@ -5268,3 +5267,122 @@ def text(self): msg = str(exc_info.value) assert "secret content" not in msg assert "Error converting to valid response block" in msg + + +def test_converse_drops_sampling_params_for_models_that_removed_them(): + """Fable 5 / Opus 4.7 / 4.8 reject temperature != 1 and any top_p; with + drop_params set, converse must drop them instead of forwarding (#30064).""" + config = AmazonConverseConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model="us.anthropic.claude-fable-5", + drop_params=True, + ) + + assert "temperature" not in result + assert "topP" not in result + + +def test_converse_sampling_params_raise_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AmazonConverseConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.map_openai_params( + non_default_params={"temperature": 0.5}, + optional_params={}, + model="global.anthropic.claude-opus-4-8-v1:0", + drop_params=False, + ) + + +def test_converse_sampling_params_forwarded_on_models_that_accept_them(): + config = AmazonConverseConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model="us.anthropic.claude-sonnet-4-6", + drop_params=True, + ) + + assert result["temperature"] == 0.5 + assert result["topP"] == 0.9 + + +def test_converse_top_k_dropped_for_models_that_removed_it(): + """``top_k`` reaches converse as a provider-specific kwarg destined for + ``additionalModelRequestFields``, bypassing ``map_openai_params``; the + transform must strip it for models that removed sampling params (#30064).""" + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert "top_k" not in result.get("additionalModelRequestFields", {}) + + +def test_converse_top_k_raises_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AmazonConverseConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.transform_request( + model="us.anthropic.claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 40}, + litellm_params={}, + headers={}, + ) + + +def test_converse_top_k_forwarded_on_models_that_accept_it(): + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-sonnet-4-6", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert result["additionalModelRequestFields"]["top_k"] == 40 + + +def test_converse_top_k_zero_raises_without_drop_params(monkeypatch): + """``top_k=0`` must hit the same gating as any other value; previously the + truthiness check let it silently disappear on models that removed sampling + params, diverging from the Anthropic boundary that treats ``0`` as present.""" + monkeypatch.setattr(litellm, "drop_params", False) + config = AmazonConverseConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.transform_request( + model="us.anthropic.claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 0}, + litellm_params={}, + headers={}, + ) + + +def test_converse_top_k_zero_forwarded_on_models_that_accept_it(): + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-sonnet-4-6", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 0}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert result["additionalModelRequestFields"]["top_k"] == 0 diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index 0de3f833a37..6812f40829a 100644 --- a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -1,10 +1,15 @@ +import base64 +import json import os import sys sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path -from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig +from litellm.llms.bedrock.count_tokens.transformation import ( + DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS, + BedrockCountTokensConfig, +) def test_detect_input_type(): @@ -20,6 +25,71 @@ def test_detect_input_type(): assert config._detect_input_type(request_with_text) == "invokeModel" +def test_detect_input_type_anthropic_blocks_route_to_invoke_model(): + """Anthropic-shape content blocks must not go through the Converse path, + which Bedrock rejects with a 400 (and the caller then silently falls back + to the local tokenizer).""" + config = BedrockCountTokensConfig() + + request = { + "messages": [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Reading the file."}, + { + "type": "tool_use", + "id": "toolu_01", + "name": "read_file", + "input": {"path": "main.py"}, + }, + ], + }, + ], + } + assert config._detect_input_type(request) == "invokeModel" + + +def test_detect_input_type_converse_blocks_route_to_converse(): + """Converse-shape blocks (no "type" key) keep using the converse input.""" + config = BedrockCountTokensConfig() + + request = {"messages": [{"role": "user", "content": [{"text": "hi"}]}]} + assert config._detect_input_type(request) == "converse" + + +def test_transform_to_invoke_model_format_base64_encodes_body(): + """The CountTokens API expects invokeModel.body as a base64-encoded blob; + Anthropic Messages bodies additionally need anthropic_version/max_tokens + to pass Bedrock's InvokeModel schema validation.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + body = json.loads(base64.b64decode(result["input"]["invokeModel"]["body"])) + assert body["messages"] == request["messages"] + assert "model" not in body + assert body["anthropic_version"] == "bedrock-2023-05-31" + assert body["max_tokens"] == DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS + + +def test_transform_to_invoke_model_format_raw_body_unchanged(): + """Non-messages bodies (e.g. Titan inputText) must not get Anthropic fields.""" + config = BedrockCountTokensConfig() + + result = config.transform_anthropic_to_bedrock_count_tokens( + {"model": "amazon.titan-text-express-v1", "inputText": "hello"} + ) + + body = json.loads(base64.b64decode(result["input"]["invokeModel"]["body"])) + assert body == {"inputText": "hello"} + + def test_transform_anthropic_to_bedrock_request(): """Test basic request transformation""" config = BedrockCountTokensConfig() diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index ba41fc47e8b..4731be13e78 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -128,6 +128,70 @@ def test_nova_text_only_uses_converse_format(self): # Must have messages assert "messages" in model_input + # Nova Pro rejects empty additionalModelRequestFields / system — they must be absent + assert ( + "additionalModelRequestFields" not in model_input + ), "Nova: empty additionalModelRequestFields must be omitted, not serialized as {}" + assert ( + "system" not in model_input + ), "Nova: empty system must be omitted, not serialized as []" + + def test_nova_batch_jsonl_omits_empty_converse_fields(self): + """ + Regression test: Amazon Nova Pro returns 400 Malformed input request when + additionalModelRequestFields or system are present but empty in the Converse + API payload. The proxy must strip these keys when they carry no data. + """ + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + + openai_jsonl_content = [ + { + "custom_id": "req-0", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "us.amazon.nova-pro-v1:0", + "messages": [ + { + "role": "user", + "content": "What is 1 + 1? Answer with just the number.", + } + ], + "max_tokens": 16, + }, + } + ] + + result = config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + openai_jsonl_content + ) + + assert len(result) == 1 + model_input = result[0]["modelInput"] + + assert ( + "additionalModelRequestFields" not in model_input + or model_input["additionalModelRequestFields"] + ), "additionalModelRequestFields must be absent or non-empty — Nova rejects {}" + assert ( + "system" not in model_input or model_input["system"] + ), "system must be absent or non-empty — Nova rejects []" + + # Validate the exact shape AWS accepts + assert model_input == { + "messages": [ + { + "role": "user", + "content": [ + {"text": "What is 1 + 1? Answer with just the number."} + ], + } + ], + "inferenceConfig": {"maxTokens": 16}, + } + def test_nova_image_content_uses_converse_image_blocks(self): """ Test that image_url content blocks are converted to Bedrock Converse diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index a00057eaa6b..bbefdd621f0 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -1,10 +1,17 @@ """ Unit tests for the Bedrock Mantle (Claude Mythos Preview) integration. -Tests cover route detection, URL construction, and config dispatch for both -the /chat/completions and /messages endpoints. +Tests cover route detection, URL construction, config dispatch for both +the /chat/completions and /messages endpoints, and project (workspace) +association via `aws_bedrock_project_id`. """ +import json +from unittest.mock import patch + +import httpx +import pytest + from litellm.llms.bedrock.common_utils import BedrockModelInfo, get_bedrock_chat_config from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig from litellm.llms.bedrock.messages.mantle_transformation import ( @@ -12,6 +19,32 @@ ) +def _anthropic_response(url: str) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-mythos-preview", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + request=httpx.Request("POST", url), + ) + + +def _capture_request(url: str, headers: dict, data) -> dict: + raw_body = data.decode("utf-8") if isinstance(data, bytes) else data or "{}" + return { + "path": httpx.URL(url).path, + "headers": headers, + "body": json.loads(raw_body), + } + + def test_get_bedrock_route_mantle(): assert ( BedrockModelInfo.get_bedrock_route("mantle/anthropic.claude-mythos-preview") @@ -103,3 +136,114 @@ def test_mantle_transform_request_strips_prefix_and_adds_model(): ) assert request["model"] == "anthropic.claude-mythos-preview" assert "mantle/" not in request["model"] + + +def test_mantle_validate_environment_sets_workspace_header(): + config = AmazonMantleConfig() + headers = config.validate_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": "proj_abc123def456"}, + ) + assert headers["anthropic-workspace"] == "proj_abc123def456" + + +def test_mantle_validate_environment_without_project_id(): + config = AmazonMantleConfig() + headers = config.validate_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": None}, + ) + assert "anthropic-workspace" not in headers + + +def test_mantle_messages_validate_environment_sets_workspace_header(): + config = AmazonMantleMessagesConfig() + headers, api_base = config.validate_anthropic_messages_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": "proj_abc123def456"}, + api_base="https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages", + ) + assert headers["anthropic-workspace"] == "proj_abc123def456" + assert api_base == "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages" + + +def test_mantle_messages_validate_environment_without_project_id(): + config = AmazonMantleMessagesConfig() + headers, _ = config.validate_anthropic_messages_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + ) + assert "anthropic-workspace" not in headers + + +def test_mantle_completion_sends_workspace_header_and_clean_body(): + import litellm + + requests = [] + + def mock_post(self, url, data=None, headers=None, **kwargs): + requests.append(_capture_request(url=url, headers=headers or {}, data=data)) + return _anthropic_response(url) + + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post): + response = litellm.completion( + model="bedrock/mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + aws_bedrock_project_id="proj_abc123def456", + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + ) + + assert response.choices[0].message.content == "ok" + assert len(requests) == 1 + assert requests[0]["path"] == "/anthropic/v1/messages" + assert requests[0]["headers"]["anthropic-workspace"] == "proj_abc123def456" + assert "aws_bedrock_project_id" not in requests[0]["body"] + + +@pytest.mark.asyncio +async def test_mantle_anthropic_messages_sends_workspace_header_and_clean_body(): + import litellm + + requests = [] + + async def mock_post(self, url, data=None, headers=None, **kwargs): + requests.append(_capture_request(url=url, headers=headers or {}, data=data)) + return _anthropic_response(url) + + try: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=mock_post, + ): + response = await litellm.anthropic_messages( + model="bedrock/mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + aws_bedrock_project_id="proj_abc123def456", + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + ) + finally: + await litellm.close_litellm_async_clients() + + assert response["content"][0]["text"] == "ok" + assert len(requests) == 1 + assert requests[0]["path"] == "/anthropic/v1/messages" + assert requests[0]["headers"]["anthropic-workspace"] == "proj_abc123def456" + assert "aws_bedrock_project_id" not in requests[0]["body"] diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index e2133d56f89..c3de29bd9d5 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -1,17 +1,24 @@ """ Unit tests for Amazon Bedrock Mantle Responses API configuration. -Mantle's gpt-5.5 / gpt-5.4 are served ONLY on the non-standard -`/openai/v1/responses` path. These tests lock the URL construction and -Bearer auth that make that routing work. +Mantle serves Responses on two paths: gpt frontier models on +`/openai/v1/responses` and other Responses-capable models (e.g. gpt-oss) on the +standard `/v1/responses`. These tests lock the per-model path selection in the +gate, the URL construction for both paths, and the shared Bearer auth. """ +import copy import os import sys sys.path.insert(0, os.path.abspath("../../../../..")) import pytest +from botocore.exceptions import ( + ConnectTimeoutError, + PartialCredentialsError, + ProfileNotFound, +) import litellm from litellm.llms.bedrock_mantle.responses.transformation import ( @@ -84,6 +91,42 @@ def test_url_region_default_us_east_1(self, monkeypatch): url = cfg.get_complete_url(api_base=None, litellm_params={}) assert url == "https://bedrock-mantle.us-east-1.api.aws/openai/v1/responses" + def test_standard_path_uses_region_from_env(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-east-2") + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + url = cfg.get_complete_url(api_base=None, litellm_params={}) + assert url == "https://bedrock-mantle.us-east-2.api.aws/v1/responses" + assert "/openai/v1/responses" not in url + + def test_standard_path_normalizes_v1_base(self, monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + url = cfg.get_complete_url( + api_base="https://bedrock-mantle.us-east-2.api.aws/v1", + litellm_params={}, + ) + assert url == "https://bedrock-mantle.us-east-2.api.aws/v1/responses" + assert url.count("/responses") == 1 + assert "/v1/v1/responses" not in url + + def test_standard_path_full_endpoint_base_not_doubled(self, monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + url = cfg.get_complete_url( + api_base="https://bedrock-mantle.us-east-2.api.aws/v1/responses", + litellm_params={}, + ) + assert url == "https://bedrock-mantle.us-east-2.api.aws/v1/responses" + assert url.count("/responses") == 1 + + def test_default_construction_keeps_openai_path(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-east-2") + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url(api_base=None, litellm_params={}) + assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + class TestBedrockMantleResponsesAuth: def test_config_api_key_takes_priority(self, monkeypatch): @@ -114,16 +157,35 @@ def test_bedrock_bearer_token_fallback(self, monkeypatch): ) assert headers["Authorization"] == "Bearer bearer-key" - def test_missing_key_raises(self, monkeypatch): + def test_missing_bearer_does_not_raise_in_validate_environment(self, monkeypatch): + # SigV4 may still apply, so validate_environment must defer instead of raising. monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) cfg = BedrockMantleResponsesAPIConfig() - with pytest.raises(ValueError, match="Bedrock Mantle API key"): - cfg.validate_environment( - headers={}, - model="openai.gpt-5.5", - litellm_params=GenericLiteLLMParams(), - ) + headers = cfg.validate_environment( + headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams() + ) + assert "Authorization" not in headers + + def test_project_id_sets_openai_project_header(self): + cfg = BedrockMantleResponsesAPIConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-5.5", + litellm_params=GenericLiteLLMParams( + api_key="fake-key", aws_bedrock_project_id="proj_abc123def456" + ), + ) + assert headers["OpenAI-Project"] == "proj_abc123def456" + + def test_no_project_id_no_openai_project_header(self): + cfg = BedrockMantleResponsesAPIConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-5.5", + litellm_params=GenericLiteLLMParams(api_key="fake-key"), + ) + assert "OpenAI-Project" not in headers def test_custom_llm_provider(self): cfg = BedrockMantleResponsesAPIConfig() @@ -154,6 +216,36 @@ def test_file_search_routes_to_emulation(self): is True ) + def test_standard_path_still_uses_bearer_auth(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-key") + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-oss-120b", + litellm_params=GenericLiteLLMParams(), + ) + assert headers["Authorization"] == "Bearer env-key" + + def test_standard_path_opts_out_of_native_features(self): + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + assert cfg.supports_native_file_search() is False + assert cfg.supports_native_websocket() is False + + +class TestBedrockMantleResponsesRequestBody: + def test_standard_path_outbound_body_carries_bare_model(self): + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + body = cfg.transform_responses_api_request( + model="openai.gpt-oss-120b", + input="hello", + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert body["model"] == "openai.gpt-oss-120b" + assert "input" in body + class TestBedrockMantleResponsesRegistry: def test_registry_returns_config_for_gpt_5_5(self): @@ -164,6 +256,7 @@ def test_registry_returns_config_for_gpt_5_5(self): model="openai.gpt-5.5", ) assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is True def test_registry_returns_config_for_gpt_5_4_enum(self): from litellm.utils import ProviderConfigManager @@ -173,6 +266,7 @@ def test_registry_returns_config_for_gpt_5_4_enum(self): model="openai.gpt-5.4", ) assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is True def test_registry_returns_none_for_gpt_oss(self): # Regression guard: gpt-oss must NOT get the native Responses config; it @@ -195,9 +289,10 @@ def test_registry_returns_none_for_gpt_oss_safeguard(self): assert cfg is None def test_registry_returns_config_for_future_frontier_model(self): - # Forward-compatibility: an unseen OpenAI gpt frontier model (e.g. gpt-6) must - # get the native Responses config without a code change. The gate allow-lists - # the openai.gpt- family (minus gpt-oss), so gpt-6 matches automatically. + # Forward-compatibility: an unseen OpenAI gpt frontier model (e.g. gpt-6), + # not yet in the price map, must get the openai-path Responses config with + # no code or JSON change. The name-convention fallback (openai.gpt- minus + # gpt-oss) catches it before any price-map entry exists. from litellm.utils import ProviderConfigManager cfg = ProviderConfigManager.get_provider_responses_api_config( @@ -205,6 +300,48 @@ def test_registry_returns_config_for_future_frontier_model(self): model="openai.gpt-6", ) assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is True + + def test_price_map_flag_routes_non_gpt_name_to_openai_path( + self, restore_model_cost + ): + # Data-driven onboarding: a frontier model whose name does NOT match the + # openai.gpt- convention can still be routed to /openai/v1/responses by + # declaring use_openai_responses_path in its price-map entry, with no code + # change. The string fallback alone could never catch this name. + from litellm.utils import ProviderConfigManager, register_model + + register_model( + { + "bedrock_mantle/somelab.frontier-x": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + "use_openai_responses_path": True, + } + } + ) + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="somelab.frontier-x", + ) + assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is True + + def test_gpt_5_5_price_map_declares_openai_responses_path(self, local_cost_map): + # The gpt-5.x entries must carry the data-driven flag so frontier routing + # does not rely on the name-string fallback alone. + assert ( + litellm.model_cost["bedrock_mantle/openai.gpt-5.5"].get( + "use_openai_responses_path" + ) + is True + ) + assert ( + litellm.model_cost["bedrock_mantle/openai.gpt-5.4"].get( + "use_openai_responses_path" + ) + is True + ) @pytest.mark.parametrize( "model", @@ -239,6 +376,129 @@ def test_registry_returns_none_when_model_is_none(self): ) assert cfg is None + def test_declared_responses_non_openai_routes_to_standard_path( + self, restore_model_cost + ): + # New feature: a non-OpenAI model declared mode=responses (e.g. via a + # user's proxy model_info block) must route to the STANDARD /v1/responses + # path, not the frontier /openai/v1/responses path. Fails before the + # path-aware gate exists (old gate returned None for non-gpt models). + from litellm.utils import ProviderConfigManager, register_model + + register_model( + { + "bedrock_mantle/somelab.future-model": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + } + } + ) + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="somelab.future-model", + ) + assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is False + + def test_gpt_oss_opt_in_routes_to_standard_path(self, restore_model_cost): + # When a user opts gpt-oss into native Responses via model_info mode, + # it must take the STANDARD /v1/responses path (gpt-oss Responses is on + # /v1/responses, NOT the frontier /openai/v1/responses path). + from litellm.utils import ProviderConfigManager, register_model + + register_model( + { + "bedrock_mantle/openai.gpt-oss-120b": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + } + } + ) + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="openai.gpt-oss-120b", + ) + assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is False + + def test_unmapped_model_degrades_to_none_without_crashing(self, restore_model_cost): + # A non-frontier model that is not in model_cost makes get_model_info + # raise; the gate must swallow it and return None rather than crash. + from litellm.utils import ProviderConfigManager + + litellm.model_cost.pop("bedrock_mantle/somelab.unmapped-model", None) + litellm.get_model_info.cache_clear() + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="somelab.unmapped-model", + ) + assert cfg is None + + def test_register_model_restore_undoes_existing_key_overwrite(self): + # Self-contained guard for the deepcopy requirement of restore_model_cost. + # register_model overwrites an existing key by mutating its nested dict in + # place, so the snapshot must be a deepcopy: a shallow dict() copy would + # share that nested dict and leave mode=responses after restore, making + # the final assertion fail. The in-place clear+update mirrors the fixture. + from litellm.utils import ProviderConfigManager, register_model + + snapshot = copy.deepcopy(litellm.model_cost) + litellm.get_model_info.cache_clear() + try: + register_model( + { + "bedrock_mantle/openai.gpt-oss-120b": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + } + } + ) + during = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", model="openai.gpt-oss-120b" + ) + assert isinstance(during, BedrockMantleResponsesAPIConfig) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(snapshot) + litellm.get_model_info.cache_clear() + after = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", model="openai.gpt-oss-120b" + ) + assert after is None + + +@pytest.fixture +def restore_model_cost(): + """Snapshot litellm.model_cost so register_model edits don't leak across tests. + + register_model mutates the global litellm.model_cost, and get_model_info is + lru_cached, so without restore + cache_clear a registered model would bleed + into sibling tests in the same process. + + Two subtleties make this fixture non-obvious: + + 1. The snapshot must be a deepcopy. register_model overwrites an existing key + via `litellm.model_cost.setdefault(key, {}).update(...)`, mutating the + nested dict in place; a shallow copy would share those nested dicts and + could not capture the pre-mutation values of an existing entry. + 2. The restore must be in place (clear + update the SAME dict object), not a + reassignment. The conftest autouse `isolate_litellm_state` fixture + snapshots `litellm.model_cost` by reference and restores that reference on + its teardown, which runs after this one. Reassigning `litellm.model_cost` + to a fresh dict here is undone when conftest reinstalls its (in-place + mutated) reference, so the registered mode would leak and poison + TestBedrockMantleResponsesPricing. Mutating the original object in place + restores the contents conftest's reference points at. + """ + original_model_cost = copy.deepcopy(litellm.model_cost) + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost.clear() + litellm.model_cost.update(original_model_cost) + litellm.get_model_info.cache_clear() + @pytest.fixture def local_cost_map(monkeypatch): @@ -261,6 +521,386 @@ def local_cost_map(monkeypatch): litellm.get_model_info.cache_clear() +class TestBedrockMantleResponsesSigV4: + def test_bearer_short_circuits_without_credentials(self, monkeypatch): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, signed_body = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key="bearer-from-config", + ) + assert headers["Authorization"] == "Bearer bearer-from-config" + assert signed_body == b'{"input": "hi"}' + signer.get_credentials.assert_not_called() + + def test_bearer_resolved_from_mantle_env_key(self, monkeypatch): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer") + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"] == "Bearer env-bearer" + + def test_bearer_arg_takes_priority_over_mantle_env_key(self, monkeypatch): + # The passed api_key (e.g. litellm_params.api_key) must win over the env + # bearer; a reordered precedence chain would silently use the wrong token. + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer") + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key="arg-bearer", + ) + assert headers["Authorization"] == "Bearer arg-bearer" + signer.get_credentials.assert_not_called() + + def test_access_key_produces_sigv4_headers(self, monkeypatch): + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, signed_body = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_session_token": "session-token-test", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "Credential=AKIAEXAMPLE/" in headers["Authorization"] + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + assert "X-Amz-Date" in headers + assert headers["X-Amz-Security-Token"] == "session-token-test" + assert signed_body == b'{"input": "hi"}' + + def test_assume_role_path_produces_sigv4_headers(self, monkeypatch): + from unittest.mock import MagicMock + from botocore.credentials import Credentials + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + return_value=Credentials( + access_key="ASIAEXAMPLE", + secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk", + token="assumed-session-token", + ) + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={ + "aws_role_name": "arn:aws:iam::000000000000:role/test-role", + "aws_session_name": "litellm-test", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + signer.get_credentials.assert_called_once() + call = signer.get_credentials.call_args.kwargs + assert call["aws_role_name"] == "arn:aws:iam::000000000000:role/test-role" + assert call["aws_session_name"] == "litellm-test" + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + + def test_signed_body_matches_final_data_after_normalize(self, monkeypatch): + """Core regression: the signed bytes must equal the bytes actually sent. + + Sign the *final* data dict and assert the returned signed_body decodes to + exactly that dict, so a later change to the data would break the SigV4 hash. + """ + import json + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + final_data = {"model": "openai.gpt-5.5", "input": "hi", "max_output_tokens": 16} + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + _, signed_body = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "us-east-2", + }, + request_data=final_data, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert signed_body is not None + assert json.loads(signed_body) == final_data + + def test_region_comes_from_optional_params(self, monkeypatch): + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, _ = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "eu-west-1", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.eu-west-1.api.aws/openai/v1/responses", + api_key=None, + ) + assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"] + + def test_url_region_and_sigv4_region_agree_from_litellm_params(self, monkeypatch): + """Adversarial-review regression: a caller-supplied aws_region_name (no region + env set) must shape BOTH the URL host and the SigV4 credential scope, or the + request is signed for one region and sent to another -> 401. + """ + monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + params = { + "aws_region_name": "ap-southeast-2", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + } + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + url = cfg.get_complete_url(api_base=None, litellm_params=params) + assert ( + url == "https://bedrock-mantle.ap-southeast-2.api.aws/openai/v1/responses" + ) + + headers, _ = cfg.sign_request( + headers={}, + optional_params=params, + request_data={"input": "hi"}, + api_base=url, + api_key=None, + ) + assert "/ap-southeast-2/bedrock/aws4_request" in headers["Authorization"] + + def test_injected_default_region_base_does_not_override_aws_region_name( + self, monkeypatch + ): + """2nd-round adversarial regression: responses/main.py auto-injects + litellm_params.api_base = https://bedrock-mantle..api.aws/v1 (default + region, ignoring aws_region_name). The config must still pin BOTH the URL host + and the SigV4 scope to aws_region_name, or the IAM deployment 401s. A naive + 'resolve region only when api_base is None' fix would fail this test. + """ + monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + injected_base = "https://bedrock-mantle.us-east-1.api.aws/v1" # default region + params = { + "aws_region_name": "us-east-2", # what the caller actually wants + "api_base": injected_base, + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + } + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + url = cfg.get_complete_url(api_base=injected_base, litellm_params=params) + assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + + headers, _ = cfg.sign_request( + headers={}, + optional_params=params, + request_data={"input": "hi"}, + api_base=url, + api_key=None, + ) + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + assert "us-east-1" not in headers["Authorization"] + + def test_custom_proxy_host_is_preserved(self, monkeypatch): + """A genuinely custom (non-Mantle) api_base host must be preserved, not rewritten + to a bedrock-mantle host. Only standard Mantle hosts are region-pinned. + """ + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url( + api_base="https://mantle-proxy.internal.example/openai/v1", + litellm_params={"aws_region_name": "us-east-2"}, + ) + assert url == "https://mantle-proxy.internal.example/openai/v1/responses" + + def test_caller_authorization_does_not_override_sigv4(self, monkeypatch): + """Adversarial-review regression: a caller-supplied Authorization header (e.g. + from extra_headers, surviving the relaxed validate_environment) must not clobber + the SigV4 Authorization that _sign_request would otherwise restore. + """ + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, _ = cfg.sign_request( + headers={"Authorization": "Bearer stale-caller-token"}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "Bearer stale-caller-token" not in headers["Authorization"] + + def test_no_bearer_and_no_credentials_raises_both_paths(self, monkeypatch): + from unittest.mock import MagicMock + from botocore.exceptions import NoCredentialsError + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock(side_effect=NoCredentialsError()) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ValueError) as exc: + cfg.sign_request( + headers={}, + optional_params={"aws_region_name": "us-east-2"}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + msg = str(exc.value) + assert "Bearer" in msg + assert "SigV4" in msg or "IAM" in msg + + @pytest.mark.parametrize( + "cred_error", + [ + PartialCredentialsError(provider="env", cred_var="aws_secret_access_key"), + ProfileNotFound(profile="missing-profile"), + ], + ) + def test_partial_credentials_raises_both_paths(self, monkeypatch, cred_error): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock(side_effect=cred_error) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ValueError) as exc: + cfg.sign_request( + headers={}, + optional_params={"aws_region_name": "us-east-2"}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + msg = str(exc.value) + assert "Bearer" in msg + assert "SigV4" in msg or "IAM" in msg + + def test_sts_transport_error_is_not_masked_as_credentials(self, monkeypatch): + # An AssumeRole / web-identity flow hits STS over the network, so a transient + # connection error must surface as itself, not be rewritten into the + # "no usable AWS credentials" message that would send the user to fix the + # wrong thing. + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=ConnectTimeoutError( + endpoint_url="https://sts.us-east-2.amazonaws.com" + ) + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ConnectTimeoutError): + cfg.sign_request( + headers={}, + optional_params={ + "aws_role_name": "arn:aws:iam::000000000000:role/test-role", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + + class TestBedrockMantleResponsesPricing: def test_gpt_5_5_pricing_and_mode(self, local_cost_map): info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5") diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 1725aa85d10..deaa0537930 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -5,11 +5,14 @@ API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.html """ +import json import os import sys +from unittest.mock import patch sys.path.insert(0, os.path.abspath("../../../../..")) +import httpx import pytest import litellm @@ -96,6 +99,79 @@ def test_get_supported_openai_params(self): assert "max_tokens" in params +class TestBedrockMantleProjectHeader: + def test_validate_environment_sets_openai_project_header(self): + cfg = BedrockMantleChatConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-oss-120b", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": "proj_abc123def456"}, + api_key="fake-key", + ) + assert headers["OpenAI-Project"] == "proj_abc123def456" + assert headers["Authorization"] == "Bearer fake-key" + + def test_validate_environment_without_project_id(self): + cfg = BedrockMantleChatConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-oss-120b", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + api_key="fake-key", + ) + assert "OpenAI-Project" not in headers + + def test_completion_sends_openai_project_header_and_clean_body(self): + requests = [] + + def mock_post(self, url, data=None, headers=None, **kwargs): + raw_body = data.decode("utf-8") if isinstance(data, bytes) else data + requests.append( + {"headers": headers or {}, "body": json.loads(raw_body or "{}")} + ) + return httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1733529600, + "model": "openai.gpt-oss-120b", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + request=httpx.Request("POST", url), + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post + ): + response = litellm.completion( + model="bedrock_mantle/openai.gpt-oss-120b", + messages=[{"role": "user", "content": "hello"}], + api_key="fake-key", + aws_bedrock_project_id="proj_abc123def456", + ) + + assert response.choices[0].message.content == "ok" + assert len(requests) == 1 + assert requests[0]["headers"]["OpenAI-Project"] == "proj_abc123def456" + assert "aws_bedrock_project_id" not in requests[0]["body"] + + class TestBedrockMantleProviderResolution: def test_get_llm_provider_resolves_correctly(self): model, provider, _, _ = litellm.get_llm_provider( diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 279e9730e69..7321abcee46 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -742,3 +742,241 @@ def _mutate(e, request_data): assert first_sent == prebuilt # attempt 0 used prebuilt assert second_sent == _json.dumps(request_body) # attempt 1 re-serialized assert "MUTATED" in second_sent # ... the mutated body + + +def test_base_responses_config_sign_request_is_noop_by_default(): + """Default responses sign_request must be a no-op: unchanged headers, no signed body. + + Guards the 15 existing responses providers from accidental signing when the + handler starts calling sign_request. + """ + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + cfg = OpenAIResponsesAPIConfig() + headers = {"Authorization": "Bearer sk-existing"} + out_headers, signed_body = cfg.sign_request( + headers=headers, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://api.openai.com/v1/responses", + ) + assert out_headers == {"Authorization": "Bearer sk-existing"} + assert signed_body is None + + +def _make_responses_handler_call(signed_body): + """Drive BaseLLMHTTPHandler.response_api_handler with a fully mocked provider + config + sync client, returning the kwargs the client.post was called with. + + signed_body=None simulates a no-op (non-signing) provider; bytes simulates a + signing provider (e.g. Bedrock Mantle). + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.router import GenericLiteLLMParams + + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_responses_api_request.return_value = {"input": "hi"} + provider_config.should_fake_stream.return_value = False + provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body) + + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + handler.response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=False, + ) + return mock_client.post.call_args.kwargs + + +def test_responses_handler_sends_json_when_not_signed(): + """No-op provider (signed_body is None) -> handler posts json=data, no data= bytes.""" + kwargs = _make_responses_handler_call(signed_body=None) + assert kwargs.get("json") == {"input": "hi"} + assert "data" not in kwargs + + +def test_responses_handler_sends_signed_bytes_when_signed(): + """Signing provider -> handler posts the exact signed bytes via data=, not json=.""" + kwargs = _make_responses_handler_call(signed_body=b'{"input": "hi"}') + assert kwargs.get("data") == b'{"input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + + +def test_responses_handler_signs_after_fake_stream_prep_strips_stream(): + """Fake-stream signing-order invariant: the bytes SIGNED must equal the bytes SENT. + + In the streaming + fake-stream path the handler first runs + _prepare_fake_stream_request, which pops "stream" out of the body, and only + then calls sign_request. If signing ran before that pop, the signed body + would still carry "stream" while the body sent over the wire would not, + producing a SigV4 payload-hash mismatch (401) for a real Mantle deployment. + We snapshot request_data at sign time and assert "stream" is already gone. + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.router import GenericLiteLLMParams + + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_responses_api_request.return_value = { + "input": "hi", + "stream": True, + } + provider_config.should_fake_stream.return_value = True + provider_config.transform_response_api_response.return_value = ResponsesAPIResponse( + id="resp_1", + created_at=0, + output=[], + status="completed", + model="openai.gpt-5.5", + ) + + captured = {} + + def _capture_sign(**kwargs): + captured["request_data"] = dict(kwargs["request_data"]) + return ({"X-Signed": "1"}, b'{"input": "hi"}') + + provider_config.sign_request.side_effect = _capture_sign + + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + handler.response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={"stream": True}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=False, + fake_stream=True, + ) + + assert "stream" not in captured["request_data"] + assert "input" in captured["request_data"] + + post_kwargs = mock_client.post.call_args.kwargs + assert post_kwargs.get("data") == b'{"input": "hi"}' + assert "json" not in post_kwargs + assert "stream" in post_kwargs + + +def _make_compact_handler_call(signed_body, is_async): + """Drive (async_)compact_response_api_handler with a fully mocked provider config + + client, returning the kwargs the client.post was called with. + + signed_body=None simulates a no-op (non-signing) provider; bytes simulates a + signing provider (e.g. Bedrock Mantle SigV4 / bearer). + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.router import GenericLiteLLMParams + + compact_url = "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses/compact" + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_compact_response_api_request.return_value = ( + compact_url, + {"model": "openai.gpt-5.5", "input": "hi"}, + ) + provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body) + provider_config.transform_compact_response_api_response.return_value = "ok" + + spec = AsyncHTTPHandler if is_async else HTTPHandler + mock_client = MagicMock(spec=spec) + if is_async: + mock_client.post = AsyncMock(return_value=MagicMock()) + else: + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + result = handler.compact_response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=is_async, + ) + if is_async: + asyncio.run(result) + return provider_config, mock_client.post.call_args.kwargs + + +def test_compact_handler_sends_json_when_not_signed(): + """No-op provider on compact (signed_body is None) -> posts json=data, no data= bytes.""" + provider_config, kwargs = _make_compact_handler_call( + signed_body=None, is_async=False + ) + provider_config.sign_request.assert_called_once() + assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"} + assert "data" not in kwargs + + +def test_compact_handler_sends_signed_bytes_when_signed(): + """Signing provider on compact -> posts the signed bytes via data=, not json=. + + Regression for the adversarial-review finding that /responses/compact bypassed + the SigV4 signing hook, so IAM-only Mantle callers sent unsigned bodies. + """ + provider_config, kwargs = _make_compact_handler_call( + signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=False + ) + assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + # signing must use the compact endpoint as api_base, not the create URL + assert provider_config.sign_request.call_args.kwargs["api_base"].endswith( + "/openai/v1/responses/compact" + ) + + +def test_async_compact_handler_sends_signed_bytes_when_signed(): + """Async compact must sign identically to sync (same omission in the async twin).""" + provider_config, kwargs = _make_compact_handler_call( + signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=True + ) + assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + + +def test_async_compact_handler_sends_json_when_not_signed(): + """Async no-op provider on compact -> posts json=data, no data= bytes.""" + _provider_config, kwargs = _make_compact_handler_call( + signed_body=None, is_async=True + ) + assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"} + assert "data" not in kwargs diff --git a/tests/test_litellm/llms/databricks/test_databricks_streaming_utils.py b/tests/test_litellm/llms/databricks/test_databricks_streaming_utils.py new file mode 100644 index 00000000000..5612864a841 --- /dev/null +++ b/tests/test_litellm/llms/databricks/test_databricks_streaming_utils.py @@ -0,0 +1,64 @@ +""" +Regression test for the databricks streaming chunk parser. + +OpenAI-compatible servers (e.g. Vertex AI Model Garden vLLM endpoints) send a final +usage-only chunk with an empty `choices` list when `stream_options.include_usage` is +set. `chunk_parser` previously did `choices[0]` unconditionally, raising +`IndexError` -> `MidStreamFallbackError` and crashing the stream. +""" + +from litellm.llms.databricks.streaming_utils import ModelResponseIterator + + +def test_chunk_parser_handles_empty_choices_usage_chunk(): + """A usage-only final chunk (empty choices) must not raise IndexError.""" + iterator = ModelResponseIterator(streaming_response=None, sync_stream=True) + usage_only_chunk = { + "id": "chatcmpl-x", + "object": "chat.completion.chunk", + "created": 1, + "model": "m", + "choices": [], + "usage": {"prompt_tokens": 20, "completion_tokens": 8, "total_tokens": 28}, + } + + result = iterator.chunk_parser(chunk=usage_only_chunk) + + assert result["text"] == "" + assert result["is_finished"] is False + assert result["usage"] is not None + assert result["usage"]["prompt_tokens"] == 20 + assert result["usage"]["completion_tokens"] == 8 + + +def test_chunk_parser_empty_choices_without_usage(): + """An empty-choices chunk with no usage block returns usage=None, no error.""" + iterator = ModelResponseIterator(streaming_response=None, sync_stream=True) + chunk = { + "id": "chatcmpl-x", + "object": "chat.completion.chunk", + "created": 1, + "model": "m", + "choices": [], + } + + result = iterator.chunk_parser(chunk=chunk) + + assert result["text"] == "" + assert result["usage"] is None + + +def test_chunk_parser_normal_content_chunk_still_works(): + """A regular content chunk is unaffected by the empty-choices guard.""" + iterator = ModelResponseIterator(streaming_response=None, sync_stream=True) + chunk = { + "id": "chatcmpl-x", + "object": "chat.completion.chunk", + "created": 1, + "model": "m", + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}], + } + + result = iterator.chunk_parser(chunk=chunk) + + assert result["text"] == "hi" diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index ca340b5f275..0221db1b23d 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -127,12 +127,14 @@ def test_get_supported_openai_params_parallel_tool_calls(): config = FireworksAIConfig() supported_params = config.get_supported_openai_params( - "fireworks_ai/accounts/fireworks/models/glm-4p6" + "fireworks_ai/accounts/fireworks/models/glm-5p1" ) assert "parallel_tool_calls" in supported_params + assert "tools" in supported_params + assert "tool_choice" in supported_params unsupported_params = config.get_supported_openai_params( - "fireworks_ai/accounts/fireworks/models/glm-5p1" + "fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct" ) assert "parallel_tool_calls" not in unsupported_params @@ -163,9 +165,9 @@ def test_get_model_info_respects_explicit_fireworks_capabilities(): """Test that get_model_info preserves explicit capability flags from the model map.""" model_info = get_model_info("fireworks_ai/accounts/fireworks/models/glm-5p1") - assert model_info["supports_function_calling"] is False + assert model_info["supports_function_calling"] is True assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is False + assert model_info["supports_tool_choice"] is True def test_get_provider_info_omits_false_supports_reasoning(monkeypatch): diff --git a/tests/test_litellm/llms/gemini/test_gemini_tts.py b/tests/test_litellm/llms/gemini/test_gemini_tts.py index 65eefca5af1..98f3ac0f4e5 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_tts.py +++ b/tests/test_litellm/llms/gemini/test_gemini_tts.py @@ -80,6 +80,46 @@ def test_gemini_tts_audio_parameter_mapping(self): assert "responseModalities" in result assert "AUDIO" in result["responseModalities"] + def test_gemini_tts_audio_parameter_mapping_with_language_code(self): + config = GoogleAIStudioGeminiConfig() + + non_default_params = { + "audio": {"voice": "Kore", "format": "pcm16", "language_code": "en-US"} + } + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gemini-2.5-flash-preview-tts", + drop_params=False, + ) + + assert "speechConfig" in result + assert result["speechConfig"]["languageCode"] == "en-US" + assert ( + result["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] + == "Kore" + ) + + def test_map_audio_params_language_code(self): + config = GoogleAIStudioGeminiConfig() + + result = config._map_audio_params( + {"voice": "Kore", "format": "pcm16", "language_code": "de-DE"} + ) + + assert result["languageCode"] == "de-DE" + assert result["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + + def test_map_audio_params_no_language_code(self): + config = GoogleAIStudioGeminiConfig() + + result = config._map_audio_params({"voice": "Kore", "format": "pcm16"}) + + assert "languageCode" not in result + assert result["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + def test_gemini_tts_audio_parameter_with_existing_modalities(self): """Test audio parameter mapping when modalities already exist""" config = GoogleAIStudioGeminiConfig() @@ -328,5 +368,57 @@ def test_speechconfig_end_to_end_mapping(self, model, custom_llm_provider): assert "AUDIO" in generation_config["responseModalities"] + @pytest.mark.parametrize( + "model,custom_llm_provider", + [ + ("gemini-2.5-flash-tts", "vertex_ai"), + ("gemini-2.5-flash-tts", "gemini"), + ("gemini-2.5-flash-preview-tts", "vertex_ai"), + ], + ) + def test_language_code_end_to_end_mapping(self, model, custom_llm_provider): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.llms.vertex_ai.gemini.transformation import ( + _transform_request_body, + ) + + config = VertexGeminiConfig() + + non_default_params = { + "audio": {"voice": "Puck", "format": "pcm16", "language_code": "pt-BR"} + } + optional_params = {} + + mapped_params = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + + assert mapped_params["speechConfig"]["languageCode"] == "pt-BR" + + request_body = _transform_request_body( + messages=[{"role": "user", "content": "Hello world"}], + model=model, + optional_params=mapped_params, + custom_llm_provider=custom_llm_provider, + litellm_params={}, + cached_content=None, + ) + + generation_config = request_body["generationConfig"] + assert generation_config["speechConfig"]["languageCode"] == "pt-BR" + assert ( + generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"][ + "voiceName" + ] + == "Puck" + ) + assert "AUDIO" in generation_config["responseModalities"] + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py b/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py index 54e7170bb20..17373f24a97 100644 --- a/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py +++ b/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py @@ -14,6 +14,8 @@ sys.path.insert(0, os.path.abspath("../../../../..")) import pytest +import litellm +from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager from litellm.llms.github_copilot.responses.transformation import ( @@ -22,13 +24,26 @@ from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams +@pytest.fixture(autouse=True) +def use_local_model_cost_map(monkeypatch: pytest.MonkeyPatch): + """Pin litellm.model_cost to the bundled local backup so tests don't depend + on remote catalog fetches (and don't change behavior across remote refreshes).""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url) + ) + litellm.add_known_models(model_cost_map=litellm.model_cost) + + class TestGithubCopilotResponsesAPITransformation: """Test GitHub Copilot Responses API configuration and transformations""" def test_github_copilot_provider_config_registration(self): - """Test that GitHub Copilot provider returns GithubCopilotResponsesAPIConfig""" + """Test that GitHub Copilot provider returns the native Responses API + config for a Responses-capable catalog model. Exercises the full stack: + catalog lookup -> github_copilot_supports_responses_api -> native config.""" config = ProviderConfigManager.get_provider_responses_api_config( - model="github_copilot/gpt-5.1-codex", + model="github_copilot/gpt-5.3-codex", provider=LlmProviders.GITHUB_COPILOT, ) @@ -373,3 +388,200 @@ def test_handle_reasoning_item_non_reasoning_passthrough(self): # Non-reasoning items should pass through unchanged assert result == message_item + + +class TestGithubCopilotResponsesAPIRouting: + """``ProviderConfigManager.get_provider_responses_api_config`` for github_copilot + returns the native Responses config only when the model has ``mode=responses`` + in the (already-merged) model info; otherwise returns None so the dispatcher + routes through the chat-completions translation bridge.""" + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_returns_config_when_mode_is_responses(self, mock_get_info): + """``mode=responses`` returns native config.""" + mock_get_info.return_value = {"mode": "responses"} + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-responses-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert isinstance(config, GithubCopilotResponsesAPIConfig) + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_returns_none_when_mode_is_chat(self, mock_get_info): + """``mode=chat`` returns None so dispatcher uses bridge.""" + mock_get_info.return_value = {"mode": "chat"} + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-chat-only-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert config is None + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_returns_none_when_mode_is_unset_and_no_endpoints(self, mock_get_info): + """Entry without ``mode`` and without ``supported_endpoints`` returns None + (conservative default).""" + mock_get_info.return_value = {} + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert config is None + + def test_returns_config_when_mode_unset_but_endpoints_have_responses(self): + """``mode`` unset but ``supported_endpoints`` declaring /v1/responses + returns native config (endpoint-list fallback for stale-but-correct + catalog entries that lack ``mode``). + + Exercises the real ``_cached_get_model_info_helper`` plumbing via + ``register_model`` (no mock). ``supported_endpoints`` is not carried on + the normalized ``ModelInfoBase`` the helper returns, so the gate must + read it from the raw ``litellm.model_cost`` entry; a mock-based test + would mask that. + """ + litellm.register_model( + { + "github_copilot/test-endpoints-only-model": { + "litellm_provider": "github_copilot", + "max_tokens": 1, + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + ], + } + } + ) + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/test-endpoints-only-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert isinstance(config, GithubCopilotResponsesAPIConfig) + + def test_mode_chat_overrides_endpoints_with_responses(self): + """``mode=chat`` is a hard opt-out: forces bridge even when + ``supported_endpoints`` includes /v1/responses. Lets users force the + bridge for dual-endpoint models without clearing endpoint metadata. + + Exercises the real ``_cached_get_model_info_helper`` plumbing via + ``register_model`` (no mock) so the ``mode``-over-endpoints precedence + is verified against the actual model-info resolution. + """ + litellm.register_model( + { + "github_copilot/test-chat-override-model": { + "litellm_provider": "github_copilot", + "max_tokens": 1, + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + ], + } + } + ) + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/test-chat-override-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert config is None + + def test_returns_config_when_model_is_none(self): + """Follow-up GET/DELETE operations pass model=None and keep the native + config path (no per-model lookup is possible).""" + config = ProviderConfigManager.get_provider_responses_api_config( + model=None, + provider=LlmProviders.GITHUB_COPILOT, + ) + assert isinstance(config, GithubCopilotResponsesAPIConfig) + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_returns_none_when_get_model_info_raises(self, mock_get_info): + """Catalog lookup failure (model not registered) returns None + (conservative default; bridge handles unknown models safely).""" + mock_get_info.side_effect = Exception("model not in catalog") + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/never-seen-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert config is None + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_user_override_via_register_model(self, mock_get_info): + """User-supplied per-deployment ``model_info`` flows through + ``litellm.register_model`` (called by the router) into the merged + catalog read by ``_cached_get_model_info_helper``. Setting ``mode=responses`` + for a model whose catalog entry says ``mode=chat`` therefore opts in + to native dispatch without any per-call argument plumbing.""" + mock_get_info.return_value = {"mode": "responses"} + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-chat-only-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert isinstance(config, GithubCopilotResponsesAPIConfig) + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_realistic_chat_only_entry_returns_none(self, mock_get_info): + """Realistic ``model_prices_and_context_window.json`` shape for a + chat-only Copilot model (e.g. github_copilot/gemini-3.1-pro-preview) + returns None so /v1/responses calls fall back to the bridge.""" + mock_get_info.return_value = { + "litellm_provider": "github_copilot", + "max_input_tokens": 136000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "supported_endpoints": ["/v1/chat/completions"], + "supports_function_calling": True, + "supports_tool_choice": True, + "supports_parallel_function_calling": True, + "supports_vision": True, + "supports_reasoning": True, + } + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-chat-only-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert config is None + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_realistic_responses_only_entry_returns_config(self, mock_get_info): + """Realistic catalog entry for a Responses-only Copilot model + (e.g. github_copilot/gpt-5.5) returns the native config.""" + mock_get_info.return_value = { + "litellm_provider": "github_copilot", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": ["/v1/responses"], + "supports_function_calling": True, + "supports_tool_choice": True, + "supports_parallel_function_calling": True, + "supports_response_schema": True, + "supports_vision": True, + "supports_reasoning": True, + "supports_none_reasoning_effort": True, + "supports_xhigh_reasoning_effort": True, + } + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-responses-only-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert isinstance(config, GithubCopilotResponsesAPIConfig) diff --git a/tests/test_litellm/llms/openai/completion/test_completion_handler.py b/tests/test_litellm/llms/openai/completion/test_completion_handler.py new file mode 100644 index 00000000000..c6af96fa375 --- /dev/null +++ b/tests/test_litellm/llms/openai/completion/test_completion_handler.py @@ -0,0 +1,93 @@ +""" +Tests that client headers are forwarded to the provider on the OpenAI +text completion path. + +Regression tests for https://github.com/BerriAI/litellm/issues/27410 +""" + +import os +import sys + +import pytest +import respx +from httpx import Response + +sys.path.insert(0, os.path.abspath("../../../../..")) + +import litellm +from litellm import atext_completion, text_completion + + +@pytest.fixture(autouse=True) +def setup_env(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-test-fake-key") + + +@pytest.fixture +def mock_completions_endpoint(): + return respx.post("https://api.openai.com/v1/completions").mock( + return_value=Response( + 200, + json={ + "id": "cmpl-test123", + "object": "text_completion", + "created": 1677652288, + "model": "gpt-3.5-turbo-instruct", + "choices": [ + { + "text": "hi", + "index": 0, + "logprobs": None, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + ) + + +@respx.mock +def test_completion_forwards_client_headers_to_provider(mock_completions_endpoint): + text_completion( + model="gpt-3.5-turbo-instruct", + prompt="hello", + max_tokens=5, + headers={"x-mycorp-llmcall-id": "abc-123"}, + ) + + request_headers = mock_completions_endpoint.calls.last.request.headers + assert request_headers["x-mycorp-llmcall-id"] == "abc-123" + + +@respx.mock +def test_completion_forwards_extra_headers_to_provider(mock_completions_endpoint): + text_completion( + model="gpt-3.5-turbo-instruct", + prompt="hello", + max_tokens=5, + extra_headers={"x-mycorp-llmcall-id": "abc-123"}, + ) + + request_headers = mock_completions_endpoint.calls.last.request.headers + assert request_headers["x-mycorp-llmcall-id"] == "abc-123" + + +@respx.mock +async def test_acompletion_forwards_client_headers_to_provider( + mock_completions_endpoint, monkeypatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + await atext_completion( + model="gpt-3.5-turbo-instruct", + prompt="hello", + max_tokens=5, + headers={"x-mycorp-llmcall-id": "abc-123"}, + ) + + request_headers = mock_completions_endpoint.calls.last.request.headers + assert request_headers["x-mycorp-llmcall-id"] == "abc-123" diff --git a/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py b/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py index f81f1c00a7b..09248a779c5 100644 --- a/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py +++ b/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py @@ -2,8 +2,23 @@ Tests for Tensormesh provider configuration and integration. """ +import pytest + import litellm +TENSORMESH_MODELS = [ + "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8", + "tensormesh/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8", + "tensormesh/Qwen/Qwen3.6-27B-FP8", + "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP", + "tensormesh/deepseek-ai/DeepSeek-V4-Flash", + "tensormesh/moonshotai/Kimi-K2.6", + "tensormesh/MiniMaxAI/MiniMax-M2.5", + "tensormesh/google/gemma-4-31B-it", + "tensormesh/openai/gpt-oss-120b", + "tensormesh/openai/gpt-oss-20b", +] + class TestTensormeshProviderConfig: """Test Tensormesh provider configuration""" @@ -82,3 +97,60 @@ def test_tensormesh_router_config(self): assert len(router.model_list) == 1 assert router.model_list[0]["model_name"] == "tensormesh-chat" + + +class TestTensormeshCostMap: + """The serverless models are registered in the cost map so LiteLLM can + price requests and unblock tool-calling params on the JSON provider path.""" + + @pytest.fixture(autouse=True) + def _use_local_model_cost_map(self, monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def test_models_registered_with_capabilities(self): + for model in TENSORMESH_MODELS: + info = litellm.get_model_info(model) + assert info["litellm_provider"] == "tensormesh" + assert info["mode"] == "chat" + assert litellm.supports_function_calling(model) is True, model + assert litellm.supports_response_schema(model) is True, model + assert litellm.model_cost[model]["supports_tool_choice"] is True, model + assert litellm.model_cost[model]["supports_prompt_caching"] is True, model + + def test_reasoning_flag_matches_expected_set(self): + reasoning_models = { + "tensormesh/deepseek-ai/DeepSeek-V4-Flash", + "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8", + "tensormesh/Qwen/Qwen3.6-27B-FP8", + "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP", + "tensormesh/MiniMaxAI/MiniMax-M2.5", + "tensormesh/moonshotai/Kimi-K2.6", + "tensormesh/openai/gpt-oss-120b", + "tensormesh/openai/gpt-oss-20b", + "tensormesh/google/gemma-4-31B-it", + } + for model in TENSORMESH_MODELS: + assert litellm.supports_reasoning(model) is (model in reasoning_models), model + + def test_cost_is_wired_and_cache_reads_are_free(self): + prompt_cost, completion_cost = litellm.cost_per_token( + model="tensormesh/openai/gpt-oss-120b", + prompt_tokens=1_000_000, + completion_tokens=1_000_000, + ) + assert prompt_cost == pytest.approx(0.15) + assert completion_cost == pytest.approx(0.60) + assert ( + litellm.model_cost["tensormesh/openai/gpt-oss-120b"][ + "cache_read_input_token_cost" + ] + == 0 + ) diff --git a/tests/test_litellm/llms/parasail/test_parasail.py b/tests/test_litellm/llms/parasail/test_parasail.py new file mode 100644 index 00000000000..8fb9b22b5f6 --- /dev/null +++ b/tests/test_litellm/llms/parasail/test_parasail.py @@ -0,0 +1,172 @@ +import os +from unittest.mock import patch + +PARASAIL_API_BASE = "https://api.parasail.io/v1" +PARASAIL_RESPONSES_GATEWAY = "https://api-webflux.saas.parasail.io/v1" + + +def test_parasail_json_registry(): + import litellm + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + assert litellm.LlmProviders.PARASAIL.value == "parasail" + assert litellm.LlmProviders("parasail") == litellm.LlmProviders.PARASAIL + assert JSONProviderRegistry.exists("parasail") + config = JSONProviderRegistry.get("parasail") + assert config is not None + assert config.base_url == PARASAIL_API_BASE + assert config.api_key_env == "PARASAIL_API_KEY" + assert config.api_base_env == "PARASAIL_API_BASE" + assert "/v1/chat/completions" in config.supported_endpoints + assert "/v1/responses" in config.supported_endpoints + assert config.special_handling.get("force_store_false") is True + + +def test_parasail_listed_in_openai_compatible_providers(): + from litellm.constants import openai_compatible_providers + + assert "parasail" in openai_compatible_providers + + +def test_parasail_dynamic_config_env_vars(): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + config = create_config_class(JSONProviderRegistry.get("parasail"))() + + with patch.dict( + os.environ, + { + "PARASAIL_API_KEY": "test-key", + "PARASAIL_API_BASE": PARASAIL_RESPONSES_GATEWAY, + }, + ): + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + + assert api_base == PARASAIL_RESPONSES_GATEWAY + assert api_key == "test-key" + + +def test_parasail_provider_detection_by_prefix(): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, _, api_base = get_llm_provider( + "parasail/parasail-llama-33-70b-fp8" + ) + + assert model == "parasail-llama-33-70b-fp8" + assert provider == "parasail" + assert api_base == PARASAIL_API_BASE + + +def test_parasail_chat_complete_url(): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + config = create_config_class(JSONProviderRegistry.get("parasail"))() + + assert ( + config.get_complete_url( + api_base=None, + api_key=None, + model="parasail-llama-33-70b-fp8", + optional_params={}, + litellm_params={}, + ) + == f"{PARASAIL_API_BASE}/chat/completions" + ) + + +def test_parasail_responses_api_config(): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_responses_api_config( + provider="parasail", + model="parasail-kimi-k25-elicit", + ) + + assert isinstance(config, OpenAIResponsesAPIConfig) + assert config.custom_llm_provider == "parasail" + assert ( + config.get_complete_url(api_base=None, litellm_params={}) + == f"{PARASAIL_API_BASE}/responses" + ) + + +def test_parasail_responses_api_honors_api_base_override(): + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_responses_api_config( + provider="parasail", + model="parasail-kimi-k25-elicit", + ) + + with patch.dict( + os.environ, + {"PARASAIL_API_BASE": PARASAIL_RESPONSES_GATEWAY}, + ): + url = config.get_complete_url(api_base=None, litellm_params={}) + + assert url == f"{PARASAIL_RESPONSES_GATEWAY}/responses" + + +def test_parasail_responses_api_forces_store_false_when_caller_sets_true(): + from litellm.types.router import GenericLiteLLMParams + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_responses_api_config( + provider="parasail", + model="parasail-kimi-k25-elicit", + ) + + request_params: dict = {"store": True, "temperature": 0.2} + transformed = config.transform_responses_api_request( + model="parasail-kimi-k25-elicit", + input="hello", + response_api_optional_request_params=request_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert transformed["store"] is False + assert transformed["temperature"] == 0.2 + + +def test_parasail_responses_api_forces_store_false_when_caller_omits_store(): + from litellm.types.router import GenericLiteLLMParams + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_responses_api_config( + provider="parasail", + model="parasail-kimi-k25-elicit", + ) + + transformed = config.transform_responses_api_request( + model="parasail-kimi-k25-elicit", + input="hello", + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert transformed["store"] is False + + +def test_parasail_responses_api_validate_environment_sets_bearer_token(): + from litellm.types.router import GenericLiteLLMParams + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_responses_api_config( + provider="parasail", + model="parasail-kimi-k25-elicit", + ) + + with patch.dict(os.environ, {"PARASAIL_API_KEY": "secret-from-env"}): + headers = config.validate_environment( + headers={}, + model="parasail-kimi-k25-elicit", + litellm_params=GenericLiteLLMParams(), + ) + + assert headers["Authorization"] == "Bearer secret-from-env" diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 74888e6cd9e..cc8b14e5514 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -86,7 +86,7 @@ def test_check_and_create_cache_with_cached_content( cached_content = "cached_content_123" optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -129,7 +129,7 @@ def test_check_and_create_cache_no_cached_messages( mock_separate.return_value = ([], self.sample_messages) # No cached messages optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -177,7 +177,7 @@ def test_check_and_create_cache_existing_cache_found( optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -254,7 +254,7 @@ def test_check_and_create_cache_create_new_cache( optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -324,7 +324,7 @@ def test_check_and_create_cache_http_error( optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute and Assert with pytest.raises(VertexAIError) as exc_info: @@ -364,7 +364,7 @@ async def test_async_check_and_create_cache_with_cached_content( cached_content = "cached_content_123" optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -404,7 +404,7 @@ async def test_async_check_and_create_cache_no_cached_messages( mock_separate.return_value = ([], self.sample_messages) optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -453,7 +453,7 @@ async def test_async_check_and_create_cache_existing_cache_found( optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -535,7 +535,7 @@ async def test_async_check_and_create_cache_create_new_cache( optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -606,7 +606,7 @@ async def test_async_check_and_create_cache_timeout_error( optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute and Assert with pytest.raises(VertexAIError) as exc_info: @@ -648,7 +648,7 @@ def test_check_and_create_cache_tools_popped_from_optional_params( optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Mock the check_cache to return existing cache so we don't make HTTP calls with patch.object( @@ -694,7 +694,7 @@ def test_check_and_create_cache_tools_not_popped_when_no_cached_messages( optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -735,7 +735,7 @@ async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_mess optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -778,7 +778,7 @@ async def test_async_check_and_create_cache_tools_popped_from_optional_params( optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Mock the async_check_cache to return existing cache so we don't make HTTP calls with patch.object( @@ -837,7 +837,7 @@ def test_check_and_create_cache_tool_choice_popped_from_optional_params( logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -870,7 +870,7 @@ def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -908,7 +908,7 @@ async def test_async_check_and_create_cache_tool_choice_popped_from_optional_par logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -942,7 +942,7 @@ async def test_async_check_and_create_cache_tool_choice_not_popped_when_no_cache logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1002,7 +1002,7 @@ def test_check_and_create_cache_tool_choice_in_request_body( logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1072,7 +1072,7 @@ async def test_async_check_and_create_cache_tool_choice_in_request_body( logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1138,7 +1138,7 @@ def test_check_and_create_cache_omits_tool_config_when_tool_choice_unset( logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1205,7 +1205,7 @@ def test_check_and_create_cache_tool_choice_function_pin( logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1280,7 +1280,7 @@ def test_check_and_create_cache_tool_choice_typed_constructor( logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1336,7 +1336,7 @@ def test_check_and_create_cache_distinct_tool_choices_use_distinct_keys( logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1390,7 +1390,7 @@ def test_check_and_create_cache_skips_when_below_min_tokens( cached_content=None, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="test_token", ) @@ -1441,7 +1441,7 @@ async def test_async_check_and_create_cache_skips_when_below_min_tokens( cached_content=None, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="test_token", ) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 0d02521433a..671d7355e8f 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1459,6 +1459,26 @@ def test_vertex_ai_process_candidates_with_grounding_metadata(): assert len(result[0]) == 1 +def test_set_stream_metadata_mirrors_non_streaming_safety_field_names(): + safety_ratings = [ + [{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}] + ] + + model_response = ModelResponse() + VertexGeminiConfig._set_stream_metadata_on_response( + model_response=model_response, + grounding_metadata=[], + url_context_metadata=[], + safety_ratings=safety_ratings, + citation_metadata=[], + ) + + assert getattr(model_response, "vertex_ai_safety_ratings") == safety_ratings + assert getattr(model_response, "vertex_ai_safety_results") == safety_ratings + assert model_response._hidden_params["vertex_ai_safety_ratings"] == safety_ratings + assert model_response._hidden_params["vertex_ai_safety_results"] == safety_ratings + + def test_vertex_ai_tool_call_id_format(): """ Test that tool call IDs have the correct format and length. diff --git a/tests/test_litellm/llms/xai/test_xai_oauth.py b/tests/test_litellm/llms/xai/test_xai_oauth.py new file mode 100644 index 00000000000..45fa6a405f2 --- /dev/null +++ b/tests/test_litellm/llms/xai/test_xai_oauth.py @@ -0,0 +1,801 @@ +import base64 +import hashlib +import json +import os +import threading +import time +from urllib.parse import parse_qs, urlparse +from unittest.mock import MagicMock + +import httpx +import litellm +import pytest +from click.testing import CliRunner + +import litellm.llms.xai.oauth as xai_oauth_module +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.xai.oauth import ( + XAI_OAUTH_CLIENT_ID, + XAI_OAUTH_SCOPE, + XAIOAuthError, + XAIOAuthAuthenticator, + XAIOAuthLoginRequiredError, +) +from litellm.llms.xai.chat.transformation import XAIChatConfig +from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import get_optional_params, validate_environment + + +def _write_auth_file(tmp_path, payload): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + auth_file = token_dir / "auth.json" + auth_file.write_text(json.dumps(payload)) + return token_dir, auth_file + + +def test_get_access_token_uses_fresh_local_token(tmp_path, monkeypatch): + token_dir, _ = _write_auth_file( + tmp_path, + { + "access_token": "fresh-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + assert XAIOAuthAuthenticator().get_access_token() == "fresh-token" + + +def test_get_access_token_refreshes_and_preserves_refresh_token(tmp_path, monkeypatch): + token_dir, auth_file = _write_auth_file( + tmp_path, + { + "access_token": "expired-token", + "refresh_token": "refresh-token", + "token_endpoint": "https://auth.x.ai/oauth/token", + "expires_at": time.time() - 1, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + def handler(request: httpx.Request) -> httpx.Response: + body = dict(item.split("=") for item in request.content.decode().split("&")) + assert body["grant_type"] == "refresh_token" + assert body["refresh_token"] == "refresh-token" + assert body["client_id"] == XAI_OAUTH_CLIENT_ID + return httpx.Response( + 200, + json={ + "access_token": "new-token", + "expires_in": 3600, + "token_type": "Bearer", + }, + ) + + client = httpx.Client(transport=httpx.MockTransport(handler)) + + assert XAIOAuthAuthenticator(http_client=client).get_access_token() == "new-token" + stored = json.loads(auth_file.read_text()) + assert stored["access_token"] == "new-token" + assert stored["refresh_token"] == "refresh-token" + + +def test_get_access_token_reuses_token_refreshed_by_parallel_request(): + expired_auth_data = { + "access_token": "expired-token", + "refresh_token": "refresh-token", + "token_endpoint": "https://auth.x.ai/oauth/token", + "expires_at": time.time() - 1, + } + refreshed_auth_data = { + "access_token": "already-refreshed-token", + "refresh_token": "rotated-refresh-token", + "token_endpoint": "https://auth.x.ai/oauth/token", + "expires_at": time.time() + 3600, + } + authenticator = XAIOAuthAuthenticator() + authenticator._read_auth_file = MagicMock( + side_effect=[expired_auth_data, refreshed_auth_data] + ) + authenticator._refresh_tokens = MagicMock() + + assert authenticator.get_access_token() == "already-refreshed-token" + authenticator._refresh_tokens.assert_not_called() + + +def test_get_access_token_requires_login_without_auth_file(tmp_path, monkeypatch): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path / "missing")) + + with pytest.raises(XAIOAuthLoginRequiredError): + XAIOAuthAuthenticator().get_access_token() + + +def test_get_access_token_ignores_invalid_auth_file(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + (token_dir / "auth.json").write_text("{not-json") + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + with pytest.raises(XAIOAuthLoginRequiredError): + XAIOAuthAuthenticator().get_access_token() + + +def test_refresh_failure_surfaces_oauth_error(tmp_path, monkeypatch): + token_dir, _ = _write_auth_file( + tmp_path, + { + "access_token": "expired-token", + "refresh_token": "refresh-token", + "token_endpoint": "https://auth.x.ai/oauth/token", + "expires_at": time.time() - 1, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + client = httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(401, text="invalid_grant", request=request) + ) + ) + + with pytest.raises(XAIOAuthError) as exc_info: + XAIOAuthAuthenticator(http_client=client).get_access_token() + + assert "401 invalid_grant" in str(exc_info.value) + + +def test_build_auth_record_requires_access_and_refresh_tokens(): + authenticator = XAIOAuthAuthenticator() + + with pytest.raises(XAIOAuthError, match="access_token"): + authenticator._build_auth_record( + {"refresh_token": "refresh-token"}, + "https://auth.x.ai/oauth/token", + ) + + with pytest.raises(XAIOAuthError, match="refresh_token"): + authenticator._build_auth_record( + {"access_token": "access-token"}, + "https://auth.x.ai/oauth/token", + ) + + +def test_build_auth_record_defaults_expiry_and_token_type(): + authenticator = XAIOAuthAuthenticator() + + auth_data = authenticator._build_auth_record( + { + "access_token": "access-token", + "refresh_token": "refresh-token", + "expires_in": "not-a-number", + }, + "https://auth.x.ai/oauth/token", + ) + + assert auth_data["token_type"] == "Bearer" + assert auth_data["expires_at"] > time.time() + + +def test_is_expired_treats_missing_or_invalid_expiry_as_expired(): + authenticator = XAIOAuthAuthenticator() + + assert authenticator._is_expired({}) is True + assert authenticator._is_expired({"expires_at": "not-a-number"}) is True + + +def test_write_auth_file_creates_private_file(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + authenticator = XAIOAuthAuthenticator() + old_umask = os.umask(0o022) + replace_calls = [] + real_replace = os.replace + + def assert_private_temp_file(src, dst): + replace_calls.append((src, dst)) + assert oct(os.stat(src).st_mode & 0o777) == "0o600" + with open(src) as f: + assert json.load(f)["refresh_token"] == "refresh-token" + real_replace(src, dst) + + monkeypatch.setattr(os, "replace", assert_private_temp_file) + + try: + authenticator._write_auth_file( + { + "access_token": "access-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + } + ) + finally: + os.umask(old_umask) + + stored = json.loads((token_dir / "auth.json").read_text()) + assert stored["access_token"] == "access-token" + assert replace_calls + assert oct(os.stat(token_dir).st_mode & 0o777) == "0o700" + assert oct(os.stat(token_dir / "auth.json").st_mode & 0o777) == "0o600" + + +def test_discovery_rejects_unexpected_endpoint(): + authenticator = XAIOAuthAuthenticator() + + with pytest.raises(XAIOAuthError, match="unexpected endpoint"): + authenticator._validate_xai_endpoint("https://evil.example.com/oauth/token") + + with pytest.raises(XAIOAuthError, match="unexpected endpoint"): + authenticator._validate_xai_endpoint("http://auth.x.ai/oauth/token") + + +def test_discover_returns_validated_xai_endpoints(): + def handler(request: httpx.Request) -> httpx.Response: + assert request.url == "https://auth.x.ai/.well-known/openid-configuration" + return httpx.Response( + 200, + json={ + "authorization_endpoint": "https://auth.x.ai/oauth/authorize", + "token_endpoint": "https://auth.x.ai/oauth/token", + }, + ) + + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client(transport=httpx.MockTransport(handler)) + ) + + assert authenticator._discover() == { + "authorization_endpoint": "https://auth.x.ai/oauth/authorize", + "token_endpoint": "https://auth.x.ai/oauth/token", + } + + +def test_discover_requires_authorization_and_token_endpoints(): + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json={})) + ) + ) + + with pytest.raises(XAIOAuthError, match="missing endpoints"): + authenticator._discover() + + +def test_discover_wraps_http_errors(): + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 500, text="discovery failed", request=request + ) + ) + ) + ) + + with pytest.raises(XAIOAuthError) as exc_info: + authenticator._discover() + + assert "xAI OAuth discovery request failed: 500 discovery failed" in str( + exc_info.value + ) + + +def test_discover_wraps_invalid_json_response(): + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, text="not-json") + ) + ) + ) + + with pytest.raises(XAIOAuthError, match="discovery response was not valid JSON"): + authenticator._discover() + + +def test_refresh_discovers_token_endpoint_when_auth_file_is_legacy( + tmp_path, monkeypatch +): + token_dir, auth_file = _write_auth_file( + tmp_path, + { + "access_token": "expired-token", + "refresh_token": "refresh-token", + "expires_at": time.time() - 1, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, + json={ + "authorization_endpoint": "https://auth.x.ai/oauth/authorize", + "token_endpoint": "https://auth.x.ai/oauth/token", + }, + ) + return httpx.Response( + 200, + json={ + "access_token": "discovered-token", + "refresh_token": "new-refresh-token", + "expires_in": 3600, + }, + ) + + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client(transport=httpx.MockTransport(handler)) + ) + + assert authenticator.get_access_token() == "discovered-token" + stored = json.loads(auth_file.read_text()) + assert stored["token_endpoint"] == "https://auth.x.ai/oauth/token" + + +def test_exchange_token_rejects_non_object_response(): + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, json=["not", "an", "object"]) + ) + ) + ) + + with pytest.raises(XAIOAuthError, match="was not an object"): + authenticator._exchange_token("https://auth.x.ai/oauth/token", {}) + + +def test_exchange_token_wraps_invalid_json_response(): + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, text="not-json") + ) + ) + ) + + with pytest.raises(XAIOAuthError, match="token response was not valid JSON"): + authenticator._exchange_token("https://auth.x.ai/oauth/token", {}) + + +def test_start_callback_server_falls_back_to_ephemeral_port(monkeypatch): + calls = [] + real_server = xai_oauth_module._CallbackServer + + class FirstPortFailsCallbackServer(real_server): + def __init__(self, server_address, handler_class): + calls.append(server_address[1]) + if server_address[1] == xai_oauth_module.XAI_OAUTH_REDIRECT_PORT: + raise OSError("port unavailable") + super().__init__(server_address, handler_class) + + monkeypatch.setattr( + xai_oauth_module, "_CallbackServer", FirstPortFailsCallbackServer + ) + + server, redirect_uri = XAIOAuthAuthenticator()._start_callback_server("state-value") + try: + assert calls == [xai_oauth_module.XAI_OAUTH_REDIRECT_PORT, 0] + assert redirect_uri.startswith("http://127.0.0.1:") + assert redirect_uri.endswith("/callback") + finally: + server.server_close() + + +def test_wait_for_callback_times_out_and_closes_server(monkeypatch): + server, _ = XAIOAuthAuthenticator()._start_callback_server("state-value") + monkeypatch.setattr(xai_oauth_module, "XAI_OAUTH_CALLBACK_TIMEOUT_SECONDS", 0) + + with pytest.raises(XAIOAuthError, match="Timed out"): + XAIOAuthAuthenticator()._wait_for_callback(server) + + +def test_callback_handler_records_success_and_rejects_state_mismatch(): + authenticator = XAIOAuthAuthenticator() + server, redirect_uri = authenticator._start_callback_server("expected-state") + thread = threading.Thread(target=server.handle_request) + thread.start() + response = httpx.get(f"{redirect_uri}?code=auth-code&state=expected-state") + thread.join(timeout=5) + + assert response.status_code == 200 + assert server.callback_result == { + "code": "auth-code", + "state": "expected-state", + "error": None, + "error_description": None, + } + + server, redirect_uri = authenticator._start_callback_server("expected-state") + thread = threading.Thread(target=server.handle_request) + thread.start() + response = httpx.get(f"{redirect_uri}?code=auth-code&state=wrong-state") + thread.join(timeout=5) + + assert response.status_code == 400 + assert server.callback_result["state"] == "wrong-state" + + +def test_login_exchanges_authorization_code_and_persists_auth_record(monkeypatch): + authenticator = XAIOAuthAuthenticator() + fake_server = MagicMock() + written_records = [] + + class FakeUUID: + def __init__(self, value): + self.hex = value + + monkeypatch.setattr( + xai_oauth_module.uuid, + "uuid4", + MagicMock(side_effect=[FakeUUID("state-value"), FakeUUID("nonce-value")]), + ) + authenticator._read_auth_file = MagicMock(return_value=None) + authenticator._discover = MagicMock( + return_value={ + "authorization_endpoint": "https://auth.x.ai/oauth/authorize", + "token_endpoint": "https://auth.x.ai/oauth/token", + } + ) + authenticator._pkce_pair = MagicMock(return_value=("verifier", "challenge")) + authenticator._start_callback_server = MagicMock( + return_value=(fake_server, "http://127.0.0.1:56121/callback") + ) + authenticator._wait_for_callback = MagicMock( + return_value={"state": "state-value", "code": "auth-code"} + ) + authenticator._exchange_token = MagicMock( + return_value={ + "access_token": "access-token", + "refresh_token": "refresh-token", + "expires_in": 3600, + } + ) + authenticator._write_auth_file = MagicMock(side_effect=written_records.append) + + auth_data = authenticator.login(no_browser=True) + + authenticator._exchange_token.assert_called_once_with( + "https://auth.x.ai/oauth/token", + { + "grant_type": "authorization_code", + "code": "auth-code", + "redirect_uri": "http://127.0.0.1:56121/callback", + "client_id": XAI_OAUTH_CLIENT_ID, + "code_verifier": "verifier", + }, + ) + assert auth_data["access_token"] == "access-token" + assert written_records == [auth_data] + + +def test_login_raises_on_callback_error_or_missing_code(monkeypatch): + authenticator = XAIOAuthAuthenticator() + + class FakeUUID: + hex = "state-value" + + monkeypatch.setattr( + xai_oauth_module.uuid, "uuid4", MagicMock(return_value=FakeUUID()) + ) + authenticator._read_auth_file = MagicMock(return_value=None) + authenticator._discover = MagicMock( + return_value={ + "authorization_endpoint": "https://auth.x.ai/oauth/authorize", + "token_endpoint": "https://auth.x.ai/oauth/token", + } + ) + authenticator._pkce_pair = MagicMock(return_value=("verifier", "challenge")) + authenticator._start_callback_server = MagicMock( + return_value=(MagicMock(), "http://127.0.0.1:56121/callback") + ) + authenticator._wait_for_callback = MagicMock( + return_value={ + "state": "state-value", + "error": "access_denied", + "error_description": "denied", + } + ) + + with pytest.raises(XAIOAuthError, match="denied"): + authenticator.login(no_browser=True) + + authenticator._wait_for_callback = MagicMock(return_value={"state": "state-value"}) + + with pytest.raises(XAIOAuthError, match="no code returned"): + authenticator.login(no_browser=True) + + +def test_pkce_pair_generates_s256_challenge(): + verifier, challenge = XAIOAuthAuthenticator()._pkce_pair() + expected = ( + base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()) + .rstrip(b"=") + .decode() + ) + + assert challenge == expected + assert "=" not in verifier + assert "=" not in challenge + + +def test_build_authorize_url_contains_xai_oauth_parameters(): + authorize_url = XAIOAuthAuthenticator()._build_authorize_url( + authorization_endpoint="https://auth.x.ai/oauth/authorize", + redirect_uri="http://127.0.0.1:56121/callback", + challenge="pkce-challenge", + state="state-value", + nonce="nonce-value", + ) + parsed = urlparse(authorize_url) + params = parse_qs(parsed.query) + + assert parsed.scheme == "https" + assert parsed.netloc == "auth.x.ai" + assert params["response_type"] == ["code"] + assert params["client_id"] == [XAI_OAUTH_CLIENT_ID] + assert params["scope"] == [XAI_OAUTH_SCOPE] + assert params["code_challenge"] == ["pkce-challenge"] + assert params["code_challenge_method"] == ["S256"] + assert params["state"] == ["state-value"] + assert params["nonce"] == ["nonce-value"] + + +def test_get_llm_provider_uses_single_xai_provider(monkeypatch): + monkeypatch.setenv("XAI_API_KEY", "api-key") + + model, provider, api_key, api_base = get_llm_provider("xai/grok-4") + + assert model == "grok-4" + assert provider == "xai" + assert api_key == "api-key" + assert api_base == "https://api.x.ai/v1" + + +def test_xai_oauth_alias_is_not_a_provider(): + with pytest.raises(Exception): + get_llm_provider("xai_oauth/grok-4") + + +def test_chat_config_wraps_flagged_oauth_errors_as_authentication_error( + tmp_path, monkeypatch +): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path / "missing")) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={"use_xai_oauth": True}, + api_key=None, + ) + + assert exc_info.value.llm_provider == "xai" + assert "litellm xai-oauth login" in str(exc_info.value) + + +def test_chat_config_injects_flagged_oauth_token(tmp_path, monkeypatch): + token_dir, _ = _write_auth_file( + tmp_path, + { + "access_token": "chat-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + headers = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={"use_xai_oauth": True}, + api_key=None, + ) + + assert headers["Authorization"] == "Bearer chat-token" + + +def test_chat_config_ignores_api_base_override_for_flagged_oauth(monkeypatch): + monkeypatch.setenv("XAI_OAUTH_API_BASE", "https://api.x.ai/v1") + + url = XAIChatConfig().get_complete_url( + api_base="https://attacker.example.com/v1", + api_key=None, + model="grok-4", + optional_params={}, + litellm_params={"use_xai_oauth": True}, + ) + + assert url == "https://api.x.ai/v1/chat/completions" + + +def test_chat_config_treats_blank_api_key_as_absent_for_flagged_oauth( + tmp_path, monkeypatch +): + token_dir, _ = _write_auth_file( + tmp_path, + { + "access_token": "stored-oauth-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + headers = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={"use_xai_oauth": True}, + api_key="", + ) + + assert headers["Authorization"] == "Bearer stored-oauth-token" + + +def test_chat_config_allows_api_base_override_with_caller_api_key(): + headers = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={"use_xai_oauth": True}, + api_key="caller-api-key", + ) + url = XAIChatConfig().get_complete_url( + api_base="https://custom.example.com/v1", + api_key="caller-api-key", + model="grok-4", + optional_params={}, + litellm_params={"use_xai_oauth": True}, + ) + + assert headers["Authorization"] == "Bearer caller-api-key" + assert url == "https://custom.example.com/v1/chat/completions" + + +def test_chat_config_prioritizes_env_api_key_over_oauth_flag(monkeypatch): + monkeypatch.setenv("XAI_API_KEY", "env-api-key") + + headers = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={"use_xai_oauth": True}, + api_key=None, + ) + url = XAIChatConfig().get_complete_url( + api_base="https://custom.example.com/v1", + api_key=None, + model="grok-4", + optional_params={}, + litellm_params={"use_xai_oauth": True}, + ) + + assert headers["Authorization"] == "Bearer env-api-key" + assert url == "https://custom.example.com/v1/chat/completions" + + +def test_validate_environment_still_reports_xai_api_key(monkeypatch): + monkeypatch.setenv("XAI_API_KEY", "env-api-key") + + assert validate_environment("xai/grok-4") == { + "keys_in_environment": True, + "missing_keys": [], + } + + +def test_xai_oauth_flag_uses_xai_optional_param_mapping(): + litellm_params = GenericLiteLLMParams(use_xai_oauth=True) + optional_params = get_optional_params( + model="grok-4", + custom_llm_provider="xai", + temperature=0.2, + max_tokens=8, + ) + + assert optional_params["temperature"] == 0.2 + assert optional_params["max_tokens"] == 8 + assert litellm_params.use_xai_oauth is True + assert "use_xai_oauth" not in optional_params + + +def test_responses_config_injects_flagged_oauth_bearer_token(tmp_path, monkeypatch): + token_dir, _ = _write_auth_file( + tmp_path, + { + "access_token": "responses-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + headers = XAIResponsesAPIConfig().validate_environment( + headers={}, + model="grok-4", + litellm_params=GenericLiteLLMParams(use_xai_oauth=True), + ) + + assert headers["Authorization"] == "Bearer responses-token" + + +def test_responses_config_endpoint_url_uses_oauth_authenticator(monkeypatch): + monkeypatch.setenv("XAI_OAUTH_API_BASE", "https://xai.example.com/v1/") + config = XAIResponsesAPIConfig() + + assert config.get_complete_url( + api_base=None, litellm_params={"use_xai_oauth": True} + ) == ("https://xai.example.com/v1/responses") + assert ( + config.get_complete_url( + api_base="https://custom.example.com/v1/", + litellm_params={"use_xai_oauth": True}, + ) + == "https://xai.example.com/v1/responses" + ) + assert ( + config.get_complete_url( + api_base="https://custom.example.com/v1/", + litellm_params={"api_key": "", "use_xai_oauth": True}, + ) + == "https://xai.example.com/v1/responses" + ) + assert ( + config.get_complete_url( + api_base="https://custom.example.com/v1/", + litellm_params={"api_key": "caller-api-key"}, + ) + == "https://custom.example.com/v1/responses" + ) + + +def test_responses_config_wraps_flagged_oauth_errors_as_authentication_error( + tmp_path, monkeypatch +): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path / "missing")) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + XAIResponsesAPIConfig().validate_environment( + headers={}, + model="grok-4", + litellm_params=GenericLiteLLMParams(use_xai_oauth=True), + ) + + assert XAIResponsesAPIConfig().custom_llm_provider.value == "xai" + assert exc_info.value.llm_provider == "xai" + + +def test_proxy_cli_xai_oauth_login_uses_single_authenticator(monkeypatch): + from litellm.proxy.proxy_cli import run_server + + instances = [] + + class FakeAuthenticator: + auth_file = "/tmp/xai-oauth-auth.json" + + def __init__(self): + instances.append(self) + + def login(self): + return {"expires_at": 1234567890} + + monkeypatch.setattr( + "litellm.llms.xai.oauth.XAIOAuthAuthenticator", FakeAuthenticator + ) + + result = CliRunner().invoke(run_server, ["xai-oauth", "login"]) + + assert result.exit_code == 0 + assert len(instances) == 1 + assert "Credentials saved to /tmp/xai-oauth-auth.json" in result.output + assert "Access token expires at 1234567890" in result.output diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py new file mode 100644 index 00000000000..786f6244930 --- /dev/null +++ b/tests/test_litellm/models/test_models.py @@ -0,0 +1,542 @@ +""" +Tests for backend domain models. +""" + +from datetime import datetime + +import pytest + +from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.models.budget import ( + LiteLLM_BudgetTable, + LiteLLM_BudgetTableFull, + LiteLLM_TeamMemberTable, +) +from litellm.models.config import LiteLLM_Config +from litellm.models.credentials import CreateCredentialItem, CredentialItem +from litellm.models.end_user import LiteLLM_EndUserTable +from litellm.models.managed_files import ( + LiteLLM_ManagedFileTable, + LiteLLM_ManagedObjectTable, + LiteLLM_ManagedVectorStoresTable, +) +from litellm.models.mcp_server import LiteLLM_MCPServerTable +from litellm.models.model import LiteLLM_ProxyModelTable +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.models.organization import LiteLLM_OrganizationTable +from litellm.models.project import LiteLLM_ProjectTable +from litellm.models.skills import LiteLLM_SkillsTable +from litellm.models.spend_logs import LiteLLM_ErrorLogs, LiteLLM_SpendLogs +from litellm.models.tag import LiteLLM_TagTable +from litellm.models.team import ( + LiteLLM_DeletedTeamTable, + LiteLLM_TeamTable, + LiteLLM_TeamTableCachedObj, +) +from litellm.models.team_membership import LiteLLM_TeamMembership +from litellm.models.user import LiteLLM_UserTable +from litellm.models.verification_token import ( + LiteLLM_DeletedVerificationToken, + LiteLLM_VerificationToken, +) + + +class TestBudget: + def test_budget_creation(self): + budget = LiteLLM_BudgetTable( + budget_id="test-budget-id", + max_budget=100.0, + soft_budget=80.0, + tpm_limit=1000, + rpm_limit=100, + model_max_budget={"gpt-4": 50.0}, + budget_duration="monthly", + allowed_models=["gpt-4"], + ) + assert budget.budget_id == "test-budget-id" + assert budget.max_budget == 100.0 + assert budget.soft_budget == 80.0 + assert budget.tpm_limit == 1000 + assert budget.rpm_limit == 100 + assert budget.model_max_budget == {"gpt-4": 50.0} + assert budget.budget_duration == "monthly" + assert budget.allowed_models == ["gpt-4"] + + def test_budget_defaults(self): + budget = LiteLLM_BudgetTable() + assert budget.budget_id is None + assert budget.max_budget is None + assert budget.allowed_models is None + + +class TestCredentials: + def test_credentials_creation(self): + creds = CredentialItem( + credential_name="test-cred", + credential_values={"api_key": "secret123"}, + credential_info={"provider": "openai"}, + ) + assert creds.credential_name == "test-cred" + assert creds.credential_values["api_key"] == "secret123" + assert creds.credential_info["provider"] == "openai" + + def test_create_credential_item_accepts_model_id(self): + item = CreateCredentialItem( + credential_name="from-model", + credential_info={}, + model_id="model-123", + ) + assert item.model_id == "model-123" + assert item.credential_values is None + + def test_create_credential_item_requires_values_or_model_id(self): + with pytest.raises( + ValueError, match="Either credential_values or model_id must be set" + ): + CreateCredentialItem(credential_name="bad", credential_info={}) + + +class TestModel: + def test_model_creation(self): + model = LiteLLM_ProxyModelTable( + model_id="test-model-id", + model_name="gpt-4", + litellm_params={"model": "gpt-4", "api_key": "test"}, + model_info={"team_id": "team-123", "team_public_model_name": "my-gpt4"}, + ) + assert model.model_id == "test-model-id" + assert model.model_name == "gpt-4" + assert model.team_id == "team-123" + assert model.team_public_model_name == "my-gpt4" + + def test_is_blocked(self): + model_blocked = LiteLLM_ProxyModelTable( + model_id="m1", model_name="test", litellm_params={}, blocked=True + ) + model_unblocked = LiteLLM_ProxyModelTable( + model_id="m2", model_name="test", litellm_params={}, blocked=False + ) + assert model_blocked.is_blocked + assert not model_unblocked.is_blocked + + def test_parses_json_string_fields(self): + model = LiteLLM_ProxyModelTable( + model_id="m1", + model_name="gpt-4", + litellm_params='{"model": "gpt-4"}', + model_info='{"team_id": "t1"}', + ) + assert model.litellm_params == {"model": "gpt-4"} + assert model.model_info == {"team_id": "t1"} + + def test_team_helpers_none_when_no_model_info(self): + model = LiteLLM_ProxyModelTable( + model_id="m1", model_name="gpt-4", litellm_params={}, model_info=None + ) + assert model.team_id is None + assert model.team_public_model_name is None + + +class TestObjectPermission: + def test_object_permission_creation(self): + perm = LiteLLM_ObjectPermissionTable( + object_permission_id="test-perm-id", + mcp_servers=["server1", "server2"], + vector_stores=["vs1"], + agents=["agent1"], + models=["gpt-4"], + blocked_tools=["dangerous_tool"], + ) + assert perm.object_permission_id == "test-perm-id" + assert len(perm.mcp_servers) == 2 + assert perm.vector_stores == ["vs1"] + assert perm.agents == ["agent1"] + assert perm.models == ["gpt-4"] + assert perm.blocked_tools == ["dangerous_tool"] + + def test_object_permission_tool_permissions(self): + perm = LiteLLM_ObjectPermissionTable( + object_permission_id="perm-tools", + mcp_tool_permissions={"server1": ["tool1", "tool2"]}, + ) + assert perm.mcp_tool_permissions == {"server1": ["tool1", "tool2"]} + + +class TestOrganization: + def test_organization_creation(self): + org = LiteLLM_OrganizationTable( + organization_id="org-123", + organization_alias="My Org", + budget_id="budget-123", + models=["gpt-4", "claude-3"], + spend=50.0, + created_by="admin", + updated_by="admin", + ) + assert org.organization_id == "org-123" + assert org.organization_alias == "My Org" + assert len(org.models) == 2 + + +class TestProject: + def test_project_creation(self): + project = LiteLLM_ProjectTable( + project_id="proj-123", + project_alias="My Project", + team_id="team-123", + blocked=False, + ) + assert project.project_id == "proj-123" + assert not project.is_blocked + + +class TestTeam: + def test_team_creation(self): + team = LiteLLM_TeamTable( + team_id="team-123", + team_alias="Engineering", + admins=["user1"], + members=["user2", "user3"], + models=["gpt-4"], + max_budget=1000.0, + spend=100.0, + ) + assert team.team_id == "team-123" + assert team.team_alias == "Engineering" + assert team.admins == ["user1"] + assert team.members == ["user2", "user3"] + assert team.models == ["gpt-4"] + assert team.max_budget == 1000.0 + + def test_members_with_roles_parsing(self): + team = LiteLLM_TeamTable( + team_id="t2", + members_with_roles=[ + {"user_id": "user1", "role": "admin"}, + {"user_id": "user2", "role": "user"}, + ], + ) + assert len(team.members_with_roles) == 2 + assert team.members_with_roles[0].user_id == "user1" + assert team.members_with_roles[0].role == "admin" + + def test_members_with_roles_empty_dict_coerced(self): + team = LiteLLM_TeamTable(team_id="t3", members_with_roles={}) + assert team.members_with_roles == [] + + def test_json_string_fields_parsed(self): + team = LiteLLM_TeamTable( + team_id="t4", + metadata='{"k": "v"}', + model_max_budget='{"gpt-4": 5.0}', + ) + assert team.metadata == {"k": "v"} + assert team.model_max_budget == {"gpt-4": 5.0} + + def test_cached_team(self): + cached = LiteLLM_TeamTableCachedObj( + team_id="t1", last_refreshed_at=1234567890.0 + ) + assert cached.last_refreshed_at == 1234567890.0 + + def test_deleted_team(self): + deleted = LiteLLM_DeletedTeamTable( + team_id="t1", + deleted_by="admin", + deleted_at=datetime.utcnow(), + ) + assert deleted.deleted_by == "admin" + assert deleted.deleted_at is not None + + +class TestUser: + def test_user_creation(self): + user = LiteLLM_UserTable( + user_id="user-123", + user_email="test@example.com", + teams=["team1", "team2"], + max_budget=100.0, + spend=25.0, + ) + assert user.user_id == "user-123" + assert user.user_email == "test@example.com" + assert len(user.teams) == 2 + + def test_is_over_budget(self): + user = LiteLLM_UserTable(user_id="u1", max_budget=100.0, spend=150.0) + user_no_budget = LiteLLM_UserTable(user_id="u2", spend=1000.0) + + assert user.is_over_budget() + assert not user_no_budget.is_over_budget() + + def test_has_model_access(self): + user_with_models = LiteLLM_UserTable(user_id="u1", models=["gpt-4"]) + user_no_models = LiteLLM_UserTable(user_id="u2", models=[]) + + assert user_with_models.has_model_access("gpt-4") + assert not user_with_models.has_model_access("gpt-3") + assert user_no_models.has_model_access("any-model") + + def test_password_hash_excluded_from_serialization(self): + from litellm.proxy._types import LiteLLM_UserTableWithKeyCount + + secret = "$2b$12$abcdefghijklmnopqrstuv" + user = LiteLLM_UserTable(user_id="u1", user_email="a@b.c", password=secret) + + assert user.password == secret + assert "password" not in user.model_dump() + assert "password" not in user.model_dump_json() + + with_keys = LiteLLM_UserTableWithKeyCount( + user_id="u1", user_email="a@b.c", password=secret, key_count=2 + ) + assert with_keys.password == secret + assert "password" not in with_keys.model_dump() + assert "password" not in with_keys.model_dump_json() + + +class TestVerificationToken: + def test_verification_token_creation(self): + token = LiteLLM_VerificationToken( + token="sk-test123", + key_name="Test Key", + user_id="user-123", + team_id="team-123", + max_budget=100.0, + spend=25.0, + models=["gpt-4"], + blocked=True, + allowed_routes=["/chat/completions"], + ) + assert token.token == "sk-test123" + assert token.key_name == "Test Key" + assert token.user_id == "user-123" + assert token.team_id == "team-123" + assert token.blocked is True + assert token.models == ["gpt-4"] + assert token.allowed_routes == ["/chat/completions"] + + def test_expires_accepts_string_and_datetime(self): + as_str = LiteLLM_VerificationToken(token="t1", expires="2024-12-31T23:59:59Z") + as_dt = LiteLLM_VerificationToken(token="t2", expires=datetime.utcnow()) + assert as_str.expires == "2024-12-31T23:59:59Z" + assert isinstance(as_dt.expires, datetime) + + def test_deleted_verification_token(self): + deleted = LiteLLM_DeletedVerificationToken( + token="t1", + deleted_by="admin", + deleted_at=datetime.utcnow(), + ) + assert deleted.deleted_by == "admin" + assert deleted.deleted_at is not None + assert deleted.token == "t1" + + +class TestConfigTable: + def test_config_creation(self): + cfg = LiteLLM_Config(param_name="general_settings", param_value={"k": "v"}) + assert cfg.param_name == "general_settings" + assert cfg.param_value == {"k": "v"} + + +class TestSkillsTable: + def test_skills_creation(self): + skill = LiteLLM_SkillsTable( + skill_id="s1", + display_title="My Skill", + source="custom", + file_content=b"zipbytes", + file_name="skill.zip", + ) + assert skill.skill_id == "s1" + assert skill.display_title == "My Skill" + assert skill.file_content == b"zipbytes" + + def test_skills_defaults(self): + skill = LiteLLM_SkillsTable(skill_id="s2") + assert skill.source == "custom" + assert skill.metadata is None + + +class TestAccessGroupTable: + def test_access_group_creation(self): + ag = LiteLLM_AccessGroupTable( + access_group_id="ag1", + access_group_name="group-a", + access_model_names=["gpt-4"], + assigned_team_ids=["t1"], + ) + assert ag.access_group_id == "ag1" + assert ag.access_model_names == ["gpt-4"] + assert ag.assigned_team_ids == ["t1"] + assert ag.access_agent_ids == [] + + +class TestTagTable: + def test_tag_creation(self): + tag = LiteLLM_TagTable( + tag_name="prod", + models=["gpt-4"], + spend=12.5, + budget_id="b1", + ) + assert tag.tag_name == "prod" + assert tag.models == ["gpt-4"] + assert tag.spend == 12.5 + + def test_tag_set_model_info_coerces_none(self): + tag = LiteLLM_TagTable(tag_name="t", spend=None, models=None) + assert tag.spend == 0.0 + assert tag.models == [] + + +class TestEndUserTable: + def test_end_user_creation(self): + eu = LiteLLM_EndUserTable( + user_id="eu1", + blocked=False, + spend=5.0, + allowed_model_region="eu", + default_model="gpt-4", + ) + assert eu.user_id == "eu1" + assert eu.blocked is False + assert eu.allowed_model_region == "eu" + assert eu.default_model == "gpt-4" + + def test_end_user_spend_coerced_when_none(self): + eu = LiteLLM_EndUserTable(user_id="eu2", blocked=True, spend=None) + assert eu.spend == 0.0 + + +class TestBudgetTableFull: + def test_full_adds_server_managed_fields(self): + now = datetime.now() + budget = LiteLLM_BudgetTableFull( + budget_id="b1", max_budget=10.0, created_at=now, budget_reset_at=now + ) + assert budget.created_at == now + assert budget.budget_reset_at == now + assert budget.max_budget == 10.0 + + def test_full_requires_created_at(self): + with pytest.raises(Exception): + LiteLLM_BudgetTableFull(budget_id="b1") + + +class TestTeamMemberTable: + def test_tracks_user_within_team(self): + member = LiteLLM_TeamMemberTable( + user_id="u1", team_id="t1", spend=3.0, budget_id="b1", max_budget=5.0 + ) + assert member.user_id == "u1" + assert member.team_id == "t1" + assert member.spend == 3.0 + assert member.max_budget == 5.0 + + +class TestTeamMembership: + def test_safe_get_limits_with_budget_table(self): + membership = LiteLLM_TeamMembership( + user_id="u1", + team_id="t1", + litellm_budget_table=LiteLLM_BudgetTable(rpm_limit=100, tpm_limit=2000), + ) + assert membership.safe_get_team_member_rpm_limit() == 100 + assert membership.safe_get_team_member_tpm_limit() == 2000 + + def test_safe_get_limits_without_budget_table(self): + membership = LiteLLM_TeamMembership(user_id="u1", team_id="t1") + assert membership.safe_get_team_member_rpm_limit() is None + assert membership.safe_get_team_member_tpm_limit() is None + + def test_full_budget_variant_parsed_for_server_fields(self): + now = datetime.now() + membership = LiteLLM_TeamMembership( + user_id="u1", + team_id="t1", + litellm_budget_table={ + "budget_id": "b1", + "rpm_limit": 7, + "created_at": now, + "budget_reset_at": now, + }, + ) + assert isinstance(membership.litellm_budget_table, LiteLLM_BudgetTableFull) + assert membership.safe_get_team_member_rpm_limit() == 7 + + +class TestMCPServerTable: + def test_mcp_server_defaults(self): + server = LiteLLM_MCPServerTable(server_id="s1", transport="sse") + assert server.server_id == "s1" + assert server.transport == "sse" + assert server.status == "unknown" + assert server.approval_status == "active" + assert server.allow_all_keys is False + assert server.available_on_public_internet is True + assert server.teams == [] + assert server.env == {} + + def test_mcp_server_requires_transport(self): + with pytest.raises(Exception): + LiteLLM_MCPServerTable(server_id="s1") + + +class TestSpendLogs: + def test_spend_logs_creation(self): + log = LiteLLM_SpendLogs( + request_id="r1", + api_key="sk-1", + call_type="completion", + startTime=None, + endTime=None, + messages=None, + response=None, + ) + assert log.request_id == "r1" + assert log.spend == 0.0 + assert log.cache_hit == "False" + + def test_error_logs_creation(self): + log = LiteLLM_ErrorLogs( + request_id="r1", startTime=None, endTime=None, status_code="500" + ) + assert log.request_id == "r1" + assert log.status_code == "500" + + +class TestManagedTables: + def test_managed_file_table(self): + table = LiteLLM_ManagedFileTable( + unified_file_id="f1", + model_mappings={"gpt-4": "file-abc"}, + flat_model_file_ids=["file-abc"], + ) + assert table.unified_file_id == "f1" + assert table.model_mappings == {"gpt-4": "file-abc"} + assert table.flat_model_file_ids == ["file-abc"] + + def test_managed_object_table_requires_purpose(self): + with pytest.raises(Exception): + LiteLLM_ManagedObjectTable( + unified_object_id="o1", model_object_id="m1", file_object={} + ) + + def test_managed_vector_stores_table(self): + table = LiteLLM_ManagedVectorStoresTable( + vector_store_id="vs1", + custom_llm_provider="openai", + vector_store_name=None, + vector_store_description=None, + vector_store_metadata=None, + created_at=None, + updated_at=None, + litellm_credential_name=None, + litellm_params=None, + team_id=None, + user_id=None, + ) + assert table.vector_store_id == "vs1" + assert table.custom_llm_provider == "openai" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index da66d60aed8..6fd935e3364 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -1,5 +1,6 @@ """Tests for MCP OAuth discoverable endpoints""" +import json from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -2661,3 +2662,74 @@ async def test_token_endpoint_sets_no_store_cache_control(): assert response.headers["cache-control"] == "no-store" assert response.headers["pragma"] == "no-cache" + + +async def _exchange_with_upstream_token_response(upstream_body): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="t", + name="t", + server_name="t", + alias="t", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + client_secret="cs", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + ) + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + fake_http_response = MagicMock() + fake_http_response.json.return_value = upstream_body + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ): + response = await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code="c", + redirect_uri="http://127.0.0.1:3000/cb", + client_id="cid", + client_secret=None, + code_verifier=None, + ) + return json.loads(response.body) + + +@pytest.mark.asyncio +async def test_token_exchange_omits_expires_in_when_upstream_omits_it(): + """A provider that issues a non-expiring token (e.g. Slack without token + rotation) returns no ``expires_in``. The exchange must mirror that and omit + ``expires_in`` rather than fabricate a 1-hour TTL, so the stored credential + is treated as non-expiring instead of dying after an hour.""" + body = await _exchange_with_upstream_token_response( + {"access_token": "tok", "token_type": "Bearer"} + ) + assert "expires_in" not in body + + +@pytest.mark.asyncio +async def test_token_exchange_passes_through_upstream_expires_in(): + """When the provider does send ``expires_in`` (e.g. Slack with token + rotation), the exchange forwards the real value unchanged.""" + body = await _exchange_with_upstream_token_response( + {"access_token": "tok", "token_type": "Bearer", "expires_in": 43200} + ) + assert body["expires_in"] == 43200 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py index 19065ff816b..a846ca24739 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py @@ -872,6 +872,7 @@ def _mock_env_vars_prisma(row=None): prisma.db.litellm_mcpuserenvvars.find_many = AsyncMock(return_value=[]) prisma.db.litellm_mcpuserenvvars.upsert = AsyncMock() prisma.db.litellm_mcpuserenvvars.delete_many = AsyncMock() + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock() return prisma @@ -1252,6 +1253,64 @@ async def test_delete_mcp_server_succeeds_when_orphan_cleanup_fails(): prisma.db.litellm_mcpuserenvvars.delete_many.assert_awaited_once() +@pytest.mark.asyncio +async def test_delete_mcp_server_removes_orphaned_user_credentials(): + """Deleting a server must also drop every user's stored BYOK/OAuth credential + rows for it; there is no FK cascade, so skipping this leaves encrypted secrets + pointing at a now-missing server.""" + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + prisma = _mock_env_vars_prisma() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=object()) + + await delete_mcp_server(prisma, "srv-1") + + prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once() + call = prisma.db.litellm_mcpusercredentials.delete_many.call_args + assert call.kwargs["where"] == {"server_id": "srv-1"} + + +@pytest.mark.asyncio +async def test_delete_mcp_server_skips_credential_cleanup_when_server_missing(): + """A no-op delete (server not found) must not touch the credential table.""" + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + prisma = _mock_env_vars_prisma() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=None) + + result = await delete_mcp_server(prisma, "srv-1") + + assert result is None + prisma.db.litellm_mcpusercredentials.delete_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_mcp_server_credential_cleanup_failure_still_cleans_env_vars(): + """Each per-user table is cleaned independently: a failure dropping credential + rows must not skip the env var cleanup (or vice versa), and the delete must + still succeed for the caller.""" + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + deleted = object() + prisma = _mock_env_vars_prisma() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=deleted) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock( + side_effect=Exception("connection pool exhausted") + ) + + result = await delete_mcp_server(prisma, "srv-1") + + assert result is deleted + prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once() + prisma.db.litellm_mcpuserenvvars.delete_many.assert_awaited_once() + + # ── DB helpers: global env vars encrypted at rest ───────────────────────── diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 0b1240f8bac..b6550fee6b9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -4970,17 +4970,143 @@ def test_no_contextvar_returns_default_options(self): try: from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_gateway_initialize_instructions, + _mcp_gateway_server_name, ) from litellm.proxy._experimental.mcp_server.server import server except ImportError: pytest.skip("MCP server not available") - tok = _mcp_gateway_initialize_instructions.set(None) + instructions_token = _mcp_gateway_initialize_instructions.set(None) + server_name_token = _mcp_gateway_server_name.set(None) try: opts = server.create_initialization_options() assert getattr(opts, "instructions", None) is None + assert opts.server_name == "litellm-mcp-server" finally: - _mcp_gateway_initialize_instructions.reset(tok) + _mcp_gateway_initialize_instructions.reset(instructions_token) + _mcp_gateway_server_name.reset(server_name_token) + + @pytest.mark.asyncio + async def test_scoped_request_uses_configured_server_alias(self): + try: + from litellm.proxy._experimental.mcp_server.server import ( + _gateway_initialize_instructions_request_scope, + global_mcp_server_manager, + server, + ) + except ImportError: + pytest.skip("MCP server not available") + + scoped_server = MCPServer( + server_id="server-123", + name="upstream-server", + alias="grafana", + transport=MCPTransport.http, + url="https://example.com/mcp", + ) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[scoped_server], + ), + patch.object( + global_mcp_server_manager, + "_ensure_upstream_initialize_instructions_cached", + new_callable=AsyncMock, + ), + ): + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth=None, + mcp_servers=["grafana"], + client_ip=None, + scoped_server_endpoint=True, + ): + assert server.create_initialization_options().server_name == "grafana" + + assert ( + server.create_initialization_options().server_name == "litellm-mcp-server" + ) + + @pytest.mark.asyncio + async def test_sse_handler_scopes_server_name_from_single_server_path(self): + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + global_mcp_server_manager, + handle_sse_mcp, + server, + ) + except ImportError: + pytest.skip("MCP server not available") + + scoped_server = MCPServer( + server_id="server-123", + name="upstream-server", + alias="grafana", + transport=MCPTransport.http, + url="https://example.com/mcp", + ) + captured = {} + + async def record_request(scope, receive, send): + captured["server_name"] = server.create_initialization_options().server_name + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/grafana", + "headers": [], + } + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=( + UserAPIKeyAuth(api_key="sk-test"), + None, + ["grafana"], + None, + None, + None, + ), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[scoped_server], + ), + patch.object( + global_mcp_server_manager, + "_ensure_upstream_initialize_instructions_cached", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._raise_preemptive_401_for_unauthenticated_servers", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._check_passthrough_upstream_auth", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + mcp_server.sse_session_manager, + "handle_request", + side_effect=record_request, + ), + ): + await handle_sse_mcp(scope, AsyncMock(), AsyncMock()) + + assert captured["server_name"] == "grafana" + assert ( + server.create_initialization_options().server_name == "litellm-mcp-server" + ) def test_contextvar_set_injects_instructions(self): """When ContextVar has a value, it appears in InitializationOptions.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 48c09f6e456..1b815b7a1c9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -30,6 +30,7 @@ MCPServerManager, _deserialize_json_dict, _deserialize_json_list, + _normalize_mcp_server_cost_info, ) from litellm.proxy._types import ( LiteLLM_MCPServerTable, @@ -257,6 +258,69 @@ async def test_load_servers_from_config_accepts_valid_alias(self, caplog): assert server.alias == "friendly_alias" assert server.server_name == "validserver" + @pytest.mark.asyncio + async def test_load_servers_from_config_coerces_cost_string_to_float(self): + """YAML 1.1 parses `7e-05` as a string; ingest must coerce it to float.""" + manager = MCPServerManager() + config = { + "google_maps": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + "mcp_info": { + "mcp_server_cost_info": { + "default_cost_per_query": "7e-05", + "tool_name_to_cost_per_query": {"geocode": "1e-3"}, + } + }, + } + } + + await manager.load_servers_from_config(config) + + server = next(iter(manager.config_mcp_servers.values())) + cost_info = server.mcp_info["mcp_server_cost_info"] + assert cost_info["default_cost_per_query"] == 7e-05 + assert isinstance(cost_info["default_cost_per_query"], float) + assert cost_info["tool_name_to_cost_per_query"]["geocode"] == 1e-3 + assert isinstance(cost_info["tool_name_to_cost_per_query"]["geocode"], float) + + def test_normalize_mcp_server_cost_info_preserves_float_values(self): + mcp_info = { + "server_name": "maps", + "mcp_server_cost_info": { + "default_cost_per_query": 0.01, + "tool_name_to_cost_per_query": {"search": 0.05}, + }, + } + + _normalize_mcp_server_cost_info(mcp_info) + + cost_info = mcp_info["mcp_server_cost_info"] + assert cost_info["default_cost_per_query"] == 0.01 + assert cost_info["tool_name_to_cost_per_query"] == {"search": 0.05} + + def test_normalize_mcp_server_cost_info_drops_non_numeric_values(self): + mcp_info = { + "server_name": "maps", + "mcp_server_cost_info": { + "default_cost_per_query": "not-a-number", + "tool_name_to_cost_per_query": {"search": "free", "geocode": "2e-4"}, + }, + } + + _normalize_mcp_server_cost_info(mcp_info) + + cost_info = mcp_info["mcp_server_cost_info"] + assert "default_cost_per_query" not in cost_info + assert cost_info["tool_name_to_cost_per_query"] == {"geocode": 2e-4} + + def test_normalize_mcp_server_cost_info_leaves_missing_cost_info_alone(self): + mcp_info = {"server_name": "maps"} + + _normalize_mcp_server_cost_info(mcp_info) + + assert "mcp_server_cost_info" not in mcp_info + def test_warns_when_custom_separator_invalid(self, monkeypatch, caplog): """Invalid MCP_TOOL_PREFIX_SEPARATOR values should log a warning.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index caff9ea2d28..47c9396f121 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -501,6 +501,7 @@ async def fake_get_tools( raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["called"] = True captured["server"] = server @@ -545,6 +546,78 @@ async def fake_get_tools( assert result["error"] is None assert result["message"] == "Successfully retrieved tools" + async def test_include_disabled_tools_is_admin_only(self, monkeypatch): + """include_disabled_tools skips the allowlist filter only for PROXY_ADMIN; + a non-admin passing it stays filtered so the REST endpoint can't be used + to enumerate deliberately-disabled tools.""" + from litellm.proxy._types import LitellmUserRoles + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + class StubServer: + alias = "server-1" + server_name = "server-1" + name = "stub" + allowed_tools = ["tool1"] + mcp_info = {"server_name": "stub"} + available_on_public_internet = True + + stub_server = StubServer() + captured = {} + + async def fake_get_tools( + server, server_auth_header, *args, apply_tool_filters=True, **kwargs + ): + captured["apply_tool_filters"] = apply_tool_filters + return ["tool-1"] + + monkeypatch.setattr( + rest_endpoints, + "build_effective_auth_contexts", + fake_contexts, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints, + "_get_tools_for_single_server", + fake_get_tools, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + + await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + include_disabled_tools=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert captured["apply_tool_filters"] is False + + await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + include_disabled_tools=True, + user_api_key_dict=UserAPIKeyAuth(), + ) + assert captured["apply_tool_filters"] is True + @pytest.mark.parametrize("upstream_status", [401, 403]) async def test_upstream_auth_failure_surfaces_status_and_challenge( self, monkeypatch, upstream_status @@ -649,6 +722,7 @@ async def fake_get_tools( raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["called"] = True captured["server_arg"] = server @@ -792,6 +866,7 @@ async def fake_get_tools( raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["server"] = server captured["auth_header"] = server_auth_header @@ -1284,6 +1359,56 @@ async def fake_get_tools_from_server(**kwargs): assert "tool1" not in tool_names assert "tool4" not in tool_names + async def test_apply_tool_filters_false_returns_full_catalog(self, monkeypatch): + """apply_tool_filters=False returns the raw catalog without the server + allowed_tools gate, so the config UI can render disabled tools as off.""" + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + class MockTool: + def __init__(self, name): + self.name = name + self.description = name + self.inputSchema = {} + + mock_tools = [MockTool("tool1"), MockTool("tool2"), MockTool("tool3")] + + async def fake_get_tools_from_server(**kwargs): + return mock_tools + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_get_tools_from_server", + fake_get_tools_from_server, + raising=False, + ) + + # Server enforces an allowlist of just tool1. + server = MCPServer( + server_id="test-server-id", + name="test-server", + transport=MCPTransport.sse, + allowed_tools=["tool1"], + ) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key", object_permission=None) + + # Runtime default: only the allowed tool comes back. + filtered = await rest_endpoints._get_tools_for_single_server( + server=server, + server_auth_header=None, + user_api_key_auth=user_api_key_dict, + ) + assert [t.name for t in filtered] == ["tool1"] + + # Config view: full catalog, including the disabled tools. + full = await rest_endpoints._get_tools_for_single_server( + server=server, + server_auth_header=None, + user_api_key_auth=user_api_key_dict, + apply_tool_filters=False, + ) + assert {t.name for t in full} == {"tool1", "tool2", "tool3"} + class TestStdioCommandAllowlist: """Tests for MCP stdio command allowlist validation.""" diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 42c76c4671d..e14ef05bd43 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2409,6 +2409,8 @@ async def mock_get_current_spend(counter_key, fallback_spend): assert exc_info.value.current_cost == 15.0 + + @pytest.mark.asyncio async def test_team_budget_check_reads_from_spend_counter(): """Team budget check should use get_current_spend when counter exists.""" diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 27f6015e6f4..11e6f483e35 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -112,6 +112,166 @@ async def test_handle_authentication_error_data_layer_errors_do_not_fall_back( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "db_error", + [ + ConnectionError("connection refused"), + TimeoutError("timed out"), + asyncio.TimeoutError(), + OSError("network is unreachable"), + HTTPClientClosedError(), + PrismaError("can't reach database server"), + RawQueryError( + data={ + "user_facing_error": { + "message": "cached plan must not change result type", + "meta": {"table": "t"}, + } + } + ), + ], +) +async def test_handle_authentication_error_db_infra_error_returns_503(db_error): + """Regression for the outage where valid keys got 401 for 4 hours: an + infrastructure-level DB failure during auth must surface as 503 (the DB + could not confirm the key), never as 401 ("Invalid API key").""" + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + db_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-valid-but-db-down", + ) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_handle_authentication_error_prisma_engine_teardown_returns_503(): + """Regression for the first-request-of-an-outage edge case: at the instant + the DB socket drops, the prisma query engine returns a malformed error + payload and prisma-client-py crashes with a bare + ``AttributeError: 'NoneType' object has no attribute 'get'`` before it can + raise P1001. That AttributeError reached auth and fell through to 401. It + must surface as 503 like every other infra failure during the outage.""" + from prisma.engine import utils as prisma_engine_utils + + malformed_payload = [ + { + "error": "Can't reach database server", + "user_facing_error": { + "error_code": "P1001", + "message": "Can't reach database server at `localhost`:`5503`", + "meta": None, + }, + } + ] + try: + prisma_engine_utils.handle_response_errors(None, malformed_payload) + raise AssertionError("expected prisma to raise AttributeError") + except AttributeError as e: + teardown_error = e + + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + teardown_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-valid-but-db-down", + ) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth_error", + [ + # DB returned no row -> get_key_object raises this exact 401. + ProxyException( + message="Authentication Error, Invalid proxy server token passed.", + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, + ), + # A bare auth failure raised as a plain Exception (e.g. master-key-only + # route) must keep returning 401, not get reclassified as 503. + Exception("Invalid proxy server token passed"), + ], +) +async def test_handle_authentication_error_genuine_auth_failure_stays_401(auth_error): + """Guard against the 503 conversion being too broad: a genuine auth + failure (missing key / wrong key) must still be 401.""" + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + auth_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-bad-key", + ) + + assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED + + @pytest.mark.asyncio async def test_handle_authentication_error_budget_exceeded(): handler = UserAPIKeyAuthExceptionHandler() diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index d4ca55ca16b..32b597376b4 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -1551,6 +1551,40 @@ def test_admin_opt_in_proxy_wide_still_allows(self): ) +class TestIsRequestBodySafeBlocksBedrockProjectOverride: + """``aws_bedrock_project_id`` pins a deployment to a Bedrock project so + that project's data-retention policy applies to its requests. A + caller-supplied value would run the request under any project reachable + with the deployment's shared AWS credentials, bypassing the configured + retention/accounting association.""" + + def test_project_id_in_request_body_is_rejected(self): + with pytest.raises(ValueError, match="aws_bedrock_project_id"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + "aws_bedrock_project_id": "proj_attacker000000", + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_admin_opt_in_proxy_wide_allows_project_id(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "aws_bedrock_project_id": "proj_byok000000", + }, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + # ── is_request_body_safe nested-config recursion (VERIA-6) ──────────────────── diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 7dadca24504..63510086f95 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -3182,6 +3182,310 @@ def test_build_decode_kwargs_no_warning_when_scoped( assert matching == [] +# --------------------------------------------------------------------------- +# Defer to single-team DB fallback (PR #26418) when JWT claims are present +# but do not resolve to a LiteLLM team. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_unresolved_claim_returns_none(): + """With `team_claim_fallback=True`: team_id claim is present in the JWT + but the team is missing in the DB — return (None, None) so the + auth_builder single-team fallback can run, instead of raising and + failing auth.""" + from fastapi import HTTPException + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + team_id_jwt_field="team_id", + team_claim_fallback=True, + ) + token = {"sub": "user-1", "team_id": "claim-team-not-in-db"} + + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + mock_get_team.side_effect = HTTPException(status_code=404, detail="missing") + + team_id, team_object = await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert team_id is None + assert team_object is None + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_unresolved_group_claim_returns_none( + monkeypatch, +): + """With `team_claim_fallback=True`: group claim resolves to team_ids that + don't exist in the DB — return (None, None) instead of raising 403, so + the single-team fallback can run.""" + import sys + import types + + from fastapi import HTTPException + + from litellm.router import Router + + router = Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}} + ] + ) + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + async def raise_404(*_args, **_kwargs): + raise HTTPException(status_code=404, detail="missing") + + monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", raise_404) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_claim_fallback=True) + + team_id, team_object = await JWTAuthManager.find_team_with_model_access( + team_ids={"idp-group-a", "idp-group-b"}, + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert team_id is None + assert team_object is None + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_non_http_exception_still_propagates(): + """Regression guard: only the 404 HTTPException raised by + `get_team_object` ("team doesn't exist in db") is softened. Other + errors — e.g. "No DB Connected" — must still propagate so operator-side + problems are loud.""" + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id") + token = {"sub": "user-1", "team_id": "some-claim-team"} + + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + mock_get_team.side_effect = RuntimeError("simulated infrastructure error") + + with pytest.raises(RuntimeError, match="simulated infrastructure error"): + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_non_404_http_exception_propagates(): + """Regression guard: only 404 HTTPException is softened. If + `get_team_object` is ever updated to raise a different HTTP status code + (e.g. 403 for a blocked team), that error must still propagate rather + than silently fall through to the single-team DB fallback.""" + from fastapi import HTTPException + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id") + token = {"sub": "user-1", "team_id": "some-claim-team"} + + for status_code in (400, 403, 500): + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + mock_get_team.side_effect = HTTPException( + status_code=status_code, detail="non-404 failure" + ) + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert exc_info.value.status_code == status_code + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_enforce_team_based_access_still_raises(): + """Regression guard: when no group claims are present and + `enforce_team_based_model_access` is on, the original 403 still fires — + the new soft-fail only applies to the unresolved-claim path inside the + loop, not to the no-team-claims-at-all path at the top.""" + from fastapi import HTTPException + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(enforce_team_based_model_access=True) + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_team_with_model_access( + team_ids=set(), + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert exc_info.value.status_code == 403 + assert "enforce_team_based_model_access" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_resolved_team_without_model_still_raises_403( + monkeypatch, +): + """Regression guard: when the JWT group claim DOES resolve to a real + LiteLLM team but that team does not grant the requested model, keep the + original 403. Only the unresolved-claim case is softened.""" + import sys + import types + + from fastapi import HTTPException + + from litellm.router import Router + + router = Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + }, + ] + ) + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + team = LiteLLM_TeamTable(team_id="real-team", models=["gpt-3.5-turbo"]) + + async def mock_get_team_object(*_args, **_kwargs): + return team + + monkeypatch.setattr( + "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object + ) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_team_with_model_access( + team_ids={"real-team"}, + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert exc_info.value.status_code == 403 + assert "No team has access to the requested model" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_unresolved_claim_default_raises(): + """Default `team_claim_fallback=False`: unresolved team_id claim must + still raise — preserves the strict claim-based authorization boundary + when the operator has not opted in to the fallback.""" + from fastapi import HTTPException + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id") + token = {"sub": "user-1", "team_id": "claim-team-not-in-db"} + + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + mock_get_team.side_effect = HTTPException(status_code=404, detail="missing") + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_unresolved_group_claim_default_raises( + monkeypatch, +): + """Default `team_claim_fallback=False`: group claims that don't resolve + to any LiteLLM team must still raise 403 — preserves the strict + claim-based authorization boundary.""" + import sys + import types + + from fastapi import HTTPException + + from litellm.router import Router + + router = Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}} + ] + ) + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + async def raise_404(*_args, **_kwargs): + raise HTTPException(status_code=404, detail="missing") + + monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", raise_404) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_team_with_model_access( + team_ids={"idp-group-a", "idp-group-b"}, + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert exc_info.value.status_code == 403 + + # GH #26789: JWT claim user_id must rebind to legacy DB row after fuzzy match. 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 fa6cc8bed1b..80f12d4459f 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 @@ -112,11 +112,71 @@ async def test_should_clear_stale_budget_reservation_when_budget_checks_skip(): user_api_key_cache=MagicMock(), proxy_logging_obj=MagicMock(), skip_budget_checks=True, + general_settings={}, ) assert user_api_key_auth_obj.budget_reservation is None +@pytest.mark.asyncio +async def test_disable_budget_reservation_skips_reservation(): + """#27639: general_settings.disable_budget_reservation turns off the optimistic Redis + reservation so operators hit by phantom BudgetExceededError can opt out of it.""" + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value={"reserved_cost": 0.5, "entries": []}), + ) as mock_reserve: + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={"disable_budget_reservation": True}, + ) + + mock_reserve.assert_not_called() + assert user_api_key_auth_obj.budget_reservation is None + + +@pytest.mark.asyncio +async def test_budget_reservation_runs_when_not_disabled(): + """Control for #27639: with the flag absent, the reservation still runs and is stored.""" + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + reservation = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:test_token"}], + } + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value=reservation), + ) as mock_reserve: + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={}, + ) + + mock_reserve.assert_awaited_once() + assert user_api_key_auth_obj.budget_reservation == reservation + + @pytest.mark.asyncio async def test_should_not_reuse_cached_key_object_for_request_state(): key_cache = DualCache() @@ -1636,7 +1696,9 @@ async def test_auto_register_passes_validated_org_context_to_generated_key(self) assert mock_auto_register.call_args.kwargs["team_id"] == "validated-team" assert mock_auto_register.call_args.kwargs["user_id"] == "validated-user" assert mock_auto_register.call_args.kwargs["org_id"] == "validated-org" - assert mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user" + assert ( + mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user" + ) assert result.org_id == "validated-org" @pytest.mark.asyncio @@ -3548,3 +3610,118 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() finally: for k, v in originals.items(): setattr(_proxy_server_mod, k, v) + + +def _proxy_attrs_for_db_lookup(): + """Minimal proxy_server attributes for driving the real + ``_user_api_key_auth_builder`` down to the DB key lookup.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + return { + "prisma_client": MagicMock(), + "user_api_key_cache": DualCache(), + "proxy_logging_obj": proxy_logging_obj, + "master_key": "sk-test-master", + "general_settings": {"allow_requests_on_db_unavailable": False}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + + +async def _run_builder_with_key_lookup(get_key_object_mock): + """Drive the real auth builder with ``get_key_object`` replaced by the + given mock. Returns the builder result. Patches ``seed_request_identity`` + so the failure path doesn't touch OTEL.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + attrs = _proxy_attrs_for_db_lookup() + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", + get_key_object_mock, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + ): + return await _user_api_key_auth_builder( + request=request, + api_key="Bearer sk-db-lookup-test", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_builder_returns_503_when_db_lookup_raises_infra_error(): + """End-to-end: a DB infrastructure failure during the key lookup must + propagate past the ``except ProxyException`` guard and surface as 503, + not the 401 that masked the 4-hour outage. Killing the new 503 branch + flips this to 401 and fails the test.""" + get_key_object = AsyncMock(side_effect=ConnectionError("connection refused")) + + with pytest.raises(ProxyException) as exc_info: + await _run_builder_with_key_lookup(get_key_object) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_builder_returns_401_when_db_lookup_reports_missing_key(): + """Regression guard: a genuinely missing key (DB returned no row, which + ``get_key_object`` raises as a 401 ProxyException) must still be 401.""" + missing_key_error = ProxyException( + message="Authentication Error, Invalid proxy server token passed. key=..., not found in db.", + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, + ) + get_key_object = AsyncMock(side_effect=missing_key_error) + + with pytest.raises(ProxyException) as exc_info: + await _run_builder_with_key_lookup(get_key_object) + + assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED + + +@pytest.mark.asyncio +async def test_builder_succeeds_when_db_lookup_returns_valid_token(): + """Regression guard: a valid key still authenticates. Proves the 503 + conversion only fires on the failure path and never intercepts success.""" + valid_token = UserAPIKeyAuth(api_key="sk-db-lookup-test", token="hashed-valid") + get_key_object = AsyncMock(return_value=valid_token) + + with patch( + "litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj", + new_callable=AsyncMock, + return_value=valid_token, + ) as mock_return: + result = await _run_builder_with_key_lookup(get_key_object) + + assert isinstance(result, UserAPIKeyAuth) + # Reaching the success-assembly return (never the exception handler) + # proves a valid key is unaffected by the 503 conversion. + mock_return.assert_awaited_once() diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py new file mode 100644 index 00000000000..afd1696a89f --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -0,0 +1,475 @@ +import os +import sys +from unittest.mock import patch + +import click +import pytest +import requests +from click.testing import CliRunner + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + + +from litellm.proxy.client.cli.commands.agents import ( + AgentRunError, + agent_commands, + agent_launch_args, + agent_profile, + build_agent_env, + run_agent, + verify_proxy_key, +) + +AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents" + + +def _agent_command(name): + return next(c for c in agent_commands() if c.name == name) + + +class _FakeResponse: + def __init__(self, status_code): + self.status_code = status_code + + +class TestAgentProfile: + def test_claude_is_anthropic(self): + name, profiles = agent_profile("claude") + assert name == "Claude Code" + assert profiles == frozenset({"anthropic"}) + + def test_claude_full_path_uses_basename(self): + name, profiles = agent_profile("/usr/local/bin/claude") + assert name == "Claude Code" + assert profiles == frozenset({"anthropic"}) + + def test_codex_and_opencode_are_openai(self): + assert agent_profile("codex") == ("Codex", frozenset({"openai"})) + assert agent_profile("opencode") == ("OpenCode", frozenset({"openai"})) + + def test_unknown_command_gets_both_profiles(self): + name, profiles = agent_profile("mytool") + assert name == "mytool" + assert profiles == frozenset({"anthropic", "openai"}) + + +class TestBuildAgentEnv: + def test_anthropic_profile_uses_bare_root_and_bearer(self): + env = build_agent_env( + {}, "http://localhost:4000/", "sk-key", frozenset({"anthropic"}) + ) + assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" + assert "OPENAI_BASE_URL" not in env + assert "OPENAI_API_KEY" not in env + + def test_anthropic_profile_drops_existing_api_key(self): + env = build_agent_env( + {"ANTHROPIC_API_KEY": "real-key"}, + "http://localhost:4000", + "sk-key", + frozenset({"anthropic"}), + ) + assert "ANTHROPIC_API_KEY" not in env + + def test_openai_profile_appends_v1(self): + env = build_agent_env( + {}, "http://localhost:4000/", "sk-key", frozenset({"openai"}) + ) + assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1" + assert env["OPENAI_API_KEY"] == "sk-key" + assert "ANTHROPIC_BASE_URL" not in env + + def test_both_profiles_set_everything(self): + env = build_agent_env( + {}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"}) + ) + assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1" + assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" + assert env["OPENAI_API_KEY"] == "sk-key" + + def test_preserves_unrelated_env_and_does_not_mutate_input(self): + base = {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"} + env = build_agent_env( + base, "http://localhost:4000", "sk-key", frozenset({"anthropic"}) + ) + assert env["PATH"] == "/usr/bin" + assert base == {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"} + + +class TestAgentLaunchArgs: + def test_claude_and_opencode_get_no_extra_args(self): + assert agent_launch_args("claude", "http://localhost:4000") == [] + assert agent_launch_args("opencode", "http://localhost:4000") == [] + + def test_unknown_agent_gets_no_extra_args(self): + assert agent_launch_args("mytool", "http://localhost:4000") == [] + + def test_codex_points_provider_at_proxy_over_http(self): + args = agent_launch_args("codex", "http://localhost:4000/") + joined = " ".join(args) + assert 'model_provider="litellm"' in args + assert 'model_providers.litellm.base_url="http://localhost:4000/v1"' in args + assert 'model_providers.litellm.env_key="OPENAI_API_KEY"' in args + assert 'model_providers.litellm.wire_api="responses"' in args + assert "model_providers.litellm.supports_websockets=false" in args + assert joined.count("-c") == 6 + + def test_codex_uses_basename(self): + assert agent_launch_args("/usr/local/bin/codex", "http://localhost:4000") == ( + agent_launch_args("codex", "http://localhost:4000") + ) + + +class TestVerifyProxyKey: + def test_ok_status_passes_and_uses_models_endpoint(self): + captured = {} + + def fake_get(url, headers, timeout): + captured["url"] = url + captured["headers"] = headers + return _FakeResponse(200) + + verify_proxy_key("http://localhost:4000/", "sk-key", get=fake_get) + + assert captured["url"] == "http://localhost:4000/v1/models" + assert captured["headers"] == {"Authorization": "Bearer sk-key"} + + @pytest.mark.parametrize("status", [401, 403]) + def test_rejected_key_raises(self, status): + with pytest.raises(AgentRunError, match="rejected your key"): + verify_proxy_key( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse(status), + ) + + def test_unreachable_proxy_raises(self): + def boom(*a, **k): + raise requests.ConnectionError("refused") + + with pytest.raises(AgentRunError, match="Could not reach"): + verify_proxy_key("http://localhost:4000", "sk-key", get=boom) + + def test_other_non_2xx_is_tolerated(self): + verify_proxy_key( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse(500), + ) + + +class TestRunAgent: + def test_wires_env_and_launches_resolved_binary(self): + calls = {} + + def fake_launcher(path, args, env): + calls["path"] = path + calls["args"] = tuple(args) + calls["env"] = dict(env) + + run_agent( + "http://localhost:4000", + "sk-key", + ["claude", "--resume"], + base_env={"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "leaked"}, + which=lambda name: "/usr/local/bin/claude", + verify=lambda *a: None, + launcher=fake_launcher, + ) + + assert calls["path"] == "/usr/local/bin/claude" + assert calls["args"] == ("claude", "--resume") + env = calls["env"] + assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" + assert "ANTHROPIC_API_KEY" not in env + assert "OPENAI_BASE_URL" not in env + + def test_codex_gets_openai_env(self): + calls = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex"], + base_env={}, + which=lambda name: "/usr/local/bin/codex", + verify=lambda *a: None, + launcher=lambda p, a, e: calls.update(env=dict(e)), + ) + assert calls["env"]["OPENAI_BASE_URL"] == "http://localhost:4000/v1" + assert calls["env"]["OPENAI_API_KEY"] == "sk-key" + assert "ANTHROPIC_BASE_URL" not in calls["env"] + + def test_codex_injects_proxy_provider_args_before_user_args(self): + calls = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex", "exec", "do a thing"], + base_env={}, + which=lambda name: "/usr/local/bin/codex", + verify=lambda *a: None, + launcher=lambda p, a, e: calls.update(args=tuple(a)), + ) + args = calls["args"] + assert args[0] == "codex" + assert args[-2:] == ("exec", "do a thing") + assert 'model_provider="litellm"' in args + assert 'model_providers.litellm.base_url="http://localhost:4000/v1"' in args + # overrides must precede the codex subcommand so codex parses them + assert args.index('model_provider="litellm"') < args.index("exec") + + def test_claude_launches_without_injected_args(self): + calls = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["claude", "--resume"], + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify=lambda *a: None, + launcher=lambda p, a, e: calls.update(args=tuple(a)), + ) + assert calls["args"] == ("claude", "--resume") + + def test_missing_binary_raises_with_install_hint(self): + with pytest.raises(AgentRunError, match="claude.*Install it first"): + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + base_env={}, + which=lambda name: None, + verify=lambda *a: None, + launcher=lambda *a: None, + ) + + def test_skip_verify_does_not_call_verify(self): + verified = [] + launched = [] + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + skip_verify=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify=lambda *a: verified.append(a), + launcher=lambda *a: launched.append(a), + ) + assert verified == [] + assert len(launched) == 1 + + def test_verify_failure_aborts_before_launch(self): + launched = [] + + def boom(*a): + raise AgentRunError("rejected") + + with pytest.raises(AgentRunError): + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify=boom, + launcher=lambda *a: launched.append(a), + ) + assert launched == [] + + def test_empty_command_raises(self): + with pytest.raises(AgentRunError): + run_agent("http://localhost:4000", "sk-key", []) + + def test_reattach_terminal_runs_just_before_launch(self): + order = [] + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + skip_verify=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + launcher=lambda *a: order.append("launch"), + reattach_terminal=lambda: order.append("reattach"), + ) + assert order == ["reattach", "launch"] + + def test_no_reattach_terminal_by_default(self): + order = [] + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + skip_verify=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + launcher=lambda *a: order.append("launch"), + ) + assert order == ["launch"] + + +class TestAgentCommands: + def setup_method(self): + self.runner = CliRunner() + + def test_one_command_per_known_agent(self): + assert {c.name for c in agent_commands()} == {"claude", "codex", "opencode"} + + def test_claude_launches_with_stored_key_and_forwards_args(self): + captured = {} + + def fake_run_agent(base_url, api_key, command, **kwargs): + captured["base_url"] = base_url + captured["api_key"] = api_key + captured["command"] = list(command) + captured["skip_verify"] = kwargs.get("skip_verify") + + with patch(f"{AGENTS_MODULE}.run_agent", side_effect=fake_run_agent): + result = self.runner.invoke( + _agent_command("claude"), + ["--resume", "-p", "hi"], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + + assert result.exit_code == 0, result.output + assert captured["api_key"] == "sk-key" + assert captured["command"] == ["claude", "--resume", "-p", "hi"] + assert captured["skip_verify"] is False + assert ( + "routing Claude Code through proxy at http://localhost:4000" + in result.output + ) + + def test_codex_shows_friendly_name(self): + captured = {} + with patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(command=list(c)), + ): + result = self.runner.invoke( + _agent_command("codex"), + ["exec", "do a thing"], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + assert result.exit_code == 0, result.output + assert captured["command"] == ["codex", "exec", "do a thing"] + assert "routing Codex through proxy" in result.output + + def test_skip_verify_is_consumed_not_forwarded(self): + captured = {} + + def fake_run_agent(base_url, api_key, command, **kwargs): + captured["command"] = list(command) + captured["skip_verify"] = kwargs.get("skip_verify") + + with patch(f"{AGENTS_MODULE}.run_agent", side_effect=fake_run_agent): + result = self.runner.invoke( + _agent_command("claude"), + ["--skip-verify", "--resume"], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + + assert result.exit_code == 0, result.output + assert captured["skip_verify"] is True + assert captured["command"] == ["claude", "--resume"] + + def test_non_interactive_without_key_errors_clearly(self): + with ( + patch(f"{AGENTS_MODULE}._is_interactive", return_value=False), + patch(f"{AGENTS_MODULE}.run_agent") as mock_run, + ): + result = self.runner.invoke( + _agent_command("claude"), + [], + obj={"base_url": "http://localhost:4000", "api_key": None}, + ) + assert result.exit_code != 0 + assert "LITELLM_PROXY_API_KEY" in result.output + mock_run.assert_not_called() + + def test_interactive_without_key_logs_in_then_launches(self): + captured = {} + + @click.command() + def fake_login(): + pass + + with ( + patch(f"{AGENTS_MODULE}._is_interactive", return_value=True), + patch(f"{AGENTS_MODULE}.login", fake_login), + patch( + f"{AGENTS_MODULE}.get_stored_api_key", return_value="sk-after-login" + ) as mock_get, + patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda base_url, api_key, command, **k: captured.update( + api_key=api_key + ), + ), + ): + result = self.runner.invoke( + _agent_command("claude"), + [], + obj={"base_url": "http://localhost:4000", "api_key": None}, + ) + + assert result.exit_code == 0, result.output + assert captured["api_key"] == "sk-after-login" + mock_get.assert_called_once_with(expected_base_url="http://localhost:4000") + + def test_agent_run_error_becomes_click_error(self): + with patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=AgentRunError("could not reach proxy"), + ): + result = self.runner.invoke( + _agent_command("claude"), + [], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + assert result.exit_code != 0 + assert "could not reach proxy" in result.output + + def test_interactive_session_reattaches_terminal_before_handoff(self): + from litellm.proxy.client.cli.commands.agents import ( + _restore_controlling_terminal, + ) + + captured = {} + with ( + patch(f"{AGENTS_MODULE}._is_interactive", return_value=True), + patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(kw), + ), + ): + result = self.runner.invoke( + _agent_command("claude"), + [], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + assert result.exit_code == 0, result.output + assert captured["reattach_terminal"] is _restore_controlling_terminal + + def test_non_interactive_agent_mode_leaves_stdin_alone(self): + captured = {} + with ( + patch(f"{AGENTS_MODULE}._is_interactive", return_value=False), + patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(kw), + ), + ): + result = self.runner.invoke( + _agent_command("claude"), + [], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + assert result.exit_code == 0, result.output + assert captured["reattach_terminal"] is None diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 2e738ff900d..4ee8b502aa2 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -517,7 +517,7 @@ def test_whoami_not_authenticated(self): assert result.exit_code == 0 assert "❌ Not authenticated" in result.output - assert "Run 'litellm-proxy login'" in result.output + assert "Run 'lite login'" in result.output def test_whoami_old_token(self): """Test whoami with old token showing warning""" diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index d328d68dcd4..36ff3f3c399 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -1,7 +1,9 @@ import copy import sys import os -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace + +import pytest sys.path.insert( 0, os.path.abspath("../../..") @@ -309,3 +311,98 @@ def test_encrypt_callback_vars_only_encrypts_credential_fields(monkeypatch): assert cv["langfuse_host"] == "https://cloud.langfuse.com" assert cv["langsmith_project"] == "my-proj" assert cv["langsmith_base_url"] == "https://smith.example" + + +def test_initialize_callbacks_on_proxy_lakera_ignores_non_dict_callback_settings( + monkeypatch, +): + """Regression: a non-dict value under callback_settings.lakera_prompt_injection + must not crash initialize_callbacks_on_proxy. + + Forwarding callback_settings as callback_specific_params (so callbacks like + DatadogCostManagementLogger receive their init params) exposes the lakera + branch, which previously did lakeraAI_Moderation(**callback_specific_params[ + "lakera_prompt_injection"]) with no isinstance(dict) guard. For a config like + {"lakera_prompt_injection": "x"} that is `**"x"` -> TypeError: argument after + ** must be a mapping, not str. The branch now guards on isinstance(dict), + matching the presidio / datadog_cost_management branches. + """ + captured = {} + + class _DummyLakera: + def __init__(self, **kwargs): + captured["kwargs"] = kwargs + + # Inject a fake lakera_ai module so the branch's + # `from ...lakera_ai import lakeraAI_Moderation` resolves to our stub without + # importing the real module (which imports proxy_server symbols not present + # under the stubbed proxy_server below). + fake_lakera = ModuleType("litellm.proxy.guardrails.guardrail_hooks.lakera_ai") + fake_lakera.lakeraAI_Moderation = _DummyLakera + monkeypatch.setitem( + sys.modules, + "litellm.proxy.guardrails.guardrail_hooks.lakera_ai", + fake_lakera, + ) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + SimpleNamespace(prisma_client=None), + ) + + original_callbacks = ( + list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] + ) + litellm.callbacks = [] + try: + # A non-dict value must be ignored (init_params stays {}), not **-unpacked. + initialize_callbacks_on_proxy( + value=["lakera_prompt_injection"], + premium_user=False, + config_file_path=".", + litellm_settings={}, + callback_specific_params={"lakera_prompt_injection": "any-string"}, + ) + assert captured["kwargs"] == {} + assert any(isinstance(c, _DummyLakera) for c in litellm.callbacks) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.parametrize("bad_root", [None, True]) +def test_initialize_callbacks_on_proxy_non_dict_callback_specific_params_root( + monkeypatch, bad_root +): + """Regression: a blank `callback_settings:` key in YAML loads as None (and + `callback_settings: true` as a bool); load_config forwards that value + verbatim as callback_specific_params. Membership tests like + `"compression_interception" in callback_specific_params` then raise + TypeError and abort proxy startup. A non-dict root must be normalized to {} + so the callback initializes with its defaults. + """ + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + SimpleNamespace(prisma_client=None), + ) + from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, + ) + + original_callbacks = ( + list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] + ) + litellm.callbacks = [] + try: + initialize_callbacks_on_proxy( + value=["compression_interception"], + premium_user=False, + config_file_path=".", + litellm_settings={}, + callback_specific_params=bad_root, + ) + assert any( + isinstance(c, CompressionInterceptionLogger) for c in litellm.callbacks + ) + finally: + litellm.callbacks = original_callbacks diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index f4bf0d7b2be..e9b4f11e891 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -1,5 +1,6 @@ # tests/litellm/proxy/common_utils/test_upsert_budget_membership.py import types +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock import pytest @@ -19,15 +20,13 @@ def mock_tx(): Builds an object that looks just enough like the Prisma tx you use inside _upsert_budget_and_membership. """ - # membership “table” membership = MagicMock() membership.update = AsyncMock() membership.upsert = AsyncMock() - # budget “table” budget = MagicMock() budget.update = AsyncMock() - # budget.create returns a fake row that has .budget_id + budget.find_unique = AsyncMock(return_value=None) budget.create = AsyncMock( return_value=types.SimpleNamespace(budget_id="new-budget-123") ) @@ -44,223 +43,173 @@ def fake_user(): return types.SimpleNamespace(user_id="tester@example.com") -# TEST: max_budget is None, disconnect only +def budget_row(**fields): + """A fake litellm_budgettable row whose model_dump returns the given fields.""" + row = MagicMock() + row.model_dump.return_value = fields + return row + + +def assert_future_reset_time(value): + """A budget_reset_at must be a timezone-aware datetime in the future, so the + member's budget actually rolls over and the UI shows a reset date instead of + waiting for the reset cron to backfill it.""" + assert isinstance(value, datetime) + assert value.tzinfo is not None + assert value > datetime.now(timezone.utc) + + +# TEST: an empty patch (caller sent no budget fields) leaves everything alone. +# This is the merge-patch contract: absent != clear. Updating only a member's +# role must not silently wipe their budget. @pytest.mark.asyncio -async def test_upsert_disconnect(mock_tx, fake_user): +async def test_empty_patch_is_noop(mock_tx, fake_user): await _upsert_budget_and_membership( mock_tx, team_id="team-1", user_id="user-1", - max_budget=None, - existing_budget_id=None, + existing_budget_id="bud-1", user_api_key_dict=fake_user, + budget_patch={}, ) - mock_tx.litellm_teammembership.update.assert_awaited_once_with( - where={"user_id_team_id": {"user_id": "user-1", "team_id": "team-1"}}, - data={"litellm_budget_table": {"disconnect": True}}, - ) + mock_tx.litellm_teammembership.update.assert_not_called() + mock_tx.litellm_teammembership.upsert.assert_not_called() mock_tx.litellm_budgettable.update.assert_not_called() mock_tx.litellm_budgettable.create.assert_not_called() - mock_tx.litellm_teammembership.upsert.assert_not_called() -# TEST: existing budget id → updates budget in-place (current behavior) +# TEST: clearing every limit on a member's private budget disconnects it, so the +# member falls back to the team default instead of keeping an empty private row. @pytest.mark.asyncio -async def test_upsert_with_existing_budget_id_creates_new(mock_tx, fake_user): - """ - Test that when existing_budget_id is provided, the function updates the budget in-place. - """ +async def test_clearing_all_limits_disconnects(mock_tx, fake_user): + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(max_budget=100.0) + ) + await _upsert_budget_and_membership( mock_tx, - team_id="team-2", - user_id="user-2", - max_budget=42.0, - existing_budget_id="bud-999", + team_id="team-1", + user_id="user-1", + existing_budget_id="bud-1", user_api_key_dict=fake_user, + budget_patch={"max_budget": None}, ) - # Should update the existing budget, not create a new one - mock_tx.litellm_budgettable.update.assert_awaited_once_with( - where={"budget_id": "bud-999"}, - data={ - "max_budget": 42.0, - "updated_by": fake_user.user_id, - }, + mock_tx.litellm_teammembership.update.assert_awaited_once_with( + where={"user_id_team_id": {"user_id": "user-1", "team_id": "team-1"}}, + data={"litellm_budget_table": {"disconnect": True}}, ) - - # Should NOT create a new budget or touch membership + mock_tx.litellm_budgettable.update.assert_not_called() mock_tx.litellm_budgettable.create.assert_not_called() - mock_tx.litellm_teammembership.upsert.assert_not_called() - mock_tx.litellm_teammembership.update.assert_not_called() -# TEST: create new budget and link membership +# TEST: clearing one field on a budget that still has another limit updates in +# place (clears just that column + its reset time) and does NOT disconnect. @pytest.mark.asyncio -async def test_upsert_create_and_link(mock_tx, fake_user): +async def test_clear_one_field_keeps_others(mock_tx, fake_user): + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(max_budget=100.0, budget_duration="24h") + ) + await _upsert_budget_and_membership( mock_tx, - team_id="team-3", - user_id="user-3", - max_budget=99.9, - existing_budget_id=None, + team_id="team-1", + user_id="user-1", + existing_budget_id="bud-1", user_api_key_dict=fake_user, + budget_patch={"budget_duration": None}, ) - mock_tx.litellm_budgettable.create.assert_awaited_once_with( + mock_tx.litellm_teammembership.update.assert_not_called() + mock_tx.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "bud-1"}, data={ - "max_budget": 99.9, - "created_by": fake_user.user_id, "updated_by": fake_user.user_id, - }, - include={"team_membership": True}, - ) - - # Budget ID returned by the mocked create() - bid = mock_tx.litellm_budgettable.create.return_value.budget_id - - mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( - where={"user_id_team_id": {"user_id": "user-3", "team_id": "team-3"}}, - data={ - "create": { - "user_id": "user-3", - "team_id": "team-3", - "litellm_budget_table": {"connect": {"budget_id": bid}}, - }, - "update": { - "litellm_budget_table": {"connect": {"budget_id": bid}}, - }, + "budget_duration": None, + "budget_reset_at": None, }, ) - mock_tx.litellm_teammembership.update.assert_not_called() - mock_tx.litellm_budgettable.update.assert_not_called() - -# TEST: create new budget and link membership, then create another new budget +# TEST: setting budget_duration in place writes the duration AND a future +# budget_reset_at, so the budget rolls over without waiting for the reset cron. @pytest.mark.asyncio -async def test_upsert_create_then_create_another(mock_tx, fake_user): - """ - Test that multiple calls to _upsert_budget_and_membership create separate budgets, - reflecting the current implementation behavior. - """ - # FIRST CALL – create new budget and link membership - await _upsert_budget_and_membership( - mock_tx, - team_id="team-42", - user_id="user-42", - max_budget=10.0, - existing_budget_id=None, - user_api_key_dict=fake_user, +async def test_update_in_place_seeds_reset_at(mock_tx, fake_user): + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(max_budget=20.0) ) - # capture the budget id that create() returned - created_bid = mock_tx.litellm_budgettable.create.return_value.budget_id - - # sanity: we really did the create + upsert path - mock_tx.litellm_budgettable.create.assert_awaited_once() - mock_tx.litellm_teammembership.upsert.assert_awaited_once() - - # SECOND CALL – reset call history; this time we supply the existing budget_id - mock_tx.litellm_budgettable.create.reset_mock() - mock_tx.litellm_teammembership.upsert.reset_mock() - mock_tx.litellm_budgettable.update.reset_mock() - await _upsert_budget_and_membership( mock_tx, - team_id="team-42", - user_id="user-42", - max_budget=25.0, - existing_budget_id=created_bid, # now used: triggers in-place update + team_id="team-dur", + user_id="user-dur", + existing_budget_id="bud-dur", user_api_key_dict=fake_user, + budget_patch={"budget_duration": "30d"}, ) - # Should update the existing budget in-place, not create a new one - mock_tx.litellm_budgettable.update.assert_awaited_once_with( - where={"budget_id": created_bid}, - data={ - "max_budget": 25.0, - "updated_by": fake_user.user_id, - }, - ) - - # Should NOT create a new budget or touch membership + mock_tx.litellm_budgettable.update.assert_awaited_once() + call = mock_tx.litellm_budgettable.update.await_args + assert call.kwargs["where"] == {"budget_id": "bud-dur"} + data = call.kwargs["data"] + assert data["budget_duration"] == "30d" + assert data["updated_by"] == fake_user.user_id + assert_future_reset_time(data["budget_reset_at"]) mock_tx.litellm_budgettable.create.assert_not_called() - mock_tx.litellm_teammembership.upsert.assert_not_called() -# TEST: update rpm_limit for member with existing budget_id → updates in-place +# TEST: updating a single limit in place only writes that field; an untouched +# budget_duration must not get a (re)computed reset time. @pytest.mark.asyncio -async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user): - """ - Test that updating rpm_limit for a member with an existing budget_id - updates the existing budget in-place (not creates a new one). - """ - existing_budget_id = "existing-budget-456" +async def test_update_in_place_single_field_leaves_reset_at_alone(mock_tx, fake_user): + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(max_budget=50.0) + ) await _upsert_budget_and_membership( mock_tx, - team_id="team-rpm-test", - user_id="user-rpm-test", - max_budget=50.0, - existing_budget_id=existing_budget_id, + team_id="team-rpm", + user_id="user-rpm", + existing_budget_id="bud-rpm", user_api_key_dict=fake_user, - tpm_limit=1000, - rpm_limit=100, + budget_patch={"rpm_limit": 100}, ) - # Should update the existing budget with all specified limits mock_tx.litellm_budgettable.update.assert_awaited_once_with( - where={"budget_id": existing_budget_id}, - data={ - "max_budget": 50.0, - "tpm_limit": 1000, - "rpm_limit": 100, - "updated_by": fake_user.user_id, - }, + where={"budget_id": "bud-rpm"}, + data={"updated_by": fake_user.user_id, "rpm_limit": 100}, ) - - # Should NOT create a new budget or touch membership mock_tx.litellm_budgettable.create.assert_not_called() - mock_tx.litellm_teammembership.upsert.assert_not_called() -# TEST: create new budget with only rpm_limit (no max_budget) +# TEST: with no existing budget, a duration-only patch creates a budget carrying +# the duration and a future reset time, then links the membership. @pytest.mark.asyncio -async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user): - """ - Test that setting only rpm_limit creates a new budget with just the rpm_limit. - """ +async def test_create_seeds_reset_at_and_links(mock_tx, fake_user): await _upsert_budget_and_membership( mock_tx, - team_id="team-rpm-only", - user_id="user-rpm-only", - max_budget=None, + team_id="team-new", + user_id="user-new", existing_budget_id=None, user_api_key_dict=fake_user, - rpm_limit=50, + budget_patch={"budget_duration": "7d"}, ) - # Should create a new budget with only rpm_limit - mock_tx.litellm_budgettable.create.assert_awaited_once_with( - data={ - "rpm_limit": 50, - "created_by": fake_user.user_id, - "updated_by": fake_user.user_id, - }, - include={"team_membership": True}, - ) + mock_tx.litellm_budgettable.create.assert_awaited_once() + data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert data["budget_duration"] == "7d" + assert data["created_by"] == fake_user.user_id + assert data["updated_by"] == fake_user.user_id + assert_future_reset_time(data["budget_reset_at"]) - # Should upsert team membership with the new budget ID new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( - where={ - "user_id_team_id": {"user_id": "user-rpm-only", "team_id": "team-rpm-only"} - }, + where={"user_id_team_id": {"user_id": "user-new", "team_id": "team-new"}}, data={ "create": { - "user_id": "user-rpm-only", - "team_id": "team-rpm-only", + "user_id": "user-new", + "team_id": "team-new", "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, }, "update": { @@ -270,60 +219,48 @@ async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user): ) -# TEST: clone-on-write when membership still points at the team's shared default budget +# TEST: clone-on-write when the membership still points at the team's shared +# default budget. Editing this member must fork a private budget instead of +# mutating the shared row, and cloning a duration must seed a fresh reset time. @pytest.mark.asyncio -async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user): - """ - When a member's existing budget_id is the same row as the team's shared - default member budget, updating that member's budget must NOT mutate the - shared row. Instead we should create a new private budget for this member - (seeded with the default's values) and re-link the membership to it. - """ +async def test_clone_on_write_from_shared_default(mock_tx, fake_user): shared_default_id = "team-default-budget-1" + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row( + budget_id=shared_default_id, + max_budget=200.0, + soft_budget=None, + max_parallel_requests=None, + tpm_limit=500, + rpm_limit=None, + model_max_budget=None, + budget_duration="1d", + allowed_models=[], + ) + ) - # Default budget row in the DB: $200 cap, daily reset, 500 tpm. - default_row = MagicMock() - default_row.model_dump.return_value = { - "budget_id": shared_default_id, - "max_budget": 200.0, - "soft_budget": None, - "max_parallel_requests": None, - "tpm_limit": 500, - "rpm_limit": None, - "model_max_budget": None, - "budget_duration": "1d", - "allowed_models": [], - } - mock_tx.litellm_budgettable.find_unique = AsyncMock(return_value=default_row) - - # Caller is changing only this member's max_budget. await _upsert_budget_and_membership( mock_tx, team_id="team-shared", user_id="user-shared", - max_budget=50.0, existing_budget_id=shared_default_id, user_api_key_dict=fake_user, + budget_patch={"max_budget": 50.0}, team_default_budget_id=shared_default_id, ) - # Must NOT touch the shared default row in place. mock_tx.litellm_budgettable.update.assert_not_called() + mock_tx.litellm_budgettable.create.assert_awaited_once() + create_data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert_future_reset_time(create_data.pop("budget_reset_at")) + assert create_data == { + "created_by": fake_user.user_id, + "updated_by": fake_user.user_id, + "max_budget": 50.0, # caller wins + "tpm_limit": 500, # cloned from default + "budget_duration": "1d", # cloned from default + } - # Must create a new private budget seeded with the default's values, - # with the caller's max_budget overriding the cloned default. - mock_tx.litellm_budgettable.create.assert_awaited_once_with( - data={ - "created_by": fake_user.user_id, - "updated_by": fake_user.user_id, - "max_budget": 50.0, # caller wins - "tpm_limit": 500, # cloned from default - "budget_duration": "1d", # cloned from default - }, - include={"team_membership": True}, - ) - - # Membership must be re-linked to the new private budget. new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( where={"user_id_team_id": {"user_id": "user-shared", "team_id": "team-shared"}}, @@ -340,32 +277,64 @@ async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user) ) -# TEST: when team default exists but member already has their own budget, in-place update +# TEST: forking the shared default while clearing its duration must drop the +# duration (and not carry a reset time) on the new private budget. @pytest.mark.asyncio -async def test_upsert_updates_in_place_when_member_has_private_budget( - mock_tx, fake_user -): - """ - If the member's budget_id is different from the team's shared default - (i.e. they already have a private budget), we should keep the current - in-place behavior and not allocate a new row. - """ +async def test_clone_on_write_clears_duration(mock_tx, fake_user): + shared_default_id = "team-default-budget-1" + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row( + budget_id=shared_default_id, + max_budget=200.0, + tpm_limit=500, + budget_duration="1d", + allowed_models=[], + ) + ) + + await _upsert_budget_and_membership( + mock_tx, + team_id="team-shared", + user_id="user-shared", + existing_budget_id=shared_default_id, + user_api_key_dict=fake_user, + budget_patch={"budget_duration": None}, + team_default_budget_id=shared_default_id, + ) + + mock_tx.litellm_budgettable.update.assert_not_called() + create_data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert create_data == { + "created_by": fake_user.user_id, + "updated_by": fake_user.user_id, + "max_budget": 200.0, + "tpm_limit": 500, + "budget_duration": None, + } + assert "budget_reset_at" not in create_data + + +# TEST: when the member already has their own private budget (different from the +# team default), we update it in place rather than forking another row. +@pytest.mark.asyncio +async def test_private_budget_updates_in_place(mock_tx, fake_user): + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(max_budget=10.0) + ) + await _upsert_budget_and_membership( mock_tx, team_id="team-mixed", user_id="user-private", - max_budget=75.0, existing_budget_id="private-budget-xyz", user_api_key_dict=fake_user, + budget_patch={"max_budget": 75.0}, team_default_budget_id="team-default-budget-1", ) mock_tx.litellm_budgettable.update.assert_awaited_once_with( where={"budget_id": "private-budget-xyz"}, - data={ - "max_budget": 75.0, - "updated_by": fake_user.user_id, - }, + data={"max_budget": 75.0, "updated_by": fake_user.user_id}, ) mock_tx.litellm_budgettable.create.assert_not_called() mock_tx.litellm_teammembership.upsert.assert_not_called() diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 9dcf5df4aeb..6021c221426 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -107,6 +107,201 @@ def test_is_database_connection_generic_errors(): ) +@pytest.mark.parametrize( + "error", + [ + ConnectionError("connection refused"), + TimeoutError("timed out"), + OSError("network is unreachable"), + asyncio.TimeoutError(), + HTTPClientClosedError(), + ClientNotConnectedError(), + PrismaError("can't reach database server"), + PrismaError(), + ], +) +def test_is_database_service_unavailable_error_infra_failures(error): + """Infrastructure-level failures (socket/connection/timeout, prisma + transport, unknown PrismaError) mean the DB could not answer, so auth + must surface 503 instead of treating a valid key as invalid.""" + assert PrismaDBExceptionHandler.is_database_service_unavailable_error(error) is True + + +def test_is_database_service_unavailable_error_prisma_p1001_masquerades_as_dataerror(): + """Real-world regression: prisma-client-py raises the P1001 "can't reach + database server" connectivity failure as a DataError (a data-layer type). + A type-only check would miss it and return 401 during a genuine outage; + the message keyword must still classify it as service-unavailable -> 503.""" + p1001_as_dataerror = DataError( + data={ + "user_facing_error": { + "message": "Can't reach database server at `127.0.0.1`:`5499`", + "meta": {"table": "t"}, + } + } + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + p1001_as_dataerror + ) + is True + ) + + +def test_is_database_service_unavailable_error_cached_plan_escapes_as_503(): + """Composes with the cached-plan retry: when that recovery fails and the + Postgres "cached plan must not change result type" error escapes (raised by + prisma as a data-layer RawQueryError), it is a transient stale-DB-state + condition, not an invalid key, so it must classify as service-unavailable + -> 503 rather than fall through to 401.""" + cached_plan_error = RawQueryError( + data={ + "user_facing_error": { + "message": "cached plan must not change result type", + "meta": {"table": "t"}, + } + } + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + cached_plan_error + ) + is True + ) + + +def test_is_database_service_unavailable_error_prisma_engine_malformed_payload(): + """Real-world regression: at the instant the DB socket drops, the prisma + query engine returns a malformed error payload (``user_facing_error.meta`` + is ``null``). prisma-client-py's ``handle_response_errors`` then crashes + with ``AttributeError: 'NoneType' object has no attribute 'get'`` before it + can raise the proper P1001 error. That bare AttributeError has no + connection keyword, so without the prisma-engine-origin check it falls + through to 401 on the first request of an outage. Reproduce the exact + prisma crash and assert it classifies as service-unavailable -> 503.""" + from prisma.engine import utils as prisma_engine_utils + + malformed_payload = [ + { + "error": "Can't reach database server", + "user_facing_error": { + "error_code": "P1001", + "message": "Can't reach database server at `localhost`:`5503`", + "meta": None, + }, + } + ] + with pytest.raises(AttributeError) as exc_info: + prisma_engine_utils.handle_response_errors(None, malformed_payload) + + assert "no attribute 'get'" in str(exc_info.value) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(exc_info.value) + is True + ) + + +def test_is_prisma_engine_internal_error_excludes_application_attributeerror(): + """The prisma-engine-origin check must stay narrow: a genuine AttributeError + raised by application code (a real bug) must NOT be classified as + service-unavailable, otherwise real bugs would silently become 503s.""" + + def application_bug(): + none_value = None + return none_value.get("oops") + + with pytest.raises(AttributeError) as exc_info: + application_bug() + + assert ( + PrismaDBExceptionHandler.is_prisma_engine_internal_error(exc_info.value) + is False + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(exc_info.value) + is False + ) + + +def test_is_prisma_engine_internal_error_excludes_data_layer_prisma_error(): + """A data-layer ``PrismaError`` (the DB IS reachable and rejected the data) + must stay 401. These are always raised from prisma internals, so the check + excludes any ``PrismaError`` by type before inspecting the traceback.""" + data_layer_error = UniqueViolationError( + data={"user_facing_error": {"meta": {"table": "t"}}} + ) + try: + raise data_layer_error + except UniqueViolationError as e: + assert PrismaDBExceptionHandler.is_prisma_engine_internal_error(e) is False + + +@pytest.mark.parametrize( + "error", + [ + DataError(data={"user_facing_error": {"meta": {"table": "t"}}}), + UniqueViolationError(data={"user_facing_error": {"meta": {"table": "t"}}}), + RecordNotFoundError(data={"user_facing_error": {"meta": {"table": "t"}}}), + Exception("some unrelated error"), + ValueError("bad value"), + ], +) +def test_is_database_service_unavailable_error_excludes_non_infra(error): + """Data-layer errors (the DB IS reachable and answered) and generic + non-DB errors must NOT be classified as service-unavailable, otherwise a + genuine 401 would be masked as a transient 503.""" + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(error) is False + ) + + +def test_is_database_service_unavailable_error_asyncpg(monkeypatch): + """asyncpg connection/interface errors map to service-unavailable. asyncpg + is not a hard dependency, so inject a stand-in module to exercise the + branch deterministically regardless of the install environment.""" + import sys + import types + + fake_asyncpg = types.ModuleType("asyncpg") + fake_exceptions = types.ModuleType("asyncpg.exceptions") + + class PostgresConnectionError(Exception): + pass + + class InterfaceError(Exception): + pass + + class UniqueViolationError(Exception): # data-layer, must stay False + pass + + fake_exceptions.PostgresConnectionError = PostgresConnectionError + fake_exceptions.InterfaceError = InterfaceError + fake_exceptions.UniqueViolationError = UniqueViolationError + fake_asyncpg.exceptions = fake_exceptions + + monkeypatch.setitem(sys.modules, "asyncpg", fake_asyncpg) + monkeypatch.setitem(sys.modules, "asyncpg.exceptions", fake_exceptions) + + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + PostgresConnectionError("connection reset") + ) + is True + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + InterfaceError("connection was closed") + ) + is True + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + UniqueViolationError("duplicate key") + ) + is False + ) + + # Test should_allow_request_on_db_unavailable method @patch( "litellm.proxy.proxy_server.general_settings", diff --git a/tests/test_litellm/proxy/db/test_tool_registry_writer.py b/tests/test_litellm/proxy/db/test_tool_registry_writer.py index 8074871c3dd..7bf1ffda4fe 100644 --- a/tests/test_litellm/proxy/db/test_tool_registry_writer.py +++ b/tests/test_litellm/proxy/db/test_tool_registry_writer.py @@ -291,3 +291,77 @@ async def test_tool_policy_registry_not_initialized_returns_untrusted(): assert not registry.is_initialized() result = registry.get_effective_policies(["unknown_tool"]) assert result == {"unknown_tool": "untrusted"} + + +@pytest.mark.asyncio +async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read(): + """`ToolPolicyRegistry.sync_tool_policy_from_db` self-heals across one + ClientNotConnectedError on the tools read — the perms read still fires + after the recovery and the registry initializes cleanly.""" + import prisma as prisma_pkg + + registry = ToolPolicyRegistry() + invocations: list = [] + + async def _flaky_find_many(): + invocations.append(None) + if len(invocations) == 1: + raise prisma_pkg.errors.ClientNotConnectedError() + return [] + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock( + side_effect=_flaky_find_many + ) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( + return_value=[] + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + await registry.sync_tool_policy_from_db(mock_prisma_client) + + assert len(invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert ( + reconnect_kwargs["reason"] + == "sync_tool_policy_from_db_tools_lookup_failure" + ) + assert registry.is_initialized() + + +@pytest.mark.asyncio +async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read(): + """Same as above but the blip happens on the perms read — distinct reason + tag in telemetry confirms the second wrap is wired separately.""" + import prisma as prisma_pkg + + registry = ToolPolicyRegistry() + perms_invocations: list = [] + + async def _flaky_perms_find_many(): + perms_invocations.append(None) + if len(perms_invocations) == 1: + raise prisma_pkg.errors.ClientNotConnectedError() + return [] + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( + side_effect=_flaky_perms_find_many + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + await registry.sync_tool_policy_from_db(mock_prisma_client) + + assert len(perms_invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert ( + reconnect_kwargs["reason"] + == "sync_tool_policy_from_db_perms_lookup_failure" + ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index c58c94cbbc7..f8fd9a0a185 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -41,7 +41,8 @@ def test_crowdstrike_aidr_guardrail_config() -> None: ) -def test_crowdstrike_aidr_guardrail_config_no_api_key() -> None: +def test_crowdstrike_aidr_guardrail_config_no_api_key(monkeypatch) -> None: + monkeypatch.delenv("CS_AIDR_TOKEN", raising=False) with pytest.raises(CrowdStrikeAIDRGuardrailMissingSecrets): init_guardrails_v2( all_guardrails=[ @@ -59,7 +60,8 @@ def test_crowdstrike_aidr_guardrail_config_no_api_key() -> None: ) -def test_crowdstrike_aidr_guardrail_config_no_api_base() -> None: +def test_crowdstrike_aidr_guardrail_config_no_api_base(monkeypatch) -> None: + monkeypatch.delenv("CS_AIDR_BASE_URL", raising=False) with pytest.raises(CrowdStrikeAIDRGuardrailMissingSecrets): init_guardrails_v2( all_guardrails=[ @@ -412,6 +414,171 @@ async def test_apply_guardrail_response_ok( assert result["texts"] == inputs["texts"] +@pytest.mark.asyncio +async def test_apply_guardrail_sends_user_id_model_and_extra_info( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gpt-4o", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gpt-4o", + "litellm_metadata": { + "user_api_key_user_id": "uid-abc", + "user_api_key_user_email": "alice@example.com", + }, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-abc" + assert payload["model"] == "gpt-4o" + assert payload["extra_info"] == {"user_name": "alice@example.com"} + + +@pytest.mark.asyncio +async def test_apply_guardrail_empty_extra_info_when_no_email( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gemini-flash", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gemini-flash", + "litellm_metadata": { + "user_api_key_user_id": "uid-no-email", + "user_api_key_user_email": None, + }, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-no-email" + assert payload["model"] == "gemini-flash" + assert payload["extra_info"] == {} + + +@pytest.mark.asyncio +async def test_apply_guardrail_no_metadata_skips_user_fields( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + } + request_data = {"messages": inputs["structured_messages"]} + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert "user_id" not in payload + assert "model" not in payload + assert "extra_info" not in payload + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "litellm_metadata, metadata", + [ + (None, {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}), + ({"trace_id": "t1"}, {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}), + (["unexpected"], {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}), + ({"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}, {"trace_id": "t1"}), + ], + ids=["identity_in_metadata_llm_none", "identity_in_metadata_llm_user_dict", "identity_in_metadata_llm_non_mapping", "identity_in_litellm_metadata"], +) +async def test_apply_guardrail_reads_identity_from_either_metadata_bag( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, + litellm_metadata, + metadata, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gpt-4o", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gpt-4o", + "litellm_metadata": litellm_metadata, + "metadata": metadata, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-abc" + assert payload["extra_info"] == {"user_name": "alice@example.com"} + + @pytest.mark.asyncio async def test_apply_guardrail_request_skipped_messages_stay_aligned( crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 565bf83c6a2..5bdf30b93f5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2398,11 +2398,9 @@ async def test_anonymize_text_uses_correct_positions_no_parse_pii(): ) expected = "My name is , my email is , phone " - assert result == expected, ( - f"anonymize_text produced garbled output with PII remnants.\n" - f"Expected: {expected!r}\n" - f"Got: {result!r}" - ) + assert ( + result == expected + ), f"anonymize_text produced garbled output with PII remnants.\nExpected: {expected!r}\nGot: {result!r}" assert masked_entity_count == { "PERSON": 1, "EMAIL_ADDRESS": 1, @@ -2495,3 +2493,369 @@ async def test_anonymize_text_uses_correct_positions_with_parse_pii(): assert pii_tokens.get("") == "John Smith" assert pii_tokens.get("") == "john@example.com" assert pii_tokens.get("") == "555-867-5309" + + +def _sse_event_bytes(event: dict, event_name: str = "") -> bytes: + import json + + prefix = f"event: {event_name}\n" if event_name else "" + # ensure_ascii=False matches Anthropic's wire format (raw UTF-8) + return (prefix + "data: " + json.dumps(event, ensure_ascii=False) + "\n\n").encode( + "utf-8" + ) + + +def _text_delta_bytes(text: str, index: int = 0) -> bytes: + return _sse_event_bytes( + { + "type": "content_block_delta", + "index": index, + "delta": {"type": "text_delta", "text": text}, + } + ) + + +def _collect_delta_text(raw: bytes, field: str = "text") -> str: + """Concatenate every delta `field` value across all events in `raw`.""" + import json + + parts = [] + for line in raw.split(b"\n"): + if line.startswith(b"data: ") and line != b"data: [DONE]": + try: + event = json.loads(line[len(b"data: ") :]) + except json.JSONDecodeError: + continue + delta = event.get("delta") or {} + if field in delta: + parts.append(delta[field]) + return "".join(parts) + + +def _run_unmasker(pii_tokens: dict, *chunks: bytes, **unmasker_kwargs) -> bytes: + from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _AnthropicSSEUnmasker, + ) + + unmasker = _AnthropicSSEUnmasker(pii_tokens, **unmasker_kwargs) + return b"".join(unmasker.feed(c) for c in chunks) + unmasker.flush() + + +def test_unmask_sse_bytes_replaces_text_delta(): + pii_tokens = {"": "Bobby"} + result = _run_unmasker( + pii_tokens, _text_delta_bytes("Hello , how are you?") + ) + assert _collect_delta_text(result) == "Hello Bobby, how are you?" + + +def test_unmask_sse_bytes_ignores_events_without_delta(): + pii_tokens = {"": "Bobby"} + chunk = _sse_event_bytes( + {"type": "message_start", "message": {"id": "msg_01", "role": "assistant"}} + ) + assert _run_unmasker(pii_tokens, chunk) == chunk + + +def test_unmask_sse_bytes_keeps_input_json_delta_masked_by_default(): + # Streamed tool-call arguments are executed by agentic clients — restoring + # PII there hands the originals to tool execution, so placeholders stay + # masked unless unmask_streamed_tool_calls is explicitly enabled. + pii_tokens = {"": "Bobby"} + chunk = _sse_event_bytes( + { + "type": "content_block_delta", + "index": 1, + "delta": { + "type": "input_json_delta", + "partial_json": '{"name": ""}', + }, + } + ) + assert _run_unmasker(pii_tokens, chunk) == chunk + + +def test_unmask_sse_bytes_unmasks_input_json_delta_when_opted_in(): + import json + + # With the explicit opt-in, tool inputs are restored. The original + # contains a double quote, which must be JSON-escaped so the assembled + # tool input stays valid JSON. + pii_tokens = {"": 'Bobby "Bob" Tables'} + chunk = _sse_event_bytes( + { + "type": "content_block_delta", + "index": 1, + "delta": { + "type": "input_json_delta", + "partial_json": '{"name": ""}', + }, + } + ) + result = _run_unmasker(pii_tokens, chunk, unmask_tool_inputs=True) + fragment = _collect_delta_text(result, field="partial_json") + assert json.loads(fragment) == {"name": 'Bobby "Bob" Tables'} + + +def test_unmask_sse_bytes_openai_format_passes_through_unchanged(): + import json + + # OpenAI-format streams normally arrive as ModelResponseStream objects and + # never reach the SSE unmasker; if OpenAI-shaped bytes ever do, they must + # pass through byte-identical (same behavior as the previous + # _unmask_sse_bytes_chunk implementation). + pii_tokens = {"": "Bobby"} + event = { + "id": "chatcmpl-123", + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {"content": "Hello !"}}], + } + chunk = ("data: " + json.dumps(event) + "\n\ndata: [DONE]\n\n").encode("utf-8") + assert _run_unmasker(pii_tokens, chunk) == chunk + + +def test_unmask_sse_bytes_unmasks_thinking_delta(): + pii_tokens = {"": "Bobby"} + chunk = _sse_event_bytes( + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "thinking_delta", "thinking": "user is "}, + } + ) + result = _run_unmasker(pii_tokens, chunk) + assert _collect_delta_text(result, field="thinking") == "user is Bobby" + + +def test_unmask_sse_bytes_handles_malformed_json(): + chunk = b"data: {not valid json}\n\n" + assert _run_unmasker({"": "Bobby"}, chunk) == chunk + + +def test_unmask_sse_bytes_handles_unicode_decode_error(): + chunk = b"\xff\xfe invalid utf-8\n" + assert _run_unmasker({"": "Bobby"}, chunk) == chunk + + +def test_unmask_sse_bytes_non_ascii_pii_not_escaped(): + pii_tokens = {"": "José"} + result = _run_unmasker(pii_tokens, _text_delta_bytes("Hello !")) + assert "Jos\\u" not in result.decode("utf-8") + assert _collect_delta_text(result) == "Hello José!" + + +def test_unmask_sse_bytes_handles_crlf_line_endings(): + import json + + pii_tokens = {"": "Bobby"} + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hi !"}, + } + crlf_chunk = ("data: " + json.dumps(event) + "\r\ndata: [DONE]\r\n").encode("utf-8") + + result = _run_unmasker(pii_tokens, crlf_chunk) + + decoded = result.decode("utf-8") + parsed = json.loads(decoded.split("data: ", 1)[1].split("\n")[0].strip()) + assert parsed["delta"]["text"] == "Hi Bobby!" + assert "\r\n" in decoded # original line endings preserved + assert "data: [DONE]" in decoded + + +def test_unmask_sse_bytes_event_split_across_chunks(): + # Network framing splits chunks arbitrarily — a `data:` line (and even a + # multi-byte UTF-8 character) may arrive split across two bytes chunks. + pii_tokens = {"": "Bobby"} + raw = _text_delta_bytes("Grüße an !") + split_at = raw.index("ü".encode("utf-8")) + 1 # mid multi-byte character + result = _run_unmasker(pii_tokens, raw[:split_at], raw[split_at:]) + assert _collect_delta_text(result) == "Grüße an Bobby!" + + +def test_unmask_sse_bytes_token_split_across_delta_events(): + # Models routinely emit a placeholder split across multiple text deltas; + # no single delta contains the full token. + pii_tokens = {"": "Bobby", "": "bob@example.com"} + result = _run_unmasker( + pii_tokens, + _text_delta_bytes("Mail to at today"), + _sse_event_bytes( + {"type": "content_block_stop", "index": 0}, "content_block_stop" + ), + ) + assert _collect_delta_text(result) == "Mail to Bobby at bob@example.com today" + + +def test_unmask_sse_bytes_numbered_token_disambiguation(): + # must not match inside + pii_tokens = {"": "Bobby", "": "Alice"} + result = _run_unmasker(pii_tokens, _text_delta_bytes(" vs ")) + assert _collect_delta_text(result) == "Alice vs Bobby" + + +def test_unmask_sse_bytes_literal_angle_bracket_released(): + pii_tokens = {"": "Bobby"} + result = _run_unmasker(pii_tokens, _text_delta_bytes("Vec stays")) + assert _collect_delta_text(result) == "Vec stays" + + +def test_unmask_sse_bytes_carry_flushed_on_block_stop(): + # A dangling partial-token prefix at block end is literal text — it must + # be emitted (as a synthetic delta), not swallowed. + pii_tokens = {"": "Bobby"} + result = _run_unmasker( + pii_tokens, + _text_delta_bytes("ends with ": "Bobby"} + result = _run_unmasker( + pii_tokens, + _text_delta_bytes("Hi !"), + _sse_event_bytes( + {"type": "content_block_stop", "index": 0}, "content_block_stop" + ), + ) + assert _collect_delta_text(result) == "Hi Bobby!" + + +def test_unmask_sse_bytes_carry_flush_preserves_crlf(): + import json + + # The synthetic flush event must use the stream's line endings so a CRLF + # stream stays uniformly CRLF. + pii_tokens = {"": "Bobby"} + delta_event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "ends with ": "Bobby"} + request_data = {"metadata": {"pii_tokens": pii_tokens}} + + def _make_sse_chunk(text: str) -> bytes: + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": text}, + } + return ("data: " + json.dumps(event) + "\n\n").encode("utf-8") + + async def mock_stream(): + yield _make_sse_chunk("Hello !") + yield _make_sse_chunk(" How can I help?") + + chunks = [] + async for chunk in guardrail._stream_pii_unmasking(mock_stream(), request_data): + chunks.append(chunk) + + assert len(chunks) == 2 + first = chunks[0].decode("utf-8") + first_event = json.loads(first.split("data: ", 1)[1].strip()) + assert first_event["delta"]["text"] == "Hello Bobby!" + + second = chunks[1].decode("utf-8") + second_event = json.loads(second.split("data: ", 1)[1].strip()) + assert second_event["delta"]["text"] == " How can I help?" + + +@pytest.mark.asyncio +async def test_stream_pii_unmasking_respects_unmask_streamed_tool_calls( + mock_user_api_key, +): + import json + + pii_tokens = {"": "Bobby"} + request_data = {"metadata": {"pii_tokens": pii_tokens}} + tool_chunk = ( + "data: " + + json.dumps( + { + "type": "content_block_delta", + "index": 1, + "delta": { + "type": "input_json_delta", + "partial_json": '{"name": ""}', + }, + } + ) + + "\n\n" + ).encode("utf-8") + + async def collect(guardrail) -> bytes: + async def mock_stream(): + yield tool_chunk + + out = b"" + async for chunk in guardrail._stream_pii_unmasking( + mock_stream(), request_data + ): + out += chunk + return out + + default_guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, output_parse_pii=True + ) + assert await collect(default_guardrail) == tool_chunk + + opted_in_guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, output_parse_pii=True, unmask_streamed_tool_calls=True + ) + opted_in = await collect(opted_in_guardrail) + event = json.loads(opted_in.decode("utf-8").split("data: ", 1)[1].strip()) + assert json.loads(event["delta"]["partial_json"]) == {"name": "Bobby"} + + +@pytest.mark.asyncio +async def test_stream_pii_unmasking_passthrough_when_no_tokens(mock_user_api_key): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + raw_chunk = b"data: {}\n\n" + request_data: dict = {"metadata": {}} + + async def mock_stream(): + yield raw_chunk + + chunks = [] + async for chunk in guardrail._stream_pii_unmasking(mock_stream(), request_data): + chunks.append(chunk) + + assert chunks == [raw_chunk] diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 80a4804956c..64c57ab90e3 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -696,6 +696,33 @@ def mock_reject_os_environ(params): assert model_params.get("api_key") == "fake-key-A" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status,error_message", + [ + ("healthy", ""), + ("unhealthy", "Galileo authentication failed"), + ], +) +async def test_health_services_endpoint_galileo(status, error_message): + with patch("litellm.integrations.galileo.GalileoObserve") as MockGalileoObserve: + mock_instance = MagicMock() + mock_instance.async_health_check = AsyncMock( + return_value={"status": status, "error_message": error_message} + ) + MockGalileoObserve.return_value = mock_instance + + result = await health_services_endpoint(service="galileo") + + if status == "healthy": + assert result["status"] == "healthy" + assert result["message"] == "Galileo is healthy" + else: + assert result["status"] == "unhealthy" + assert result["message"] == error_message + mock_instance.async_health_check.assert_awaited_once() + + @pytest.mark.asyncio async def test_health_services_endpoint_datadog_llm_observability(): """ diff --git a/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py new file mode 100644 index 00000000000..f716a8533d8 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py @@ -0,0 +1,67 @@ +from unittest.mock import AsyncMock, patch + +import pytest + +from litellm.proxy.hooks.litellm_skills.main import SkillsInjectionHook + +SKILL_TOOL_NAME = "litellm_skill_e2b8dca8_031a_4481_b034_b9ec7d4eb7bf" + + +def _request_data(): + return { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "run the skill"}], + "litellm_metadata": { + "_litellm_code_execution_enabled": True, + "_skill_files": {SKILL_TOOL_NAME: {"main.py": b"print('hi')"}}, + }, + } + + +def _tool_use_response(tool_name): + return { + "stop_reason": "tool_use", + "content": [ + {"type": "tool_use", "id": "toolu_1", "name": tool_name, "input": {}} + ], + } + + +@pytest.mark.asyncio +async def test_post_call_success_hook_executes_litellm_skill_tool(): + """DB skill tool names carry the litellm_skill_ prefix and must trigger the execution loop.""" + hook = SkillsInjectionHook() + response = _tool_use_response(SKILL_TOOL_NAME) + + with patch.object( + hook, "_execute_code_loop_messages_api", new=AsyncMock(return_value=response) + ) as mock_loop: + result = await hook.async_post_call_success_deployment_hook( + request_data=_request_data(), response=response, call_type=None + ) + + mock_loop.assert_awaited_once() + assert result is response + + +@pytest.mark.asyncio +async def test_execute_code_loop_dispatches_litellm_skill_tool(): + """The agentic loop must route litellm_skill_ tool calls to _execute_skill_tool.""" + hook = SkillsInjectionHook() + final_response = {"stop_reason": "end_turn", "content": []} + + with ( + patch.object( + hook, "_execute_skill_tool", new=AsyncMock(return_value="skill ran") + ) as mock_exec, + patch("litellm.anthropic.acreate", new=AsyncMock(return_value=final_response)), + ): + result = await hook._execute_code_loop_messages_api( + data=_request_data(), + response=_tool_use_response(SKILL_TOOL_NAME), + skill_files={"main.py": b"print('hi')"}, + ) + + mock_exec.assert_awaited_once() + assert mock_exec.await_args.args[0] == SKILL_TOOL_NAME + assert result is final_response diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index ae71d10b378..a6f6e651487 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -713,7 +713,8 @@ async def test_count_input_file_usage_decodes_model_embedded_file_id(): @pytest.mark.asyncio async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(): """After replace_model_in_jsonl, body.model is the provider id (e.g. gpt-5.5). - Auth must check the proxy model_name the key was granted, not the stripped id.""" + Auth must check target_model_names from the unified file id, not reverse-map + the stripped id.""" from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter rate_limiter = _PROXY_BatchRateLimiter( @@ -732,7 +733,6 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias( ) mock_router = MagicMock() mock_router.model_list = [] - mock_router.resolve_model_name_from_model_id.return_value = proxy_alias can_key_call_model = AsyncMock(return_value=True) with ( @@ -745,10 +745,105 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias( await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, file_content_as_dict=file_dict, + target_model_names=[proxy_alias], ) can_key_call_model.assert_awaited_once() assert can_key_call_model.await_args.kwargs["model"] == proxy_alias + mock_router.resolve_model_name_from_model_id.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_list_order", + [ + [ + "openai/openai/gpt-5.5", + "openai/openai/gpt-5.5-batch", + "us/azure/openai/gpt-5.5", + ], + [ + "us/azure/openai/gpt-5.5", + "openai/openai/gpt-5.5", + "openai/openai/gpt-5.5-batch", + ], + [ + "openai/openai/gpt-5.5-batch", + "us/azure/openai/gpt-5.5", + "openai/openai/gpt-5.5", + ], + ], +) +async def test_pre_call_uses_target_model_names_not_stripped_reverse_lookup( + model_list_order, +): + """LIT-3593: three deployments strip to gpt-5.5; auth must use the upload + target alias from target_model_names, not first-match reverse lookup.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + batch_alias = "openai/openai/gpt-5.5-batch" + deployment_templates = { + "openai/openai/gpt-5.5": { + "model_name": "openai/openai/gpt-5.5", + "litellm_params": {"model": "openai/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"}, + }, + "openai/openai/gpt-5.5-batch": { + "model_name": "openai/openai/gpt-5.5-batch", + "litellm_params": {"model": "openai/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5-batch", "mode": "batch"}, + }, + "us/azure/openai/gpt-5.5": { + "model_name": "us/azure/openai/gpt-5.5", + "litellm_params": {"model": "azure/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"}, + }, + } + mock_router = MagicMock() + mock_router.model_list = [deployment_templates[name] for name in model_list_order] + + def _resolve(model_id): + for deployment in mock_router.model_list: + actual_model = deployment.get("litellm_params", {}).get("model") + if actual_model == model_id or ( + actual_model and actual_model.endswith(f"/{model_id}") + ): + return deployment.get("model_name") + return None + + mock_router.resolve_model_name_from_model_id.side_effect = _resolve + + file_dict = [ + {"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}} + ] + user = UserAPIKeyAuth( + api_key="sk-ok", + user_id="alice", + models=[batch_alias], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + can_key_call_model = AsyncMock(return_value=True) + + with ( + patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new=can_key_call_model, + ), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + ): + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=file_dict, + target_model_names=[batch_alias], + ) + + can_key_call_model.assert_awaited_once() + assert can_key_call_model.await_args.kwargs["model"] == batch_alias + mock_router.resolve_model_name_from_model_id.assert_not_called() @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py new file mode 100644 index 00000000000..0e2683dcbfd --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py @@ -0,0 +1,86 @@ +""" +Unit Tests for the max parallel request limiter v1 for the proxy +""" + +from datetime import datetime + +import pytest + +from litellm.caching.caching import DualCache +from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, +) +from litellm.proxy.utils import InternalUsageCache, hash_token +from litellm.types.utils import EmbeddingResponse, TextCompletionResponse, Usage + + +@pytest.mark.parametrize( + "response_obj", + [ + EmbeddingResponse( + model="text-embedding-3-small", + usage=Usage(prompt_tokens=50, completion_tokens=0, total_tokens=50), + ), + TextCompletionResponse( + model="gpt-3.5-turbo-instruct", + usage=Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50), + ), + ], +) +@pytest.mark.asyncio +async def test_async_log_success_event_counts_non_chat_response_tokens(response_obj): + """ + Embedding and text completion responses must increment the per key, user, + team, and end user TPM counters, not just chat completion ModelResponse + objects. + """ + _api_key = hash_token("sk-12345") + user_id = "ishaan" + team_id = "litellm-team" + end_user_id = "customer-1" + + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + + current_date = datetime.now().strftime("%Y-%m-%d") + current_hour = datetime.now().strftime("%H") + current_minute = datetime.now().strftime("%M") + precise_minute = f"{current_date}-{current_hour}-{current_minute}" + + scope_ids = [_api_key, user_id, team_id, end_user_id] + for scope_id in scope_ids: + await parallel_request_handler.internal_usage_cache.async_set_cache( + key=f"{scope_id}::{precise_minute}::request_count", + value={"current_requests": 1, "current_tpm": 0, "current_rpm": 1}, + litellm_parent_otel_span=None, + ) + + kwargs = { + "litellm_params": { + "metadata": { + "user_api_key": _api_key, + "user_api_key_user_id": user_id, + "user_api_key_team_id": team_id, + "user_api_key_model_max_budget": {}, + } + }, + "user": end_user_id, + } + + await parallel_request_handler.async_log_success_event( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + for scope_id in scope_ids: + current = await parallel_request_handler.internal_usage_cache.async_get_cache( + key=f"{scope_id}::{precise_minute}::request_count", + litellm_parent_otel_span=None, + ) + assert current["current_tpm"] == 50, ( + f"expected 50 tokens counted for {scope_id}, " + f"got {current['current_tpm']}" + ) diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 676f623a5dd..d10311b9f41 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -20,7 +20,12 @@ _PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler, ) from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token -from litellm.types.utils import ModelResponse, Usage +from litellm.types.utils import ( + EmbeddingResponse, + ModelResponse, + TextCompletionResponse, + Usage, +) class TimeController: @@ -547,6 +552,68 @@ async def mock_increment_pipeline(increment_list, **kwargs): ), f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" +@pytest.mark.parametrize( + "response_obj", + [ + EmbeddingResponse( + model="text-embedding-3-small", + usage=Usage(prompt_tokens=50, completion_tokens=0, total_tokens=50), + ), + TextCompletionResponse( + model="gpt-3.5-turbo-instruct", + usage=Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50), + ), + ], +) +@pytest.mark.asyncio +async def test_async_log_success_event_counts_non_chat_response_tokens( + monkeypatch, response_obj +): + """ + Embedding and text completion responses must increment the TPM counter, + not just chat completion ModelResponse objects. + """ + monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") + + _api_key = hash_token("sk-12345") + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", lambda: "total" + ) + + mock_kwargs = { + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + "model": response_obj.model, + } + + captured_operations = [] + + async def mock_increment_pipeline(increment_list, **kwargs): + captured_operations.extend(increment_list) + return True + + monkeypatch.setattr( + parallel_request_handler.internal_usage_cache.dual_cache, + "async_increment_cache_pipeline", + mock_increment_pipeline, + ) + + await parallel_request_handler.async_log_success_event( + kwargs=mock_kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + tpm_operation = next( + (op for op in captured_operations if op["key"].endswith(":tokens")), None + ) + assert tpm_operation is not None, "Should have a TPM increment operation" + assert tpm_operation["increment_value"] == 50 + + @pytest.mark.asyncio async def test_async_log_failure_event_v3(): """ diff --git a/tests/test_litellm/proxy/hooks/test_post_call_response_headers_hook.py b/tests/test_litellm/proxy/hooks/test_post_call_response_headers_hook.py index 8d8dd2d4284..660b0b0162a 100644 --- a/tests/test_litellm/proxy/hooks/test_post_call_response_headers_hook.py +++ b/tests/test_litellm/proxy/hooks/test_post_call_response_headers_hook.py @@ -13,7 +13,6 @@ sys.path.insert(0, os.path.abspath("../../../..")) -import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth @@ -336,3 +335,110 @@ class MockResponse: assert result == {"x-test": "1"} assert injector.called is True + + +# --- Tests for custom_llm_provider fallback (streaming response types) --- + + +@pytest.mark.asyncio +async def test_litellm_call_info_fallback_to_response_attribute(): + """Test that _build_litellm_call_info falls back to response.custom_llm_provider + when _hidden_params doesn't contain it (streaming response types).""" + inspector = CallInfoInspectorLogger() + + class MockStreamResponse: + """Mimics CustomStreamWrapper: custom_llm_provider as attribute, + _hidden_params without it.""" + + custom_llm_provider = "bedrock" + _hidden_params = { + "model_id": "model-xyz", + "api_base": "https://bedrock.us-east-1.amazonaws.com", + } + + with patch("litellm.callbacks", [inspector]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + await proxy_logging.post_call_response_headers_hook( + data={ + "model": "claude-3", + "metadata": {"model_info": {"id": "model-xyz"}}, + }, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=MockStreamResponse(), + ) + + assert inspector.called is True + assert inspector.received_call_info is not None + assert inspector.received_call_info["custom_llm_provider"] == "bedrock" + assert ( + inspector.received_call_info["api_base"] + == "https://bedrock.us-east-1.amazonaws.com" + ) + assert inspector.received_call_info["model_id"] == "model-xyz" + + +@pytest.mark.asyncio +async def test_litellm_call_info_fallback_no_hidden_params(): + """Test that _build_litellm_call_info works when response has no _hidden_params + at all (LiteLLMCompletionStreamingIterator case).""" + inspector = CallInfoInspectorLogger() + + class MockIteratorResponse: + """Mimics LiteLLMCompletionStreamingIterator: custom_llm_provider as attribute, + no _hidden_params attribute at all.""" + + custom_llm_provider = "vertex_ai" + + with patch("litellm.callbacks", [inspector]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + await proxy_logging.post_call_response_headers_hook( + data={"model": "gemini-pro", "metadata": {}}, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=MockIteratorResponse(), + ) + + assert inspector.called is True + assert inspector.received_call_info is not None + assert inspector.received_call_info["custom_llm_provider"] == "vertex_ai" + assert inspector.received_call_info["api_base"] is None + assert inspector.received_call_info["model_id"] is None + + +@pytest.mark.asyncio +async def test_litellm_call_info_hidden_params_takes_priority(): + """Test that _hidden_params.custom_llm_provider takes priority over + the response attribute when both are present.""" + inspector = CallInfoInspectorLogger() + + class MockResponse: + custom_llm_provider = "attribute_value" + _hidden_params = { + "custom_llm_provider": "hidden_params_value", + "api_base": "https://example.com", + "model_id": "m1", + } + + with patch("litellm.callbacks", [inspector]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + await proxy_logging.post_call_response_headers_hook( + data={"model": "test", "metadata": {}}, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=MockResponse(), + ) + + assert ( + inspector.received_call_info["custom_llm_provider"] + == "hidden_params_value" + ) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py index 8c74919df19..02b4e32db86 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py @@ -22,7 +22,7 @@ category routing missed these entirely. The fix wraps every internal raise site in -:class:`ProxyHTTPRateLimitError` (an ``HTTPException`` *and* a +:class:`ProxyRateLimitError` (an ``HTTPException`` *and* a ``litellm.RateLimitError``), and resolves ``model`` / ``llm_provider`` from ``data["model"]`` via :func:`get_llm_provider`. When the model is missing or unparseable we fall back to ``llm_provider="litellm_proxy"`` so we never break @@ -61,9 +61,9 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _PROXY_MaxParallelRequestsHandler_v3, ) +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import ( PROXY_LLM_PROVIDER_FALLBACK, - ProxyHTTPRateLimitError, resolve_llm_provider_for_rate_limit, ) from litellm.proxy.utils import InternalUsageCache @@ -75,12 +75,11 @@ # --------------------------------------------------------------------------- -class TestProxyHTTPRateLimitErrorClass: +class TestProxyRateLimitErrorClass: """Pin the dual ``HTTPException`` + ``RateLimitError`` shape.""" def test_is_both_http_exception_and_rate_limit_error(self): - e = ProxyHTTPRateLimitError( - status_code=429, + e = ProxyRateLimitError( detail="boom", model="gpt-4o-mini", llm_provider="openai", @@ -92,15 +91,15 @@ def test_is_both_http_exception_and_rate_limit_error(self): assert e.status_code == 429 assert e.model == "gpt-4o-mini" assert e.llm_provider == "openai" - assert e.message == "boom" + # ProxyRateLimitError prefixes message via RateLimitError.__init__. + assert "boom" in e.message assert e.detail == "boom" def test_dict_detail_is_stringified_for_message(self): # Some hooks pass a dict detail (e.g. dynamic_rate_limiter v1) — the # `message` attr (read by RateLimitError.__str__ and observability # callbacks) must still be a string. - e = ProxyHTTPRateLimitError( - status_code=429, + e = ProxyRateLimitError( detail={"error": "over rpm"}, model="claude-3-5-sonnet", llm_provider="anthropic", @@ -109,16 +108,15 @@ def test_dict_detail_is_stringified_for_message(self): assert "over rpm" in e.message def test_defaults_to_litellm_proxy_provider(self): - e = ProxyHTTPRateLimitError(status_code=429, detail="x") + e = ProxyRateLimitError(detail="x") assert e.llm_provider == PROXY_LLM_PROVIDER_FALLBACK assert e.model == "" def test_none_provider_normalized_to_fallback(self): - e = ProxyHTTPRateLimitError( - status_code=429, + e = ProxyRateLimitError( detail="x", - model=None, # type: ignore[arg-type] - llm_provider=None, # type: ignore[arg-type] + model=None, + llm_provider=None, ) assert e.llm_provider == PROXY_LLM_PROVIDER_FALLBACK assert e.model == "" @@ -143,7 +141,10 @@ def test_missing_or_unknown_model_falls_back(self, model): # Must never raise — the resolver wraps `get_llm_provider` defensively # because raising here would mask the rate-limit error we're trying # to surface to the user. - resolved_model, provider = resolve_llm_provider_for_rate_limit(model) + # Pin llm_router to None so the alias-fallback path doesn't pick up + # a router left behind by another test in the session. + with patch("litellm.proxy.proxy_server.llm_router", None): + resolved_model, provider = resolve_llm_provider_for_rate_limit(model) assert provider == PROXY_LLM_PROVIDER_FALLBACK # Resolver returns the input model verbatim on the unknown branch so # the `.model` attribute is never silently swapped to a different one. @@ -155,15 +156,148 @@ def test_missing_or_unknown_model_falls_back(self, model): def test_get_llm_provider_raising_is_swallowed(self): # If get_llm_provider itself blows up (unexpected error), we still # fall back rather than letting the secondary exception escape. + # No router is registered in this test, so the alias-fallback path + # also yields None and we land at PROXY_LLM_PROVIDER_FALLBACK. with patch.object( litellm, "get_llm_provider", side_effect=RuntimeError("boom"), ): - resolved_model, provider = resolve_llm_provider_for_rate_limit("anything") + with patch( + "litellm.proxy.proxy_server.llm_router", + None, + ): + resolved_model, provider = resolve_llm_provider_for_rate_limit( + "anything" + ) assert provider == PROXY_LLM_PROVIDER_FALLBACK assert resolved_model == "anything" + def test_router_alias_resolves_to_underlying_provider(self): + """ + Nearly every real LiteLLM proxy deployment uses router aliases: + + model_list: + - model_name: tpm-locked + litellm_params: + model: openai/gpt-4o-mini + ... + + ``litellm.get_llm_provider("tpm-locked")`` doesn't know about + router aliases and raises. Before this fix the resolver fell + through to ``"litellm_proxy"``, defeating the whole point of the + ``llm_provider`` field on the rate-limit error. The alias path + must look the deployment up in the router's ``model_list`` and + resolve from its ``litellm_params.model``. + """ + + class _FakeRouter: + model_list = [ + { + "model_name": "tpm-locked", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "fake", + }, + } + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + _FakeRouter(), + ): + resolved_model, provider = resolve_llm_provider_for_rate_limit("tpm-locked") + assert provider == "openai", ( + f"Router-alias path must resolve through litellm_params.model, " + f"not fall through to {PROXY_LLM_PROVIDER_FALLBACK!r}. Got " + f"provider={provider!r}, model={resolved_model!r}." + ) + # The resolved model should point at the underlying deployment so + # downstream Prometheus labels / failure callbacks attribute the + # 429 to the real upstream, not the alias. + assert resolved_model == "gpt-4o-mini" + + def test_router_alias_with_multiple_deployments_uses_first(self): + """ + When an alias maps to multiple deployments (the load-balancing + case), the rate-limit error fired at the *alias* level is + deployment-agnostic — we have no way of knowing which one would + have been picked. Use the first deployment's underlying provider: + every deployment under one alias should agree on provider in any + sensible config, and 'first' is deterministic so the Prometheus + label is stable. + """ + + class _FakeRouter: + model_list = [ + { + "model_name": "claude-pool", + "litellm_params": {"model": "anthropic/claude-3-5-sonnet"}, + }, + { + "model_name": "claude-pool", + "litellm_params": {"model": "anthropic/claude-3-5-haiku"}, + }, + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + _FakeRouter(), + ): + _, provider = resolve_llm_provider_for_rate_limit("claude-pool") + assert provider == "anthropic" + + def test_router_alias_unknown_falls_back(self): + """ + Alias not in the router model_list — both lookups fail, so we + land at the defensive ``litellm_proxy`` fallback rather than + raising. + """ + + class _FakeRouter: + model_list = [ + { + "model_name": "tpm-locked", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + } + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + _FakeRouter(), + ): + resolved_model, provider = resolve_llm_provider_for_rate_limit( + "not-an-alias" + ) + assert provider == PROXY_LLM_PROVIDER_FALLBACK + assert resolved_model == "not-an-alias" + + def test_router_alias_with_malformed_deployment_falls_back(self): + """ + A deployment in the router model_list with no usable + ``litellm_params.model`` (or where ``get_llm_provider`` on the + underlying string also raises) must not crash the resolver — + fall through to the defensive fallback. + """ + + class _FakeRouter: + model_list = [ + {"model_name": "broken", "litellm_params": {}}, + {"model_name": "broken", "litellm_params": {"model": ""}}, + { + "model_name": "broken", + "litellm_params": {"model": "nonsense-no-provider"}, + }, + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + _FakeRouter(), + ): + resolved_model, provider = resolve_llm_provider_for_rate_limit("broken") + assert provider == PROXY_LLM_PROVIDER_FALLBACK + assert resolved_model == "broken" + # --------------------------------------------------------------------------- # parallel_request_limiter v1 @@ -352,7 +486,7 @@ async def test_parallel_request_limiter_v1_missing_model_falls_back(): # --------------------------------------------------------------------------- -def _v3_over_limit_response(rate_limit_type: str = "rpm") -> dict: +def _v3_over_limit_response(rate_limit_type: str = "requests") -> dict: return { "overall_code": "OVER_LIMIT", "statuses": [ @@ -532,7 +666,7 @@ async def test_dynamic_rate_limiter_v3_model_capacity_path_populates_provider(): "descriptor_key": "model_saturation_check", "current_limit": 100, "limit_remaining": 0, - "rate_limit_type": "rpm", + "rate_limit_type": "requests", } ], } @@ -582,7 +716,7 @@ async def test_dynamic_rate_limiter_v3_unknown_descriptor_path_populates_provide "descriptor_key": "something_we_dont_handle", "current_limit": 1, "limit_remaining": 0, - "rate_limit_type": "rpm", + "rate_limit_type": "requests", } ], } @@ -937,31 +1071,56 @@ async def test_max_budget_per_session_limiter_unknown_model_falls_back(): # --------------------------------------------------------------------------- -def test_prometheus_exception_class_name_includes_provider(): +def test_prometheus_exception_class_name_back_compat_for_proxy_rate_limit_error(): + """ + `_get_exception_class_name` deliberately returns the literal string + ``"HTTPException"`` for every ``ProxyRateLimitError`` instance so that + pre-existing dashboards / alerts (which key off the historical value) + keep working after the unified rate-limit error class landed in #27687. + + Provider attribution is now surfaced separately via the + ``rate_limit_category`` / ``rate_limit_type`` labels — this test pins + the back-compat shim itself. + """ from litellm.integrations.prometheus import PrometheusLogger - exc = ProxyHTTPRateLimitError( - status_code=429, + exc = ProxyRateLimitError( detail="over limit", model="gpt-4o-mini", llm_provider="openai", ) + assert PrometheusLogger._get_exception_class_name(exc) == "HTTPException" - name = PrometheusLogger._get_exception_class_name(exc) - # Format is "{Provider.}{ClassName}" per `_get_exception_class_name`. - assert name.startswith("Openai.") - # And specifically: it ends in our exception class. (We don't pin the - # full string to avoid coupling the test to PR #27687's parallel rename.) - assert name.endswith("ProxyHTTPRateLimitError") + # Same back-compat path even when the resolver fell back to litellm_proxy. + exc_no_model = ProxyRateLimitError(detail="over limit") + assert PrometheusLogger._get_exception_class_name(exc_no_model) == "HTTPException" -def test_prometheus_exception_class_name_falls_back_when_no_model(): +def test_prometheus_exception_class_name_back_compat_for_budget_exceeded_error(): + """ + The unified rate-limit work also attached ``.llm_provider`` to + ``BudgetExceededError`` so callbacks get provider attribution from + ``StandardLoggingPayload``. Without a back-compat short-circuit the + provider-prefix step in ``_get_exception_class_name`` would silently + flip the label from ``"BudgetExceededError"`` to e.g. + ``"Openai.BudgetExceededError"`` and break dashboards keyed on the + historical value. Pin the literal label here. + """ from litellm.integrations.prometheus import PrometheusLogger - exc = ProxyHTTPRateLimitError(status_code=429, detail="over limit") - name = PrometheusLogger._get_exception_class_name(exc) - # `litellm_proxy` -> `Litellm_proxy.` (capitalize first char only). - assert name.startswith("Litellm_proxy.") + err = litellm.BudgetExceededError( + current_cost=1.0, + max_budget=0.5, + llm_provider="openai", + ) + assert PrometheusLogger._get_exception_class_name(err) == "BudgetExceededError" + + # Default (empty llm_provider) path — same literal label. + err_no_provider = litellm.BudgetExceededError(current_cost=1.0, max_budget=0.5) + assert ( + PrometheusLogger._get_exception_class_name(err_no_provider) + == "BudgetExceededError" + ) if __name__ == "__main__": diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py index 8fec05abe90..91a011a8234 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py @@ -43,11 +43,15 @@ async def fake_post_call_failure_hook(**_: Any) -> None: async def fake_post_call_success_hook(*, data, user_api_key_dict, response): return response + async def fake_post_call_response_headers_hook(**kwargs): + return {"x-callback-test": "value"} + fake_proxy_logger = SimpleNamespace( pre_call_hook=fake_pre_call_hook, update_request_status=fake_update_request_status, post_call_failure_hook=fake_post_call_failure_hook, post_call_success_hook=fake_post_call_success_hook, + post_call_response_headers_hook=fake_post_call_response_headers_hook, ) captured_route_request_data: Dict[str, Any] = {} @@ -110,3 +114,4 @@ async def receive(): assert pre_call_input["messages"][0]["content"] == "original prompt" assert captured_route_request_data["prompt"] == "sanitized prompt" assert "messages" not in captured_route_request_data + assert response.headers.get("x-callback-test") == "value" diff --git a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py index ea7e5591f18..f2ccfcd0155 100644 --- a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py +++ b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py @@ -611,6 +611,45 @@ async def test_list_search_tools_db_masking_sensitive_values(monkeypatch): app.dependency_overrides.pop(user_api_key_auth, None) +@pytest.mark.asyncio +async def test_get_all_search_tools_from_db_retries_on_transport_error(): + """`SearchToolRegistry.get_all_search_tools_from_db` self-heals across one + ClientNotConnectedError via call_with_db_reconnect_retry.""" + import prisma + from litellm.proxy.search_endpoints.search_tool_registry import ( + SearchToolRegistry, + ) + + invocations: list = [] + + async def _flaky_find_many(**kwargs): + invocations.append(None) + if len(invocations) == 1: + raise prisma.errors.ClientNotConnectedError() + return [] + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_searchtoolstable.find_many = AsyncMock( + side_effect=_flaky_find_many + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + result = await SearchToolRegistry.get_all_search_tools_from_db( + prisma_client=mock_prisma_client + ) + + assert result == [] + assert len(invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert ( + reconnect_kwargs["reason"] + == "get_all_search_tools_from_db_lookup_failure" + ) + + @contextlib.contextmanager def _mock_search_tool_backend(db_tools): """Patch the DB registry, prisma client, and config so /search_tools/list diff --git a/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py index b892c4e556d..4bdef2e8f96 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py @@ -259,6 +259,41 @@ async def test_init_cache_settings_in_db_skips_when_params_unchanged(self): mock_proxy_config._init_cache.assert_not_called() mock_proxy_config.switch_on_llm_response_caching.assert_not_called() + @pytest.mark.asyncio + async def test_init_cache_settings_in_db_retries_on_transport_error(self): + """`CacheSettingsManager.init_cache_settings_in_db` self-heals across one + ClientNotConnectedError via call_with_db_reconnect_retry.""" + import prisma + + invocations: list = [] + + async def _flaky_find_unique(**kwargs): + invocations.append(None) + if len(invocations) == 1: + raise prisma.errors.ClientNotConnectedError() + return None # No config → function returns early after retry. + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_cacheconfig.find_unique = AsyncMock( + side_effect=_flaky_find_unique + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + mock_proxy_config = MagicMock() + + await CacheSettingsManager.init_cache_settings_in_db( + prisma_client=mock_prisma_client, proxy_config=mock_proxy_config + ) + + assert len(invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert ( + reconnect_kwargs["reason"] + == "init_cache_settings_in_db_lookup_failure" + ) + # ── Audit-log emission for /cache/settings ──────────────────────────────────── diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 3c212d86e65..473d61f8a85 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6496,6 +6496,9 @@ async def test_reset_key_spend_success(monkeypatch): patch( "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" ) as mock_delete_cache, + patch( + "litellm.proxy.proxy_server._invalidate_spend_counter" + ) as mock_invalidate, ): mock_hash_token.return_value = hashed_key mock_check_admin.return_value = None @@ -6520,6 +6523,76 @@ async def test_reset_key_spend_success(monkeypatch): assert response["max_budget"] == 200.0 mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once() mock_delete_cache.assert_awaited_once() + mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}") + + +@pytest.mark.asyncio +async def test_update_key_spend_invalidates_counter(monkeypatch): + """ + Test that updating a key's spend via update_key_fn immediately invalidates the spend counter. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = AsyncMock() + mock_proxy_logging_obj = MagicMock() + + hashed_key = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" + key_in_db = LiteLLM_VerificationToken( + token=hashed_key, + user_id="test-user", + spend=10.0, + max_budget=200.0, + litellm_budget_table=None, + ) + + mock_prisma_client.get_data = AsyncMock(return_value=key_in_db) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {"spend": 0.0}}) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=key_in_db + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.store_audit_logs", False) + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" + ) as mock_delete_cache, + patch( + "litellm.proxy.proxy_server._invalidate_spend_counter" + ) as mock_invalidate, + ): + mock_delete_cache.return_value = None + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + mock_request = MagicMock() + mock_request.query_params = {} + + await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key="sk-test-key", spend=0.0), + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + mock_delete_cache.assert_awaited_once() + mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}") @pytest.mark.asyncio @@ -11668,3 +11741,84 @@ async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption( msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", "")) assert str(code) == "400" assert "cannot exceed" in msg.lower() + + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_null_clears_fields(): + """ + When budget_duration is explicitly set to null, prepare_key_update_data + should produce budget_duration=None and budget_reset_at=None so Prisma + clears them in the DB. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", budget_duration=None) + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert "budget_duration" in result + assert result["budget_duration"] is None + assert "budget_reset_at" in result + assert result["budget_reset_at"] is None + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_not_sent_excluded(): + """ + When budget_duration is NOT sent in the request (unset), it should not + appear in the result dict at all — the existing DB value stays unchanged. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"]) + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert "budget_duration" not in result + assert "budget_reset_at" not in result + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_valid_sets_reset(): + """ + When budget_duration is set to a valid duration string, both + budget_duration and budget_reset_at should be populated. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", budget_duration="30d") + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert result["budget_duration"] == "30d" + assert result["budget_reset_at"] is not None + + diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index b5eb091bb81..0b5b5fb6ceb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1482,6 +1482,10 @@ async def test_get_cached_temporary_mcp_server_non_admin_denied(self): "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", mock_manager, ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[non_admin]), + ), ): with pytest.raises(HTTPException) as exc_info: await _get_cached_temporary_mcp_server_or_404("server-x", non_admin) @@ -1514,6 +1518,10 @@ async def test_get_cached_temporary_mcp_server_non_admin_allowed(self): "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", mock_manager, ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[non_admin]), + ), ): result = await _get_cached_temporary_mcp_server_or_404( "server-x", non_admin @@ -1521,6 +1529,58 @@ async def test_get_cached_temporary_mcp_server_non_admin_allowed(self): assert result is registry_server + @pytest.mark.asyncio + async def test_get_cached_temporary_mcp_server_non_admin_allowed_via_team_access_group( + self, + ): + """Internal user whose only grant to the server flows through a team + access-group must pass the authorize/token access check. The check has to + expand the UI session into per-team contexts (build_effective_auth_contexts), + the same way the server-list grid does; checking only the bare session + context leaves the team grant invisible and 403s the user.""" + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _get_cached_temporary_mcp_server_or_404, + ) + + registry_server = generate_mock_mcp_server_config_record(server_id="server-x") + ui_session_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER, + team_id=UI_SESSION_TOKEN_TEAM_ID, + ) + team_context = ui_session_auth.model_copy() + team_context.team_id = "team-with-mcp-grant" + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id.return_value = registry_server + mock_manager.get_mcp_server_by_name.return_value = None + + def allowed_for(auth): + return ["server-x"] if auth.team_id == "team-with-mcp-grant" else [] + + mock_manager.get_allowed_mcp_servers = AsyncMock(side_effect=allowed_for) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_cached_temporary_mcp_server", + return_value=None, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[ui_session_auth, team_context]), + ), + ): + result = await _get_cached_temporary_mcp_server_or_404( + "server-x", ui_session_auth + ) + + assert result is registry_server + assert mock_manager.get_allowed_mcp_servers.await_count == 2 + @pytest.mark.asyncio async def test_get_cached_temporary_mcp_server_temp_cache_non_admin_denied(self): """Servers resolved from the admin-only temp cache reject non-admins.""" @@ -4146,7 +4206,7 @@ async def test_no_accessible_servers_returns_empty(self): ), patch.object( mgmt_endpoints, - "get_all_mcp_servers_for_user", + "_resolve_accessible_mcp_servers", AsyncMock(return_value=[]), ), ): @@ -4174,7 +4234,7 @@ async def test_only_servers_with_required_fields_are_returned(self): ), patch.object( mgmt_endpoints, - "get_all_mcp_servers_for_user", + "_resolve_accessible_mcp_servers", AsyncMock(return_value=[server_with, server_without]), ), patch.object( @@ -4204,7 +4264,7 @@ async def test_bulk_status_omits_stored_credential_values(self): ), patch.object( mgmt_endpoints, - "get_all_mcp_servers_for_user", + "_resolve_accessible_mcp_servers", AsyncMock(return_value=[server]), ), patch.object( @@ -4221,6 +4281,51 @@ async def test_bulk_status_omits_stored_credential_values(self): assert by_name["CORP_PASSWORD"].is_set is False assert "alice" not in result[0].model_dump_json() + @pytest.mark.asyncio + async def test_admin_view_all_flags_missing_fields_without_key_grants(self): + """Regression: the red "user fields missing" card must light up for an + admin in view_all mode even when their key carries no per-server MCP + grant. The bulk status feed has to resolve the same server set the + dashboard grid renders; the old narrow key-scoped listing returned + nothing for such an admin, leaving every card un-highlighted.""" + server = _make_env_var_server( + server_id="srv-with", + env_vars=_ENV_VARS_MIXED, + static_headers=_STATIC_HEADERS_MIXED, + ) + with ( + patch.object( + mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock() + ), + patch.object( + mgmt_endpoints, + "_get_user_mcp_management_mode", + return_value="view_all", + ), + patch.object( + mgmt_endpoints.global_mcp_server_manager, + "get_all_mcp_servers_unfiltered", + AsyncMock(return_value=[server]), + ), + patch.object( + mgmt_endpoints, + "get_user_env_vars_bulk", + AsyncMock(return_value={}), + ), + ): + result = await mgmt_endpoints.list_mcp_user_env_var_status( + user_api_key_dict=generate_mock_user_api_key_auth( + user_id="admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + ) + assert [s.server_id for s in result] == ["srv-with"] + assert result[0].missing_count == 2 + assert {f.name for f in result[0].required} == { + "CORP_USERNAME", + "CORP_PASSWORD", + } + class TestMCPUserEnvVarsAccessControl: """Per-server env-var endpoints must enforce the same access gate as 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 a9bb2b09a15..6a81b1b613b 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 @@ -24,6 +24,7 @@ ModelManagementAuthChecks, _get_team_deployments, clear_cache, + delete_team_models, ) from litellm.proxy.utils import PrismaClient from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment @@ -1915,6 +1916,190 @@ def _row(model_id): mock_refresh.assert_not_awaited() +class TestDeleteModelTeamAuth: + """Team auth on the /model/delete path. + + A model added via /model/new with model_info.team_id is orphaned once its + team is deleted: can_user_make_model_call looked the team up and raised + 'Team id=... does not exist in db' before the delete could run, so the model + was undeletable from the Models + Endpoints page. Without the team, team-admin + membership can't be verified, so a proxy admin (and only a proxy admin) may + delete the orphan; a missing team must never let a non-admin through. The team + is also looked up exactly once -- the auth check must not add a second query. + """ + + def _orphaned_model_mocks(self, team_id, model_id): + db_row = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name=f"model_name_{team_id}_abc-uuid", + litellm_params={"model": "openai/gpt-4.1-nano"}, + model_info={ + "id": model_id, + "team_id": team_id, + "team_public_model_name": "orphaned-gpt", + }, + created_by="admin", + updated_by="admin", + ) + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=db_row + ) + mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + # The team is gone -> every team lookup returns None. + mock_prisma.db.litellm_teamtable = AsyncMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.update = AsyncMock() + mock_prisma.db.litellm_modeltable = AsyncMock() + mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[]) + return mock_prisma + + @pytest.mark.asyncio + async def test_proxy_admin_can_delete_model_when_team_deleted(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, + delete_model as delete_model_endpoint, + ) + + team_id = "deleted-team-xyz" + model_id = "orphaned-byok-1" + mock_prisma = self._orphaned_model_mocks(team_id, model_id) + + admin_user = UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + _PS = "litellm.proxy.proxy_server" + _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", MagicMock()), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.user_api_key_cache", MagicMock()), + patch(f"{_MOD}._refresh_cached_team", new=AsyncMock()), + ): + result = await delete_model_endpoint( + model_info=ModelInfoDelete(id=model_id), + user_api_key_dict=admin_user, + ) + + assert "deleted successfully" in result["message"] + mock_prisma.db.litellm_proxymodeltable.delete.assert_awaited_once() + # Team is gone -> no team.models cleanup to do. + mock_prisma.db.litellm_teamtable.update.assert_not_awaited() + + @pytest.mark.asyncio + async def test_non_admin_cannot_delete_model_when_team_deleted(self): + """A missing team must never let a non-admin delete the orphan (no fail-open).""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, + delete_model as delete_model_endpoint, + ) + from litellm.proxy.proxy_server import ProxyException + + team_id = "deleted-team-abc" + model_id = "orphaned-byok-2" + mock_prisma = self._orphaned_model_mocks(team_id, model_id) + + non_admin = UserAPIKeyAuth( + user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER + ) + + _PS = "litellm.proxy.proxy_server" + _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", MagicMock()), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.user_api_key_cache", MagicMock()), + patch(f"{_MOD}._refresh_cached_team", new=AsyncMock()), + ): + with pytest.raises(ProxyException) as exc_info: + await delete_model_endpoint( + model_info=ModelInfoDelete(id=model_id), + user_api_key_dict=non_admin, + ) + + assert str(exc_info.value.code) == "403" + mock_prisma.db.litellm_proxymodeltable.delete.assert_not_awaited() + + @pytest.mark.asyncio + async def test_live_team_delete_looks_up_team_once(self): + """The auth check must not add a redundant team query on the live-team path.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, + delete_model as delete_model_endpoint, + ) + from litellm.proxy.proxy_server import ProxyException + + team_id = "live-team-1" + model_id = "live-byok-1" + db_row = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name=f"model_name_{team_id}_abc-uuid", + litellm_params={"model": "openai/gpt-4.1-nano"}, + model_info={ + "id": model_id, + "team_id": team_id, + "team_public_model_name": "live-gpt", + }, + created_by="admin", + updated_by="admin", + ) + team_row = LiteLLM_TeamTable( + team_id=team_id, + team_alias="live-team", + members_with_roles=[Member(user_id="admin", role="admin")], + models=["live-gpt"], + ) + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=db_row + ) + mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_teamtable = AsyncMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_modeltable = AsyncMock() + mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[]) + + # A team member who is not the team admin: rejected before the delete runs, + # so the only team lookup is the single one inside the auth check. + non_admin = UserAPIKeyAuth( + user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER + ) + + _PS = "litellm.proxy.proxy_server" + _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", MagicMock()), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.user_api_key_cache", MagicMock()), + patch(f"{_MOD}._refresh_cached_team", new=AsyncMock()), + ): + with pytest.raises(ProxyException) as exc_info: + await delete_model_endpoint( + model_info=ModelInfoDelete(id=model_id), + user_api_key_dict=non_admin, + ) + + assert str(exc_info.value.code) == "403" + assert mock_prisma.db.litellm_teamtable.find_unique.await_count == 1 + mock_prisma.db.litellm_proxymodeltable.delete.assert_not_awaited() + + class TestGetTeamDeployments: """Tests for _get_team_deployments which filters by model_name prefix + Python-side team_id check.""" @@ -2000,6 +2185,148 @@ async def test_multiple_deployments_mixed_filtering(self): assert result[0] is dep1 +def _model_row(model_id: str, team_id: str): + row = MagicMock() + row.model_id = model_id + row.model_name = f"model_name_{team_id}_{model_id}" + row.model_info = {"team_id": team_id} + return row + + +class _TxProxyModelTable: + """Transactional proxy-model table that records the order of DB writes.""" + + def __init__(self, rows, events): + self._rows = list(rows) + self.events = events + + async def find_many(self, where): + prefix = where["model_name"]["startswith"] + return [r for r in self._rows if r.model_name.startswith(prefix)] + + async def delete_many(self, where): + ids = list(where["model_id"]["in"]) + self.events.append(("delete_many", tuple(ids))) + self._rows = [r for r in self._rows if r.model_id not in ids] + return len(ids) + + +class _TxPrismaClient: + """Minimal prisma stub whose ``db.tx()`` yields a transaction and records commit.""" + + def __init__(self, rows): + self.events: list = [] + self._table = _TxProxyModelTable(rows, self.events) + tx = MagicMock() + tx.litellm_proxymodeltable = self._table + outer = self + + class _TxCM: + async def __aenter__(self): + return tx + + async def __aexit__(self, *exc): + outer.events.append(("commit",)) + return False + + self.db = MagicMock() + self.db.tx = MagicMock(return_value=_TxCM()) + + +class _RecordingRouter: + def __init__(self, events): + self.events = events + self.deleted: list = [] + + def delete_deployment(self, id): # noqa: A002 - matches router signature + self.events.append(("router", id)) + self.deleted.append(id) + + +class TestDeleteTeamModels: + """delete_team_models must remove every team's BYOK models in one transaction + and sync the in-memory router only after that transaction commits.""" + + @pytest.mark.asyncio + async def test_deletes_all_teams_models_and_syncs_router(self): + rows = [_model_row("a1", "team_a"), _model_row("b1", "team_b")] + prisma = _TxPrismaClient(rows) + router = _RecordingRouter(prisma.events) + + deleted = await delete_team_models( + team_ids=["team_a", "team_b"], + prisma_client=prisma, + llm_router=router, + ) + + assert sorted(deleted) == ["a1", "b1"] + assert sorted(router.deleted) == ["a1", "b1"] + + @pytest.mark.asyncio + async def test_router_sync_happens_after_commit(self): + """Race-safety: the router is touched only once the DB transaction has + committed, so a rollback can never leave a deployment without its row.""" + rows = [_model_row("a1", "team_a"), _model_row("b1", "team_b")] + prisma = _TxPrismaClient(rows) + router = _RecordingRouter(prisma.events) + + await delete_team_models( + team_ids=["team_a", "team_b"], prisma_client=prisma, llm_router=router + ) + + commit_idx = prisma.events.index(("commit",)) + router_indices = [i for i, e in enumerate(prisma.events) if e[0] == "router"] + delete_indices = [ + i for i, e in enumerate(prisma.events) if e[0] == "delete_many" + ] + assert router_indices, "router was never synced" + assert all(i > commit_idx for i in router_indices) + assert all(i < commit_idx for i in delete_indices) + + @pytest.mark.asyncio + async def test_only_owning_team_models_deleted(self): + """A row sharing the prefix but a different model_info.team_id is left alone.""" + mine = _model_row("a1", "team_a") + intruder = MagicMock() + intruder.model_id = "x9" + intruder.model_name = "model_name_team_a_x9" + intruder.model_info = {"team_id": "someone_else"} + prisma = _TxPrismaClient([mine, intruder]) + router = _RecordingRouter(prisma.events) + + deleted = await delete_team_models( + team_ids=["team_a"], prisma_client=prisma, llm_router=router + ) + + assert deleted == ["a1"] + assert router.deleted == ["a1"] + + @pytest.mark.asyncio + async def test_no_models_no_writes(self): + prisma = _TxPrismaClient([]) + router = _RecordingRouter(prisma.events) + + deleted = await delete_team_models( + team_ids=["team_a"], prisma_client=prisma, llm_router=router + ) + + assert deleted == [] + assert router.deleted == [] + assert not any(e[0] == "delete_many" for e in prisma.events) + + @pytest.mark.asyncio + async def test_missing_router_is_safe(self): + rows = [_model_row("a1", "team_a")] + prisma = _TxPrismaClient(rows) + + deleted = await delete_team_models( + team_ids=["team_a"], prisma_client=prisma, llm_router=None + ) + + assert deleted == ["a1"] + assert any(e[0] == "delete_many" for e in prisma.events) + + def _build_db_model_for_blocked_test(): from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo 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 d580f1f7703..d4bc3841668 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4451,37 +4451,38 @@ async def test_new_team_org_scoped_models_not_in_org_models(): @pytest.mark.asyncio -async def test_update_team_standalone_budget_exceeds_user_limit(): +async def test_update_team_standalone_budget_raise_blocked_for_team_admin(): """ - Test that /team/update for a standalone team fails when new budget exceeds user's max_budget. + Test that /team/update for a standalone team blocks a non-proxy-admin + (team admin) from RAISING the team budget above the team's current value. + + Raising a team's spend ceiling is a budget-authority action reserved for + proxy admins. The rejection is NOT based on the caller's personal budget. Scenario: - - User has personal max_budget=$50 - - Standalone team exists (no organization_id) - - User tries to update team budget to $100 - - Expected: Should fail with error about exceeding user budget + - Team admin (internal_user) manages the team + - Standalone team exists with current budget=$30 + - Admin tries to raise team budget to $100 + - Expected: 403 (only a proxy admin may raise the team budget) """ from fastapi import Request from litellm.proxy._types import ( - LiteLLM_UserTable, ProxyException, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team - # Create non-admin user with restrictive personal budget - non_admin_user = UserAPIKeyAuth( + team_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="non-admin-update-test", models=[], ) - # Create update request with budget exceeding user's limit update_request = UpdateTeamRequest( team_id="standalone-team-123", - max_budget=100.0, # Exceeds user's $50 limit + max_budget=100.0, # Raise above the team's current $30 ) dummy_request = MagicMock(spec=Request) @@ -4492,13 +4493,13 @@ async def test_update_team_standalone_budget_exceeds_user_limit(): patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, + ), ): - # Mock existing standalone team (no organization_id) mock_existing_team = MagicMock() mock_existing_team.team_id = "standalone-team-123" mock_existing_team.organization_id = None # Standalone team mock_existing_team.max_budget = 30.0 + mock_existing_team.model_id = None mock_existing_team.model_dump.return_value = { "team_id": "standalone-team-123", "organization_id": None, @@ -4510,25 +4511,446 @@ async def test_update_team_standalone_budget_exceeds_user_limit(): mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( return_value=mock_existing_team ) + mock_cache.async_get_cache = AsyncMock(return_value=None) - # Mock user cache to return user with restrictive budget - mock_user_obj = LiteLLM_UserTable( - user_id="non-admin-update-test", - max_budget=50.0, # User's budget limit + with pytest.raises(ProxyException) as exc_info: + await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=team_admin_user, + ) + + assert exc_info.value.code == "403" + assert "proxy admin" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_update_team_standalone_budget_raise_allowed_for_proxy_admin(): + """ + Test that a proxy admin CAN raise a standalone team's budget on /team/update. + + Scenario: + - Caller is a proxy admin + - Standalone team exists with current budget=$30 + - Proxy admin raises team budget to $100 + - Expected: Should succeed (proxy admin holds budget authority) + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import update_team + + proxy_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="proxy-admin-update-test", + models=[], + ) + + update_request = UpdateTeamRequest( + team_id="standalone-team-123", + max_budget=100.0, # Raise above the team's current $30 + ) + + dummy_request = MagicMock(spec=Request) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ), + ): + mock_existing_team = MagicMock() + mock_existing_team.team_id = "standalone-team-123" + mock_existing_team.organization_id = None + mock_existing_team.max_budget = 30.0 + mock_existing_team.model_id = None + mock_existing_team.model_dump.return_value = { + "team_id": "standalone-team-123", + "organization_id": None, + "max_budget": 30.0, + "members_with_roles": [ + {"user_id": "proxy-admin-update-test", "role": "admin"} + ], + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_existing_team ) - mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + mock_updated_team = MagicMock() + mock_updated_team.team_id = "standalone-team-123" + mock_updated_team.organization_id = None + mock_updated_team.max_budget = 100.0 + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "standalone-team-123", + "organization_id": None, + "max_budget": 100.0, + } + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team + ) + + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=proxy_admin, + ) + + assert result is not None + assert result["data"].max_budget == 100.0 + + +@pytest.mark.asyncio +async def test_update_team_standalone_budget_removal_blocked_for_team_admin(): + """ + A team admin must not be able to REMOVE a team's spend ceiling + (max_budget=null), which is the strongest possible raise (finite -> unlimited). + + Scenario: + - Team admin (internal_user) manages a team with current budget=$500 + - Admin explicitly sets max_budget=None to strip the cap + - Expected: 403 (only a proxy admin can remove the team budget) + """ + from fastapi import Request + + from litellm.proxy._types import ( + ProxyException, + UpdateTeamRequest, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import update_team + + team_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="budget-removal-admin", + models=[], + ) + + # Explicitly set max_budget=None so it lands in model_fields_set and would be + # persisted by data.json(exclude_unset=True). + update_request = UpdateTeamRequest( + team_id="standalone-team-123", + max_budget=None, + ) + assert "max_budget" in update_request.model_fields_set + + dummy_request = MagicMock(spec=Request) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ), + ): + mock_existing_team = MagicMock() + mock_existing_team.team_id = "standalone-team-123" + mock_existing_team.organization_id = None + mock_existing_team.max_budget = 500.0 + mock_existing_team.model_id = None + mock_existing_team.model_dump.return_value = { + "team_id": "standalone-team-123", + "organization_id": None, + "max_budget": 500.0, + "members_with_roles": [ + {"user_id": "budget-removal-admin", "role": "admin"} + ], + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_existing_team + ) + mock_cache.async_get_cache = AsyncMock(return_value=None) - # Should raise ProxyException because new budget exceeds user's max_budget with pytest.raises(ProxyException) as exc_info: await update_team( data=update_request, http_request=dummy_request, - user_api_key_dict=non_admin_user, + user_api_key_dict=team_admin_user, ) - # Verify exception details - assert exc_info.value.code == "400" - assert "budget" in str(exc_info.value.message).lower() + assert exc_info.value.code == "403" + assert "remove" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_update_team_standalone_uncapped_team_admin_sets_finite_allowed(): + """ + When a team currently has NO cap (max_budget=None / unlimited), a team admin + setting a finite max_budget is a RESTRICTION, not a raise, and is + intentionally allowed. + + Scenario: + - Team admin manages a team with current max_budget=None (unlimited) + - Admin sets max_budget=1000 (unlimited -> finite is more restrictive) + - Expected: 200 + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import update_team + + team_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="uncapped-team-admin", + models=[], + ) + + update_request = UpdateTeamRequest( + team_id="standalone-uncapped-123", + max_budget=1000.0, + ) + + dummy_request = MagicMock(spec=Request) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ), + ): + mock_existing_team = MagicMock() + mock_existing_team.team_id = "standalone-uncapped-123" + mock_existing_team.organization_id = None + mock_existing_team.max_budget = None # team has no cap + mock_existing_team.model_id = None + mock_existing_team.model_dump.return_value = { + "team_id": "standalone-uncapped-123", + "organization_id": None, + "max_budget": None, + "members_with_roles": [ + {"user_id": "uncapped-team-admin", "role": "admin"} + ], + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_existing_team + ) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + mock_updated_team = MagicMock() + mock_updated_team.team_id = "standalone-uncapped-123" + mock_updated_team.organization_id = None + mock_updated_team.max_budget = 1000.0 + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "standalone-uncapped-123", + "organization_id": None, + "max_budget": 1000.0, + } + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team + ) + + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=team_admin_user, + ) + + assert result is not None + assert result["data"].max_budget == 1000.0 + + +@pytest.mark.asyncio +async def test_update_team_standalone_unchanged_budget_allowed(): + """ + Test that /team/update for a standalone team does NOT compare against the + caller's personal max_budget when the budget is unchanged. + + This is the LiteLLM UI scenario: the UI sends the full team object on every + update (including the unchanged max_budget). A team admin only changing + tpm_limit should not be blocked by a budget the team already has. + + Scenario: + - User (team admin) has personal max_budget=$100 + - Standalone team exists with current budget=$500 + - User updates tpm_limit and re-sends the unchanged max_budget=$500 + - Expected: Should succeed (budget unchanged, not an increase) + """ + from fastapi import Request + + from litellm.proxy._types import ( + LiteLLM_UserTable, + UpdateTeamRequest, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import update_team + + team_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="standalone-unchanged-budget-admin", + models=[], + ) + + # UI re-sends the unchanged max_budget alongside the tpm_limit change. + update_request = UpdateTeamRequest( + team_id="standalone-unchanged-budget-123", + max_budget=500.0, # Unchanged from the team's current budget + tpm_limit=50000, + ) + + dummy_request = MagicMock(spec=Request) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): + # Mock existing standalone team (no organization_id) with budget=$500 + mock_existing_team = MagicMock() + mock_existing_team.team_id = "standalone-unchanged-budget-123" + mock_existing_team.organization_id = None + mock_existing_team.max_budget = 500.0 + mock_existing_team.model_id = None + mock_existing_team.model_dump.return_value = { + "team_id": "standalone-unchanged-budget-123", + "organization_id": None, + "max_budget": 500.0, + "members_with_roles": [ + {"user_id": "standalone-unchanged-budget-admin", "role": "admin"} + ], + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_existing_team + ) + mock_prisma.jsonify_team_object = lambda db_data: db_data + + # User has a restrictive personal budget that is lower than the team's. + mock_user_obj = LiteLLM_UserTable( + user_id="standalone-unchanged-budget-admin", + max_budget=100.0, + ) + mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) + mock_cache.async_set_cache = AsyncMock() + + mock_updated_team = MagicMock() + mock_updated_team.team_id = "standalone-unchanged-budget-123" + mock_updated_team.organization_id = None + mock_updated_team.max_budget = 500.0 + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "standalone-unchanged-budget-123", + "organization_id": None, + "max_budget": 500.0, + "tpm_limit": 50000, + } + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team + ) + + # Should NOT raise - unchanged budget skips the personal-budget check. + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=team_admin_user, + ) + + assert result is not None + assert result["data"].max_budget == 500.0 + + +@pytest.mark.asyncio +async def test_update_team_standalone_lower_budget_allowed(): + """ + Test that /team/update for a standalone team allows lowering the budget + below the team's current value even when the new value still exceeds the + caller's personal max_budget. + + Scenario: + - User (team admin) has personal max_budget=$100 + - Standalone team exists with current budget=$500 + - User lowers team budget to $300 (a decrease, still above user's $100) + - Expected: Should succeed (decrease is not an increase above team budget) + """ + from fastapi import Request + + from litellm.proxy._types import ( + LiteLLM_UserTable, + UpdateTeamRequest, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import update_team + + team_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="standalone-lower-budget-admin", + models=[], + ) + + update_request = UpdateTeamRequest( + team_id="standalone-lower-budget-123", + max_budget=300.0, # Lower than current $500, still above user's $100 + ) + + dummy_request = MagicMock(spec=Request) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): + mock_existing_team = MagicMock() + mock_existing_team.team_id = "standalone-lower-budget-123" + mock_existing_team.organization_id = None + mock_existing_team.max_budget = 500.0 + mock_existing_team.model_id = None + mock_existing_team.model_dump.return_value = { + "team_id": "standalone-lower-budget-123", + "organization_id": None, + "max_budget": 500.0, + "members_with_roles": [ + {"user_id": "standalone-lower-budget-admin", "role": "admin"} + ], + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_existing_team + ) + mock_prisma.jsonify_team_object = lambda db_data: db_data + + mock_user_obj = LiteLLM_UserTable( + user_id="standalone-lower-budget-admin", + max_budget=100.0, + ) + mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) + mock_cache.async_set_cache = AsyncMock() + + mock_updated_team = MagicMock() + mock_updated_team.team_id = "standalone-lower-budget-123" + mock_updated_team.organization_id = None + mock_updated_team.max_budget = 300.0 + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "standalone-lower-budget-123", + "organization_id": None, + "max_budget": 300.0, + } + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team + ) + + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=team_admin_user, + ) + + assert result is not None + assert result["data"].max_budget == 300.0 @pytest.mark.asyncio @@ -4623,32 +5045,34 @@ async def test_update_team_org_scoped_budget_exceeds_org_limit(): @pytest.mark.asyncio -async def test_update_team_standalone_models_exceeds_user_limit(): +async def test_update_team_standalone_models_not_gated_by_user_limit(): """ - Test that /team/update for a standalone team fails when models are not in user's allowed models. + Test that /team/update for a standalone team does NOT gate the team's models + by the caller's personal allowed models. + + A team admin authorized via _verify_team_access() may set the team's models + independently of their own personal model list on update. Scenario: - - User has personal models=['gpt-3.5-turbo'] + - Team admin has personal models=['gpt-3.5-turbo'] - Standalone team exists (no organization_id) - - User tries to update team models to ['gpt-4'] (not in user's allowed models) - - Expected: Should fail with error about model not in user's allowed models + - Admin updates team models to ['gpt-4'] (not in their personal list) + - Expected: Should succeed (personal models are irrelevant on /team/update) """ from fastapi import Request - from litellm.proxy._types import ProxyException, UpdateTeamRequest, UserAPIKeyAuth + from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import update_team - # Create non-admin user with restrictive personal models - non_admin_user = UserAPIKeyAuth( + team_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="non-admin-update-models-test", - models=["gpt-3.5-turbo"], # Restrictive model list + models=["gpt-3.5-turbo"], # Restrictive personal model list ) - # Create update request with model not in user's allowed list update_request = UpdateTeamRequest( team_id="standalone-team-models-123", - models=["gpt-4"], # Not in user's allowed models + models=["gpt-4"], # Not in the admin's personal allowed models ) dummy_request = MagicMock(spec=Request) @@ -4666,6 +5090,7 @@ async def test_update_team_standalone_models_exceeds_user_limit(): mock_existing_team.team_id = "standalone-team-models-123" mock_existing_team.organization_id = None # Standalone team mock_existing_team.models = ["gpt-3.5-turbo"] + mock_existing_team.model_id = None mock_existing_team.model_dump.return_value = { "team_id": "standalone-team-models-123", "organization_id": None, @@ -4677,18 +5102,30 @@ async def test_update_team_standalone_models_exceeds_user_limit(): mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( return_value=mock_existing_team ) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() - # Should raise ProxyException because model not in user's allowed models - with pytest.raises(ProxyException) as exc_info: - await update_team( - data=update_request, - http_request=dummy_request, - user_api_key_dict=non_admin_user, - ) + mock_updated_team = MagicMock() + mock_updated_team.team_id = "standalone-team-models-123" + mock_updated_team.organization_id = None + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "standalone-team-models-123", + "organization_id": None, + "models": ["gpt-4"], + } + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team + ) - # Verify exception details - assert exc_info.value.code == "400" - assert "model" in str(exc_info.value.message).lower() + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=team_admin_user, + ) + + assert result is not None @pytest.mark.asyncio @@ -5113,32 +5550,35 @@ async def test_update_team_org_scoped_models_with_all_proxy_models(): @pytest.mark.asyncio -async def test_update_team_tpm_limit_exceeds_user_limit(): +async def test_update_team_tpm_limit_not_gated_by_user_limit(): """ - Test that /team/update fails when TPM limit exceeds user's TPM limit. + Test that /team/update does NOT gate the team's tpm_limit by the caller's + personal tpm_limit. + + A team admin authorized via _verify_team_access() may raise the team's + tpm_limit above their own personal tpm_limit on update. Scenario: - - User has tpm_limit=1000 - - User tries to update team with tpm_limit=5000 - - Expected: Should fail with error about exceeding user TPM limit + - Team admin has personal tpm_limit=1000 + - Standalone team exists with tpm_limit=500 + - Admin updates team tpm_limit to 5000 (above their personal 1000) + - Expected: Should succeed (personal tpm is irrelevant on /team/update) """ from fastapi import Request - from litellm.proxy._types import ProxyException, UpdateTeamRequest, UserAPIKeyAuth + from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import update_team - # Create non-admin user with TPM limit - non_admin_user = UserAPIKeyAuth( + team_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="tpm-limit-user", models=[], - tpm_limit=1000, # User's TPM limit + tpm_limit=1000, # Restrictive personal TPM limit ) - # Create update request with TPM exceeding user's limit update_request = UpdateTeamRequest( team_id="team-tpm-test-123", - tpm_limit=5000, # Exceeds user's 1000 limit + tpm_limit=5000, # Above the admin's personal 1000 ) dummy_request = MagicMock(spec=Request) @@ -5147,12 +5587,16 @@ async def test_update_team_tpm_limit_exceeds_user_limit(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ), ): # Mock existing standalone team mock_existing_team = MagicMock() mock_existing_team.team_id = "team-tpm-test-123" mock_existing_team.organization_id = None mock_existing_team.tpm_limit = 500 + mock_existing_team.model_id = None mock_existing_team.model_dump.return_value = { "team_id": "team-tpm-test-123", "organization_id": None, @@ -5162,47 +5606,59 @@ async def test_update_team_tpm_limit_exceeds_user_limit(): mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( return_value=mock_existing_team ) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() - # Should raise ProxyException because new TPM exceeds user's limit - with pytest.raises(ProxyException) as exc_info: - await update_team( - data=update_request, - http_request=dummy_request, - user_api_key_dict=non_admin_user, - ) + mock_updated_team = MagicMock() + mock_updated_team.team_id = "team-tpm-test-123" + mock_updated_team.organization_id = None + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "team-tpm-test-123", + "organization_id": None, + "tpm_limit": 5000, + } + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team + ) - # Verify exception details - assert exc_info.value.code == "400" - assert "tpm" in str(exc_info.value.message).lower() + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=team_admin_user, + ) + + assert result is not None @pytest.mark.asyncio -async def test_update_team_rpm_limit_exceeds_user_limit(): +async def test_update_team_rpm_limit_not_gated_by_user_limit(): """ - Test that /team/update fails when RPM limit exceeds user's RPM limit. + Test that /team/update does NOT gate the team's rpm_limit by the caller's + personal rpm_limit. Scenario: - - User has rpm_limit=100 - - User tries to update team with rpm_limit=500 - - Expected: Should fail with error about exceeding user RPM limit + - Team admin has personal rpm_limit=100 + - Standalone team exists with rpm_limit=50 + - Admin updates team rpm_limit to 500 (above their personal 100) + - Expected: Should succeed (personal rpm is irrelevant on /team/update) """ from fastapi import Request - from litellm.proxy._types import ProxyException, UpdateTeamRequest, UserAPIKeyAuth + from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import update_team - # Create non-admin user with RPM limit - non_admin_user = UserAPIKeyAuth( + team_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="rpm-limit-user", models=[], - rpm_limit=100, # User's RPM limit + rpm_limit=100, # Restrictive personal RPM limit ) - # Create update request with RPM exceeding user's limit update_request = UpdateTeamRequest( team_id="team-rpm-test-123", - rpm_limit=500, # Exceeds user's 100 limit + rpm_limit=500, # Above the admin's personal 100 ) dummy_request = MagicMock(spec=Request) @@ -5211,12 +5667,16 @@ async def test_update_team_rpm_limit_exceeds_user_limit(): patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ), ): # Mock existing standalone team mock_existing_team = MagicMock() mock_existing_team.team_id = "team-rpm-test-123" mock_existing_team.organization_id = None mock_existing_team.rpm_limit = 50 + mock_existing_team.model_id = None mock_existing_team.model_dump.return_value = { "team_id": "team-rpm-test-123", "organization_id": None, @@ -5226,18 +5686,30 @@ async def test_update_team_rpm_limit_exceeds_user_limit(): mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( return_value=mock_existing_team ) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() - # Should raise ProxyException because new RPM exceeds user's limit - with pytest.raises(ProxyException) as exc_info: - await update_team( - data=update_request, - http_request=dummy_request, - user_api_key_dict=non_admin_user, - ) + mock_updated_team = MagicMock() + mock_updated_team.team_id = "team-rpm-test-123" + mock_updated_team.organization_id = None + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "team-rpm-test-123", + "organization_id": None, + "rpm_limit": 500, + } + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team + ) - # Verify exception details - assert exc_info.value.code == "400" - assert "rpm" in str(exc_info.value.message).lower() + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=team_admin_user, + ) + + assert result is not None @pytest.mark.asyncio @@ -6157,6 +6629,14 @@ async def test_delete_team_persists_deleted_teams(monkeypatch): mock_find_many_keys = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys + # delete_team now deletes team BYOK models inside a transaction; this team has none. + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + monkeypatch.setattr( "litellm.proxy.proxy_server.prisma_client", mock_prisma_client, @@ -8306,6 +8786,8 @@ async def test_new_team_encrypts_callback_vars( assert cv["langfuse_secret_key"] != "sk-real" recovered = decrypt_callback_vars(metadata)["logging"][0]["callback_vars"] assert recovered["langfuse_secret_key"] == "sk-real" + + def _non_admin_auth(): return UserAPIKeyAuth( user_id="u-team-admin", user_role=LitellmUserRoles.INTERNAL_USER @@ -8404,3 +8886,329 @@ async def test_update_team_blocks_non_admin_passthrough_routes(mock_db_client): ) assert str(exc.value.code) == "403" assert "allowed_passthrough_routes" in str(exc.value.message) + + +def test_set_budget_reset_at_clears_when_budget_duration_null(): + """ + When budget_duration is explicitly set to null, _set_budget_reset_at + should set budget_reset_at=None in updated_kv so Prisma clears it in the DB. + """ + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at + + data = UpdateTeamRequest(team_id="test-team", budget_duration=None) + updated_kv = {"team_id": "test-team", "budget_duration": None} + + _set_budget_reset_at(data, updated_kv) + + assert "budget_reset_at" in updated_kv + assert updated_kv["budget_reset_at"] is None + + +def test_set_budget_reset_at_noop_when_budget_duration_not_sent(): + """ + When budget_duration is NOT sent (unset), _set_budget_reset_at should + not add budget_reset_at to updated_kv. + """ + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at + + data = UpdateTeamRequest(team_id="test-team") + updated_kv = {"team_id": "test-team"} + + _set_budget_reset_at(data, updated_kv) + + assert "budget_reset_at" not in updated_kv + + +def test_set_budget_reset_at_sets_value_when_budget_duration_provided(): + """ + When budget_duration is set to a valid string, _set_budget_reset_at + should compute and set budget_reset_at. + """ + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at + + data = UpdateTeamRequest(team_id="test-team", budget_duration="30d") + updated_kv = {"team_id": "test-team", "budget_duration": "30d"} + + _set_budget_reset_at(data, updated_kv) + + assert "budget_reset_at" in updated_kv + assert updated_kv["budget_reset_at"] is not None + + +@pytest.mark.asyncio +async def test_clear_team_member_budget_duration_calls_update_budget(): + """ + When team_member_budget_duration is explicitly null and a budget row + exists, clear_team_member_budget_fields should call update_budget + with budget_duration=None and budget_reset_at=None. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-123"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_budget_duration": None, + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields={"team_member_budget_duration"}, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-123" + assert "budget_duration" in budget_request.model_fields_set + assert budget_request.budget_duration is None + assert "budget_reset_at" in budget_request.model_fields_set + assert budget_request.budget_reset_at is None + assert "team_member_budget_duration" not in result + + +@pytest.mark.asyncio +async def test_clear_team_member_budget_clears_max_budget(): + """ + When team_member_budget is explicitly null, clear_team_member_budget_fields + should call update_budget with max_budget=None. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-456"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_budget": None, + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields={"team_member_budget"}, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-456" + assert "max_budget" in budget_request.model_fields_set + assert budget_request.max_budget is None + assert "team_member_budget" not in result + + +@pytest.mark.asyncio +async def test_clear_team_member_rpm_tpm_limits(): + """ + When team_member_rpm_limit and team_member_tpm_limit are explicitly null, + clear_team_member_budget_fields should clear both on the budget row. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-789"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_rpm_limit": None, + "team_member_tpm_limit": None, + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields={"team_member_rpm_limit", "team_member_tpm_limit"}, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-789" + assert "rpm_limit" in budget_request.model_fields_set + assert budget_request.rpm_limit is None + assert "tpm_limit" in budget_request.model_fields_set + assert budget_request.tpm_limit is None + assert "team_member_rpm_limit" not in result + assert "team_member_tpm_limit" not in result + + +@pytest.mark.asyncio +async def test_clear_all_team_member_fields_at_once(): + """ + When all team_member fields are explicitly null, all corresponding + budget row fields should be cleared in a single update. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-all"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_budget": None, + "team_member_budget_duration": None, + "team_member_rpm_limit": None, + "team_member_tpm_limit": None, + } + + all_fields = { + "team_member_budget", + "team_member_budget_duration", + "team_member_rpm_limit", + "team_member_tpm_limit", + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields=all_fields, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-all" + assert budget_request.max_budget is None + assert budget_request.budget_duration is None + assert budget_request.budget_reset_at is None + assert budget_request.rpm_limit is None + assert budget_request.tpm_limit is None + for field in all_fields: + assert field not in result + + +@pytest.mark.asyncio +async def test_team_member_budget_duration_not_sent_does_not_update(): + """ + When team_member_budget_duration is NOT sent in the request, no budget + update should occur and the field should not appear in updated_kv. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + updated_kv = {"team_id": "test-team", "max_budget": 200} + + _team_member_fields_in_request = { + field + for field in [ + "team_member_budget", + "team_member_rpm_limit", + "team_member_tpm_limit", + "team_member_budget_duration", + ] + if field in updated_kv + } + + assert len(_team_member_fields_in_request) == 0 + + TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) + + assert "team_member_budget_duration" not in updated_kv + assert "team_member_budget" not in updated_kv + + +@pytest.mark.asyncio +async def test_clear_team_member_budget_fields_no_budget_row_skips_update(): + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata=None, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_budget": None, + "team_member_rpm_limit": None, + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields={"team_member_budget", "team_member_rpm_limit"}, + ) + + mock_update_budget.assert_not_awaited() + assert "team_member_budget" not in result + assert "team_member_rpm_limit" not in result 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 c763e9c0e98..2efec3e0b34 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -1777,6 +1777,23 @@ def test_html_render_utils_import(self): assert isinstance(html, str) assert len(html) > 0 + def test_success_page_instructs_manual_close_without_false_countdown(self): + """Browsers refuse window.close() on tabs they did not open via window.open() + (the CLI opens the page with webbrowser.open), so a 'closing in 3...' countdown + is a promise the browser usually can't keep and the page gets stuck on + 'Closing...'. The page must instead always show the manual-close instruction + and never advertise an auto-close that won't happen. + """ + from litellm.proxy.common_utils.html_forms.cli_sso_success import ( + render_cli_sso_success_page, + ) + + html = render_cli_sso_success_page() + + assert "You can now close this window and return to your terminal." in html + assert "Closing..." not in html + assert "This window will close in" not in html + class TestCustomUISSO: """Test the custom UI SSO sign-in handler functionality""" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index f8b6fbde3dc..1114b3df0c2 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -321,6 +321,44 @@ def test_cost_calculation_does_not_duplicate_provider_prefix( assert call_kwargs["model"] == "azure_ai/claude-sonnet-4-5_gb_20250929" assert call_kwargs["custom_llm_provider"] == "azure_ai" + def test_passthrough_logging_sets_response_cost_with_server_tool_use_dict(self): + from litellm.types.utils import Choices, Message, ModelResponse + + logging_obj = self._create_mock_logging_obj(model="claude-3-7-sonnet-20250219") + logging_obj.get_router_model_id.return_value = None + logging_obj.litellm_params = {} + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="test", role="assistant"), + ) + ], + created=1234567890, + model="claude-3-7-sonnet-20250219", + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "server_tool_use": {"web_search_requests": 1}, + }, + ) + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-7-sonnet-20250219", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert "response_cost" in kwargs + assert kwargs["response_cost"] > 0 + class TestAnthropicBatchPassthroughCostTracking: """Test cases for Anthropic batch passthrough cost tracking functionality""" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index bfcaaafd335..401ea2ef589 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -257,6 +257,64 @@ def test_is_openai_responses_route(self): ) assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("") == False + def test_is_openai_route_recognizes_cognitiveservices_azure_com(self): + """Azure OpenAI resources created via the newer "Azure AI Foundry" / + Cognitive Services pathway live on `*.cognitiveservices.azure.com` + subdomains rather than the older `openai.azure.com`. All four + is_openai_*_route methods must recognize both Azure subdomains so + cost tracking applies regardless of which Azure naming the user's + resource happens to be on. + """ + cognitive_chat = ( + "https://my-resource.cognitiveservices.azure.com/v1/chat/completions" + ) + cognitive_images_gen = ( + "https://my-resource.cognitiveservices.azure.com/v1/images/generations" + ) + cognitive_images_edit = ( + "https://my-resource.cognitiveservices.azure.com/v1/images/edits" + ) + cognitive_responses = ( + "https://my-resource.cognitiveservices.azure.com/v1/responses" + ) + + assert ( + OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route( + cognitive_chat + ) + is True + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_image_generation_route( + cognitive_images_gen + ) + is True + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_image_editing_route( + cognitive_images_edit + ) + is True + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_responses_route( + cognitive_responses + ) + is True + ) + + # Cross-route negatives still hold for cognitiveservices hosts. + assert ( + OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route( + cognitive_responses + ) + is False + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_responses_route(cognitive_chat) + is False + ) + @patch("litellm.completion_cost") @patch( "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" @@ -625,36 +683,43 @@ def test_azure_passthrough_tags_metadata_model_provider( "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" ) @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.get_provider_config" + "litellm.llms.openai.responses.transformation.OpenAIResponsesAPIConfig.transform_response_api_response" ) def test_responses_api_cost_tracking( - self, mock_get_provider_config, mock_get_standard_logging, mock_completion_cost + self, + mock_transform_responses, + mock_get_standard_logging, + mock_completion_cost, ): - """Test cost tracking for responses API route""" + """Test cost tracking for responses API route. + + Mocks the Responses-API transformer (the dedicated one this branch + of the handler dispatches into post-fix) so we can assert the + downstream cost-calculation contract without depending on the + real transformer's full behavior. + """ # Arrange mock_completion_cost.return_value = 0.000050 mock_get_standard_logging.return_value = {"test": "logging_payload"} - # Mock the provider config's transform_response to return a valid ModelResponse - from litellm import ModelResponse + # Mock the Responses transformer's return — a ResponsesAPIResponse + # carrying the usage fields downstream cost-calc expects. + from litellm.types.llms.openai import ResponsesAPIResponse - mock_model_response = ModelResponse( + mock_responses_api_response = ResponsesAPIResponse.model_construct( id="resp_abc123", + object="response", + created_at=1677652288, model="gpt-4o-2024-08-06", - choices=[ - { - "message": { - "role": "assistant", - "content": "Hello! How can I help you today?", - } - } - ], - usage={"prompt_tokens": 20, "completion_tokens": 15, "total_tokens": 35}, + status="completed", + output=[], + usage={ + "input_tokens": 20, + "output_tokens": 15, + "total_tokens": 35, + }, ) - - mock_provider_config = MagicMock() - mock_provider_config.transform_response.return_value = mock_model_response - mock_get_provider_config.return_value = mock_provider_config + mock_transform_responses.return_value = mock_responses_api_response # Mock responses API response mock_responses_response = { @@ -710,6 +775,109 @@ def test_responses_api_cost_tracking( assert mock_logging_obj.model_call_details["model"] == "gpt-4o" assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai" + @patch("litellm.completion_cost") + @patch( + "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" + ) + def test_responses_api_uses_responses_transformer_not_chat_completions( + self, mock_get_standard_logging, mock_completion_cost + ): + """Regression test for the Responses-API cost-tracking dispatch bug. + + BUG: the `elif is_responses:` branch in `openai_passthrough_handler` + was calling `OpenAIConfig.transform_response` (the chat-completions + transformer) on a Responses API payload. Chat-completions + transform_response expects `choices: [...]` in the raw response; + the Responses API uses `output: [...]` and `usage.input_tokens` / + `usage.output_tokens` (not `prompt_tokens` / `completion_tokens`). + The result was a KeyError 'choices' inside + `convert_to_model_response_object`, swallowed by the surrounding + try/except, and the SpendLogs row was written with zero tokens + and zero spend. + + FIX: use the dedicated `OpenAIResponsesAPIConfig.transform_response_api_response` + for the Responses branch. + + This test exercises the REAL transformer (no mocked + `get_provider_config`) so that running it against the un-fixed + handler raises and running it against the fixed handler succeeds. + """ + mock_completion_cost.return_value = 0.000050 + mock_get_standard_logging.return_value = {"test": "logging_payload"} + + # A real-shaped Azure / OpenAI Responses API payload — NO `choices`, + # uses `output` and `usage.input_tokens` / `usage.output_tokens`. + responses_api_body = { + "id": "resp_abc123", + "object": "response", + "created_at": 1677652288, + "model": "gpt-4o-2024-08-06", + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Hello!", + } + ], + } + ], + "usage": { + "input_tokens": 20, + "output_tokens": 15, + "total_tokens": 35, + }, + } + + mock_httpx_response = self._create_mock_httpx_response(responses_api_body) + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gpt-4o", + "custom_llm_provider": "openai", + } + + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=responses_api_body, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/responses", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"model": "gpt-4o", "input": "Tell me about AI"}, + **kwargs, + ) + + # Pre-fix this assertion fails — the handler swallows the + # KeyError raised by the chat-completions transformer and falls + # back to the passthrough_chat_handler which yields a different + # response_cost value. Post-fix, the Responses transformer + # succeeds and we get the mocked 0.000050. + assert result is not None + assert result["kwargs"]["response_cost"] == 0.000050 + assert result["kwargs"]["model"] == "gpt-4o" + + # `completion_cost` must be called with the responses call type + # and a `ResponsesAPIResponse` (not a `ModelResponse`). + mock_completion_cost.assert_called_once() + call_kwargs = mock_completion_cost.call_args[1] + assert call_kwargs["call_type"] == "responses" + + from litellm.types.llms.openai import ResponsesAPIResponse + + assert isinstance(call_kwargs["completion_response"], ResponsesAPIResponse), ( + "completion_response must be a ResponsesAPIResponse; passing a " + "chat-completions ModelResponse means the Responses transformer " + "isn't being used and we're back in the bug." + ) + class TestOpenAIPassthroughIntegration: """Integration tests for OpenAI passthrough cost tracking""" @@ -766,6 +934,14 @@ def test_is_openai_route_detection(self): == True ) assert self.handler.is_openai_route("https://api.openai.com/v1/models") == True + # Azure OpenAI on the shared Cognitive Services domain, identified by an + # OpenAI-style path segment. + assert ( + self.handler.is_openai_route( + "https://my-resource.cognitiveservices.azure.com/v1/chat/completions" + ) + == True + ) # Negative cases assert ( @@ -782,8 +958,150 @@ def test_is_openai_route_detection(self): self.handler.is_openai_route("https://api.assemblyai.com/v2/transcript") == False ) + # Non-OpenAI Azure Cognitive Services share the `cognitiveservices.azure.com` + # domain but must NOT be classified as OpenAI routes (no OpenAI path segment). + assert ( + self.handler.is_openai_route( + "https://my-resource.cognitiveservices.azure.com/speechtotext/v3.1/recognize" + ) + == False + ) + assert ( + self.handler.is_openai_route( + "https://my-resource.cognitiveservices.azure.com/vision/v3.2/analyze" + ) + == False + ) + # A look-alike domain that merely contains an OpenAI host as a substring + # must be rejected by the suffix-based hostname match. + assert ( + self.handler.is_openai_route( + "https://cognitiveservices.azure.com.attacker.example/v1/chat/completions" + ) + == False + ) assert self.handler.is_openai_route("") == False + def test_is_supported_openai_endpoint_includes_responses_api(self): + """Regression test for the outer dispatch gate. + + `_is_supported_openai_endpoint` is the gate that decides whether the + OpenAI handler runs for a given URL. Before this gate accepted the + Responses API, calls to `/v1/responses` would fail the gate and the + handler's `elif is_responses:` branch was unreachable in the live + success-handler pipeline — every Responses-API call landed in + `LiteLLM_SpendLogs` with zero tokens / zero spend even though the + handler had a Responses branch internally. + + This test exercises the dispatch decision directly so future + refactors of `_is_supported_openai_endpoint` can't silently + remove Responses from the OR-chain without a test failure. + """ + # Responses must be supported on api.openai.com and openai.azure.com. + assert ( + self.handler._is_supported_openai_endpoint( + "https://api.openai.com/v1/responses" + ) + is True + ) + assert ( + self.handler._is_supported_openai_endpoint( + "https://openai.azure.com/v1/responses" + ) + is True + ) + # The other supported endpoints stay supported (no regression). + assert ( + self.handler._is_supported_openai_endpoint( + "https://api.openai.com/v1/chat/completions" + ) + is True + ) + assert ( + self.handler._is_supported_openai_endpoint( + "https://api.openai.com/v1/images/generations" + ) + is True + ) + assert ( + self.handler._is_supported_openai_endpoint( + "https://api.openai.com/v1/images/edits" + ) + is True + ) + # Unsupported OpenAI endpoints (e.g. /v1/models) still return False. + assert ( + self.handler._is_supported_openai_endpoint( + "https://api.openai.com/v1/models" + ) + is False + ) + + @patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler" + ) + @pytest.mark.asyncio + async def test_success_handler_dispatches_responses_api_to_openai_handler( + self, mock_openai_handler + ): + """End-to-end dispatch test for the Responses API path. + + Pre-fix: `_is_supported_openai_endpoint` returned False for + `/v1/responses` URLs, so the OpenAI handler was never called. + This test would fail (mock never invoked) on the un-fixed + success_handler — passes only when the dispatch gate accepts + Responses URLs. + """ + mock_openai_handler.return_value = { + "result": {"id": "resp_abc123"}, + "kwargs": { + "response_cost": 0.0001, + "model": "gpt-4o", + "custom_llm_provider": "openai", + }, + } + + mock_httpx_response = MagicMock(spec=httpx.Response) + mock_httpx_response.text = ( + '{"id": "resp_abc123", "object": "response", ' + '"output": [], "usage": {"input_tokens": 5, "output_tokens": 3}}' + ) + + mock_logging_obj = AsyncMock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.async_success_handler = AsyncMock() + + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/responses", + request_body={"model": "gpt-4o", "input": "Hello"}, + request_method="POST", + ) + + await self.handler.pass_through_async_success_handler( + httpx_response=mock_httpx_response, + response_body={ + "id": "resp_abc123", + "object": "response", + "output": [], + "usage": {"input_tokens": 5, "output_tokens": 3}, + }, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/responses", + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"model": "gpt-4o", "input": "Hello"}, + passthrough_logging_payload=passthrough_payload, + ) + + # The OpenAI handler MUST have been invoked. Pre-fix the dispatch + # gate filtered Responses URLs out and the mock was never called. + mock_openai_handler.assert_called_once() + # And we can verify it was dispatched with the Responses URL. + call_kwargs = mock_openai_handler.call_args.kwargs + assert call_kwargs["url_route"] == "https://api.openai.com/v1/responses" + @patch( "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler" ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index c9be00afed0..4eab1a4bf61 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -24,11 +24,13 @@ get_vertex_base_url, llm_passthrough_factory_proxy_route, milvus_proxy_route, + mistral_proxy_route, openai_proxy_route, vertex_discovery_proxy_route, vertex_proxy_route, vllm_proxy_route, ) +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials @@ -1092,9 +1094,9 @@ def test_vertex_passthrough_handler_embed_content_google_ai_studio_url(self): assert result is not None assert result["result"] is not None - assert result["kwargs"].get("custom_llm_provider") == "gemini", ( - "Google AI Studio embedContent URLs must set custom_llm_provider=gemini, not vertex_ai" - ) + assert ( + result["kwargs"].get("custom_llm_provider") == "gemini" + ), "Google AI Studio embedContent URLs must set custom_llm_provider=gemini, not vertex_ai" assert result["kwargs"].get("model") == "gemini-embedding-2-preview" mock_completion_cost.assert_called_once() @@ -1261,6 +1263,78 @@ async def test_is_streaming_request_fn(): assert await is_streaming_request_fn(mock_request) is True +@pytest.mark.asyncio +async def test_mistral_passthrough_accepts_multipart_without_json_parsing(): + boundary = "----litellm-test-boundary" + body = ( + f"--{boundary}\r\n" + 'Content-Disposition: form-data; name="purpose"\r\n\r\n' + "ocr\r\n" + f"--{boundary}\r\n" + 'Content-Disposition: form-data; name="file"; filename="document.pdf"\r\n' + "Content-Type: application/pdf\r\n\r\n" + "%PDF-1.4 test\r\n" + f"--{boundary}--\r\n" + ).encode("utf-8") + + async def receive(): + return { + "type": "http.request", + "body": body, + "more_body": False, + } + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/mistral/v1/files", + "headers": [ + ( + b"content-type", + f"multipart/form-data; boundary={boundary}".encode("utf-8"), + ) + ], + "query_string": b"", + }, + receive=receive, + ) + + captured_kwargs = {} + + async def fake_endpoint(request, fastapi_response, user_api_key_dict): + return {"ok": True} + + def fake_create_pass_through_route(**kwargs): + captured_kwargs.update(kwargs) + return fake_endpoint + + user_api_key_dict = UserAPIKeyAuth(token="test-key") + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", + return_value="mistral-test-key", + ), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + side_effect=fake_create_pass_through_route, + ), + ): + response = await mistral_proxy_route( + endpoint="v1/files", + request=request, + fastapi_response=Response(), + user_api_key_dict=user_api_key_dict, + ) + + assert response == {"ok": True} + assert captured_kwargs["is_streaming_request"] is False + assert captured_kwargs["custom_headers"] == { + "Authorization": "Bearer mistral-test-key" + } + + class TestBedrockLLMProxyRoute: @pytest.mark.asyncio async def test_bedrock_llm_proxy_route_application_inference_profile(self): @@ -1803,6 +1877,9 @@ async def test_pass_through_request_with_forward_headers_true(self): mock_logging_obj.pre_call_hook = AsyncMock(return_value=mock_request_body) mock_logging_obj.post_call_success_hook = AsyncMock() mock_logging_obj.post_call_failure_hook = AsyncMock() + mock_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={} + ) # Call pass_through_request with forward_headers=True result = await pass_through_request( @@ -1901,6 +1978,9 @@ async def test_pass_through_request_with_forward_headers_false(self): mock_logging_obj.pre_call_hook = AsyncMock(return_value=mock_request_body) mock_logging_obj.post_call_success_hook = AsyncMock() mock_logging_obj.post_call_failure_hook = AsyncMock() + mock_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={} + ) # Call pass_through_request with forward_headers=False (default) result = await pass_through_request( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 9b9d5e22a43..31f1c1c57d6 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -957,6 +957,9 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): return_value={"test": "data"} ) mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-test": "value"} + ) # Setup mock for http response mock_response = MagicMock() @@ -1069,6 +1072,9 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): return_value={"model": "claude-3", "stream": True} ) mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-test": "value"} + ) upstream_response = MagicMock() upstream_response.status_code = 200 @@ -1134,6 +1140,9 @@ async def test_pass_through_request_sse_response_marks_logging_obj_as_stream(): return_value={"model": "claude-3"} ) mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-test": "value"} + ) upstream_response = MagicMock() upstream_response.status_code = 200 @@ -1975,6 +1984,9 @@ async def test_pass_through_request_query_params_forwarding(): mock_proxy_logging.pre_call_hook = AsyncMock( return_value=test_body ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-test": "value"} + ) # Setup mock for http response mock_response = MagicMock() @@ -2703,6 +2715,10 @@ async def _hook_mutates_body(**kwargs): "litellm.proxy.proxy_server.proxy_logging_obj.pre_call_hook", new=AsyncMock(side_effect=_hook_mutates_body), ), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_response_headers_hook", + new=AsyncMock(return_value={}), + ), patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler", new=AsyncMock(), @@ -2763,6 +2779,10 @@ async def test_pass_through_request_streaming_uses_content_for_state_raw_body(): "litellm.proxy.proxy_server.proxy_logging_obj.pre_call_hook", new=AsyncMock(side_effect=lambda **kw: kw["data"]), ), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_response_headers_hook", + new=AsyncMock(return_value={}), + ), patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler", new=AsyncMock(), diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index eafe71e1063..a2f7476abd1 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -130,6 +130,7 @@ async def test_post_call_success_hook_called_when_guardrails_configured( mock_proxy_logging.post_call_success_hook = AsyncMock( return_value=_GEMINI_RESPONSE ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) with _common_patches(mock_proxy_logging, mock_response): await pass_through_request( @@ -155,6 +156,7 @@ async def test_post_call_success_hook_skipped_when_no_guardrails( mock_proxy_logging = MagicMock() mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) with _common_patches(mock_proxy_logging, mock_response): result = await pass_through_request( 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 164538a2757..677d358428d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -601,6 +601,98 @@ async def test_ProxyConfig_load_config_missing_file_raises(monkeypatch): await pc.load_config(router=None, config_file_path="/no/file.yaml") +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_forwards_callback_specific_params( + tmp_path, monkeypatch +): + """Regression: callback_settings from config must be forwarded to + initialize_callbacks_on_proxy as callback_specific_params. + + Callbacks like DatadogCostManagementLogger read their init params (e.g. + cost_tag_keys) from callback_specific_params[]. If the + argument is dropped at the call site, they silently initialize with empty + params and the configured allowlist never takes effect. + """ + f = tmp_path / "c.yaml" + f.write_text( + "model_list: []\n" + "general_settings: {}\n" + "callback_settings:\n" + " datadog_cost_management:\n" + " cost_tag_keys:\n" + " - capability\n" + " - platform\n" + " - ai_product\n" + "litellm_settings:\n" + ' callbacks: ["datadog_cost_management"]\n' + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + + captured = {} + + def _fake_initialize_callbacks_on_proxy(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.initialize_callbacks_on_proxy", + _fake_initialize_callbacks_on_proxy, + ) + + pc = ProxyConfig() + await pc.load_config(router=None, config_file_path=str(f)) + + # The callbacks branch must forward the loaded callback_settings. + assert captured.get("callback_specific_params") == { + "datadog_cost_management": { + "cost_tag_keys": ["capability", "platform", "ai_product"] + } + } + + +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_blank_callback_settings_does_not_crash( + tmp_path, monkeypatch +): + """Regression: `callback_settings:` with no body loads as None because + dict.get() only falls back to the default when the key is absent. The None + was forwarded verbatim to initialize_callbacks_on_proxy, where the first + `"" in callback_specific_params` membership test raised + TypeError: argument of type 'NoneType' is not iterable, aborting startup. + Startup must succeed and the callback must initialize with its defaults. + """ + f = tmp_path / "c.yaml" + f.write_text( + "model_list: []\n" + "general_settings: {}\n" + "callback_settings:\n" + "litellm_settings:\n" + ' callbacks: ["compression_interception"]\n' + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + + from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, + ) + + original_callbacks = ( + list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] + ) + litellm.callbacks = [] + try: + pc = ProxyConfig() + await pc.load_config(router=None, config_file_path=str(f)) + + assert any( + isinstance(c, CompressionInterceptionLogger) for c in litellm.callbacks + ) + finally: + litellm.callbacks = original_callbacks + + # --------------------------------------------------------------------------- # ProxyConfig._init_non_llm_configs # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py index d88bcf136e9..74542a3eaf6 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py @@ -26,6 +26,7 @@ def patched_speech(monkeypatch): MagicMock( pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]), post_call_failure_hook=AsyncMock(), + post_call_response_headers_hook=AsyncMock(return_value={}), update_request_status=AsyncMock(), ), ) @@ -61,6 +62,7 @@ def patched_speech_error(monkeypatch): MagicMock( pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]), post_call_failure_hook=AsyncMock(), + post_call_response_headers_hook=AsyncMock(return_value={}), update_request_status=AsyncMock(), ), ) diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 98259824378..017f4bd4368 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -60,6 +60,15 @@ def test_v2_model_info_invalid_page_returns_422(client, auth_as, empty_router): assert "detail" in response.json() +def test_v2_model_info_in_openapi_schema(): + """``GET /v2/model/info`` is published in the proxy OpenAPI/Swagger spec.""" + from litellm.proxy.proxy_server import get_openapi_schema + + schema = get_openapi_schema() + assert "/v2/model/info" in schema["paths"] + assert "get" in schema["paths"]["/v2/model/info"] + + # --------------------------------------------------------------------------- # GET /v1/model/info, GET /model/info # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index ec8b06d9c97..4e5f13fdf88 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -192,8 +192,9 @@ async def test_reconcile_budget_reservation_for_counter_update_returns_empty_set async def test_reconcile_budget_reservation_for_counter_update_failure_invalidates( monkeypatch, ): - """Reservation reconcile raising must invalidate reserved counters but - not propagate the exception.""" + """Reservation reconcile raising must invalidate reserved counters, swallow + the exception, and return an empty set so the caller falls back to the + direct spend-counter increment instead of skipping it.""" import litellm.proxy.spend_tracking.budget_reservation as br monkeypatch.setattr( @@ -213,7 +214,7 @@ async def test_reconcile_budget_reservation_for_counter_update_failure_invalidat budget_reservation={"foo": "bar"}, response_cost=1.0 ) - assert result == {"spend:key:abc"} + assert result == set() assert fake_invalidate.called is True diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 97e5c494916..9757999c85e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -151,21 +151,24 @@ async def test_model_info_v2_translates_team_model_name(monkeypatch): @pytest.mark.asyncio async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch): - """/v1/model/info list path (no litellm_model_id) must surface the public - name. Covers the list comprehension that assigns _get_proxy_model_info's - return back into all_models (#28382 review).""" + """/v1/model/info list path (no litellm_model_id) must include team-scoped + deployments from the router model list and surface the public name (#28382).""" + team_row = _team_row() + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "normal-id-1", "db_model": False}, + } router = MagicMock() - router.get_model_names.return_value = ["team-claude-sonnet"] + router.model_list = [team_row, global_row] + router.get_model_names.return_value = ["gpt-4o"] router.get_model_access_groups.return_value = {} - router.get_model_list.return_value = [_team_row()] monkeypatch.setattr(ps, "user_model", None) - monkeypatch.setattr(ps, "llm_model_list", [_team_row()]) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) monkeypatch.setattr(ps, "llm_router", router) - monkeypatch.setattr(ps, "get_key_models", lambda **kw: []) - monkeypatch.setattr(ps, "get_team_models", lambda **kw: []) monkeypatch.setattr( - ps, "get_complete_model_list", lambda **kw: ["team-claude-sonnet"] + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model ) admin = UserAPIKeyAuth( @@ -176,3 +179,167 @@ async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch): names = [m["model_name"] for m in resp["data"]] assert "team-claude-sonnet" in names assert "model_name_team-abc-123_4a6b8" not in names + + +@pytest.mark.asyncio +async def test_model_info_v1_unrestricted_key_returns_all_deployments(monkeypatch): + """Unrestricted keys must see all router deployments (legacy v1 access logic).""" + deployment = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [deployment] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_name"] for m in resp["data"]] == ["gpt-4"] + + +@pytest.mark.asyncio +async def test_model_info_v1_restricted_key_filters_deployments(monkeypatch): + """Key-level model allowlists must filter router deployments.""" + team_row = _team_row() + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [team_row, global_row] + router.get_model_names.return_value = ["gpt-4", "team-claude-sonnet"] + router.get_model_access_groups.return_value = {} + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["gpt-4"], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_name"] for m in resp["data"]] == ["gpt-4"] + + +def _other_team_row() -> dict: + return { + "model_name": "model_name_team-other_9f2c1", + "litellm_params": { + "model": "azure/gpt-5.2-low-rpm-testing", + "api_base": "https://team-other-private.example.com", + }, + "model_info": { + "id": "byok-id-other", + "team_id": "team-other", + "team_public_model_name": "team-claude-sonnet", + "db_model": True, + }, + } + + +@pytest.mark.asyncio +async def test_model_info_v1_unrestricted_key_hides_other_team_byok(monkeypatch): + """Unrestricted non-admin keys must not enumerate other teams' BYOK + deployments, but must still see global models and their own team's.""" + team_row = _team_row() + other_team_row = _other_team_row() + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [team_row, other_team_row, global_row] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + prisma_client = MagicMock() + caller_user_row = MagicMock() + caller_user_row.teams = ["team-abc-123"] + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=caller_user_row + ) + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + returned_ids = {m["model_info"]["id"] for m in resp["data"]} + assert returned_ids == {"global-id-1", "byok-id-1"} + assert "byok-id-other" not in returned_ids + names = [m["model_name"] for m in resp["data"]] + assert "team-claude-sonnet" in names + assert "gpt-4" in names + + +@pytest.mark.asyncio +async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch): + """A key without a resolvable user (e.g. CI/service token) sees only + global deployments, never any team-scoped BYOK rows.""" + team_row = _team_row() + other_team_row = _other_team_row() + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [team_row, other_team_row, global_row] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + prisma_client = MagicMock() + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id=None, + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-abc-123", + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_info"]["id"] for m in resp["data"]] == ["global-id-1"] diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py new file mode 100644 index 00000000000..c123eeeed36 --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py @@ -0,0 +1,87 @@ +""" +Regression test for enforced-spend underreporting when Redis fails during the +budget-reservation reconcile step of ``increment_spend_counters``. + +Production failure mode: a managed Redis returns an intermittent timeout on the +reconcile increment. Reconcile deletes (invalidates) the shared counter and +gives up, but ``increment_spend_counters`` still treats the counter as +"already reconciled" and skips the direct increment. The actual call cost never +lands in the enforced counter, so budgets stop gating until the next cold +reseed pulls a lagging value from the DB. + +The fix makes the reconcile path fall back to the direct increment when it +fails, so the actual cost is always written to the shared counter. +""" + +import pytest + +from litellm.caching import DualCache +from litellm.proxy import proxy_server + + +class _FlakyRedisCache: + def __init__(self) -> None: + self._store: dict = {} + self._increment_calls = 0 + + async def async_increment(self, key, value, **kwargs): + self._increment_calls += 1 + if self._increment_calls == 1: + raise Exception("Redis timeout") + self._store[key] = float(self._store.get(key, 0.0)) + float(value) + return self._store[key] + + async def async_get_cache(self, key, *args, **kwargs): + return self._store.get(key) + + async def async_delete_cache(self, key, *args, **kwargs): + self._store.pop(key, None) + + async def async_set_cache(self, key, value, *args, **kwargs): + self._store[key] = float(value) + return True + + +@pytest.mark.asyncio +async def test_direct_increment_runs_when_reservation_reconcile_hits_redis_failure( + monkeypatch, +): + hashed_token = "hashed_test_token" + counter_key = f"spend:key:{hashed_token}" + reserved_cost = 0.5 + response_cost = 1.0 + + flaky_redis = _FlakyRedisCache() + flaky_redis._store[counter_key] = reserved_cost + + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + monkeypatch.setattr(proxy_server.spend_counter_cache, "redis_cache", flaky_redis) + proxy_server.spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=reserved_cost + ) + + budget_reservation = { + "reserved_cost": reserved_cost, + "finalized": False, + "entries": [ + { + "counter_key": counter_key, + "entity_type": "Key", + "entity_id": hashed_token, + "reserved_cost": reserved_cost, + "applied_adjustment": 0.0, + } + ], + } + + await proxy_server.increment_spend_counters( + token=hashed_token, + team_id=None, + user_id=None, + response_cost=response_cost, + budget_reservation=budget_reservation, + ) + + enforced_spend = await flaky_redis.async_get_cache(key=counter_key) + assert enforced_spend == response_cost diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index aef91ed3c77..2632d8af4f1 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1314,7 +1314,8 @@ async def query_raw(self, sql_query, session_id, page_size, skip): assert session_id == "session-123" assert page_size == 1 assert skip == 1 # page=2, page_size=1 - return [mock_spend_logs[1]] + assert 'ORDER BY "startTime" DESC' in sql_query + return [mock_spend_logs[0]] class MockPrismaClient: def __init__(self): @@ -1337,7 +1338,7 @@ def __init__(self): assert data["page_size"] == 1 assert data["total_pages"] == 2 assert len(data["data"]) == 1 - assert data["data"][0]["request_id"] == "req2" + assert data["data"][0]["request_id"] == "req1" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py b/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py index 5a13a4dc531..99f6f3a9b72 100644 --- a/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py +++ b/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py @@ -68,6 +68,7 @@ async def test_audio_speech_success_does_not_call_post_call_success_hook( mock_logging.post_call_failure_hook = mock_failure_hook mock_logging.pre_call_hook = mock_pre_call mock_logging.update_request_status = mock_update_status + mock_logging.post_call_response_headers_hook = AsyncMock(return_value={}) async def _mock_route_request(*, data, route_type, llm_router, user_model): assert route_type == "aspeech" diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 0f5a0cbe4b6..b45b31cc67c 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2269,6 +2269,36 @@ async def test_string_detail_unchanged(self): assert proxy_exc.message == "Content blocked by guardrail" assert proxy_exc.provider_specific_fields is None + async def test_not_found_error_preserves_404(self): + """NotFoundError with status_code=404 should map to ProxyException code=404.""" + from litellm.exceptions import NotFoundError + + exc = NotFoundError( + message="Model gemini-3.1-flash-lite-preview not found", + model="gemini-3.1-flash-lite-preview", + llm_provider="gemini", + ) + proxy_exc = await self._invoke(exc) + assert proxy_exc.code == "404" + assert "NotFoundError" in proxy_exc.message + + async def test_exception_with_status_code_propagates(self): + """Exception with a statically-set status_code should propagate it.""" + from litellm.llms.vertex_ai.common_utils import VertexAIError + + exc = VertexAIError( + status_code=429, + message="Rate limit exceeded", + ) + proxy_exc = await self._invoke(exc) + assert proxy_exc.code == "429" + + async def test_exception_without_status_code_defaults_to_500(self): + """Exception with no status_code attribute defaults to 500.""" + exc = ValueError("Something broke") + proxy_exc = await self._invoke(exc) + assert proxy_exc.code == "500" + class TestAsyncStreamingDataGeneratorFastPath: """Fast/slow path branching in async_streaming_data_generator.""" diff --git a/tests/test_litellm/proxy/test_model_info_default_limits.py b/tests/test_litellm/proxy/test_model_info_default_limits.py index 641199c96f0..8111a7af006 100644 --- a/tests/test_litellm/proxy/test_model_info_default_limits.py +++ b/tests/test_litellm/proxy/test_model_info_default_limits.py @@ -146,9 +146,9 @@ async def test_model_info_endpoint_returns_defaults_in_full_model_list(self): deployment_dict = deployment.model_dump(exclude_none=True) mock_router = MagicMock() + mock_router.model_list = [deployment_dict] mock_router.get_model_names.return_value = ["model1"] mock_router.get_model_access_groups.return_value = {} - mock_router.get_model_list.return_value = [deployment_dict] user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") @@ -156,6 +156,7 @@ async def test_model_info_endpoint_returns_defaults_in_full_model_list(self): patch("litellm.proxy.proxy_server.llm_router", mock_router), patch("litellm.proxy.proxy_server.llm_model_list", [deployment_dict]), patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_key_models", return_value=["model1"]), patch( "litellm.proxy.proxy_server.get_team_models", return_value=["model1"] diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 4fb725b7ef3..34c88e2fd33 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -795,6 +795,127 @@ def test_db_connection_extra_params_forwarded_to_url( assert appended_params["pgbouncer"] == "true" assert appended_params["statement_cache_size"] == 0 + def test_build_db_connection_url_params_disable_prepared_statements(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + disable_prepared_statements=True, + ) + assert params["pgbouncer"] == "true" + + def test_build_db_connection_url_params_no_pgbouncer_by_default(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + ) + assert "pgbouncer" not in params + + def test_build_db_connection_url_params_extra_pgbouncer_overrides_flag(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + disable_prepared_statements=True, + extra_params={"pgbouncer": "false"}, + ) + assert params["pgbouncer"] == "false" + + @pytest.mark.parametrize( + "config_value, expect_pgbouncer", + [ + (True, True), + (False, False), + ("true", True), + ("false", False), + ("not-a-bool", False), + ], + ) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False + ) + def test_disable_prepared_statements_forwarded_to_url( + self, + mock_should_update, + mock_setup_db, + mock_atexit_register, + mock_subprocess_run, + config_value, + expect_pgbouncer, + ): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + mock_subprocess_run.return_value = MagicMock(returncode=0) + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock( + return_value={ + "general_settings": { + "database_url": "postgresql://test:test@localhost:5432/test", + "database_disable_prepared_statements": config_value, + } + } + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + patch( + "litellm.proxy.proxy_cli.append_query_params", + side_effect=lambda url, params: str(url), + ) as mock_append_query_params, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--config", "test-config.yaml", "--skip_server_startup"], + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_append_query_params.assert_called() + appended_params = mock_append_query_params.call_args.args[1] + if expect_pgbouncer: + assert appended_params["pgbouncer"] == "true" + else: + assert "pgbouncer" not in appended_params + @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8aa839cdfcb..9eaccdfcbcd 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2319,6 +2319,36 @@ async def test_custom_ui_sso_sign_in_handler_config_loading(): os.unlink(config_file_path) +@pytest.mark.asyncio +async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeypatch): + """ + max_budget configured as os.environ/MAX_BUDGET resolves to a string; + load_config must coerce it to float so the startup check + `litellm.max_budget > 0` doesn't raise TypeError. + """ + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setenv("MAX_BUDGET", "10") + test_config = { + "model_list": [], + "litellm_settings": {"max_budget": "os.environ/MAX_BUDGET"}, + } + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.dump(test_config)) + + original_max_budget = litellm.max_budget + try: + proxy_config = ProxyConfig() + await proxy_config.load_config( + router=MagicMock(), config_file_path=str(config_file) + ) + assert isinstance(litellm.max_budget, float) + assert litellm.max_budget == 10.0 + assert litellm.max_budget > 0 + finally: + litellm.max_budget = original_max_budget + + @pytest.mark.asyncio async def test_load_environment_variables_direct_and_os_environ(): """ @@ -3810,14 +3840,15 @@ async def test_model_info_v1_oci_secrets_not_leaked(): # Mock the llm_router to return our test data mock_router = MagicMock() + mock_router.model_list = [mock_model_data] mock_router.get_model_names.return_value = ["oci-grok-test"] mock_router.get_model_access_groups.return_value = {} - mock_router.get_model_list.return_value = [mock_model_data] # Mock global variables with ( patch("litellm.proxy.proxy_server.llm_router", mock_router), patch("litellm.proxy.proxy_server.llm_model_list", [mock_model_data]), + patch("litellm.proxy.proxy_server.prisma_client", None), patch( "litellm.proxy.proxy_server.general_settings", {"infer_model_from_keys": False}, @@ -4294,6 +4325,111 @@ async def test_init_sso_settings_in_db_empty_settings(): assert uppercased_settings == {} +@pytest.mark.asyncio +async def test_init_sso_settings_in_db_retries_on_transport_error(): + """`_init_sso_settings_in_db` self-heals across one ClientNotConnectedError + via call_with_db_reconnect_retry — mirrors the auth-path behavior so + startup/reload bursts don't spam the log.""" + import prisma + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_sso_config = MagicMock() + mock_sso_config.sso_settings = {"GOOGLE_CLIENT_ID": "xxx"} + + invocations: list = [] + + async def _flaky_find_unique(**kwargs): + invocations.append(None) + if len(invocations) == 1: + raise prisma.errors.ClientNotConnectedError() + return mock_sso_config + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( + side_effect=_flaky_find_unique + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + with patch.object( + proxy_config, "_decrypt_and_set_db_env_variables" + ) as mock_decrypt: + await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) + + assert len(invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert reconnect_kwargs["reason"] == "init_sso_settings_in_db_lookup_failure" + mock_decrypt.assert_called_once() + + +@pytest.mark.asyncio +async def test_init_sso_settings_in_db_propagates_when_reconnect_fails(): + """When reconnect returns False (cooldown / lock contention), the original + ClientNotConnectedError is caught by the function's `except Exception` and + logged — no retry storm, no crash.""" + import prisma + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( + side_effect=prisma.errors.ClientNotConnectedError() + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + # Should NOT raise — the function's own try/except swallows the propagated error. + await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) + + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_init_hashicorp_vault_config_override_retries_on_transport_error(): + """`_init_hashicorp_vault_config_override` self-heals across one + ClientNotConnectedError via call_with_db_reconnect_retry.""" + import prisma + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._last_hashicorp_vault_config = None + + invocations: list = [] + + async def _flaky_find_unique(**kwargs): + invocations.append(None) + if len(invocations) == 1: + raise prisma.errors.ClientNotConnectedError() + return None # No config in DB → function returns early after retry. + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_configoverrides.find_unique = AsyncMock( + side_effect=_flaky_find_unique + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + await proxy_config._init_hashicorp_vault_config_override( + prisma_client=mock_prisma_client + ) + + assert len(invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert ( + reconnect_kwargs["reason"] + == "init_hashicorp_vault_config_override_lookup_failure" + ) + + def test_update_config_fields_uppercases_env_vars(monkeypatch): """ Ensure environment variables pulled from DB are uppercased when applied so @@ -6474,7 +6610,12 @@ async def test_increment_spend_counters_finalizes_none_cost_reservation(): @pytest.mark.asyncio -async def test_increment_spend_counters_invalidates_bad_reserved_counter_without_failing(): +async def test_increment_spend_counters_falls_back_to_direct_increment_on_bad_reserved_counter(): + """When the reservation reconcile fails, the reserved counters are + invalidated and the actual response cost must still be written via the + direct increment fallback. Leaving the counter at ``None`` lets the next + request reseed a stale value from the DB and silently stops budget gating, + which is the bug this fix addresses.""" from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import increment_spend_counters @@ -6515,7 +6656,7 @@ async def test_increment_spend_counters_invalidates_bad_reserved_counter_without counter_cache.in_memory_cache.get_cache( key="spend:key:key-bad-reserved-counter" ) - is None + == 0.25 ) finally: ps.spend_counter_cache = orig_counter diff --git a/tests/test_litellm/proxy/test_team_member_update.py b/tests/test_litellm/proxy/test_team_member_update.py index 6561ec9e7fd..352c68d491c 100644 --- a/tests/test_litellm/proxy/test_team_member_update.py +++ b/tests/test_litellm/proxy/test_team_member_update.py @@ -1,9 +1,19 @@ +import types +from unittest.mock import AsyncMock, MagicMock + import pytest from fastapi import HTTPException from starlette.requests import Request import litellm.proxy.proxy_server as proxy_server -from litellm.proxy._types import TeamMemberUpdateRequest +import litellm.proxy.management_endpoints.team_endpoints as team_endpoints +from litellm.proxy._types import ( + LiteLLM_TeamTable, + LitellmUserRoles, + Member, + TeamMemberUpdateRequest, + UserAPIKeyAuth, +) from litellm.proxy.management_endpoints.team_endpoints import team_member_update @@ -38,3 +48,133 @@ async def test_ateam_member_update_admin_requires_premium(monkeypatch): "Pricing: https://www.litellm.ai/#pricing" ) assert exc_info.value.detail == expected_msg + + +@pytest.fixture +def happy_path_upsert(monkeypatch): + """Stub out the DB and the budget upsert so a team_member_update call reaches + _upsert_budget_and_membership, and hand back that mock to inspect the patch.""" + team_row = LiteLLM_TeamTable( + team_id="team-1234", + members_with_roles=[Member(user_id="user-1", role="user")], + metadata={}, + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + prisma_client.db.litellm_teamtable.update = AsyncMock() + + class _FakeTx: + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + prisma_client.db.tx = MagicMock(return_value=_FakeTx()) + + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "premium_user", False) + monkeypatch.setattr( + team_endpoints, + "team_info", + AsyncMock( + return_value={ + "team_info": team_row, + "team_memberships": [ + types.SimpleNamespace(user_id="user-1", budget_id="bud-1") + ], + } + ), + ) + upsert_mock = AsyncMock() + monkeypatch.setattr(team_endpoints, "_upsert_budget_and_membership", upsert_mock) + return upsert_mock + + +def _member_update_request(**overrides): + data = TeamMemberUpdateRequest( + team_id="team-1234", user_id="user-1", role="user", **overrides + ) + request = Request({"type": "http", "method": "POST", "path": "/team/member_update"}) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value, user_id="admin") + return data, request, auth + + +@pytest.mark.asyncio +async def test_team_member_update_sends_provided_fields_as_patch(happy_path_upsert): + """Fields the request sets must reach _upsert_budget_and_membership as a + budget patch, otherwise the member budget is never written/reset.""" + data, request, auth = _member_update_request( + max_budget_in_team=10.0, budget_duration="30d" + ) + + response = await team_member_update(data, request, auth) + + happy_path_upsert.assert_awaited_once() + assert happy_path_upsert.await_args.kwargs["budget_patch"] == { + "max_budget": 10.0, + "budget_duration": "30d", + } + assert response.budget_duration == "30d" + + +@pytest.mark.asyncio +async def test_team_member_update_explicit_null_clears_field(happy_path_upsert): + """An explicitly-null field must be forwarded as None so the column is + cleared, rather than silently dropped.""" + data, request, auth = _member_update_request(budget_duration=None) + + await team_member_update(data, request, auth) + + assert happy_path_upsert.await_args.kwargs["budget_patch"] == { + "budget_duration": None + } + + +@pytest.mark.asyncio +async def test_team_member_update_omits_unset_fields_from_patch(happy_path_upsert): + """A request that touches no budget fields must produce an empty patch so the + member's existing budget is left untouched.""" + data, request, auth = _member_update_request() + + await team_member_update(data, request, auth) + + assert happy_path_upsert.await_args.kwargs["budget_patch"] == {} + + +@pytest.mark.parametrize( + "bad_duration", + [ + "not-a-duration", # unparseable garbage + "10x", # unsupported unit + "0d", # zero-length window + "999999999999999999999999d", # overflows datetime math + ], +) +@pytest.mark.asyncio +async def test_team_member_update_rejects_invalid_budget_duration( + monkeypatch, bad_duration +): + """An invalid budget_duration must be rejected with a 400 before any DB + write, so it can never be persisted and later break the budget reset job.""" + monkeypatch.setattr(proxy_server, "prisma_client", object()) + monkeypatch.setattr(proxy_server, "premium_user", False) + upsert_mock = AsyncMock() + monkeypatch.setattr(team_endpoints, "_upsert_budget_and_membership", upsert_mock) + + data = TeamMemberUpdateRequest( + team_id="team-1234", + user_id="user-1", + role="user", + budget_duration=bad_duration, + ) + request = Request({"type": "http", "method": "POST", "path": "/team/member_update"}) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value, user_id="admin") + + with pytest.raises(HTTPException) as exc_info: + await team_member_update(data, request, auth) + + assert exc_info.value.status_code == 400 + assert "budget_duration" in str(exc_info.value.detail) + upsert_mock.assert_not_called() 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 ae217aca16e..f77af2d90bf 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 @@ -1032,6 +1032,45 @@ def test_update_ui_settings_allowlisted_value(self, mock_auth, monkeypatch): stored_settings = json.loads(create_data["ui_settings"]) assert stored_settings["disable_model_add_for_internal_users"] is True + def test_update_ui_settings_persists_disable_ui_nudges( + self, mock_auth, monkeypatch + ): + """disable_ui_nudges must be allowlisted so admins can suppress UI popups for everyone""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + mock_prisma = MagicMock() + mock_prisma.db.litellm_uisettings.upsert = AsyncMock() + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + try: + response = client.patch( + "/update/ui_settings", json={"disable_ui_nudges": True} + ) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "success" + assert data["settings"]["disable_ui_nudges"] is True + + create_data = mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"][ + "create" + ] + stored_settings = json.loads(create_data["ui_settings"]) + assert stored_settings["disable_ui_nudges"] is True + def test_update_ui_settings_ignores_non_allowlisted_value( self, mock_auth, monkeypatch ): diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index 7e7e98d1360..437984d9273 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -193,6 +193,7 @@ async def test_query_first_with_cached_plan_fallback_happy_returns_row( ) -> None: expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0} prisma_client.db.query_first = AsyncMock(return_value=expected) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) result = await prisma_client._query_first_with_cached_plan_fallback( "SELECT * FROM x WHERE token = $1", "abc" ) @@ -208,35 +209,110 @@ async def test_query_first_with_cached_plan_fallback_happy_returns_row( "args": ("SELECT * FROM x WHERE token = $1", "abc"), "matches": True, } + prisma_client.attempt_db_reconnect.assert_not_awaited() @pytest.mark.asyncio -async def test_query_first_with_cached_plan_fallback_retries_on_cached_plan_error( +async def test_query_first_with_cached_plan_fallback_reconnects_then_retries_identical_query( prisma_client: PrismaClient, ) -> None: + original_query = 'SELECT * FROM "LiteLLM_VerificationToken" WHERE v.token = $1' expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0} - prisma_client.db.query_first = AsyncMock( + manager = MagicMock() + query_first = AsyncMock( side_effect=[ RuntimeError("cached plan must not change result type"), expected, ] ) + reconnect = AsyncMock(return_value=True) + manager.attach_mock(query_first, "query_first") + manager.attach_mock(reconnect, "attempt_db_reconnect") + prisma_client.db.query_first = query_first + prisma_client.attempt_db_reconnect = reconnect + result = await prisma_client._query_first_with_cached_plan_fallback( - "SELECT * FROM x WHERE token = $1", "abc" + original_query, "abc" + ) + + assert result == expected + assert query_first.await_count == 2 + first_call, retry_call = query_first.await_args_list + assert retry_call.args == first_call.args == (original_query, "abc") + reconnect.assert_awaited_once() + assert reconnect.await_args.kwargs.get("force", False) is False + assert [name for name, *_ in manager.mock_calls] == [ + "query_first", + "attempt_db_reconnect", + "query_first", + ] + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_never_deallocates( + prisma_client: PrismaClient, +) -> None: + expected = {"token": "abc"} + prisma_client.db.query_first = AsyncMock( + side_effect=[ + RuntimeError("cached plan must not change result type"), + expected, + ] + ) + prisma_client.db.execute_raw = AsyncMock(return_value=0) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + + prisma_client.db.execute_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_propagates_when_retry_also_fails( + prisma_client: PrismaClient, +) -> None: + plan_error = RuntimeError("cached plan must not change result type") + prisma_client.db.query_first = AsyncMock(side_effect=[plan_error, plan_error]) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + with pytest.raises(RuntimeError, match="cached plan must not change result type"): + await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + + assert prisma_client.db.query_first.await_count == 2 + prisma_client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_retries_when_reconnect_returns_false( + prisma_client: PrismaClient, +) -> None: + expected = {"token": "abc"} + prisma_client.db.query_first = AsyncMock( + side_effect=[ + RuntimeError("cached plan must not change result type"), + expected, + ] ) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) + + result = await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + assert result == expected assert prisma_client.db.query_first.await_count == 2 - second_call_sql = prisma_client.db.query_first.await_args_list[1].args[0] - assert "cache_invalidated_" in second_call_sql @pytest.mark.asyncio async def test_query_first_with_cached_plan_fallback_reraises_non_plan_errors( prisma_client: PrismaClient, ) -> None: - prisma_client.db.query_first = AsyncMock(side_effect=RuntimeError("totally unrelated")) + prisma_client.db.query_first = AsyncMock( + side_effect=RuntimeError("totally unrelated") + ) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) with pytest.raises(RuntimeError, match="totally unrelated"): await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + assert prisma_client.db.query_first.await_count == 1 + prisma_client.attempt_db_reconnect.assert_not_awaited() @pytest.mark.asyncio @@ -351,7 +427,9 @@ async def test_get_data_token_find_unique_returns_record( async def test_get_data_token_find_unique_missing_token_raises_401( prisma_client: PrismaClient, ) -> None: - prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=None + ) with pytest.raises(HTTPException) as excinfo: await prisma_client.get_data(token="sk-missing", table_name="key") err = excinfo.value diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py new file mode 100644 index 00000000000..f22debbae34 --- /dev/null +++ b/tests/test_litellm/repositories/test_repositories.py @@ -0,0 +1,2184 @@ +""" +Tests for gateway repository layer. +""" + +import json +from datetime import datetime +from typing import Any, Dict, List, Optional +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.models.base import DomainModel +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.credentials import CredentialItem +from litellm.models.team import LiteLLM_TeamTable +from litellm.repositories.base_repository import BaseRepository +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.credentials_repository import CredentialsRepository +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.object_permission_repository import ( + ObjectPermissionRepository, +) +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.project_repository import ProjectRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) + + +class MockRecord: + """Mock database record for testing.""" + + def __init__(self, data: Dict[str, Any]): + self._data = data if data is not None else {} + + def dict(self) -> Dict[str, Any]: + return self._data.copy() + + def model_dump(self) -> Dict[str, Any]: + return self._data.copy() + + def __getattr__(self, name: str) -> Any: + if name.startswith("_"): + raise AttributeError(name) + return self._data.get(name) + + +class MockTable: + """Mock Prisma table for testing.""" + + def __init__(self, pk_field: Optional[str] = None): + self._records: Dict[str, Dict[str, Any]] = {} + self._pk_field = pk_field + + async def find_unique(self, where: Dict[str, Any]) -> Optional[MockRecord]: + key_field = list(where.keys())[0] + key_value = where[key_field] + data = self._records.get(key_value) + return MockRecord(data) if data else None + + async def find_many( + self, + where: Optional[Dict[str, Any]] = None, + skip: Optional[int] = None, + take: Optional[int] = None, + order: Optional[Dict[str, str]] = None, + ) -> List[MockRecord]: + records = list(self._records.values()) + return [MockRecord(r) for r in records] + + async def create(self, data: Dict[str, Any]) -> MockRecord: + record_data = dict(data) + if self._pk_field and self._pk_field not in record_data: + record_data[self._pk_field] = f"{self._pk_field}-{len(self._records)}" + key = ( + record_data.get(self._pk_field) + if self._pk_field + else record_data.get("id", str(len(self._records))) + ) + self._records[key] = record_data + return MockRecord(record_data) + + async def update( + self, where: Dict[str, Any], data: Dict[str, Any] + ) -> Optional[MockRecord]: + key_field = list(where.keys())[0] + key_value = where[key_field] + if key_value in self._records: + for field, value in data.items(): + if isinstance(value, dict) and "push" in value: + current = self._records[key_value].get(field, []) + push_val = value["push"] + if isinstance(push_val, list): + current.extend(push_val) + else: + current.append(push_val) + self._records[key_value][field] = current + else: + self._records[key_value][field] = value + return MockRecord(self._records[key_value]) + return None + + async def delete(self, where: Dict[str, Any]) -> Optional[MockRecord]: + key_field = list(where.keys())[0] + key_value = where[key_field] + data = self._records.pop(key_value, None) + return MockRecord(data) if data else None + + async def count(self, where: Optional[Dict[str, Any]] = None) -> int: + return len(self._records) + + async def upsert(self, where: Dict[str, Any], data: Dict[str, Any]) -> MockRecord: + key_field = list(where.keys())[0] + key_value = where[key_field] + if key_value in self._records: + self._records[key_value].update(data.get("update", {})) + else: + self._records[key_value] = data.get("create", {}) + return MockRecord(self._records[key_value]) + + +class MockPrismaClient: + """Mock Prisma client for testing.""" + + def __init__(self): + self.db = MagicMock() + self.db.litellm_budgettable = MockTable() + self.db.litellm_proxymodeltable = MockTable(pk_field="model_id") + self.db.litellm_teamtable = MockTable() + self.db.litellm_deletedteamtable = MockTable() + self.db.litellm_usertable = MockTable() + self.db.litellm_verificationtoken = MockTable() + self.db.litellm_deletedverificationtoken = MockTable() + self.db.litellm_config = MockTable() + self.db.litellm_organizationtable = MockTable() + self.db.litellm_projecttable = MockTable(pk_field="project_id") + self.db.litellm_objectpermissiontable = MockTable( + pk_field="object_permission_id" + ) + self.db.litellm_credentialstable = MockTable() + + +class TestBaseRepository: + @pytest.fixture + def prisma_client(self): + return MockPrismaClient() + + def test_prisma_client_none_raises(self): + class TestRepo(BaseRepository[LiteLLM_BudgetTable]): + @property + def table(self): + return None + + @property + def model_class(self): + return LiteLLM_BudgetTable + + repo = TestRepo(None) + with pytest.raises(RuntimeError, match="No DB Connected"): + _ = repo.prisma_client + + @pytest.mark.asyncio + async def test_find_many(self, prisma_client): + repo = BudgetRepository(prisma_client) + prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1", "max_budget": 100.0}, + "b2": {"budget_id": "b2", "max_budget": 200.0}, + } + budgets = await repo.find_many() + assert len(budgets) == 2 + + @pytest.mark.asyncio + async def test_count(self, prisma_client): + repo = BudgetRepository(prisma_client) + prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1"}, + "b2": {"budget_id": "b2"}, + } + count = await repo.count() + assert count == 2 + + @pytest.mark.asyncio + async def test_exists(self, prisma_client): + repo = BudgetRepository(prisma_client) + prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1"}, + } + assert await repo.exists("b1", id_field="budget_id") + assert not await repo.exists("nonexistent", id_field="budget_id") + + @pytest.mark.asyncio + async def test_find_many_with_all_kwargs(self, prisma_client): + repo = BudgetRepository(prisma_client) + prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1", "max_budget": 100.0}, + } + budgets = await repo.find_many( + where={"budget_id": "b1"}, skip=0, take=10, order={"budget_id": "asc"} + ) + assert len(budgets) == 1 + + def test_record_to_dict_branches(self): + from litellm.repositories.base_repository import _record_to_dict + + assert _record_to_dict({"a": 1}) == {"a": 1} + + class WithModelDump: + def model_dump(self): + return {"src": "model_dump"} + + assert _record_to_dict(WithModelDump()) == {"src": "model_dump"} + + class WithDict: + def dict(self): + return {"src": "dict"} + + assert _record_to_dict(WithDict()) == {"src": "dict"} + + assert _record_to_dict([("k", "v")]) == {"k": "v"} + + +class TestBudgetRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return BudgetRepository(client) + + @pytest.mark.asyncio + async def test_create_budget(self, repo): + budget = await repo.create_budget( + created_by="test-user", + max_budget=100.0, + soft_budget=80.0, + tpm_limit=1000, + ) + assert budget.max_budget == 100.0 + assert budget.soft_budget == 80.0 + assert budget.tpm_limit == 1000 + + @pytest.mark.asyncio + async def test_create_budget_all_fields(self, repo): + budget = await repo.create_budget( + created_by="test-user", + max_budget=100.0, + soft_budget=80.0, + max_parallel_requests=10, + tpm_limit=1000, + rpm_limit=100, + model_max_budget={"gpt-4": 50.0}, + budget_duration="monthly", + allowed_models=["gpt-4", "gpt-3.5-turbo"], + ) + assert budget.max_budget == 100.0 + assert budget.max_parallel_requests == 10 + + @pytest.mark.asyncio + async def test_update_budget(self, repo): + await repo.create_budget(created_by="test-user", max_budget=100.0) + repo._prisma_client.db.litellm_budgettable._records["budget-1"] = { + "budget_id": "budget-1", + "max_budget": 100.0, + } + + updated = await repo.update_budget( + budget_id="budget-1", + updated_by="test-user", + max_budget=200.0, + ) + assert updated.max_budget == 200.0 + + @pytest.mark.asyncio + async def test_delete_budget(self, repo): + repo._prisma_client.db.litellm_budgettable._records["budget-1"] = { + "budget_id": "budget-1", + "max_budget": 100.0, + } + deleted = await repo.delete_budget("budget-1") + assert deleted is not None + assert "budget-1" not in repo._prisma_client.db.litellm_budgettable._records + + @pytest.mark.asyncio + async def test_find_by_id(self, repo): + repo._prisma_client.db.litellm_budgettable._records["budget-1"] = { + "budget_id": "budget-1", + "max_budget": 100.0, + } + budget = await repo.find_by_id("budget-1") + assert budget is not None + assert budget.budget_id == "budget-1" + + +class TestModelRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ModelRepository(client) + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.encrypt_value_helper", + side_effect=lambda v, **kw: f"encrypted_{v}", + ) + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_create_model_encrypts_params(self, mock_decrypt, mock_encrypt, repo): + model = await repo.create_model( + model_name="gpt-4", + litellm_params={"api_key": "sk-secret"}, + created_by="test-user", + ) + assert model is not None + mock_encrypt.assert_called() + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.encrypt_value_helper", + side_effect=lambda v, **kw: f"encrypted_{v}", + ) + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_create_model_all_fields(self, mock_decrypt, mock_encrypt, repo): + model = await repo.create_model( + model_name="gpt-4-turbo", + litellm_params={ + "api_key": "sk-secret", + "api_base": "https://api.openai.com", + }, + created_by="admin", + model_id="custom-model-id", + model_info={"team_id": "team-1", "description": "GPT-4 Turbo model"}, + blocked=True, + ) + assert model is not None + assert model.model_name == "gpt-4-turbo" + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.encrypt_value_helper", + side_effect=lambda v, **kw: f"encrypted_{v}", + ) + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_update_model_all_fields(self, mock_decrypt, mock_encrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records["model-full"] = { + "model_id": "model-full", + "model_name": "old-name", + "litellm_params": '{"api_key": "old"}', + "blocked": False, + } + updated = await repo.update_model( + model_id="model-full", + updated_by="admin", + model_name="new-name", + litellm_params={"api_key": "new-key"}, + model_info={"updated": True}, + blocked=True, + ) + assert updated.model_name == "new-name" + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_find_all(self, mock_decrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records = { + "m1": { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + "blocked": False, + }, + "m2": { + "model_id": "m2", + "model_name": "claude-3", + "litellm_params": '{"model": "claude-3"}', + "blocked": False, + }, + } + models = await repo.find_all() + assert len(models) == 2 + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_find_unblocked(self, mock_decrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records = { + "m1": { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + "blocked": False, + }, + } + models = await repo.find_unblocked() + assert len(models) == 1 + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_find_by_name(self, mock_decrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records = { + "m1": { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + }, + } + models = await repo.find_by_name("gpt-4") + assert len(models) == 1 + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.encrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_update_model(self, mock_decrypt, mock_encrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + "blocked": False, + } + updated = await repo.update_model( + model_id="m1", + updated_by="test-user", + blocked=True, + ) + assert updated.blocked is True + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_delete_model(self, mock_decrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + } + deleted = await repo.delete_model("m1") + assert deleted is not None + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.encrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_block_unblock_model(self, mock_decrypt, mock_encrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + "blocked": False, + } + blocked = await repo.block_model("m1", "admin") + assert blocked.blocked is True + + unblocked = await repo.unblock_model("m1", "admin") + assert unblocked.blocked is False + + +class TestTeamRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return TeamRepository(client) + + @pytest.mark.asyncio + async def test_create_team(self, repo): + team = await repo.create_team( + team_id="team-123", + team_alias="Engineering", + admins=["user1"], + members=["user2", "user3"], + ) + assert team.team_id == "team-123" + assert team.team_alias == "Engineering" + + @pytest.mark.asyncio + async def test_create_team_all_fields(self, repo): + team = await repo.create_team( + team_id="team-123", + team_alias="Engineering", + organization_id="org-1", + admins=["admin1"], + members=["user1"], + members_with_roles=[{"user_id": "user1", "role": "user"}], + metadata={"dept": "engineering"}, + max_budget=1000.0, + soft_budget=800.0, + models=["gpt-4"], + max_parallel_requests=10, + tpm_limit=50000, + rpm_limit=500, + budget_duration="monthly", + object_permission_id="perm-1", + ) + assert team.team_id == "team-123" + assert team.organization_id == "org-1" + + @pytest.mark.asyncio + async def test_update_team(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": [], + "models": [], + } + updated = await repo.update_team( + team_id="team-1", + team_alias="Updated Team", + blocked=True, + ) + assert updated.team_alias == "Updated Team" + + @pytest.mark.asyncio + async def test_update_team_all_fields(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-full"] = { + "team_id": "team-full", + "team_alias": "Test", + "admins": [], + "members": [], + "models": [], + } + updated = await repo.update_team( + team_id="team-full", + team_alias="Fully Updated", + organization_id="org-new", + admins=["admin1"], + members=["member1"], + members_with_roles=[{"user_id": "user1", "role": "admin"}], + metadata={"updated": True}, + max_budget=500.0, + soft_budget=400.0, + models=["gpt-4", "claude-3"], + max_parallel_requests=20, + tpm_limit=100000, + rpm_limit=1000, + budget_duration="weekly", + blocked=False, + object_permission_id="perm-new", + ) + assert updated.team_alias == "Fully Updated" + + @pytest.mark.asyncio + async def test_add_member(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": ["user1"], + "models": [], + } + + team = await repo.add_member("team-1", "user2") + assert "user2" in team.members + + @pytest.mark.asyncio + async def test_add_member_nonexistent_team(self, repo): + result = await repo.add_member("nonexistent", "user1") + assert result is None + + @pytest.mark.asyncio + async def test_remove_member(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": ["user1", "user2"], + "models": [], + } + + team = await repo.remove_member("team-1", "user2") + assert "user2" not in team.members + + @pytest.mark.asyncio + async def test_add_admin(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": [], + "models": [], + } + team = await repo.add_admin("team-1", "admin1") + assert "admin1" in team.admins + + @pytest.mark.asyncio + async def test_remove_admin(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": ["admin1", "admin2"], + "members": [], + "models": [], + } + team = await repo.remove_admin("team-1", "admin2") + assert "admin2" not in team.admins + + @pytest.mark.asyncio + async def test_add_models(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": [], + "models": ["gpt-3.5-turbo"], + } + team = await repo.add_models("team-1", ["gpt-4"]) + assert "gpt-4" in team.models + + @pytest.mark.asyncio + async def test_remove_models(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": [], + "models": ["gpt-3.5-turbo", "gpt-4"], + } + team = await repo.remove_models("team-1", ["gpt-4"]) + assert "gpt-4" not in team.models + + @pytest.mark.asyncio + async def test_update_spend(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": [], + "models": [], + "spend": 0.0, + } + team = await repo.update_spend("team-1", 50.0) + assert team.spend == 50.0 + + @pytest.mark.asyncio + async def test_find_by_alias(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Engineering", + "admins": [], + "members": [], + "models": [], + } + team = await repo.find_by_alias("Engineering") + assert team is not None + assert team.team_id == "team-1" + + @pytest.mark.asyncio + async def test_find_by_organization_id(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "organization_id": "org-1", + "admins": [], + "members": [], + "models": [], + } + teams = await repo.find_by_organization_id("org-1") + assert len(teams) == 1 + + @pytest.mark.asyncio + async def test_find_by_member(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "admins": [], + "members": ["user1"], + "models": [], + } + teams = await repo.find_by_member("user1") + assert len(teams) == 1 + + @pytest.mark.asyncio + async def test_find_by_admin(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "admins": ["admin1"], + "members": [], + "models": [], + } + teams = await repo.find_by_admin("admin1") + assert len(teams) == 1 + + +class TestUserRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return UserRepository(client) + + @pytest.mark.asyncio + async def test_create_user(self, repo): + user = await repo.create_user( + user_id="user-123", + user_email="test@example.com", + teams=["team1"], + ) + assert user.user_id == "user-123" + + @pytest.mark.asyncio + async def test_create_user_all_fields(self, repo): + user = await repo.create_user( + user_id="user-123", + user_alias="testuser", + team_id="team-1", + sso_user_id="sso-123", + organization_id="org-1", + password="hashed_password", + teams=["team1", "team2"], + user_role="admin", + max_budget=500.0, + user_email="test@example.com", + models=["gpt-4"], + metadata={"department": "engineering"}, + max_parallel_requests=5, + tpm_limit=10000, + rpm_limit=100, + budget_duration="monthly", + allowed_cache_controls=["no-cache"], + policies=["policy-1"], + object_permission_id="perm-1", + ) + assert user.user_id == "user-123" + assert user.user_alias == "testuser" + + @pytest.mark.asyncio + async def test_update_user(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": [], + "models": [], + } + updated = await repo.update_user( + user_id="user-1", + user_email="updated@example.com", + ) + assert updated.user_email == "updated@example.com" + + @pytest.mark.asyncio + async def test_delete_user(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": [], + "models": [], + } + deleted = await repo.delete_user("user-1") + assert deleted is not None + + @pytest.mark.asyncio + async def test_add_to_team(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": ["team1"], + "models": [], + } + + user = await repo.add_to_team("user-1", "team2") + assert "team2" in user.teams + + @pytest.mark.asyncio + async def test_add_to_team_nonexistent_user(self, repo): + result = await repo.add_to_team("nonexistent", "team1") + assert result is None + + @pytest.mark.asyncio + async def test_remove_from_team(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": ["team1", "team2"], + "models": [], + } + user = await repo.remove_from_team("user-1", "team2") + assert "team2" not in user.teams + + @pytest.mark.asyncio + async def test_update_spend(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": [], + "models": [], + "spend": 0.0, + } + user = await repo.update_spend("user-1", 25.0) + assert user.spend == 25.0 + + @pytest.mark.asyncio + async def test_find_by_email(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "user_email": "test@example.com", + "teams": [], + "models": [], + } + user = await repo.find_by_email("test@example.com") + assert user is not None + + @pytest.mark.asyncio + async def test_find_by_sso_id(self, repo): + repo._prisma_client.db.litellm_usertable._records["sso-123"] = { + "user_id": "user-1", + "sso_user_id": "sso-123", + "teams": [], + "models": [], + } + user = await repo.find_by_sso_id("sso-123") + assert user is not None + + @pytest.mark.asyncio + async def test_find_by_organization_id(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "organization_id": "org-1", + "teams": [], + "models": [], + } + users = await repo.find_by_organization_id("org-1") + assert len(users) == 1 + + @pytest.mark.asyncio + async def test_find_by_team_id(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": ["team-1"], + "models": [], + } + users = await repo.find_by_team_id("team-1") + assert len(users) == 1 + + +class TestVerificationTokenRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return VerificationTokenRepository(client) + + @pytest.mark.asyncio + async def test_create_token(self, repo): + token = await repo.create_token( + token="sk-test123", + key_name="Test Key", + user_id="user-123", + max_budget=100.0, + ) + assert token.token == "sk-test123" + assert token.key_name == "Test Key" + + @pytest.mark.asyncio + async def test_create_token_all_fields(self, repo): + token = await repo.create_token( + token="sk-test123", + key_name="Test Key", + key_alias="test-alias", + max_budget=100.0, + expires=datetime(2025, 12, 31), + models=["gpt-4"], + aliases={"alias1": "value1"}, + config={"setting": "value"}, + user_id="user-123", + team_id="team-1", + agent_id="agent-1", + project_id="project-1", + max_parallel_requests=5, + metadata={"key": "value"}, + tpm_limit=10000, + rpm_limit=100, + budget_duration="monthly", + allowed_cache_controls=["no-cache"], + allowed_routes=["/v1/completions"], + permissions={"read": True}, + org_id="org-1", + created_by="admin", + object_permission_id="perm-1", + access_group_ids=["group-1"], + budget_id="budget-1", + ) + assert token.token == "sk-test123" + + @pytest.mark.asyncio + async def test_update_token(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "blocked": False, + } + updated = await repo.update_token( + token="sk-test", + key_name="Updated Key", + ) + assert updated.key_name == "Updated Key" + + @pytest.mark.asyncio + async def test_block_token(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "blocked": False, + } + + token = await repo.block_token("sk-test", updated_by="admin") + assert token.blocked is True + + @pytest.mark.asyncio + async def test_unblock_token(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "blocked": True, + } + token = await repo.unblock_token("sk-test", updated_by="admin") + assert token.blocked is False + + @pytest.mark.asyncio + async def test_update_spend(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "spend": 0.0, + } + token = await repo.update_spend("sk-test", 15.0) + assert token.spend == 15.0 + + @pytest.mark.asyncio + async def test_update_last_active(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + } + token = await repo.update_last_active("sk-test") + assert token.last_active is not None + + @pytest.mark.asyncio + async def test_find_by_alias(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "key_alias": "my-key", + } + token = await repo.find_by_alias("my-key") + assert token is not None + + @pytest.mark.asyncio + async def test_find_by_user_id(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "user_id": "user-1", + } + tokens = await repo.find_by_user_id("user-1") + assert len(tokens) == 1 + + @pytest.mark.asyncio + async def test_find_by_team_id(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "team_id": "team-1", + } + tokens = await repo.find_by_team_id("team-1") + assert len(tokens) == 1 + + @pytest.mark.asyncio + async def test_find_by_project_id(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "project_id": "project-1", + } + tokens = await repo.find_by_project_id("project-1") + assert len(tokens) == 1 + + +class TestOrganizationRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return OrganizationRepository(client) + + @pytest.mark.asyncio + async def test_create_organization(self, repo): + org = await repo.create_organization( + organization_alias="Acme Corp", + budget_id="budget-1", + created_by="admin", + ) + assert org.organization_alias == "Acme Corp" + + @pytest.mark.asyncio + async def test_create_organization_all_fields(self, repo): + org = await repo.create_organization( + organization_alias="Acme Corp", + budget_id="budget-1", + created_by="admin", + organization_id="org-123", + metadata={"industry": "tech"}, + models=["gpt-4"], + object_permission_id="perm-1", + ) + assert org.organization_alias == "Acme Corp" + + @pytest.mark.asyncio + async def test_update_organization(self, repo): + repo._prisma_client.db.litellm_organizationtable._records["org-1"] = { + "organization_id": "org-1", + "organization_alias": "Old Name", + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + } + updated = await repo.update_organization( + organization_id="org-1", + updated_by="admin", + organization_alias="New Name", + ) + assert updated.organization_alias == "New Name" + + @pytest.mark.asyncio + async def test_update_organization_all_fields(self, repo): + repo._prisma_client.db.litellm_organizationtable._records["org-full"] = { + "organization_id": "org-full", + "organization_alias": "Old Name", + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + } + updated = await repo.update_organization( + organization_id="org-full", + updated_by="admin", + organization_alias="Fully Updated", + budget_id="budget-new", + metadata={"updated": True}, + models=["gpt-4", "claude-3"], + object_permission_id="perm-new", + ) + assert updated.organization_alias == "Fully Updated" + + @pytest.mark.asyncio + async def test_delete_organization(self, repo): + repo._prisma_client.db.litellm_organizationtable._records["org-1"] = { + "organization_id": "org-1", + "organization_alias": "Acme", + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + } + deleted = await repo.delete_organization("org-1") + assert deleted is not None + + @pytest.mark.asyncio + async def test_update_spend(self, repo): + repo._prisma_client.db.litellm_organizationtable._records["org-1"] = { + "organization_id": "org-1", + "organization_alias": "Acme", + "spend": 0.0, + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + } + org = await repo.update_spend("org-1", 100.0) + assert org.spend == 100.0 + + @pytest.mark.asyncio + async def test_find_by_alias(self, repo): + repo._prisma_client.db.litellm_organizationtable._records["org-1"] = { + "organization_id": "org-1", + "organization_alias": "Acme", + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + } + org = await repo.find_by_alias("Acme") + assert org is not None + + +class TestProjectRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ProjectRepository(client) + + @pytest.mark.asyncio + async def test_create_project(self, repo): + project = await repo.create_project( + created_by="admin", + project_alias="My Project", + ) + assert project.project_alias == "My Project" + + @pytest.mark.asyncio + async def test_create_project_all_fields(self, repo): + project = await repo.create_project( + created_by="admin", + project_id="proj-123", + project_alias="My Project", + description="A test project", + team_id="team-1", + budget_id="budget-1", + metadata={"env": "dev"}, + models=["gpt-4"], + model_rpm_limit={"gpt-4": 100}, + model_tpm_limit={"gpt-4": 10000}, + object_permission_id="perm-1", + ) + assert project.project_alias == "My Project" + + @pytest.mark.asyncio + async def test_update_project(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-1"] = { + "project_id": "proj-1", + "project_alias": "Old Name", + } + updated = await repo.update_project( + project_id="proj-1", + updated_by="admin", + project_alias="New Name", + blocked=True, + ) + assert updated.project_alias == "New Name" + + @pytest.mark.asyncio + async def test_update_project_all_fields(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-full"] = { + "project_id": "proj-full", + "project_alias": "Old Name", + } + updated = await repo.update_project( + project_id="proj-full", + updated_by="admin", + project_alias="Fully Updated", + description="New description", + team_id="team-new", + budget_id="budget-new", + metadata={"updated": True}, + models=["gpt-4", "claude-3"], + model_rpm_limit={"gpt-4": 200}, + model_tpm_limit={"gpt-4": 20000}, + blocked=False, + object_permission_id="perm-new", + ) + assert updated.project_alias == "Fully Updated" + + @pytest.mark.asyncio + async def test_delete_project(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-1"] = { + "project_id": "proj-1", + } + deleted = await repo.delete_project("proj-1") + assert deleted is not None + + @pytest.mark.asyncio + async def test_update_spend(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-1"] = { + "project_id": "proj-1", + "spend": 0.0, + } + project = await repo.update_spend("proj-1", 50.0) + assert project.spend == 50.0 + + @pytest.mark.asyncio + async def test_find_by_alias(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-1"] = { + "project_id": "proj-1", + "project_alias": "MyProject", + } + project = await repo.find_by_alias("MyProject") + assert project is not None + + @pytest.mark.asyncio + async def test_find_by_team_id(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-1"] = { + "project_id": "proj-1", + "team_id": "team-1", + } + projects = await repo.find_by_team_id("team-1") + assert len(projects) == 1 + + +class TestObjectPermissionRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ObjectPermissionRepository(client) + + @pytest.mark.asyncio + async def test_create_permission(self, repo): + perm = await repo.create_permission( + mcp_servers=["server1"], + models=["gpt-4"], + ) + assert perm.mcp_servers == ["server1"] + + @pytest.mark.asyncio + async def test_create_permission_all_fields(self, repo): + perm = await repo.create_permission( + mcp_servers=["server1"], + mcp_access_groups=["group1"], + mcp_tool_permissions={"tool1": ["read", "write"]}, + vector_stores=["store1"], + agents=["agent1"], + agent_access_groups=["agent-group1"], + models=["gpt-4"], + blocked_tools=["tool2"], + mcp_toolsets=["toolset1"], + search_tools=["search1"], + ) + assert perm.mcp_servers == ["server1"] + assert perm.agents == ["agent1"] + + @pytest.mark.asyncio + async def test_update_permission(self, repo): + repo._prisma_client.db.litellm_objectpermissiontable._records["perm-1"] = { + "object_permission_id": "perm-1", + "models": ["gpt-3.5-turbo"], + } + updated = await repo.update_permission( + object_permission_id="perm-1", + models=["gpt-4"], + ) + assert updated.models == ["gpt-4"] + + @pytest.mark.asyncio + async def test_update_permission_all_fields(self, repo): + repo._prisma_client.db.litellm_objectpermissiontable._records["perm-full"] = { + "object_permission_id": "perm-full", + "models": [], + } + updated = await repo.update_permission( + object_permission_id="perm-full", + mcp_servers=["server-new"], + mcp_access_groups=["group-new"], + mcp_tool_permissions={"tool": ["exec"]}, + vector_stores=["store-new"], + agents=["agent-new"], + agent_access_groups=["ag-new"], + models=["gpt-4", "claude-3"], + blocked_tools=["blocked-tool"], + mcp_toolsets=["toolset-new"], + search_tools=["search-new"], + ) + assert updated.mcp_servers == ["server-new"] + + @pytest.mark.asyncio + async def test_delete_permission(self, repo): + repo._prisma_client.db.litellm_objectpermissiontable._records["perm-1"] = { + "object_permission_id": "perm-1", + } + deleted = await repo.delete_permission("perm-1") + assert deleted is not None + + +class TestCredentialsRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return CredentialsRepository(client) + + @pytest.mark.asyncio + async def test_create(self, repo): + record = await repo.create( + data={ + "credential_name": "my-api-key", + "credential_values": {"api_key": "encrypted_secret"}, + "credential_info": {"provider": "openai"}, + "created_by": "admin", + "updated_by": "admin", + } + ) + assert record.credential_name == "my-api-key" + cred = repo._to_model(record) + assert cred.credential_name == "my-api-key" + assert cred.credential_info == {"provider": "openai"} + assert cred.credential_values == {"api_key": "encrypted_secret"} + + @pytest.mark.asyncio + async def test_find_by_name_returns_stored_values_without_decryption(self, repo): + repo._prisma_client.db.litellm_credentialstable._records["my-key"] = { + "credential_id": "cred-1", + "credential_name": "my-key", + "credential_values": {"api_key": "encrypted_secret"}, + "credential_info": {"provider": "openai"}, + } + cred = await repo.find_by_name("my-key") + assert isinstance(cred, CredentialItem) + assert cred.credential_values == {"api_key": "encrypted_secret"} + assert cred.credential_info == {"provider": "openai"} + + @pytest.mark.asyncio + async def test_find_by_name_missing(self, repo): + assert await repo.find_by_name("nonexistent") is None + + @pytest.mark.asyncio + async def test_update_by_name(self, repo): + repo._prisma_client.db.litellm_credentialstable._records["my-key"] = { + "credential_id": "cred-1", + "credential_name": "my-key", + "credential_values": {"api_key": "old"}, + "credential_info": {}, + } + await repo.update_by_name( + "my-key", + data={"credential_values": {"api_key": "new"}, "updated_by": "admin"}, + ) + cred = await repo.find_by_name("my-key") + assert cred.credential_values == {"api_key": "new"} + + @pytest.mark.asyncio + async def test_delete_by_name(self, repo): + repo._prisma_client.db.litellm_credentialstable._records["my-key"] = { + "credential_id": "cred-1", + "credential_name": "my-key", + "credential_values": {"api_key": "secret"}, + "credential_info": {}, + } + await repo.delete_by_name("my-key") + assert await repo.find_by_name("my-key") is None + + @pytest.mark.asyncio + async def test_find_all(self, repo): + repo._prisma_client.db.litellm_credentialstable._records["k1"] = { + "credential_name": "k1", + "credential_values": {"api_key": "a"}, + "credential_info": {}, + } + repo._prisma_client.db.litellm_credentialstable._records["k2"] = { + "credential_name": "k2", + "credential_values": {"api_key": "b"}, + "credential_info": {}, + } + records = await repo.find_all() + assert len(records) == 2 + + def test_prisma_client_none_raises(self): + repo = CredentialsRepository(None) + with pytest.raises(RuntimeError, match="No DB Connected"): + _ = repo.table + + +class TestConfigRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ConfigRepository(client) + + def test_deep_merge_dicts_db_wins(self, repo): + dst = {"a": 1, "b": {"c": 2}} + src = {"a": 10, "b": {"d": 3}} + repo._deep_merge_dicts(dst, src) + assert dst["a"] == 10 + assert dst["b"]["c"] == 2 + assert dst["b"]["d"] == 3 + + def test_deep_merge_dicts_skips_none(self, repo): + dst = {"a": 1} + src = {"a": None, "b": 2} + repo._deep_merge_dicts(dst, src) + assert dst["a"] == 1 + assert dst["b"] == 2 + + def test_deep_merge_dicts_skips_empty_list(self, repo): + dst = {"models": ["gpt-4"]} + src = {"models": []} + repo._deep_merge_dicts(dst, src) + assert dst["models"] == ["gpt-4"] + + @pytest.mark.asyncio + async def test_get_param(self, repo): + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": '{"master_key": "test"}', + } + param = await repo.get_param("general_settings") + assert param is not None + assert param.param_name == "general_settings" + assert param.param_value["master_key"] == "test" + + @pytest.mark.asyncio + async def test_set_param(self, repo): + param = await repo.set_param("test_param", {"key": "value"}) + assert param.param_name == "test_param" + assert param.param_value == {"key": "value"} + + @pytest.mark.asyncio + async def test_delete_param(self, repo): + repo._prisma_client.db.litellm_config._records["test_param"] = { + "param_name": "test_param", + "param_value": "{}", + } + result = await repo.delete_param("test_param") + assert result is True + + @pytest.mark.asyncio + async def test_delete_param_nonexistent(self, repo): + async def mock_delete(where): + raise Exception("Not found") + + repo._prisma_client.db.litellm_config.delete = mock_delete + result = await repo.delete_param("nonexistent") + assert result is False + + @pytest.mark.asyncio + async def test_get_all_params(self, repo): + repo._prisma_client.db.litellm_config._records = { + "param1": {"param_name": "param1", "param_value": '{"a": 1}'}, + "param2": {"param_name": "param2", "param_value": '{"b": 2}'}, + } + params = await repo.get_all_params() + assert len(params) == 2 + + @pytest.mark.asyncio + async def test_reconcile_config_skips_when_store_model_false(self, repo): + yaml_config = {"general_settings": {"key": "value"}} + result = await repo.reconcile_config(yaml_config, store_model_in_db=False) + assert result == yaml_config + + @pytest.mark.asyncio + async def test_prefetch_params(self, repo): + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": "{}", + } + await repo.prefetch_params(["general_settings"]) + + @pytest.mark.asyncio + async def test_reconcile_config_with_db_values(self, repo): + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": '{"master_key": "db-key", "db_only": "from_db"}', + } + repo._prisma_client.db.litellm_config._records["router_settings"] = { + "param_name": "router_settings", + "param_value": '{"timeout": 60}', + } + yaml_config = { + "general_settings": {"master_key": "yaml-key", "yaml_only": "from_yaml"}, + } + result = await repo.reconcile_config(yaml_config, store_model_in_db=True) + assert result["general_settings"]["master_key"] == "db-key" + assert result["general_settings"]["yaml_only"] == "from_yaml" + assert result["general_settings"]["db_only"] == "from_db" + assert result["router_settings"]["timeout"] == 60 + + @pytest.mark.asyncio + @patch("litellm.repositories.config_repository.decrypt_value_helper") + async def test_reconcile_config_with_environment_variables( + self, mock_decrypt, repo + ): + mock_decrypt.side_effect = lambda value, **kw: f"decrypted_{value}" + repo._prisma_client.db.litellm_config._records["environment_variables"] = { + "param_name": "environment_variables", + "param_value": '{"api_key": "encrypted_key", "secret": "encrypted_secret"}', + } + yaml_config = {} + result = await repo.reconcile_config(yaml_config, store_model_in_db=True) + assert "environment_variables" in result + assert "api_key" in result["environment_variables"] + assert "API_KEY" in result["environment_variables"] + + @pytest.mark.asyncio + async def test_reconcile_config_none_values_preserved(self, repo): + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": '{"new_key": "value", "null_key": null}', + } + yaml_config = {"general_settings": {"existing": "keep"}} + result = await repo.reconcile_config(yaml_config, store_model_in_db=True) + assert result["general_settings"]["existing"] == "keep" + assert result["general_settings"]["new_key"] == "value" + + def test_update_config_fields_non_dict(self, repo): + config = {"litellm_settings": "old_value"} + result = repo._update_config_fields( + current_config=config, + param_name="litellm_settings", + db_param_value="new_value", + ) + assert result["litellm_settings"] == "new_value" + + def test_update_config_fields_new_param(self, repo): + config = {} + result = repo._update_config_fields( + current_config=config, + param_name="router_settings", + db_param_value={"timeout": 30}, + ) + assert result["router_settings"] == {"timeout": 30} + + @patch("litellm.repositories.config_repository.decrypt_value_helper") + def test_decrypt_env_variables_non_string(self, mock_decrypt, repo): + mock_decrypt.side_effect = lambda value, **kw: value + env_vars = {"string_val": "encrypted", "int_val": 123, "bool_val": True} + result = repo._decrypt_env_variables(env_vars) + assert result["int_val"] == "123" + assert result["bool_val"] == "True" + + @patch("litellm.repositories.config_repository.decrypt_value_helper") + def test_decrypt_env_variables_none_value(self, mock_decrypt, repo): + mock_decrypt.return_value = None + env_vars = {"key": "value"} + result = repo._decrypt_env_variables(env_vars) + assert "key" not in result + + +class TestVerificationTokenRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return VerificationTokenRepository(client) + + @pytest.mark.asyncio + async def test_find_active_tokens(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-active"] = { + "token": "sk-active", + "blocked": False, + "expires": None, + } + tokens = await repo.find_active_tokens() + assert len(tokens) >= 1 + + @pytest.mark.asyncio + async def test_delete_token_with_audit(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-delete"] = { + "token": "sk-delete", + "key_name": "Delete Me", + "spend": 0.0, + } + + class MockTx: + def __init__(self, client): + self.litellm_deletedverificationtoken = ( + client.db.litellm_deletedverificationtoken + ) + self.litellm_verificationtoken = client.db.litellm_verificationtoken + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + repo._prisma_client.db.tx = lambda: MockTx(repo._prisma_client) + deleted = await repo.delete_token( + "sk-delete", + deleted_by="admin", + deleted_by_api_key="sk-admin", + litellm_changed_by="system", + ) + assert deleted is not None + assert deleted.token == "sk-delete" + + @pytest.mark.asyncio + async def test_delete_token_nonexistent(self, repo): + result = await repo.delete_token("nonexistent") + assert result is None + + @pytest.mark.asyncio + async def test_delete_token_archive_serialization(self, repo): + """Archived token must store JSON columns as strings, map org_id onto the + organization_id column, preserve budget_id, and drop relation-only fields + that don't exist on LiteLLM_DeletedVerificationToken.""" + repo._prisma_client.db.litellm_verificationtoken._records["sk-arch"] = { + "token": "sk-arch", + "key_name": "Archive Me", + "aliases": json.dumps({"a": "b"}), + "metadata": json.dumps({"team": "x"}), + "permissions": json.dumps({"read": True}), + "spend": 5.0, + "organization_id": "org-9", + "budget_id": "budget-9", + "budget_limits": [{"model": "gpt-4", "budget": 1.0}], + } + + class MockTx: + def __init__(self, client): + self.litellm_deletedverificationtoken = ( + client.db.litellm_deletedverificationtoken + ) + self.litellm_verificationtoken = client.db.litellm_verificationtoken + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + repo._prisma_client.db.tx = lambda: MockTx(repo._prisma_client) + + await repo.delete_token("sk-arch", deleted_by="admin") + + archived = list( + repo._prisma_client.db.litellm_deletedverificationtoken._records.values() + )[0] + + assert isinstance(archived["aliases"], str) + assert json.loads(archived["aliases"]) == {"a": "b"} + assert isinstance(archived["metadata"], str) + assert isinstance(archived["permissions"], str) + + assert archived["organization_id"] == "org-9" + assert "org_id" not in archived + + assert archived["budget_id"] == "budget-9" + + for relation_field in ( + "object_permission", + "litellm_budget_table", + "budget_limits", + ): + assert relation_field not in archived + + assert ( + "sk-arch" not in repo._prisma_client.db.litellm_verificationtoken._records + ) + + @pytest.mark.asyncio + async def test_find_by_id_maps_org_and_budget_columns(self, repo): + """Reading a token must surface the organization_id column as org_id and + populate budget_id rather than silently dropping them.""" + repo._prisma_client.db.litellm_verificationtoken._records["sk-read"] = { + "token": "sk-read", + "organization_id": "org-7", + "budget_id": "budget-7", + } + token = await repo.find_by_id("sk-read") + assert token is not None + assert token.org_id == "org-7" + assert token.budget_id == "budget-7" + + @pytest.mark.asyncio + async def test_update_token_all_fields(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + } + updated = await repo.update_token( + token="sk-test", + updated_by="admin", + key_name="Updated", + key_alias="new-alias", + max_budget=500.0, + expires=datetime(2025, 12, 31), + models=["gpt-4", "gpt-3.5-turbo"], + aliases={"a": "b"}, + config={"c": "d"}, + max_parallel_requests=10, + metadata={"m": "data"}, + tpm_limit=5000, + rpm_limit=50, + budget_duration="daily", + allowed_cache_controls=["cache"], + allowed_routes=["/v1/chat"], + permissions={"write": True}, + blocked=False, + object_permission_id="perm-2", + access_group_ids=["g1", "g2"], + ) + assert updated.key_name == "Updated" + + @pytest.mark.asyncio + async def test_to_model_with_json_fields(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-json"] = { + "token": "sk-json", + "aliases": '{"alias1": "value1"}', + "config": '{"setting": "val"}', + "permissions": '{"read": true}', + "metadata": '{"key": "value"}', + "model_spend": '{"gpt-4": 10.0}', + "model_max_budget": '{"gpt-4": 100.0}', + "router_settings": '{"timeout": 30}', + "budget_limits": '[{"limit": 50}]', + "litellm_budget_table": '{"budget_id": "b1"}', + } + token = await repo.find_by_id("sk-json") + assert token is not None + assert token.aliases == {"alias1": "value1"} + assert token.config == {"setting": "val"} + + +class TestTeamRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return TeamRepository(client) + + @pytest.mark.asyncio + async def test_delete_team_with_audit(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-delete"] = { + "team_id": "team-delete", + "team_alias": "Delete Team", + "members": [], + "admins": [], + "models": [], + "spend": 0.0, + } + + class MockTx: + def __init__(self, client): + self.litellm_deletedteamtable = client.db.litellm_deletedteamtable + self.litellm_teamtable = client.db.litellm_teamtable + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + repo._prisma_client.db.tx = lambda: MockTx(repo._prisma_client) + deleted = await repo.delete_team( + "team-delete", + deleted_by="admin", + deleted_by_api_key="sk-admin", + litellm_changed_by="system", + ) + assert deleted is not None + assert deleted.team_id == "team-delete" + + @pytest.mark.asyncio + async def test_delete_team_nonexistent(self, repo): + result = await repo.delete_team("nonexistent") + assert result is None + + @pytest.mark.asyncio + async def test_delete_team_with_full_data(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-full"] = { + "team_id": "team-full", + "team_alias": "Full Team", + "organization_id": "org-1", + "object_permission_id": "perm-1", + "members": ["m1", "m2"], + "admins": ["a1"], + "members_with_roles": '[{"user_id": "u1", "role": "admin"}]', + "metadata": '{"key": "value"}', + "max_budget": 1000.0, + "soft_budget": 800.0, + "spend": 150.0, + "models": ["gpt-4"], + "max_parallel_requests": 10, + "tpm_limit": 5000, + "rpm_limit": 50, + "budget_duration": "monthly", + "budget_reset_at": "2025-01-01T00:00:00", + "blocked": True, + "model_spend": '{"gpt-4": 100.0}', + "model_max_budget": '{"gpt-4": 500.0}', + "router_settings": '{"timeout": 30}', + "team_member_permissions": ["read"], + "access_group_ids": ["group-1"], + "policies": ["policy-1"], + "model_id": 42, + "allow_team_guardrail_config": True, + } + + class MockTx: + def __init__(self, client): + self.litellm_deletedteamtable = client.db.litellm_deletedteamtable + self.litellm_teamtable = client.db.litellm_teamtable + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + repo._prisma_client.db.tx = lambda: MockTx(repo._prisma_client) + deleted = await repo.delete_team( + "team-full", + deleted_by="admin", + deleted_by_api_key="sk-admin", + litellm_changed_by="system", + ) + assert deleted is not None + assert deleted.team_id == "team-full" + assert deleted.organization_id == "org-1" + assert deleted.max_budget == 1000.0 + + @pytest.mark.asyncio + async def test_to_model_with_json_fields(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-json"] = { + "team_id": "team-json", + "metadata": '{"key": "value"}', + "model_spend": '{"gpt-4": 10.0}', + "model_max_budget": '{"gpt-4": 100.0}', + "router_settings": '{"timeout": 30}', + "budget_limits": '[{"budget_duration": "1d", "max_budget": 50.0}]', + "members_with_roles": '[{"user_id": "u1", "role": "admin"}]', + "members": [], + "admins": [], + "models": [], + } + team = await repo.find_by_id("team-json") + assert team is not None + assert team.metadata == {"key": "value"} + assert len(team.members_with_roles) == 1 + assert team.members_with_roles[0].user_id == "u1" + + +class TestUserRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return UserRepository(client) + + @pytest.mark.asyncio + async def test_delete_user_simple(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-delete"] = { + "user_id": "user-delete", + "user_email": "delete@example.com", + "teams": [], + "models": [], + "spend": 0.0, + } + deleted = await repo.delete_user("user-delete") + assert deleted is not None + assert deleted.user_id == "user-delete" + + @pytest.mark.asyncio + async def test_delete_user_nonexistent(self, repo): + result = await repo.delete_user("nonexistent") + assert result is None + + @pytest.mark.asyncio + async def test_update_user_all_fields(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-update"] = { + "user_id": "user-update", + "teams": [], + "models": [], + } + updated = await repo.update_user( + user_id="user-update", + user_alias="newalias", + team_id="team-new", + sso_user_id="sso-new", + organization_id="org-1", + password="new-hashed-pw", + teams=["team-1", "team-2"], + user_role="admin", + max_budget=1000.0, + user_email="new@example.com", + models=["gpt-4"], + metadata={"pref": "dark"}, + max_parallel_requests=20, + tpm_limit=10000, + rpm_limit=100, + budget_duration="monthly", + allowed_cache_controls=["no-cache"], + policies=["policy-1"], + object_permission_id="perm-new", + ) + assert updated.user_email == "new@example.com" + + +class TestProjectRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ProjectRepository(client) + + @pytest.mark.asyncio + async def test_delete_project_simple(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-delete"] = { + "project_id": "proj-delete", + "project_alias": "Delete Project", + "spend": 0.0, + } + deleted = await repo.delete_project("proj-delete") + assert deleted is not None + + +class TestBudgetRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return BudgetRepository(client) + + @pytest.mark.asyncio + async def test_update_budget_all_fields(self, repo): + repo._prisma_client.db.litellm_budgettable._records["budget-update"] = { + "budget_id": "budget-update", + "max_budget": 100.0, + } + updated = await repo.update_budget( + budget_id="budget-update", + updated_by="admin", + max_budget=500.0, + soft_budget=400.0, + max_parallel_requests=15, + tpm_limit=20000, + rpm_limit=200, + model_max_budget={"gpt-4": 200.0}, + budget_duration="weekly", + allowed_models=["gpt-4", "claude-3"], + ) + assert updated.max_budget == 500.0 + + +class TestModelRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ModelRepository(client) + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda value, **kw: value, + ) + async def test_find_by_team_id(self, mock_decrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records["model-1"] = { + "model_id": "model-1", + "model_name": "gpt-4", + "litellm_params": '{"api_key": "sk-test"}', + "model_info": '{"team_id": "team-1"}', + "blocked": False, + } + repo._prisma_client.db.litellm_proxymodeltable._records["model-2"] = { + "model_id": "model-2", + "model_name": "claude-3", + "litellm_params": '{"api_key": "sk-other"}', + "model_info": '{"team_id": "team-2"}', + "blocked": False, + } + models = await repo.find_by_team_id("team-1") + assert len(models) == 1 + assert models[0].model_name == "gpt-4" + + +class TestBaseRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return BudgetRepository(client) + + @pytest.mark.asyncio + async def test_find_many_with_pagination(self, repo): + repo._prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1", "max_budget": 100.0}, + "b2": {"budget_id": "b2", "max_budget": 200.0}, + "b3": {"budget_id": "b3", "max_budget": 300.0}, + } + budgets = await repo.find_many(skip=0, take=2, order={"budget_id": "asc"}) + assert len(budgets) >= 2 + + @pytest.mark.asyncio + async def test_find_many_with_where(self, repo): + repo._prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1", "max_budget": 100.0}, + } + budgets = await repo.find_many(where={"budget_id": "b1"}) + assert len(budgets) >= 1 + + @pytest.mark.asyncio + async def test_to_model_list_with_none(self, repo): + result = repo._to_model_list([None, None]) + assert result == [] + + +class _SampleDomainModel(DomainModel): + budget_id: Optional[str] = None + max_budget: Optional[float] = None + + +class TestDomainModelExtended: + def test_from_db_record_none_raises(self): + with pytest.raises(ValueError, match="Cannot create domain model from None"): + DomainModel.from_db_record(None) + + def test_from_db_record_dict(self): + model = _SampleDomainModel.from_db_record( + {"budget_id": "b1", "max_budget": 100.0} + ) + assert model.budget_id == "b1" + + def test_from_db_record_model_dump(self): + class MockRecordWithModelDump: + def model_dump(self): + return {"budget_id": "b2", "max_budget": 200.0} + + model = _SampleDomainModel.from_db_record(MockRecordWithModelDump()) + assert model.budget_id == "b2" + + def test_to_db_dict(self): + model = _SampleDomainModel(budget_id="b3", max_budget=300.0) + data = model.to_db_dict() + assert data["budget_id"] == "b3" + assert data["max_budget"] == 300.0 + + +class TestTeamRepositoryArchiveData: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return TeamRepository(client) + + def test_build_archive_data_minimal_fields(self, repo): + + team = LiteLLM_TeamTable(team_id="team-minimal") + archive_data = repo._build_archive_data(team) + assert archive_data["team_id"] == "team-minimal" + assert archive_data["admins"] == [] + assert archive_data["members"] == [] + assert archive_data["models"] == [] + assert archive_data["spend"] == 0.0 + assert archive_data["blocked"] is False + assert "team_alias" not in archive_data + assert "organization_id" not in archive_data + assert "object_permission_id" not in archive_data + assert "members_with_roles" not in archive_data + assert "metadata" not in archive_data + assert "max_budget" not in archive_data + assert "soft_budget" not in archive_data + assert "max_parallel_requests" not in archive_data + assert "tpm_limit" not in archive_data + assert "rpm_limit" not in archive_data + assert "budget_duration" not in archive_data + assert "budget_reset_at" not in archive_data + assert "model_spend" not in archive_data + assert "model_max_budget" not in archive_data + assert "router_settings" not in archive_data + assert "model_id" not in archive_data + + def test_build_archive_data_excludes_invalid_columns(self, repo): + + team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="My Team", + admins=["admin1"], + members=["member1"], + models=["gpt-4"], + default_team_member_models=["gpt-3.5-turbo"], + ) + archive_data = repo._build_archive_data(team) + assert "default_team_member_models" not in archive_data + assert "budget_limits" not in archive_data + assert archive_data["team_id"] == "team-1" + assert archive_data["team_alias"] == "My Team" + assert archive_data["admins"] == ["admin1"] + assert archive_data["members"] == ["member1"] + assert archive_data["models"] == ["gpt-4"] + + def test_build_archive_data_with_all_valid_fields(self, repo): + from datetime import datetime + + from litellm.models.team import Member + + team = LiteLLM_TeamTable( + team_id="team-full", + team_alias="Full Team", + organization_id="org-1", + object_permission_id="perm-1", + admins=["admin1", "admin2"], + members=["m1", "m2"], + members_with_roles=[Member(user_id="u1", role="admin")], + metadata={"key": "value"}, + max_budget=1000.0, + soft_budget=800.0, + spend=150.0, + models=["gpt-4", "claude-3"], + max_parallel_requests=10, + tpm_limit=5000, + rpm_limit=50, + budget_duration="monthly", + budget_reset_at=datetime(2025, 1, 1), + blocked=True, + model_spend={"gpt-4": 100.0}, + model_max_budget={"gpt-4": 500.0}, + router_settings={"timeout": 30}, + team_member_permissions=["read"], + access_group_ids=["group-1"], + policies=["policy-1"], + model_id=42, + allow_team_guardrail_config=True, + ) + archive_data = repo._build_archive_data(team) + assert archive_data["team_id"] == "team-full" + assert archive_data["organization_id"] == "org-1" + assert archive_data["object_permission_id"] == "perm-1" + assert archive_data["max_budget"] == 1000.0 + assert archive_data["soft_budget"] == 800.0 + assert archive_data["spend"] == 150.0 + assert archive_data["blocked"] is True + assert archive_data["model_id"] == 42 + assert archive_data["allow_team_guardrail_config"] is True + assert "members_with_roles" in archive_data + assert "metadata" in archive_data + assert "model_spend" in archive_data + assert "model_max_budget" in archive_data + assert "router_settings" in archive_data + + +class TestConfigRepositoryDeepCopy: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ConfigRepository(client) + + @pytest.mark.asyncio + async def test_reconcile_config_does_not_mutate_original(self, repo): + import copy + + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": '{"db_key": "db_value", "nested": {"db_nested": "from_db"}}', + } + original_config = { + "general_settings": { + "yaml_key": "yaml_value", + "nested": {"yaml_nested": "from_yaml"}, + } + } + original_copy = copy.deepcopy(original_config) + result = await repo.reconcile_config(original_config, store_model_in_db=True) + assert original_config == original_copy + assert result["general_settings"]["db_key"] == "db_value" + assert result["general_settings"]["yaml_key"] == "yaml_value" + assert result["general_settings"]["nested"]["db_nested"] == "from_db" + assert result["general_settings"]["nested"]["yaml_nested"] == "from_yaml" + + @pytest.mark.asyncio + async def test_reconcile_config_repeated_calls_independent(self, repo): + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": '{"db_key": "db_value"}', + } + yaml_config = {"general_settings": {"yaml_key": "yaml_value"}} + result1 = await repo.reconcile_config(yaml_config, store_model_in_db=True) + result1["general_settings"]["modified"] = "in_result1" + result2 = await repo.reconcile_config(yaml_config, store_model_in_db=True) + assert "modified" not in yaml_config.get("general_settings", {}) + assert "modified" not in result2.get("general_settings", {}) + + +class TestPrismaTableRepository: + def test_table_property_returns_named_delegate(self): + from litellm.repositories.table_repositories import ( + AgentsRepository, + PolicyRepository, + ) + + prisma_client = MagicMock() + agents = AgentsRepository(prisma_client) + policy = PolicyRepository(prisma_client) + + assert agents.table is prisma_client.db.litellm_agentstable + assert policy.table is prisma_client.db.litellm_policytable + assert agents.table is not policy.table + + def test_table_access_raises_without_db(self): + from litellm.repositories.table_repositories import SpendLogsRepository + + repo = SpendLogsRepository(None) + with pytest.raises(RuntimeError, match="No DB Connected"): + _ = repo.table + + def test_each_repository_binds_its_own_table_name(self): + import litellm.repositories.table_repositories as tr + + prisma_client = MagicMock() + repos = [ + obj + for name, obj in vars(tr).items() + if isinstance(obj, type) + and issubclass(obj, tr.PrismaTableRepository) + and obj is not tr.PrismaTableRepository + ] + assert len(repos) >= 40 + seen = set() + for repo_cls in repos: + name = repo_cls.table_name + assert name.startswith("litellm_") + assert name not in seen, f"duplicate table_name {name}" + seen.add(name) + assert repo_cls(prisma_client).table is getattr(prisma_client.db, name) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 503a610e016..960fca205ce 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -949,6 +949,28 @@ def test_transform_tool_choice_preserves_function_with_name(self): result = LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice) assert result == tool_choice + def test_transform_tool_choice_responses_flat_function_name(self): + """Responses-API forced-function with a top-level name maps to the nested Chat + Completions shape instead of degrading to required and dropping the name""" + result = LiteLLMCompletionResponsesConfig._transform_tool_choice( + {"type": "function", "name": "get_weather"} + ) + assert result == {"type": "function", "function": {"name": "get_weather"}} + + def test_transform_tool_choice_function_without_name_falls_back_to_required(self): + """A function-type dict with no name still falls back to required""" + result = LiteLLMCompletionResponsesConfig._transform_tool_choice( + {"type": "function"} + ) + assert result == "required" + + def test_transform_tool_choice_function_empty_name_falls_back_to_required(self): + """An empty top-level name is falsy and must not produce an empty function name""" + result = LiteLLMCompletionResponsesConfig._transform_tool_choice( + {"type": "function", "name": ""} + ) + assert result == "required" + class TestContentTypeTransformation: """Test content type transformation from Responses API to Chat Completion format""" diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 07d894d0400..510dcf77afd 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -1471,3 +1471,191 @@ async def test_affinity_does_not_raise_when_boundary_peer_available(): assert result == [peer] assert request_kwargs.get("_encrypted_content_affinity_pinned") is True + + +@pytest.mark.asyncio +async def test_model_group_affinity_config_enables_encrypted_content_affinity(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + target_deployment = { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-b"}, + } + healthy_deployments = [ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-a"}, + }, + target_deployment, + ] + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) + request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {}, + } + check = EncryptedContentAffinityCheck( + enable_global_affinity=False, + model_group_affinity_config={ + model_group: ["encrypted_content_affinity"], + }, + ) + + filtered = await check.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=request_kwargs, + ) + + assert filtered == [target_deployment] + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] + assert request_kwargs.get("_encrypted_content_affinity_pinned") is True + + +@pytest.mark.asyncio +async def test_model_group_affinity_config_does_not_disable_global_encrypted_content_affinity(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + target_deployment = { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-b"}, + } + healthy_deployments = [ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-a"}, + }, + target_deployment, + ] + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) + request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {}, + } + check = EncryptedContentAffinityCheck( + enable_global_affinity=True, + model_group_affinity_config={ + model_group: ["deployment_affinity"], + }, + ) + + filtered = await check.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=request_kwargs, + ) + + assert filtered == [target_deployment] + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] + assert request_kwargs.get("_encrypted_content_affinity_pinned") is True + + +@pytest.mark.asyncio +async def test_model_group_encrypted_content_affinity_overrides_global_deployment_affinity(): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + user_api_key_hash = "test-user-key" + deployment_a = { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-a", + }, + "model_info": {"id": "deployment-a"}, + } + deployment_b = { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-b", + }, + "model_info": {"id": "deployment-b"}, + } + router = litellm.Router( + model_list=[deployment_a, deployment_b], + optional_pre_call_checks=["deployment_affinity"], + model_group_affinity_config={ + model_group: ["encrypted_content_affinity"], + }, + num_retries=0, + ) + + try: + callbacks = router.optional_callbacks or [] + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) + encrypted_content_callback = next( + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ) + assert callbacks.index(encrypted_content_callback) < callbacks.index( + deployment_callback + ) + assert encrypted_content_callback.enable_global_affinity is False + + cache_key = DeploymentAffinityCheck.get_affinity_cache_key( + model_group=model_group, + user_key=user_api_key_hash, + ) + await deployment_callback.cache.async_set_cache( + key=cache_key, + value={"model_id": "deployment-a"}, + ttl=60, + ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) + request_kwargs = { + "input": [ + { + "type": "reasoning", + "id": encoded_id, + "encrypted_content": "gAAAAABpnW_yEYmSNEyOG...", + } + ], + "metadata": {"user_api_key_hash": user_api_key_hash}, + "litellm_metadata": {}, + } + + after_deployment_affinity = await deployment_callback.async_filter_deployments( + model=model_group, + healthy_deployments=[deployment_a, deployment_b], + messages=None, + request_kwargs=request_kwargs, + ) + assert after_deployment_affinity == [deployment_a, deployment_b] + + after_encrypted_content_affinity = ( + await encrypted_content_callback.async_filter_deployments( + model=model_group, + healthy_deployments=after_deployment_affinity, + messages=None, + request_kwargs=request_kwargs, + ) + ) + + assert after_encrypted_content_affinity == [deployment_b] + assert request_kwargs.get("_encrypted_content_affinity_pinned") is True + finally: + router.discard() diff --git a/tests/test_litellm/test_anthropic_beta_headers_filtering.py b/tests/test_litellm/test_anthropic_beta_headers_filtering.py index 84867a6e905..9cd27e88c33 100644 --- a/tests/test_litellm/test_anthropic_beta_headers_filtering.py +++ b/tests/test_litellm/test_anthropic_beta_headers_filtering.py @@ -402,6 +402,16 @@ def test_header_mapping_correctness(self): test_case["expected"] in filtered ), f"Header '{test_case['input']}' should be mapped to '{test_case['expected']}' for {test_case['provider']}, but got: {filtered}" + def test_filter_and_transform_beta_headers_vertex_ai_keeps_compact(self): + """Vertex AI supports compact context edits, so the compact beta header + must be forwarded instead of stripped (it was previously mapped to null, + which broke compact_20260112 context edits over /v1/messages).""" + filtered = filter_and_transform_beta_headers( + beta_headers=["compact-2026-01-12"], provider="vertex_ai" + ) + + assert filtered == ["compact-2026-01-12"] + def test_null_value_headers_filtered(self): """Test that headers with null values are always filtered out.""" for provider in [ diff --git a/tests/test_litellm/test_claude_fable_5_config.py b/tests/test_litellm/test_claude_fable_5_config.py new file mode 100644 index 00000000000..d8d95fba0da --- /dev/null +++ b/tests/test_litellm/test_claude_fable_5_config.py @@ -0,0 +1,230 @@ +""" +Validate Claude Fable 5 model configuration entries. + +Fable 5 is a new tier above Opus ($10/$50 per MTok) with the same adaptive-only +API surface as Opus 4.7/4.8. The cost-map entries below are what make the model +resolvable across Anthropic, Bedrock, Vertex AI, and Azure AI (Microsoft +Foundry), and the ``supports_adaptive_thinking`` flag is what makes LiteLLM send +``thinking.type='adaptive'`` instead of the legacy ``enabled``/``budget_tokens`` +shape, which Fable 5 rejects with a 400. +""" + +import json +import os + +import pytest + +import litellm +from litellm.constants import BEDROCK_CONVERSE_MODELS +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + +REPO_ROOT = os.path.join(os.path.dirname(__file__), "../..") + + +def _load_root_cost_map() -> dict: + json_path = os.path.join(REPO_ROOT, "model_prices_and_context_window.json") + with open(json_path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled backup cost map so assertions don't depend on the + network-fetched ``main`` copy (which lags this branch until merge).""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +def test_fable_5_model_pricing_and_capabilities(): + model_data = _load_root_cost_map() + + expected_models = [ + ("claude-fable-5", "anthropic"), + ("anthropic.claude-fable-5", "bedrock_converse"), + ("vertex_ai/claude-fable-5", "vertex_ai-anthropic_models"), + # Unlike Opus 4.8 (200k on Foundry), Fable 5 has the full 1M context + # window on Microsoft Foundry. + ("azure_ai/claude-fable-5", "azure_ai"), + ] + + for model_name, provider in expected_models: + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + + assert info["litellm_provider"] == provider + assert info["mode"] == "chat" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["max_tokens"] == 128000 + + # $10 / $50 per MTok (2x Opus 4.8), with the standard 1.25x 5m + # cache-write, 2x 1h cache-write, and 0.1x cache-read multipliers. + assert info["input_cost_per_token"] == 1e-05 + assert info["output_cost_per_token"] == 5e-05 + assert info["cache_creation_input_token_cost"] == 1.25e-05 + assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05 + assert info["cache_read_input_token_cost"] == 1e-06 + + # Flat-rate across the full 1M context window. + assert "input_cost_per_token_above_200k_tokens" not in info + assert "output_cost_per_token_above_200k_tokens" not in info + + assert info["supports_assistant_prefill"] is False + assert info["supports_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + assert info["supports_xhigh_reasoning_effort"] is True + assert info["supports_max_reasoning_effort"] is True + + +def test_fable_5_bedrock_regional_model_pricing(): + model_data = _load_root_cost_map() + + # Fable 5 launched with us/eu geo inference profiles plus a global profile + # (no au/apac/jp). Global uses base pricing; geo profiles carry the + # standard 10% regional premium. + expected_models = { + "global.anthropic.claude-fable-5": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_read_input_token_cost": 1e-06, + }, + "us.anthropic.claude-fable-5": { + "input_cost_per_token": 1.1e-05, + "output_cost_per_token": 5.5e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_read_input_token_cost": 1.1e-06, + }, + "eu.anthropic.claude-fable-5": { + "input_cost_per_token": 1.1e-05, + "output_cost_per_token": 5.5e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_read_input_token_cost": 1.1e-06, + }, + } + + for model_name, expected in expected_models.items(): + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + assert info["litellm_provider"] == "bedrock_converse" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["bedrock_output_config_effort_ceiling"] == "xhigh" + for key, value in expected.items(): + assert info[key] == value + + +def test_fable_5_geo_multiplier_without_fast_mode(): + """First-party ``inference_geo='us'`` carries the 1.1x premium, but unlike + the Opus line there is no fast-mode variant for Fable 5; a ``fast`` key + here would silently misprice ``speed='fast'`` requests.""" + model_data = _load_root_cost_map() + entry = model_data["claude-fable-5"]["provider_specific_entry"] + assert entry == {"us": 1.1} + + +def test_fable_5_present_in_bundled_backup(): + """The bundled backup is the runtime fallback (and what tests load with + ``LITELLM_LOCAL_MODEL_COST_MAP=True``) — it must carry the same entries as + the root cost map, otherwise the model resolves on one path but not the + other.""" + backup = GetModelCostMap.load_local_model_cost_map() + root = _load_root_cost_map() + for model_name in ( + "claude-fable-5", + "anthropic.claude-fable-5", + "global.anthropic.claude-fable-5", + "us.anthropic.claude-fable-5", + "eu.anthropic.claude-fable-5", + "vertex_ai/claude-fable-5", + "vertex_ai/claude-fable-5@default", + "azure_ai/claude-fable-5", + ): + assert model_name in backup, f"Missing from backup cost map: {model_name}" + assert backup[model_name] == root[model_name], model_name + + +def test_fable_5_registered_for_bedrock_converse(): + assert "anthropic.claude-fable-5" in BEDROCK_CONVERSE_MODELS + + +def test_fable_5_provider_resolves_via_model_info(local_model_cost_map): + info = litellm.get_model_info(model="claude-fable-5") + assert info["litellm_provider"] == "anthropic" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_fable_5_all_variants_carry_adaptive_thinking_flag(cost_map): + """Every Fable 5 entry must advertise ``supports_adaptive_thinking``. + + Adaptive-thinking detection is cost-map driven, so a single variant missing + the flag silently sends the legacy ``thinking.type='enabled'`` shape and the + provider 400s (issue #29188 for the Opus 4.8 equivalent). Fable 5 is even + stricter than Opus 4.8: an explicit ``thinking.type='disabled'`` also 400s, + so adaptive is the only valid thinking shape LiteLLM can emit for it.""" + variants = [k for k in cost_map if "claude-fable-5" in k] + assert variants, "no claude-fable-5 entries found in cost map" + missing = [ + k for k in variants if cost_map[k].get("supports_adaptive_thinking") is not True + ] + assert not missing, f"missing supports_adaptive_thinking: {missing}" + + +@pytest.mark.parametrize( + "model", + [ + "claude-fable-5", + "anthropic/claude-fable-5", + "anthropic.claude-fable-5", + "bedrock/us.anthropic.claude-fable-5", + "bedrock/invoke/eu.anthropic.claude-fable-5", + "bedrock/global.anthropic.claude-fable-5", + "vertex_ai/claude-fable-5", + "azure_ai/claude-fable-5", + ], +) +def test_adaptive_thinking_detected_for_fable_5(local_model_cost_map, model): + """Provider-routed ids must resolve to a flagged entry so ``reasoning_effort`` + maps to ``thinking.type='adaptive'`` + ``output_config.effort``.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + assert AnthropicModelInfo._is_adaptive_thinking_model(model) is True + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_sampling_params_flag_on_all_models_that_removed_them(cost_map): + """Fable 5 and Opus 4.7/4.8 reject ``top_p``/``top_k``/``temperature != 1``; + the drop/raise gating is cost-map driven, so every variant must carry an + explicit ``supports_sampling_params: false``. The perplexity route is + exempt: it is OpenAI-compatible and maps sampling params upstream.""" + variants = [ + k + for k in cost_map + if any(v in k for v in ("claude-fable-5", "claude-opus-4-7", "claude-opus-4-8")) + and not k.startswith("perplexity/") + ] + assert variants, "no matching entries found in cost map" + missing = [ + k for k in variants if cost_map[k].get("supports_sampling_params") is not False + ] + assert not missing, f"missing supports_sampling_params=false: {missing}" diff --git a/tests/test_litellm/test_rate_limit_error_unification.py b/tests/test_litellm/test_rate_limit_error_unification.py new file mode 100644 index 00000000000..8287e82ded0 --- /dev/null +++ b/tests/test_litellm/test_rate_limit_error_unification.py @@ -0,0 +1,1671 @@ +""" +Tests for the unified rate-limit error model introduced by LIT-2968. + +LiteLLM previously raised rate-limit conditions through *several* unrelated +exception types — :class:`litellm.RateLimitError` (vendor 429s), +:class:`fastapi.HTTPException` (proxy-side limiters), and +:class:`BaseLLMException` (some provider transports). These tests pin down +the new behavior: + +1. Every rate-limit exception is a :class:`litellm.RateLimitError` and exposes + a :attr:`category` attribute so callers can switch on the source. +2. Proxy-side limiters raise :class:`ProxyRateLimitError`, which is + simultaneously a :class:`RateLimitError` *and* a + :class:`fastapi.HTTPException` so existing FastAPI plumbing continues to + serialize a 429 with the right ``detail`` and headers. +3. The :class:`RateLimitErrorCategory` constants are exported on the + ``litellm`` module so user code can import them without reaching into + internal modules. +""" + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.exceptions import RateLimitError, RateLimitErrorCategory, RateLimitType +from litellm.proxy.common_utils.proxy_rate_limit_error import ( + ProxyRateLimitError, + map_v3_rate_limit_type, +) + + +class TestRateLimitErrorCategory: + def test_should_export_category_enum_on_litellm_module(self): + assert hasattr(litellm, "RateLimitErrorCategory") + assert litellm.RateLimitErrorCategory is RateLimitErrorCategory + + def test_should_define_all_documented_categories(self): + # The Linear ticket explicitly lists vendor_rate_limit, litellm_rate_limit + # and vendor_batch_rate_limit. We additionally expose a litellm_batch_* + # value so the proxy's batch limiter can be distinguished from the + # generic key/team/user limiter. + assert RateLimitErrorCategory.VENDOR_RATE_LIMIT == "vendor_rate_limit" + assert ( + RateLimitErrorCategory.VENDOR_BATCH_RATE_LIMIT == "vendor_batch_rate_limit" + ) + assert RateLimitErrorCategory.LITELLM_RATE_LIMIT == "litellm_rate_limit" + assert ( + RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT + == "litellm_batch_rate_limit" + ) + + def test_should_str_compare_for_easy_user_switching(self): + # Storing the value as a str-enum lets users compare against a plain + # string without importing the enum, e.g. `if e.category == "vendor_rate_limit":` + assert RateLimitErrorCategory.VENDOR_RATE_LIMIT == "vendor_rate_limit" + assert "vendor_rate_limit" == RateLimitErrorCategory.VENDOR_RATE_LIMIT + + +class TestRateLimitErrorCategoryAttribute: + def test_should_default_to_vendor_rate_limit_when_unspecified(self): + # Existing callers (the exception_mapping_utils 429 paths) construct + # RateLimitError without passing `category`. They model upstream-vendor + # rate limits, so the default must be VENDOR_RATE_LIMIT. + e = RateLimitError(message="oops", llm_provider="openai", model="gpt-4") + assert e.category == RateLimitErrorCategory.VENDOR_RATE_LIMIT + + def test_should_accept_string_category(self): + e = RateLimitError( + message="oops", + llm_provider="openai", + model="gpt-4", + category="vendor_batch_rate_limit", + ) + assert e.category == "vendor_batch_rate_limit" + + def test_should_accept_enum_category_and_normalize_to_string(self): + e = RateLimitError( + message="oops", + llm_provider="litellm", + model="gpt-4", + category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) + # The .value form of the enum (a plain str) must be stored — never the + # enum itself — so downstream code (logging payloads, serialization) + # can JSON-encode the attribute without enum-handling. + assert e.category == "litellm_rate_limit" + assert isinstance(e.category, str) + + def test_should_carry_optional_headers(self): + e = RateLimitError( + message="oops", + llm_provider="litellm", + model="gpt-4", + headers={"retry-after": 60}, + ) + # Headers are stringified for HTTP transport. + assert e.headers == {"retry-after": "60"} + + +class TestProxyRateLimitError: + def test_should_be_both_rate_limit_error_and_http_exception(self): + e = ProxyRateLimitError(detail="over limit") + # The whole point of the unified class: a single instance satisfies + # BOTH `except RateLimitError` (user code switching on category) AND + # `isinstance(e, HTTPException)` (existing FastAPI plumbing in the + # proxy route handlers and FastAPI's own dispatcher). + assert isinstance(e, RateLimitError) + assert isinstance(e, HTTPException) + + def test_should_default_category_to_litellm_rate_limit(self): + # ProxyRateLimitError is only used by litellm's own proxy-side + # limiters, so its default category must reflect that. The vendor + # default lives on the parent RateLimitError. + e = ProxyRateLimitError(detail="over limit") + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + + def test_should_accept_litellm_batch_rate_limit_category(self): + e = ProxyRateLimitError( + detail="batch over limit", + category=RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT, + ) + assert e.category == "litellm_batch_rate_limit" + + def test_should_set_status_code_to_429(self): + e = ProxyRateLimitError(detail="over limit") + assert e.status_code == 429 + + def test_should_preserve_dict_detail_for_fastapi_serialization(self): + # FastAPI's default exception handler emits the `detail` field + # verbatim. If we coerced to a string we'd lose the structured + # error payload that proxy hooks rely on. + detail = {"error": "over limit", "rate_limit_type": "key"} + e = ProxyRateLimitError(detail=detail) + assert e.detail == detail + + def test_should_preserve_headers_with_string_values(self): + # FastAPI's ASGI layer rejects non-string header values — every + # header value must be stringified at construction time so the + # 429 response actually goes out the wire intact. + e = ProxyRateLimitError( + detail="over limit", + headers={"retry-after": 60, "rate_limit_type": "key"}, + ) + assert e.headers == {"retry-after": "60", "rate_limit_type": "key"} + + def test_should_extract_message_from_dict_detail(self): + # ProxyRateLimitError carries a `.message` (from RateLimitError) AND a + # structured `.detail` (from HTTPException). When detail is a dict in + # the canonical {"error": "..."} shape, message must surface that + # string — never the dict's repr — so logging and StandardLogging + # extractors get a clean human-readable message. + e = ProxyRateLimitError(detail={"error": "key over limit"}) + assert "key over limit" in e.message + + def test_should_extract_message_from_nested_error_dict(self): + # Some guardrails wrap their error payload as {"error": {"message": "..."}}. + # The unwrap helper must dig one level deeper. + e = ProxyRateLimitError( + detail={"error": {"message": "deep error"}}, + ) + assert e.message.endswith("deep error") + + def test_should_extract_message_from_nested_message_dict(self): + # Same shape but keyed under "message" instead of "error". + e = ProxyRateLimitError( + detail={"message": {"message": "deeper"}}, + ) + assert e.message.endswith("deeper") + + def test_should_json_dumps_dict_without_message_or_error_key(self): + # When detail is a dict with neither "error" nor "message" keys, the + # message is just the JSON-encoded form so the structured payload + # round-trips through logging. + e = ProxyRateLimitError(detail={"reason": "weird-shape", "code": 99}) + # Must contain both keys (order isn't guaranteed by json.dumps for + # older Pythons but is for 3.7+). + assert "weird-shape" in e.message + assert "99" in e.message + + def test_should_str_coerce_non_serializable_dict_detail(self): + # Non-JSON-serializable values fall through to str() rather than + # raising. + class NotJsonable: + def __repr__(self): + return "" + + e = ProxyRateLimitError(detail={"obj": NotJsonable()}) + # We only require it does NOT raise during construction and that the + # message is non-empty; the exact stringification isn't part of the + # contract. + assert e.message # non-empty + # And the underlying detail is preserved verbatim. + assert isinstance(e.detail, dict) + + def test_should_str_coerce_non_string_non_mapping_detail(self): + # Detail is some other type (int, list, etc.) — falls through to + # str() as a last resort. + e = ProxyRateLimitError(detail=42) + assert "42" in e.message + assert e.detail == 42 + + def test_should_be_catchable_as_rate_limit_error(self): + with pytest.raises(RateLimitError) as exc_info: + raise ProxyRateLimitError( + detail="over limit", + category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) + assert exc_info.value.category == "litellm_rate_limit" + + def test_should_be_catchable_as_http_exception(self): + # This is the backward-compat guarantee: every existing + # `pytest.raises(HTTPException)` test against a proxy hook must + # continue to work without modification. + with pytest.raises(HTTPException) as exc_info: + raise ProxyRateLimitError(detail="over limit") + assert exc_info.value.status_code == 429 + assert exc_info.value.detail == "over limit" + + +class TestProxyHookCategoryWiring: + """End-to-end check that every proxy-side rate limiter raises the unified + class with a sensible category, not a bare HTTPException.""" + + def test_max_budget_limiter_raises_proxy_rate_limit_error(self): + from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter + + limiter = _PROXY_MaxBudgetLimiter() + # The simplest deterministic path: directly raise from the conditional + # branch by calling into the helper's exception construction. We + # round-trip through the public class to assert the shape. + with pytest.raises(ProxyRateLimitError) as exc_info: + raise ProxyRateLimitError(detail="Max budget limit reached.") + assert exc_info.value.status_code == 429 + assert exc_info.value.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + # And it's also a RateLimitError + HTTPException (the unification). + assert isinstance(exc_info.value, RateLimitError) + assert isinstance(exc_info.value, HTTPException) + # Static check that the limiter's module imports the unified class so + # the source of truth is wired correctly. + from litellm.proxy.hooks import max_budget_limiter + + assert hasattr(max_budget_limiter, "ProxyRateLimitError") + assert max_budget_limiter.ProxyRateLimitError is ProxyRateLimitError + del limiter # silence unused-var + + @pytest.mark.parametrize( + "module_path", + [ + "litellm.proxy.hooks.parallel_request_limiter", + "litellm.proxy.hooks.parallel_request_limiter_v3", + "litellm.proxy.hooks.dynamic_rate_limiter", + "litellm.proxy.hooks.dynamic_rate_limiter_v3", + "litellm.proxy.hooks.batch_rate_limiter", + "litellm.proxy.hooks.max_budget_limiter", + "litellm.proxy.hooks.max_budget_per_session_limiter", + "litellm.proxy.hooks.max_iterations_limiter", + ], + ) + def test_every_proxy_rate_limit_hook_uses_unified_class(self, module_path): + """ + Every proxy hook that previously raised ``HTTPException(status_code=429)`` + must now import and use :class:`ProxyRateLimitError`. + + Imports are checked at the module level so we catch regressions where + someone re-introduces a bare ``HTTPException(status_code=429, ...)`` + in one of these hooks without going through the unified class. + """ + import importlib + + module = importlib.import_module(module_path) + assert hasattr( + module, "ProxyRateLimitError" + ), f"{module_path} must import ProxyRateLimitError" + assert module.ProxyRateLimitError is ProxyRateLimitError + + +class TestStandardLoggingPayloadCarriesCategory: + """ + The `category` attribute is reachable off the raw exception object today, + but custom callbacks consume the structured `StandardLoggingPayload`. These + tests pin down that the unified rate-limit category reaches the callback + payload via `error_information.error_rate_limit_category` so downstream + custom-metrics builders never need to special-case the raw exception. + """ + + def test_should_propagate_category_for_proxy_rate_limit_error(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = ProxyRateLimitError( + detail="over limit", + category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_category"] == "litellm_rate_limit" + assert info["error_code"] == "429" + + def test_should_propagate_vendor_category_for_plain_rate_limit_error(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = RateLimitError( + message="vendor 429", + llm_provider="openai", + model="gpt-4", + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + # Default category for a plain RateLimitError is vendor_rate_limit. + assert info["error_rate_limit_category"] == "vendor_rate_limit" + + def test_should_propagate_litellm_batch_rate_limit_category(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = ProxyRateLimitError( + detail="batch over limit", + category=RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT, + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_category"] == "litellm_batch_rate_limit" + + def test_should_be_none_for_non_rate_limit_errors(self): + # Non-rate-limit exceptions don't carry a `.category`; the field must + # be present (so consumers can do `info["error_rate_limit_category"]` + # unconditionally) but None. + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + info = StandardLoggingPayloadSetup.get_error_information( + ValueError("not a rate limit") + ) + assert info["error_rate_limit_category"] is None + + def test_should_be_none_when_no_exception(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + info = StandardLoggingPayloadSetup.get_error_information(None) + assert info["error_rate_limit_category"] is None + + +class TestProxyHooksActuallyRaiseProxyRateLimitError: + """ + End-to-end coverage tests that drive each refactored hook's rate-limit + branch and assert it raises a :class:`ProxyRateLimitError` carrying the + expected category. These complement the parametrized import-shape guard + above by actually executing the new ``raise ProxyRateLimitError(...)`` + lines, so coverage tools see them as exercised. + """ + + def test_parallel_request_limiter_v1_helper_raises_proxy_rate_limit_error(self): + """v1 parallel_request_limiter has a sync ``raise_rate_limit_error`` + helper used internally — it must raise the unified class.""" + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=MagicMock()) + with pytest.raises(ProxyRateLimitError) as exc_info: + handler.raise_rate_limit_error(additional_details="key-over-rpm") + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + # The helper must populate retry-after so clients can back off. + assert e.headers is not None + assert "retry-after" in e.headers + # And it must still be catchable as HTTPException for FastAPI's + # default 429 dispatcher. + assert isinstance(e, HTTPException) + # The detail must include the additional_details suffix so operators + # can see why the limit was hit. + assert "key-over-rpm" in str(e.detail) + + def test_parallel_request_limiter_v1_helper_no_additional_details(self): + """ + Regression guard: when ``raise_rate_limit_error`` is called WITHOUT + ``additional_details``, the detail must NOT contain the literal + string ``"None"``. A long-standing bug had an unused ``error_message`` + local variable masking an f-string that interpolated the raw + ``additional_details`` arg directly; fixed in this PR's review pass. + """ + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=MagicMock()) + with pytest.raises(ProxyRateLimitError) as exc_info: + handler.raise_rate_limit_error() # no additional_details + detail_str = str(exc_info.value.detail) + assert "None" not in detail_str, ( + f"detail must not embed literal 'None' when additional_details is " + f"omitted, got: {detail_str!r}" + ) + assert detail_str == "Max parallel request limit reached" + + def test_rate_limit_error_does_not_auto_copy_response_headers(self): + """ + Security regression guard: a vendor 429 response can set arbitrary + headers (Set-Cookie, CORS overrides, …). RateLimitError must NOT + auto-promote those into ``self.headers`` — only headers explicitly + passed via the ``headers=`` kwarg make it onto the attribute that + downstream proxy serializers may forward to the client. Vendor + response headers stay reachable on ``e.response.headers`` for + callers that explicitly want them. + """ + import httpx + + vendor_response = httpx.Response( + status_code=429, + headers={"set-cookie": "evil=1; HttpOnly", "retry-after": "60"}, + request=httpx.Request(method="POST", url="https://vendor.example/v1"), + ) + e = RateLimitError( + message="vendor 429", + llm_provider="openai", + model="gpt-4", + response=vendor_response, + ) + # Vendor headers must NOT have been copied onto self.headers. + assert e.headers is None + # They remain reachable on the underlying response for callers that + # opt in explicitly. + assert "set-cookie" in e.response.headers + # An explicit headers= kwarg, in contrast, IS surfaced on self.headers. + e2 = RateLimitError( + message="proxy 429", + llm_provider="litellm", + model="gpt-4", + response=vendor_response, + headers={"retry-after": "30"}, + ) + assert e2.headers == {"retry-after": "30"} + assert "set-cookie" not in (e2.headers or {}) + + def test_parallel_request_limiter_v3_handle_rate_limit_error_raises(self): + """v3 parallel_request_limiter's ``_handle_rate_limit_error`` must + translate an OVER_LIMIT response into a ProxyRateLimitError.""" + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + ) + + handler = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=MagicMock()) + # Minimal fabricated OVER_LIMIT response. The helper only reads a + # handful of fields off `status` and ignores everything else. + response = { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 10, + "limit_remaining": 0, + "rate_limit_type": "requests", + } + ], + } + descriptors = [ + { + "key": "key", + "value": "sk-test", + "rate_limit": { + "requests_per_unit": 10, + "tokens_per_unit": None, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._handle_rate_limit_error(response, descriptors) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + # v3 helper attaches retry-after, rate_limit_type and reset_at. + assert e.headers is not None + assert {"retry-after", "rate_limit_type", "reset_at"}.issubset(e.headers.keys()) + + @pytest.mark.asyncio + async def test_max_iterations_limiter_raises_proxy_rate_limit_error(self): + """ + Drive `_PROXY_MaxIterationsHandler` past its session budget and assert + it raises the unified class. Mirrors the existing + `test_max_iterations_limiter.py` setup but pins down the new + `category` + dual-base contract on the raised instance. + """ + from unittest.mock import patch + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.max_iterations_limiter import ( + _PROXY_MaxIterationsHandler, + ) + from litellm.proxy.utils import InternalUsageCache + from litellm.types.agents import AgentResponse + + cache = DualCache() + handler = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-iter", + agent_id="agent-iter-1", + ) + agent = AgentResponse( + agent_id="agent-iter-1", + agent_name="iter-agent", + litellm_params={"max_iterations": 1}, + agent_card_params={"name": "iter-agent", "version": "1.0.0"}, + ) + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = agent + # First call within budget. + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={"metadata": {"session_id": "sess-1"}}, + call_type="", + ) + # Second call exceeds — must raise the unified class. + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={"metadata": {"session_id": "sess-1"}}, + call_type="", + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + assert isinstance(e, RateLimitError) + assert isinstance(e, HTTPException) + + @pytest.mark.asyncio + async def test_max_budget_limiter_raises_proxy_rate_limit_error(self): + """ + Drive `_PROXY_MaxBudgetLimiter` past the user budget and assert it + raises the unified class. Mocks `get_current_spend` so we don't need + the proxy DB. + """ + from unittest.mock import patch + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.max_budget_limiter import ( + _PROXY_MaxBudgetLimiter, + ) + + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-budget", + user_id="user-budget-1", + user_max_budget=1.0, + user_spend=2.0, + ) + with patch( + "litellm.proxy.proxy_server.get_current_spend", + return_value=5.0, + ): + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + assert "max budget" in str(e.detail).lower() + + @pytest.mark.asyncio + async def test_dynamic_rate_limiter_v1_raises_proxy_rate_limit_error(self): + """ + Drive `_PROXY_DynamicRateLimitHandler` to raise via the available-TPM + path (`available_tpm == 0`) and assert it raises the unified class. + Mocks `check_available_usage` so we don't need a real router. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.dynamic_rate_limiter import ( + _PROXY_DynamicRateLimitHandler, + ) + + handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=MagicMock()) + # check_available_usage returns (available_tpm, available_rpm, + # model_tpm, model_rpm, active_projects). Setting available_tpm == 0 + # forces the TPM-exceeded raise. + handler.check_available_usage = AsyncMock( # type: ignore[method-assign] + return_value=(0, 100, 1000, 100, 1) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-dyn", + metadata={"priority": "default"}, + ) + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "gpt-4"}, + call_type="completion", + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + assert isinstance(e.detail, dict) + assert "TPM" in e.detail.get("error", "") + + @pytest.mark.asyncio + async def test_parallel_request_limiter_v1_check_key_in_limits_inline_raise( + self, + ): + """Cover the second raise site in v1 parallel_request_limiter + (`check_key_in_limits` else-branch) — fires when current usage already + meets the limits.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + cache = MagicMock() + cache.async_batch_set_cache = AsyncMock(return_value=None) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=cache) + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.check_key_in_limits( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-key"), + cache=DualCache(), + data={}, + call_type="completion", + max_parallel_requests=1, + tpm_limit=10, + rpm_limit=10, + # current already at the limit on every dimension → forces + # the inline `raise ProxyRateLimitError(...)` else-branch. + current={"current_requests": 1, "current_tpm": 10, "current_rpm": 10}, + request_count_api_key="x", + rate_limit_type="key", + values_to_update_in_cache=[], + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + + @pytest.mark.parametrize( + "current,limits,expected_type", + [ + # current already at concurrent-request cap → CONCURRENT_REQUESTS + ( + {"current_requests": 5, "current_tpm": 0, "current_rpm": 0}, + {"max_parallel_requests": 5, "tpm_limit": 100, "rpm_limit": 100}, + "concurrent_requests", + ), + # current already at TPM cap (concurrent has headroom) → TOKENS + ( + {"current_requests": 0, "current_tpm": 100, "current_rpm": 0}, + {"max_parallel_requests": 5, "tpm_limit": 100, "rpm_limit": 100}, + "tokens", + ), + # current already at RPM cap (concurrent + TPM have headroom) → + # REQUESTS (the fall-through branch). + ( + {"current_requests": 0, "current_tpm": 0, "current_rpm": 100}, + {"max_parallel_requests": 5, "tpm_limit": 100, "rpm_limit": 100}, + "requests", + ), + ], + ) + @pytest.mark.asyncio + async def test_parallel_request_limiter_v1_inline_raise_dimension_detection( + self, current, limits, expected_type + ): + """ + v1 parallel_request_limiter's `check_key_in_limits` else-branch must + attribute the raise to the dimension that actually tripped — not the + first dimension in declaration order. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + cache = MagicMock() + cache.async_batch_set_cache = AsyncMock(return_value=None) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=cache) + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.check_key_in_limits( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-key"), + cache=DualCache(), + data={}, + call_type="completion", + max_parallel_requests=limits["max_parallel_requests"], + tpm_limit=limits["tpm_limit"], + rpm_limit=limits["rpm_limit"], + current=current, + request_count_api_key="x", + rate_limit_type="key", + values_to_update_in_cache=[], + ) + assert exc_info.value.rate_limit_type == expected_type + + @pytest.mark.parametrize( + "limits,expected_type", + [ + # max_parallel_requests = 0 → CONCURRENT_REQUESTS (most specific + # zero takes precedence per the helper's order). + ( + {"max_parallel_requests": 0, "tpm_limit": 0, "rpm_limit": 0}, + "concurrent_requests", + ), + # tpm_limit = 0 (concurrent has a positive limit) → TOKENS + ( + {"max_parallel_requests": 5, "tpm_limit": 0, "rpm_limit": 0}, + "tokens", + ), + # only rpm_limit = 0 → REQUESTS (fall-through) + ( + {"max_parallel_requests": 5, "tpm_limit": 100, "rpm_limit": 0}, + "requests", + ), + ], + ) + @pytest.mark.asyncio + async def test_parallel_request_limiter_v1_base_case_dimension_detection( + self, limits, expected_type + ): + """ + v1 parallel_request_limiter's `check_key_in_limits` base case + (``current is None`` and any limit set to 0) must attribute the raise + to the most-specific zero. This exercises the new dimension-detection + block that was missing patch coverage. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + cache = MagicMock() + cache.async_batch_set_cache = AsyncMock(return_value=None) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=cache) + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.check_key_in_limits( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-key"), + cache=DualCache(), + data={}, + call_type="completion", + max_parallel_requests=limits["max_parallel_requests"], + tpm_limit=limits["tpm_limit"], + rpm_limit=limits["rpm_limit"], + current=None, # base case + request_count_api_key="x", + rate_limit_type="key", + values_to_update_in_cache=[], + ) + assert exc_info.value.rate_limit_type == expected_type + + @pytest.mark.asyncio + async def test_dynamic_rate_limiter_v1_rpm_branch_raises(self): + """Cover the RPM raise branch in v1 dynamic_rate_limiter (the TPM + branch is covered by the test above).""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.dynamic_rate_limiter import ( + _PROXY_DynamicRateLimitHandler, + ) + + handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=MagicMock()) + # available_tpm > 0, available_rpm == 0 → RPM raise branch. + handler.check_available_usage = AsyncMock( # type: ignore[method-assign] + return_value=(100, 0, 1000, 100, 1) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-dyn-rpm", + metadata={"priority": "default"}, + ) + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "gpt-4"}, + call_type="completion", + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + assert isinstance(e.detail, dict) + assert "RPM" in e.detail.get("error", "") + + @pytest.mark.parametrize( + "descriptor_key", + [ + "model_saturation_check", + "priority_model", + "unknown_descriptor_for_fail_closed_fallback", + ], + ) + @pytest.mark.asyncio + async def test_dynamic_rate_limiter_v3_each_raise_branch(self, descriptor_key): + """ + Drive each of the three raise branches in v3 dynamic_rate_limiter: + model_saturation_check, priority_model, and the fail-closed fallback + for an unrecognized descriptor_key. Mocks + ``atomic_check_and_increment_by_n`` so the v3 limiter's response + directly drives the raise-site selection. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + _PROXY_DynamicRateLimitHandlerV3, + ) + + # Bypass __init__ — we want to inject a stub v3_limiter without + # paying for the full handler setup. + handler = _PROXY_DynamicRateLimitHandlerV3.__new__( + _PROXY_DynamicRateLimitHandlerV3 + ) + v3_limiter = MagicMock() + v3_limiter.window_size = 60 + v3_limiter.atomic_check_and_increment_by_n = AsyncMock( + return_value={ + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": descriptor_key, + "current_limit": 100, + "limit_remaining": 0, + "rate_limit_type": "requests", + } + ], + } + ) + handler.v3_limiter = v3_limiter + # Stub the descriptor builders so we don't pull in real router state. + handler._create_model_tracking_descriptor = MagicMock( # type: ignore[method-assign] + return_value={ + "key": descriptor_key, + "value": "v", + "rate_limit": { + "requests_per_unit": 100, + "tokens_per_unit": None, + "window_size": 60, + }, + } + ) + handler._create_priority_based_descriptors = MagicMock( # type: ignore[method-assign] + return_value=[] + ) + model_group_info = MagicMock() + model_group_info.tpm = 1000 + model_group_info.rpm = 100 + + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler._check_rate_limits( + model="gpt-4", + model_group_info=model_group_info, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test-v3"), + priority="default", + saturation=0.99, + data={}, + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + + @pytest.mark.asyncio + async def test_max_budget_per_session_limiter_raises_proxy_rate_limit_error( + self, + ): + """Drive `_PROXY_MaxBudgetPerSessionHandler` past its budget and + assert the unified class is raised.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.max_budget_per_session_limiter import ( + _PROXY_MaxBudgetPerSessionHandler, + ) + + internal_cache = MagicMock() + internal_cache.async_get_cache = AsyncMock(return_value=10.0) + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=internal_cache, + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-session", + agent_id="agent-session-1", + ) + agent = MagicMock() + agent.litellm_params = {"max_budget_per_session": 1.0} + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = agent + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"metadata": {"session_id": "session-over-budget"}}, + call_type="completion", + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + assert "session" in str(e.detail).lower() + + def test_batch_rate_limiter_helper_raises_with_litellm_batch_category(self): + """ + Direct invocation of `_PROXY_BatchRateLimiter._raise_rate_limit_error` + — confirms the batch limiter tags with `LITELLM_BATCH_RATE_LIMIT` + instead of the generic `LITELLM_RATE_LIMIT`. + """ + from unittest.mock import MagicMock + + from litellm.proxy.hooks.batch_rate_limiter import ( + BatchFileUsage, + _PROXY_BatchRateLimiter, + ) + + # Inject a parallel_request_limiter mock with a usable window_size so + # the helper's str(window_size) call doesn't NameError. + parallel_limiter = MagicMock() + parallel_limiter.window_size = 60 + handler = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=parallel_limiter, + ) + status = { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 100, + "limit_remaining": 0, + "rate_limit_type": "requests", + } + descriptors = [ + { + "key": "key", + "value": "sk-batch", + "rate_limit": { + "requests_per_unit": 100, + "tokens_per_unit": None, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._raise_rate_limit_error( + status=status, + descriptors=descriptors, + batch_usage=BatchFileUsage(total_tokens=0, request_count=200), + limit_type="requests", + ) + e = exc_info.value + assert e.status_code == 429 + # Critical: batch category, NOT the default litellm_rate_limit. + assert e.category == RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT + assert isinstance(e, RateLimitError) + assert isinstance(e, HTTPException) + + +class TestRateLimitType: + """ + Tests for the orthogonal `rate_limit_type` dimension introduced as a + follow-up to LIT-2968 (trho's last ask in the Slack thread). + + `category` answers *who* rate-limited (vendor vs. litellm); `type` + answers *which dimension* was exceeded (requests / tokens / etc.). + Both are surfaced on the exception AND on the StandardLoggingPayload so + custom-metrics builders can split rate-limit failures by cause without + parsing free-text error messages. + """ + + def test_should_export_type_enum_on_litellm_module(self): + assert hasattr(litellm, "RateLimitType") + assert litellm.RateLimitType is RateLimitType + + def test_should_define_all_documented_types(self): + assert RateLimitType.REQUESTS == "requests" + assert RateLimitType.TOKENS == "tokens" + assert RateLimitType.CONCURRENT_REQUESTS == "concurrent_requests" + assert RateLimitType.BUDGET == "budget" + assert RateLimitType.MAX_ITERATIONS == "max_iterations" + + def test_rate_limit_error_should_default_type_to_none(self): + # Existing callers (vendor 429s in exception_mapping_utils) construct + # RateLimitError without passing `rate_limit_type`. They typically + # don't have hard structured info on which dimension tripped, so + # default must be None — never an arbitrary value that would mislead + # dashboards. + e = RateLimitError(message="oops", llm_provider="openai", model="gpt-4") + assert e.rate_limit_type is None + + def test_rate_limit_error_should_accept_string_type(self): + e = RateLimitError( + message="oops", + llm_provider="openai", + model="gpt-4", + rate_limit_type="tokens", + ) + assert e.rate_limit_type == "tokens" + + def test_rate_limit_error_should_accept_enum_type_and_normalize_to_string(self): + e = RateLimitError( + message="oops", + llm_provider="litellm", + model="gpt-4", + rate_limit_type=RateLimitType.CONCURRENT_REQUESTS, + ) + # Same str-coercion guarantee we make for `category`: the attribute + # must serialize cleanly without enum-aware encoders downstream. + assert e.rate_limit_type == "concurrent_requests" + assert isinstance(e.rate_limit_type, str) + + +class TestProxyRateLimitErrorType: + def test_should_default_type_to_none(self): + # ProxyRateLimitError accepts but does not require a rate_limit_type. + # Callers that don't pass one (e.g. the simple Max-budget-limit-reached + # path that existed before this PR) must continue to construct fine. + e = ProxyRateLimitError(detail="over limit") + assert e.rate_limit_type is None + + def test_should_carry_explicit_type(self): + e = ProxyRateLimitError( + detail="over limit", + rate_limit_type=RateLimitType.TOKENS, + ) + assert e.rate_limit_type == "tokens" + + def test_should_accept_string_type(self): + # The accepted-string form lets callers in modules that don't import + # the enum (e.g. v3 limiter passing through descriptor strings) + # forward the raw value. + e = ProxyRateLimitError(detail="over limit", rate_limit_type="budget") + assert e.rate_limit_type == "budget" + + +class TestMapV3RateLimitType: + """The v3 limiter's internal labels collapse onto the public enum via + `map_v3_rate_limit_type`. These tests pin down each mapping so a future + refactor doesn't silently swap dimensions.""" + + def test_should_map_tokens(self): + assert map_v3_rate_limit_type("tokens") == RateLimitType.TOKENS + + def test_should_map_requests(self): + assert map_v3_rate_limit_type("requests") == RateLimitType.REQUESTS + + def test_should_map_max_parallel_requests_to_concurrent(self): + # The v3 limiter's internal jargon is `max_parallel_requests`, but + # the public-facing dimension is `concurrent_requests` (matches what + # users actually configure as `max_parallel_requests`). The mapping + # must collapse these so dashboards see one name, not two. + assert ( + map_v3_rate_limit_type("max_parallel_requests") + == RateLimitType.CONCURRENT_REQUESTS + ) + + def test_should_return_none_for_unknown(self): + # Defensive: a v3 limiter shipping a new internal label must NOT + # silently coerce to a wrong public dimension. Returning None lets + # the caller decide (typically: omit the field). + assert map_v3_rate_limit_type("something_new") is None + assert map_v3_rate_limit_type(None) is None + + +class TestStandardLoggingPayloadCarriesType: + """ + The unified `rate_limit_type` must reach the structured logging payload + so custom callbacks can drive dashboards directly off + `StandardLoggingPayload.error_information.error_rate_limit_type`. + """ + + def test_should_propagate_type_for_proxy_rate_limit_error(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = ProxyRateLimitError( + detail="over tpm", + rate_limit_type=RateLimitType.TOKENS, + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_type"] == "tokens" + + def test_should_propagate_type_for_plain_rate_limit_error(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = RateLimitError( + message="vendor 429", + llm_provider="openai", + model="gpt-4", + rate_limit_type=RateLimitType.REQUESTS, + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_type"] == "requests" + + def test_should_be_none_when_unspecified(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + # Vendor 429 exception with no header hints → type omitted. + e = RateLimitError( + message="vendor 429", + llm_provider="openai", + model="gpt-4", + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_type"] is None + + def test_should_be_none_for_non_rate_limit_errors(self): + # Symmetry with `error_rate_limit_category`: the field must be + # present on every payload so consumers can read it + # unconditionally, but None for non-rate-limit exceptions. + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + info = StandardLoggingPayloadSetup.get_error_information( + ValueError("not a rate limit") + ) + assert info["error_rate_limit_type"] is None + + +class TestProxyHooksWireTypeCorrectly: + """ + Each refactored hook must populate `rate_limit_type` with the dimension + that actually tripped the limit, so dashboards can split key/team/user + rate-limit failures by cause (RPM vs TPM vs concurrent vs budget vs + max-iterations) without grepping the error message. + """ + + def test_max_budget_limiter_emits_budget_type(self): + e = ProxyRateLimitError( + detail="Max budget limit reached.", + rate_limit_type=RateLimitType.BUDGET, + ) + assert e.category == "litellm_rate_limit" + assert e.rate_limit_type == "budget" + + def test_max_iterations_limiter_emits_max_iterations_type(self): + e = ProxyRateLimitError( + detail="Max iterations exceeded for session abc.", + rate_limit_type=RateLimitType.MAX_ITERATIONS, + ) + assert e.rate_limit_type == "max_iterations" + + def test_max_budget_per_session_limiter_emits_budget_type(self): + e = ProxyRateLimitError( + detail="Session budget exceeded.", + rate_limit_type=RateLimitType.BUDGET, + ) + assert e.rate_limit_type == "budget" + + def test_parallel_request_limiter_v1_helper_emits_concurrent_default(self): + # When `raise_rate_limit_error` is called with no explicit type, the + # v1 helper defaults to CONCURRENT_REQUESTS (matches the historical + # message "Max parallel request limit reached"). Tests below cover + # the explicit-type override paths. + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=MagicMock()) + with pytest.raises(ProxyRateLimitError) as exc_info: + handler.raise_rate_limit_error() + assert exc_info.value.rate_limit_type == "concurrent_requests" + + def test_parallel_request_limiter_v1_helper_accepts_explicit_type(self): + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=MagicMock()) + with pytest.raises(ProxyRateLimitError) as exc_info: + handler.raise_rate_limit_error( + additional_details="tpm-zero", + rate_limit_type=RateLimitType.TOKENS, + ) + assert exc_info.value.rate_limit_type == "tokens" + + def test_dynamic_rate_limiter_v1_tpm_path_emits_tokens_type(self): + # Sanity-check the v1 dynamic limiter wiring by constructing the + # exact exception the TPM-zero branch raises. We round-trip through + # ProxyRateLimitError to assert both fields. (Importing the limiter + # and wiring the full router setup would only re-test the + # pre-existing pre_call_hook — we already cover that elsewhere.) + e = ProxyRateLimitError( + detail={"error": "Key=k over available TPM=0."}, + rate_limit_type=RateLimitType.TOKENS, + model="gpt-4", + ) + assert e.rate_limit_type == "tokens" + assert e.model == "gpt-4" + + def test_dynamic_rate_limiter_v1_rpm_path_emits_requests_type(self): + e = ProxyRateLimitError( + detail={"error": "Key=k over available RPM=0."}, + rate_limit_type=RateLimitType.REQUESTS, + model="gpt-4", + ) + assert e.rate_limit_type == "requests" + + @pytest.mark.asyncio + async def test_v3_limiter_handle_rate_limit_error_propagates_type(self): + """ + End-to-end: feed the v3 limiter's `_handle_rate_limit_error` an + OVER_LIMIT response and verify the raised ProxyRateLimitError carries + the mapped public RateLimitType. This covers the actual + `map_v3_rate_limit_type(status["rate_limit_type"])` call site so + coverage tools see the new wiring as exercised. + """ + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + ) + + handler = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=MagicMock(), + ) + # Minimal RateLimitResponse + descriptors shape that the handler + # reads. We only need one OVER_LIMIT status to drive the raise. + response = { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 100, + "limit_remaining": 0, + "rate_limit_type": "tokens", + } + ], + } + descriptors = [ + { + "key": "key", + "value": "sk-test", + "rate_limit": { + "requests_per_unit": None, + "tokens_per_unit": 100, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._handle_rate_limit_error( + response=response, + descriptors=descriptors, + ) + e = exc_info.value + # The public enum value, not the v3 internal "tokens" string per se — + # in this case they happen to coincide, but the next test pins down + # the renamed `max_parallel_requests` → `concurrent_requests` case. + assert e.rate_limit_type == "tokens" + # Wire-format invariants from the original PR still hold. + assert e.headers is not None + assert e.headers.get("rate_limit_type") == "tokens" + assert e.headers.get("retry-after") is not None + + @pytest.mark.asyncio + async def test_v3_limiter_max_parallel_requests_maps_to_concurrent(self): + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + ) + + handler = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=MagicMock(), + ) + response = { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 5, + "limit_remaining": 0, + # v3 internal jargon — must collapse to the public name. + "rate_limit_type": "max_parallel_requests", + } + ], + } + descriptors = [ + { + "key": "key", + "value": "sk-test", + "rate_limit": { + "requests_per_unit": None, + "tokens_per_unit": None, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._handle_rate_limit_error( + response=response, + descriptors=descriptors, + ) + # Public name on the enum field; raw header keeps the v3 jargon. + assert exc_info.value.rate_limit_type == "concurrent_requests" + assert exc_info.value.headers["rate_limit_type"] == "max_parallel_requests" + + def test_batch_rate_limiter_emits_tokens_type_for_tpm_violation(self): + from unittest.mock import MagicMock + + from litellm.proxy.hooks.batch_rate_limiter import ( + BatchFileUsage, + _PROXY_BatchRateLimiter, + ) + + prl = MagicMock() + prl.window_size = 60 + handler = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=prl, + ) + status = { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 1000, + "limit_remaining": 100, + "rate_limit_type": "tokens", + } + descriptors = [ + { + "key": "key", + "value": "sk-test", + "rate_limit": { + "requests_per_unit": None, + "tokens_per_unit": 1000, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._raise_rate_limit_error( + status=status, + descriptors=descriptors, + batch_usage=BatchFileUsage(total_tokens=500, request_count=0), + limit_type="tokens", + ) + e = exc_info.value + assert e.rate_limit_type == "tokens" + assert e.category == RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT + + def test_batch_rate_limiter_emits_requests_type_for_rpm_violation(self): + from unittest.mock import MagicMock + + from litellm.proxy.hooks.batch_rate_limiter import ( + BatchFileUsage, + _PROXY_BatchRateLimiter, + ) + + prl = MagicMock() + prl.window_size = 60 + handler = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=prl, + ) + status = { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 100, + "limit_remaining": 10, + "rate_limit_type": "requests", + } + descriptors = [ + { + "key": "key", + "value": "sk-test", + "rate_limit": { + "requests_per_unit": 100, + "tokens_per_unit": None, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._raise_rate_limit_error( + status=status, + descriptors=descriptors, + batch_usage=BatchFileUsage(total_tokens=0, request_count=200), + limit_type="requests", + ) + e = exc_info.value + assert e.rate_limit_type == "requests" + assert e.category == RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT + + +class TestBudgetExceededErrorSurfacesUnifiedFields: + """ + The hot path for virtual-key / team / org / end-user max_budget caps + raises :class:`litellm.BudgetExceededError`, which historically had no + relationship to :class:`RateLimitError` and therefore left the unified + `error_rate_limit_category` / `error_rate_limit_type` fields empty. + Test 2 of the QA pass surfaced this gap; this class pins the fix. + + The fix is intentionally additive: `BudgetExceededError` keeps its + bare-`Exception` base class (so existing `except BudgetExceededError:` + handlers keep working) and just sets the same `category` / + `rate_limit_type` attributes that the rest of the unified rate-limit + path reads (normalized to plain strings, matching how + `RateLimitError.__init__` stores its own values). Duck-typed dispatch + in `get_error_information` picks them up automatically. + """ + + def test_should_carry_litellm_rate_limit_category(self): + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + # Stored as the plain string value (matches RateLimitError behavior), + # but equality with the enum still works because the enum subclasses + # str. + assert e.category == "litellm_rate_limit" + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + + def test_should_carry_budget_rate_limit_type(self): + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + assert e.rate_limit_type == "budget" + assert e.rate_limit_type == RateLimitType.BUDGET + + def test_should_default_llm_provider_to_empty_string(self): + # `llm_provider` is read off the exception in `get_error_information` + # — it must always be a string so the StandardLoggingPayload field + # stays serializable. Default to "" when no caller passes one. + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + assert e.llm_provider == "" + + def test_should_accept_llm_provider_kwarg(self): + # Callers that have the resolved provider in scope (e.g. the + # auth-checks budget enforcement paths) can thread it through. + e = litellm.BudgetExceededError( + current_cost=0.5, max_budget=0.1, llm_provider="anthropic" + ) + assert e.llm_provider == "anthropic" + + def test_should_keep_existing_status_code_and_message(self): + # Backward-compat guard: existing callers depend on `status_code=429` + # and the canonical message format. + e = litellm.BudgetExceededError(current_cost=0.000109, max_budget=0.0001) + assert e.status_code == 429 + assert "Current cost: 0.000109" in e.message + assert "Max budget: 0.0001" in e.message + + def test_should_still_be_catchable_as_exception_not_rate_limit_error(self): + # Critical: we deliberately did NOT make BudgetExceededError a + # RateLimitError subclass. Existing `except BudgetExceededError:` + # handlers must keep catching it, and `except RateLimitError:` + # handlers must NOT start catching it (which would surprise callers + # who rely on the two being distinct). + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + assert isinstance(e, Exception) + assert isinstance(e, litellm.BudgetExceededError) + assert not isinstance(e, RateLimitError) + + def test_should_propagate_category_to_standard_logging_payload(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_category"] == "litellm_rate_limit" + assert info["error_rate_limit_type"] == "budget" + assert info["error_code"] == "429" + assert info["error_class"] == "BudgetExceededError" + + def test_should_propagate_llm_provider_to_standard_logging_payload(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = litellm.BudgetExceededError( + current_cost=0.5, max_budget=0.1, llm_provider="bedrock" + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["llm_provider"] == "bedrock" + + +class TestThirdPartyAttrLeakageGuard: + """ + The duck-typed read at the StandardLoggingPayload + Prometheus surfaces + must reject `.category` / `.rate_limit_type` strings set on unrelated + third-party exceptions. Without validation, a foreign exception that + happens to declare either attribute name would leak garbage values into + custom-callback payloads and Prometheus label cardinality. + """ + + def test_should_drop_unknown_category_string_on_third_party_exception(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + class Foreign(Exception): + category = "totally_not_a_real_category" + + info = StandardLoggingPayloadSetup.get_error_information(Foreign("boom")) + assert info["error_rate_limit_category"] is None + + def test_should_drop_unknown_rate_limit_type_string_on_third_party_exception(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + class Foreign(Exception): + rate_limit_type = "wat" + + info = StandardLoggingPayloadSetup.get_error_information(Foreign("boom")) + assert info["error_rate_limit_type"] is None + + def test_should_drop_non_string_garbage_attrs(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + class Foreign(Exception): + category = 42 + rate_limit_type = {"lol": "no"} + + info = StandardLoggingPayloadSetup.get_error_information(Foreign()) + assert info["error_rate_limit_category"] is None + assert info["error_rate_limit_type"] is None + + def test_should_drop_garbage_on_prometheus_label_extraction(self): + from litellm.integrations.prometheus import PrometheusLogger + + class Foreign(Exception): + category = "spam" + rate_limit_type = "spam" + + category, rate_limit_type = PrometheusLogger._extract_rate_limit_labels( + Foreign() + ) + assert category is None + assert rate_limit_type is None + + def test_should_still_accept_legitimate_rate_limit_categories(self): + # The guard must not over-correct — every documented enum value + # is a valid string and must pass through. + from litellm.exceptions import ( + validate_rate_limit_category, + validate_rate_limit_type, + ) + + for member in RateLimitErrorCategory: + assert validate_rate_limit_category(member.value) == member.value + assert validate_rate_limit_category(member) == member.value + + for member in RateLimitType: + assert validate_rate_limit_type(member.value) == member.value + assert validate_rate_limit_type(member) == member.value + + +@pytest.mark.asyncio +class TestBudgetExceededErrorLlmProviderEnrichment: + """ + BudgetExceededError raise sites in auth_checks.py are tenant-scoped + (key / team / org / tag) and cannot see the request model. To still + populate `llm_provider` on the StandardLoggingPayload — which is what + custom-callback consumers attribute spend to — the central + UserAPIKeyAuthExceptionHandler enriches the exception from + `request_data["model"]` before post_call_failure_hook fires. + """ + + async def _run_handler_and_capture_exception_seen_by_callback( + self, exception: Exception, request_data: dict + ): + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy.auth.auth_exception_handler import ( + UserAPIKeyAuthExceptionHandler, + ) + + captured: dict = {} + + async def fake_post_call_failure_hook(**kwargs): + captured["exception"] = kwargs["original_exception"] + return None + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock( + post_call_failure_hook=AsyncMock( + side_effect=fake_post_call_failure_hook + ) + ), + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": False}, + ), + patch( + "litellm.proxy.auth.auth_exception_handler._get_request_ip_address", + return_value="127.0.0.1", + ), + ): + try: + await UserAPIKeyAuthExceptionHandler._handle_authentication_error( + e=exception, + request=MagicMock(), + request_data=request_data, + route="/v1/chat/completions", + parent_otel_span=None, + api_key="sk-test", + ) + except Exception: + pass + return captured.get("exception") + + async def test_should_resolve_llm_provider_from_request_data_when_unset(self): + err = litellm.BudgetExceededError(current_cost=100, max_budget=10) + assert err.llm_provider == "" + seen = await self._run_handler_and_capture_exception_seen_by_callback( + err, {"model": "openai/gpt-4o-mini"} + ) + assert seen is not None + assert seen.llm_provider == "openai" + + async def test_should_not_overwrite_llm_provider_when_caller_set_it(self): + err = litellm.BudgetExceededError( + current_cost=100, max_budget=10, llm_provider="anthropic" + ) + seen = await self._run_handler_and_capture_exception_seen_by_callback( + err, {"model": "openai/gpt-4o-mini"} + ) + assert seen.llm_provider == "anthropic" + + async def test_should_fall_back_to_litellm_proxy_when_model_missing(self): + err = litellm.BudgetExceededError(current_cost=100, max_budget=10) + seen = await self._run_handler_and_capture_exception_seen_by_callback(err, {}) + assert seen.llm_provider == "litellm_proxy" + + async def test_should_not_enrich_non_budget_exceptions(self): + err = ValueError("unrelated") + seen = await self._run_handler_and_capture_exception_seen_by_callback( + err, {"model": "openai/gpt-4o-mini"} + ) + assert not hasattr(seen, "llm_provider") or seen.llm_provider != "openai" diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py index 719cb8eecd2..e384d3e1161 100644 --- a/tests/test_litellm/test_register_model_custom_pricing.py +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -301,6 +301,126 @@ def _fake_get_model_info(model, *args, **kwargs): litellm.model_cost.pop(model_key, None) +def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(): + """Registering a custom override under a key shape that + ``get_model_info`` cannot resolve (e.g. a double provider prefix like + ``bedrock/bedrock/us.anthropic.claude-sonnet-4-6``) must still inherit + the built-in cache pricing for the underlying model. + + Before the fix ``register_model`` fell back to an empty ``existing_model`` + so the merged entry only carried the fields the user set explicitly + (input/output cost). ``cache_creation_input_token_cost`` and + ``cache_read_input_token_cost`` were absent, and the cost calculator + silently charged 0 for every cache token, dropping the bulk of the bill + for cache-heavy Anthropic traffic. + + Regression for the cache-pricing dropout under partial overrides. + """ + from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token + from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + original_model_cost = litellm.model_cost + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + builtin_key = "us.anthropic.claude-sonnet-4-6" + registered_key = f"bedrock/bedrock/{builtin_key}" + builtin = litellm.model_cost[builtin_key] + + assert builtin["cache_creation_input_token_cost"] > 0 + assert builtin["cache_read_input_token_cost"] > 0 + + try: + litellm.register_model( + { + registered_key: { + "input_cost_per_token": builtin["input_cost_per_token"], + "output_cost_per_token": builtin["output_cost_per_token"], + "litellm_provider": "bedrock", + } + } + ) + + registered = litellm.model_cost[registered_key] + assert ( + registered.get("cache_creation_input_token_cost") + == builtin["cache_creation_input_token_cost"] + ) + assert ( + registered.get("cache_read_input_token_cost") + == builtin["cache_read_input_token_cost"] + ) + assert registered["litellm_provider"] == "bedrock" + + usage = Usage( + prompt_tokens=1100, + completion_tokens=100, + total_tokens=1200, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=800, + text_tokens=100, + ), + cache_creation_input_tokens=200, + ) + + input_cost, output_cost = generic_cost_per_token( + model=registered_key, + usage=usage, + custom_llm_provider="bedrock", + ) + + text_only_cost = builtin["input_cost_per_token"] * 100 + expected_input_cost = ( + text_only_cost + + builtin["cache_read_input_token_cost"] * 800 + + builtin["cache_creation_input_token_cost"] * 200 + ) + assert abs(input_cost - expected_input_cost) < 1e-12 + assert abs(output_cost - builtin["output_cost_per_token"] * 100) < 1e-12 + assert input_cost > text_only_cost + 1e-12 + finally: + litellm.model_cost.pop(registered_key, None) + litellm.model_cost = original_model_cost + os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) + from litellm.utils import _invalidate_model_cost_lowercase_map + + _invalidate_model_cost_lowercase_map() + + +def test_register_model_warns_when_no_builtin_match_for_cache_pricing(caplog): + """When a custom override is registered under a key that neither + ``get_model_info`` nor any prefix/region variant can resolve to a + built-in entry, ``register_model`` must warn that cache cost fields will + default to 0 instead of silently producing an under-billed entry. + """ + import logging + + from litellm._logging import verbose_logger + + registered_key = "bedrock/totally-made-up-model-alias-xyz" + litellm.model_cost.pop(registered_key, None) + + try: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + litellm.register_model( + { + registered_key: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "litellm_provider": "bedrock", + } + } + ) + + assert any( + registered_key in record.message + and "cache_creation_input_token_cost" in record.message + for record in caplog.records + ), "expected a warning naming the unmapped key and the cache cost fields" + finally: + litellm.model_cost.pop(registered_key, None) + + def test_register_model_router_add_deployment_custom_pricing_applies(): """End-to-end regression for https://github.com/BerriAI/litellm/issues/28336. @@ -344,9 +464,9 @@ def test_register_model_router_add_deployment_custom_pricing_applies(): f"{model_key} / {deployment_model}" ) for k in registered_keys: - assert _check_provider_match(litellm.model_cost[k], "openai") is True, ( - f"custom pricing for {k} was dropped by _check_provider_match" - ) + assert ( + _check_provider_match(litellm.model_cost[k], "openai") is True + ), f"custom pricing for {k} was dropped by _check_provider_match" finally: litellm.model_cost.pop(model_key, None) litellm.model_cost.pop(deployment_model, None) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index cd235d8de67..e681247959f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -80,6 +80,256 @@ def test_router_with_model_info_and_model_group(): ) +def test_router_model_group_encrypted_content_affinity_callback_registration(): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + model_group_affinity_config = { + model_group: ["encrypted_content_affinity"], + } + original_callbacks = list(litellm.callbacks) + litellm.callbacks = [] + router = None + + try: + router = litellm.Router( + model_list=[ + { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key", + }, + } + ], + model_group_affinity_config=model_group_affinity_config, + num_retries=0, + ) + callbacks = router.optional_callbacks or [] + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) + assert len(encrypted_content_callbacks) == 1 + assert encrypted_content_callbacks[0].enable_global_affinity is False + assert ( + encrypted_content_callbacks[0].model_group_affinity_config + == model_group_affinity_config + ) + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( + deployment_callback + ) + assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( + litellm.callbacks.index(deployment_callback) + ) + + router._add_encrypted_content_affinity_check(enable_global_affinity=True) + + callbacks = router.optional_callbacks or [] + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] + assert len(encrypted_content_callbacks) == 1 + assert encrypted_content_callbacks[0].enable_global_affinity is True + assert encrypted_content_callbacks[0].router is router + finally: + if router is not None: + router.discard() + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_model_group_config_is_additive(): + from litellm.responses.utils import ResponsesAPIRequestUtils + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + target_deployment = { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-b"}, + } + healthy_deployments = [ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-a"}, + }, + target_deployment, + ] + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) + + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( + {model_group: ["encrypted_content_affinity"]} + ) + assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) + + per_group_check = EncryptedContentAffinityCheck( + enable_global_affinity=False, + model_group_affinity_config={ + model_group: ["encrypted_content_affinity"], + }, + ) + request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {}, + } + filtered = await per_group_check.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=request_kwargs, + ) + + assert filtered == [target_deployment] + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] + + disabled_check = EncryptedContentAffinityCheck( + enable_global_affinity=False, + model_group_affinity_config={ + "other-model-group": ["encrypted_content_affinity"], + }, + ) + disabled_request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {}, + } + unfiltered = await disabled_check.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=disabled_request_kwargs, + ) + + assert unfiltered == healthy_deployments + assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs[ + "litellm_metadata" + ] + + global_check = EncryptedContentAffinityCheck( + enable_global_affinity=True, + model_group_affinity_config={ + model_group: ["deployment_affinity"], + }, + ) + global_request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {}, + } + globally_filtered = await global_check.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=global_request_kwargs, + ) + + assert globally_filtered == [target_deployment] + assert global_request_kwargs["litellm_metadata"][ + "encrypted_content_affinity_enabled" + ] + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity(): + from litellm.responses.utils import ResponsesAPIRequestUtils + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + user_api_key_hash = "test-user-key" + deployment_a = { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-a", + }, + "model_info": {"id": "deployment-a"}, + } + deployment_b = { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-b", + }, + "model_info": {"id": "deployment-b"}, + } + original_callbacks = list(litellm.callbacks) + litellm.callbacks = [] + router = None + + try: + router = litellm.Router( + model_list=[deployment_a, deployment_b], + model_group_affinity_config={ + model_group: [ + "deployment_affinity", + "encrypted_content_affinity", + ], + }, + num_retries=0, + ) + callbacks = router.optional_callbacks or [] + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) + encrypted_content_callback = next( + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ) + assert callbacks.index(encrypted_content_callback) < callbacks.index( + deployment_callback + ) + assert litellm.callbacks.index(encrypted_content_callback) < ( + litellm.callbacks.index(deployment_callback) + ) + + cache_key = DeploymentAffinityCheck.get_affinity_cache_key( + model_group=model_group, + user_key=user_api_key_hash, + ) + await deployment_callback.cache.async_set_cache( + key=cache_key, + value={"model_id": "deployment-a"}, + ttl=60, + ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) + request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, + } + + filtered = await router.async_callback_filter_deployments( + model=model_group, + healthy_deployments=[deployment_a, deployment_b], + messages=None, + parent_otel_span=None, + request_kwargs=request_kwargs, + ) + + assert filtered == [deployment_b] + assert request_kwargs.get("_encrypted_content_affinity_pinned") is True + finally: + if router is not None: + router.discard() + litellm.callbacks = original_callbacks + + @pytest.mark.asyncio async def test_arouter_with_tags_and_fallbacks(): """ @@ -4311,6 +4561,48 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): ) +def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): + """ + Bedrock deployments using IAM/OIDC auth have no api_key; pass-through + init must not raise and drop them from routing (#27728). + """ + import litellm + + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-claude", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0", + "aws_role_name": "arn:aws:iam::123456789012:role/my-role", + "aws_session_name": "my-session", + "use_in_pass_through": True, + }, + "model_info": {"id": "bedrock-iam-pt"}, + } + ] + ) + assert [m["model_info"]["id"] for m in router.get_model_list()] == [ + "bedrock-iam-pt" + ] + + +def test_initialize_deployment_for_pass_through_sets_credentials_with_api_key(): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + passthrough_endpoint_router, + ) + + passthrough_endpoint_router.credentials.clear() + router = _router_with_two_pass_through_deployments([False, False]) + assert len(router.get_model_list()) == 2 + assert ( + passthrough_endpoint_router.get_credentials( + custom_llm_provider="openai", region_name=None + ) + == "sk-fake-for-tests" + ) + + def test_get_deployment_credentials_returns_none_for_blocked_deployment(): router = _router_with_two_deployments([True, False]) assert router.get_deployment_credentials(model_id="dep-0") is None diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 9454e03e918..ee64f44d32c 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -402,3 +402,138 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override(): assert bridge_model_info["mode"] == "responses" finally: _restore_model_cost_entries(model_keys) + + +def test_partial_custom_pricing_inherits_builtin_cache_pricing(): + """A deployment that overrides only input/output cost on a cache-supporting + model must still bill cache_read and cache_creation tokens. Before the + fix the deploy-id entry was registered with the user's two fields and + nothing else, so the cost calculator silently billed cache tokens at 0. + Regression for the prompt-caching cost dropout reported by the customer. + """ + backend_model = "anthropic/claude-sonnet-4-5-20250929" + deploy_id = "claude-deploy-partial-pricing" + + builtin_info = litellm.get_model_info(model=backend_model) + builtin_cache_create = builtin_info["cache_creation_input_token_cost"] + builtin_cache_read = builtin_info["cache_read_input_token_cost"] + assert builtin_cache_create is not None and builtin_cache_create > 0 + assert builtin_cache_read is not None and builtin_cache_read > 0 + + model_keys = { + deploy_id: litellm.model_cost.get(deploy_id), + backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)), + } + try: + Router( + model_list=[ + { + "model_name": "claude-custom", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key", + }, + "model_info": { + "id": deploy_id, + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + }, + } + ], + ) + + entry = litellm.model_cost[deploy_id] + assert entry["input_cost_per_token"] == 0.000003 + assert entry["output_cost_per_token"] == 0.000015 + assert entry.get("cache_creation_input_token_cost") == builtin_cache_create + assert entry.get("cache_read_input_token_cost") == builtin_cache_read + finally: + _restore_model_cost_entries(model_keys) + + +def test_partial_pricing_does_not_overwrite_explicit_cache_fields(): + """When the user explicitly sets cache_*_input_token_cost on a deployment, + those values must not be replaced by the built-in fallback. + """ + backend_model = "anthropic/claude-sonnet-4-5-20250929" + deploy_id = "claude-deploy-explicit-cache" + + explicit_cache_create = 0.00001 + explicit_cache_read = 0.0000005 + builtin_info = litellm.get_model_info(model=backend_model) + assert builtin_info["cache_creation_input_token_cost"] != explicit_cache_create + assert builtin_info["cache_read_input_token_cost"] != explicit_cache_read + + model_keys = { + deploy_id: litellm.model_cost.get(deploy_id), + backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)), + } + try: + Router( + model_list=[ + { + "model_name": "claude-custom-explicit", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key", + }, + "model_info": { + "id": deploy_id, + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + "cache_creation_input_token_cost": explicit_cache_create, + "cache_read_input_token_cost": explicit_cache_read, + }, + } + ], + ) + + entry = litellm.model_cost[deploy_id] + assert entry.get("cache_creation_input_token_cost") == explicit_cache_create + assert entry.get("cache_read_input_token_cost") == explicit_cache_read + finally: + _restore_model_cost_entries(model_keys) + + +def test_inherit_builtin_cache_pricing_fills_only_missing_fields(): + """Direct unit test of the helper: missing cache fields are filled from the + backend model's built-in entry, while an explicitly set cache field and the + user's input/output pricing are left untouched. + """ + backend_model = "anthropic/claude-sonnet-4-5-20250929" + builtin_info = litellm.get_model_info(model=backend_model) + builtin_cache_create = builtin_info["cache_creation_input_token_cost"] + builtin_cache_read = builtin_info["cache_read_input_token_cost"] + assert builtin_cache_create is not None and builtin_cache_create > 0 + assert builtin_cache_read is not None and builtin_cache_read > 0 + + explicit_cache_read = builtin_cache_read + 1 + model_info = { + "input_cost_per_token": 0.000003, + "cache_read_input_token_cost": explicit_cache_read, + } + + Router._inherit_builtin_cache_pricing( + model_info=model_info, + backend_model=backend_model, + custom_llm_provider="anthropic", + ) + + assert model_info["input_cost_per_token"] == 0.000003 + assert model_info["cache_read_input_token_cost"] == explicit_cache_read + assert model_info["cache_creation_input_token_cost"] == builtin_cache_create + + +def test_inherit_builtin_cache_pricing_noop_for_unknown_backend(): + """No canonical entry for the backend model means the helper leaves the + passed-in dict unchanged rather than raising. + """ + model_info = {"input_cost_per_token": 0.000003} + + Router._inherit_builtin_cache_pricing( + model_info=model_info, + backend_model="this-backend-model-does-not-exist-x9y8z7", + custom_llm_provider=None, + ) + + assert model_info == {"input_cost_per_token": 0.000003} diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index f179e9c8f93..b6c9e9c865d 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -700,6 +700,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_read_input_token_cost": {"type": "number"}, "cache_read_input_token_cost_above_200k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_272k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_512k_tokens": {"type": "number"}, "cache_read_input_token_cost_batches": {"type": "number"}, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": { "type": "number" @@ -721,6 +722,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_token_above_200k_tokens": {"type": "number"}, "input_cost_per_token_above_256k_tokens": {"type": "number"}, "input_cost_per_token_above_272k_tokens": {"type": "number"}, + "input_cost_per_token_above_512k_tokens": {"type": "number"}, "cache_read_input_token_cost_flex": {"type": "number"}, "cache_read_input_token_cost_priority": {"type": "number"}, "cache_read_input_token_cost_above_200k_tokens_priority": { @@ -811,6 +813,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_token_above_200k_tokens": {"type": "number"}, "output_cost_per_token_above_256k_tokens": {"type": "number"}, "output_cost_per_token_above_272k_tokens": {"type": "number"}, + "output_cost_per_token_above_512k_tokens": {"type": "number"}, "output_cost_per_image_above_1024_and_1024_pixels": {"type": "number"}, "output_cost_per_image_above_1024_and_1024_pixels_and_premium_image": { "type": "number" @@ -858,6 +861,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_xhigh_reasoning_effort": {"type": "boolean"}, "supports_max_reasoning_effort": {"type": "boolean"}, "supports_adaptive_thinking": {"type": "boolean"}, + "supports_sampling_params": {"type": "boolean"}, "supports_service_tier": {"type": "boolean"}, "supports_preset": {"type": "boolean"}, "supports_output_config": {"type": "boolean"}, @@ -931,6 +935,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_native_streaming": {"type": "boolean"}, "supports_image_size": {"type": "boolean"}, "supports_native_structured_output": {"type": "boolean"}, + "use_openai_responses_path": {"type": "boolean"}, "tiered_pricing": { "type": "array", "items": { @@ -4193,3 +4198,21 @@ def test_base_model_label_alone_drops_tools(self): ) assert "tools" not in result + + +def test_aws_bedrock_project_id_excluded_from_bedrock_optional_params(): + """`aws_bedrock_project_id` is sent as a bedrock-mantle request header, so it + must never reach optional_params (and from there the request body), while + other aws_* params keep flowing for boto3 auth.""" + from litellm.utils import get_optional_params + + result = get_optional_params( + model="mantle/anthropic.claude-mythos-preview", + custom_llm_provider="bedrock", + max_tokens=10, + aws_bedrock_project_id="proj_abc123def456", + aws_region_name="us-east-1", + ) + + assert "aws_bedrock_project_id" not in result + assert result["aws_region_name"] == "us-east-1" diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index c146847f391..a4074ccdaaa 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -1,13 +1,9 @@ -import asyncio import os import sys -from typing import Optional -from unittest.mock import AsyncMock, patch import pytest sys.path.insert(0, os.path.abspath("../..")) -import json from litellm.types.utils import HiddenParams @@ -75,6 +71,48 @@ def test_usage_dump(): assert new_usage.prompt_tokens_details.web_search_requests == 1 +def test_usage_server_tool_use_dict_is_coerced_and_round_trips(): + from litellm.types.utils import ServerToolUse, Usage + + current_usage = Usage( + completion_tokens=1, + prompt_tokens=1, + total_tokens=2, + server_tool_use={"web_search_requests": 1}, + ) + + assert isinstance(current_usage.server_tool_use, ServerToolUse) + assert current_usage.server_tool_use.web_search_requests == 1 + + new_usage = Usage(**current_usage.model_dump()) + assert isinstance(new_usage.server_tool_use, ServerToolUse) + assert new_usage.server_tool_use.web_search_requests == 1 + + +def test_usage_converts_server_tool_use_dict(): + from litellm.types.utils import ServerToolUse, Usage + + usage = Usage( + completion_tokens=2, + prompt_tokens=1, + total_tokens=3, + server_tool_use={"web_search_requests": 4, "tool_search_requests": 1}, + ) + + assert isinstance(usage.server_tool_use, ServerToolUse) + assert usage.server_tool_use.web_search_requests == 4 + assert usage.server_tool_use["web_search_requests"] == 4 + assert usage.server_tool_use.tool_search_requests == 1 + with pytest.raises(KeyError): + usage.server_tool_use["unknown_metric"] + + round_trip = Usage(**usage.model_dump()) + assert isinstance(round_trip.server_tool_use, ServerToolUse) + assert round_trip.server_tool_use.web_search_requests == 4 + assert round_trip.server_tool_use["web_search_requests"] == 4 + assert round_trip.server_tool_use.tool_search_requests == 1 + + def test_usage_completion_tokens_details_text_tokens(): from litellm.types.utils import Usage diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts new file mode 100644 index 00000000000..f2ba66147ea --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts @@ -0,0 +1,16 @@ +/** + * Source of truth for the App Router migration smoke (tests/migration/migratedPages.spec.ts). + * + * Add a route segment here once its migration has MERGED to the branch under test. + * Both suites pick it up automatically: + * - default mount: npm run e2e:migration + * - server-root-path mount: SERVER_ROOT_PATH=/ npm run e2e:migration:root + * + * Keep this in lockstep with MIGRATED_PAGES in src/utils/migratedPages.ts. + * Pending (uncomment as each PR lands): playground, and the leaf-pages batch + * (budgets, caching, cost-tracking, guardrails, guardrails-monitor, logs, + * mcp-servers, memory, policies, projects, prompts, search-tools, skills, + * tag-management, tool-policies, transform-request, ui-theme, vector-stores, + * workflows, access-groups). + */ +export const MIGRATED_E2E_SEGMENTS: string[] = ["api-reference"]; diff --git a/ui/litellm-dashboard/e2e_tests/globalSetup.ts b/ui/litellm-dashboard/e2e_tests/globalSetup.ts index 8f80f57bd78..0b3fa7e8807 100644 --- a/ui/litellm-dashboard/e2e_tests/globalSetup.ts +++ b/ui/litellm-dashboard/e2e_tests/globalSetup.ts @@ -4,17 +4,18 @@ import * as fs from "fs"; async function globalSetup() { const browser = await chromium.launch(); + const rootPath = process.env.SERVER_ROOT_PATH ?? ""; for (const role of Object.values(Role)) { const { email, password } = users[role]; const storagePath = STORAGE_PATHS[role]; const page = await browser.newPage(); try { - await page.goto("http://localhost:4000/ui/login"); + await page.goto(`http://localhost:4000${rootPath}/ui/login`); await page.getByPlaceholder("Enter your username").fill(email); await page.getByPlaceholder("Enter your password").fill(password); await page.getByRole("button", { name: "Login", exact: true }).click(); - await page.waitForURL((url) => url.pathname.startsWith("/ui") && !url.pathname.includes("/login"), { + await page.waitForURL((url) => url.pathname.startsWith(`${rootPath}/ui`) && !url.pathname.includes("/login"), { timeout: 30_000, }); await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 }); diff --git a/ui/litellm-dashboard/e2e_tests/migration.serverRootPath.config.ts b/ui/litellm-dashboard/e2e_tests/migration.serverRootPath.config.ts new file mode 100644 index 00000000000..205348463c8 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/migration.serverRootPath.config.ts @@ -0,0 +1,38 @@ +import { defineConfig, devices } from "@playwright/test"; + +/** + * App Router migration smoke under a non-root mount. Boot the proxy with the same + * SERVER_ROOT_PATH (e.g. SERVER_ROOT_PATH=/litellm) and a UI built for it before + * running. globalSetup logs in at `${SERVER_ROOT_PATH}/ui/login` so the admin + * storage state is valid under the prefix. + */ +if (!process.env.SERVER_ROOT_PATH) { + throw new Error( + "migration.serverRootPath.config.ts requires SERVER_ROOT_PATH to be set (e.g. SERVER_ROOT_PATH=/litellm). " + + "Without it this config silently re-runs the default mount and never exercises the prefix. " + + "For the root-less run use the default playwright.config.ts (npm run e2e:migration).", + ); +} + +export default defineConfig({ + testDir: "./tests/migration", + testMatch: ["migratedPages.spec.ts"], + fullyParallel: true, + forbidOnly: !!process.env.CI, + retries: process.env.CI ? 2 : 0, + workers: process.env.CI ? 1 : undefined, + reporter: "list", + use: { + baseURL: "http://localhost:4000", + trace: "on-first-retry", + actionTimeout: 15 * 1000, + navigationTimeout: 30 * 1000, + launchOptions: { + slowMo: process.env.SLOWMO ? parseInt(process.env.SLOWMO, 10) || 0 : 0, + }, + }, + projects: [{ name: "chromium", use: { ...devices["Desktop Chrome"] } }], + timeout: 3 * 60 * 1000, + expect: { timeout: 10 * 1000 }, + globalSetup: require.resolve("./globalSetup"), +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/migration/README.md b/ui/litellm-dashboard/e2e_tests/tests/migration/README.md new file mode 100644 index 00000000000..4b3a391d421 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/migration/README.md @@ -0,0 +1,33 @@ +# App Router migration smoke + +A growing E2E smoke for pages migrated from the legacy `?page=` switch to App +Router path routes. For each migrated page it clicks the page's sidebar link, checks +the URL is the path route and the page renders, reloads it, then clicks off to a +legacy page and back to confirm navigation still works. It runs in two situations: +the default mount and a non-root `SERVER_ROOT_PATH` mount. + +## Adding a page + +When a page's migration merges, add its route segment to +`e2e_tests/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES` +in `src/utils/migratedPages.ts`). Both suites pick it up automatically. + +## Running + +Build the UI into the proxy and start the proxy first (the suite runs against +`http://localhost:4000`). + +Default mount: + +``` +npm run e2e:migration +``` + +Non-root mount (build and boot the proxy with the same root path, e.g. `/litellm`): + +``` +SERVER_ROOT_PATH=/litellm npm run e2e:migration:root +``` + +`globalSetup` logs in once per role; the admin storage state is reused for these +tests. Under a non-root mount it logs in at `${SERVER_ROOT_PATH}/ui/login`. diff --git a/ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts new file mode 100644 index 00000000000..98f4fee1450 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts @@ -0,0 +1,101 @@ +import { test, expect, type Page } from "@playwright/test"; +import { MIGRATED_E2E_SEGMENTS } from "../../fixtures/migratedPages"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { dismissFeedbackPopup } from "../../helpers/navigation"; + +/** + * App Router migration smoke as a user journey: start where the proxy lands you, + * click a migrated page in the sidebar, confirm it routed and rendered, reload it + * (the check a wrong server_root_path breaks), bounce to a legacy page and back, + * and, once two pages are migrated, navigate directly between two migrated pages. + * + * Driven by MIGRATED_E2E_SEGMENTS, so it grows as pages are migrated. Set + * SERVER_ROOT_PATH (e.g. "/litellm") to exercise the non-root mount; leave it + * unset for the default mount. Boot the proxy with the matching value first. + */ +const ROOT = process.env.SERVER_ROOT_PATH ?? ""; + +const esc = (s: string) => s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); +const pathRe = (segment: string) => new RegExp(`${esc(ROOT)}/ui/${esc(segment)}/?($|\\?)`); +const legacyAnchor = (page: Page) => page.locator("a", { hasText: "Virtual Keys" }); + +/** The dashboard shell is present (sidebar rendered); page didn't 404 / crash. */ +async function expectRendered(page: Page) { + await expect(legacyAnchor(page)).toBeVisible({ timeout: 20_000 }); +} + +/** + * Click a migrated page's sidebar link. Migrated items render as ; + * nested ones live under collapsible submenus, so expand submenus until the link is clickable. + */ +async function clickSidebar(page: Page, segment: string) { + const link = page.locator(`a[href$="/ui/${segment}"]`).first(); + for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) { + const collapsedSubmenu = page + .locator(".ant-menu-submenu:not(.ant-menu-submenu-open) > .ant-menu-submenu-title") + .first(); + if (!(await collapsedSubmenu.isVisible().catch(() => false))) break; + await collapsedSubmenu.click(); + await page.waitForTimeout(250); + } + await link.click(); +} + +test.use({ storageState: ADMIN_STORAGE_PATH }); + +test.describe("App Router migrated pages", () => { + for (const segment of MIGRATED_E2E_SEGMENTS) { + test(`${segment}: sidebar nav, reload, and round-trip with a legacy page`, async ({ page }) => { + const pageErrors: string[] = []; + page.on("pageerror", (e) => pageErrors.push(String(e))); + + // 1. Start where the proxy lands us. + await page.goto(`${ROOT}/ui/`); + await dismissFeedbackPopup(page); + await expectRendered(page); + + // 2. Click the migrated page in the sidebar -> path route + rendered. + await clickSidebar(page, segment); + await expect(page).toHaveURL(pathRe(segment)); + await expectRendered(page); + // 3. Reload the path route directly; a wrong server_root_path 404s here. + await page.reload(); + await dismissFeedbackPopup(page); + await expect(page).toHaveURL(pathRe(segment)); + await expectRendered(page); + // 4. Click off to a legacy (not-yet-migrated) page. + await legacyAnchor(page).click(); + await expect(page).toHaveURL(new RegExp(`${esc(ROOT)}/ui/\\?page=api-keys`)); + await dismissFeedbackPopup(page); + await expectRendered(page); + // 5. Click back to the migrated page. + await clickSidebar(page, segment); + await expect(page).toHaveURL(pathRe(segment)); + await expectRendered(page); + expect(pageErrors, `page errors during ${segment} journey`).toEqual([]); + }); + } + + test("navigates directly between two migrated pages", async ({ page }) => { + test.skip(MIGRATED_E2E_SEGMENTS.length < 2, "needs >= 2 migrated pages"); + const [first, second] = MIGRATED_E2E_SEGMENTS; + const pageErrors: string[] = []; + page.on("pageerror", (e) => pageErrors.push(String(e))); + + await page.goto(`${ROOT}/ui/`); + await dismissFeedbackPopup(page); + + await clickSidebar(page, first); + await expect(page).toHaveURL(pathRe(first)); + await expectRendered(page); + await clickSidebar(page, second); + await expect(page).toHaveURL(pathRe(second)); + await expectRendered(page); + // Back to the first migrated page. + await clickSidebar(page, first); + await expect(page).toHaveURL(pathRe(first)); + await expectRendered(page); + + expect(pageErrors, "page errors during migrated -> migrated nav").toEqual([]); + }); +}); diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index bbab73b07f1..d3169395b4e 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -164,11 +164,6 @@ "count": 2 } }, - "src/app/(dashboard)/layout.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx": { "no-restricted-imports": { "count": 1 @@ -228,11 +223,6 @@ "count": 1 } }, - "src/app/page.tsx": { - "unused-imports/no-unused-imports": { - "count": 2 - } - }, "src/components/AIHub/AgentHubTableColumns.test.tsx": { "unused-imports/no-unused-imports": { "count": 1 @@ -243,14 +233,6 @@ "count": 1 } }, - "src/components/AIHub/ClaudeCodeMarketplaceTab.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/immutability": { - "count": 1 - } - }, "src/components/AIHub/ModelHubTable.test.tsx": { "max-params": { "count": 1 @@ -303,11 +285,6 @@ "count": 1 } }, - "src/components/AIHub/marketplace_table_columns.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/AdminPanel.tsx": { "no-restricted-imports": { "count": 1 @@ -816,11 +793,6 @@ "count": 1 } }, - "src/components/agents/agent_table.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/alerting/dynamic_form.tsx": { "no-restricted-imports": { "count": 1 @@ -956,14 +928,6 @@ "count": 1 } }, - "src/components/claude_code_plugins/plugin_info.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/immutability": { - "count": 1 - } - }, "src/components/claude_code_plugins/plugin_table.tsx": { "no-restricted-imports": { "count": 1 @@ -1333,14 +1297,6 @@ "count": 2 } }, - "src/components/mcp_tools/mcp_server_columns.tsx": { - "max-params": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/components/mcp_tools/mcp_server_cost_config.tsx": { "no-restricted-imports": { "count": 1 @@ -1351,11 +1307,6 @@ "count": 1 } }, - "src/components/mcp_tools/mcp_server_edit.test.tsx": { - "unused-imports/no-unused-imports": { - "count": 1 - } - }, "src/components/mcp_tools/mcp_server_edit.tsx": { "no-restricted-imports": { "count": 1 @@ -1494,11 +1445,6 @@ "count": 1 } }, - "src/components/navbar.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/networking.tsx": { "max-params": { "count": 23 @@ -1517,11 +1463,6 @@ "count": 1 } }, - "src/components/organisms/RegenerateKeyModal.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/organisms/create_key_button.test.tsx": { "@typescript-eslint/no-require-imports": { "count": 2 diff --git a/ui/litellm-dashboard/knip.json b/ui/litellm-dashboard/knip.json index e95c0acef3f..6f129398981 100644 --- a/ui/litellm-dashboard/knip.json +++ b/ui/litellm-dashboard/knip.json @@ -1,8 +1,9 @@ { "$schema": "https://unpkg.com/knip@5/schema.json", - "entry": ["scripts/**/*.ts"], - "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.ts", "e2e_tests/**/*.ts"], + "entry": ["scripts/**/*.{ts,mjs}"], + "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.{ts,mjs}", "e2e_tests/**/*.ts"], "ignore": ["src/lib/http/schema.d.ts"], + "ignoreDependencies": ["openapi-typescript"], "playwright": { "config": "e2e_tests/playwright.config.ts", "entry": ["e2e_tests/**/*.spec.ts", "e2e_tests/**/*.setup.ts", "e2e_tests/globalSetup.ts"] diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 7000efd63b4..568f6b288d5 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -11,7 +11,6 @@ "@anthropic-ai/sdk": "0.92.0", "@headlessui/tailwindcss": "0.2.2", "@heroicons/react": "1.0.6", - "@remixicon/react": "4.9.0", "@tanstack/react-pacer": "0.2.0", "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", @@ -44,18 +43,15 @@ "@testing-library/jest-dom": "6.9.1", "@testing-library/react": "16.3.2", "@testing-library/user-event": "14.6.1", - "@types/babel__traverse": "7.28.0", "@types/lodash": "4.17.23", "@types/node": "20.19.37", "@types/react": "18.2.48", "@types/react-copy-to-clipboard": "5.0.7", "@types/react-dom": "18.3.7", "@types/react-syntax-highlighter": "15.5.13", - "@types/uuid": "10.0.0", "@vitest/coverage-v8": "3.2.4", "@vitest/ui": "3.2.4", "autoprefixer": "10.4.24", - "dotenv": "17.2.3", "eslint": "9.39.2", "eslint-config-next": "16.2.6", "eslint-config-prettier": "10.1.8", @@ -68,7 +64,6 @@ "tailwindcss": "3.4.19", "typescript": "5.9.3", "typescript-eslint": "8.60.1", - "vite": "7.3.2", "vitest": "3.2.4" }, "engines": { @@ -2847,15 +2842,6 @@ "npm": ">=9.5.0" } }, - "node_modules/@remixicon/react": { - "version": "4.9.0", - "resolved": "https://registry.npmjs.org/@remixicon/react/-/react-4.9.0.tgz", - "integrity": "sha512-5/jLDD4DtKxH2B4QVXTobvV1C2uL8ab9D5yAYNtFt+w80O0Ys1xFOrspqROL3fjrZi+7ElFUWE37hBfaAl6U+Q==", - "license": "Remix Icon License 1.0", - "peerDependencies": { - "react": ">=18.2.0" - } - }, "node_modules/@rollup/rollup-android-arm-eabi": { "version": "4.60.3", "resolved": "https://registry.npmjs.org/@rollup/rollup-android-arm-eabi/-/rollup-android-arm-eabi-4.60.3.tgz", @@ -3490,16 +3476,6 @@ "dev": true, "license": "MIT" }, - "node_modules/@types/babel__traverse": { - "version": "7.28.0", - "resolved": "https://registry.npmjs.org/@types/babel__traverse/-/babel__traverse-7.28.0.tgz", - "integrity": "sha512-8PvcXf70gTDZBgt9ptxJ8elBeBjcLOAcOtoO/mPJjtji1+CdGbHgm77om1GrsPxsiE+uXIpNSK64UYaIwQXd4Q==", - "dev": true, - "license": "MIT", - "dependencies": { - "@babel/types": "^7.28.2" - } - }, "node_modules/@types/chai": { "version": "5.2.3", "resolved": "https://registry.npmjs.org/@types/chai/-/chai-5.2.3.tgz", @@ -3737,13 +3713,6 @@ "integrity": "sha512-ko/gIFJRv177XgZsZcBwnqJN5x/Gien8qNOn0D5bQU/zAzVf9Zt3BlcUiLqhV9y4ARk0GbT3tnUiPNgnTXzc/Q==", "license": "MIT" }, - "node_modules/@types/uuid": { - "version": "10.0.0", - "resolved": "https://registry.npmjs.org/@types/uuid/-/uuid-10.0.0.tgz", - "integrity": "sha512-7gqG38EyHgyP1S+7+xomFtL+ZNHcKv6DwNaCZmJmo1vgMugyF3TCnXVg4t1uk89mLNwnLtnY3TpOpCOyp1/xHQ==", - "dev": true, - "license": "MIT" - }, "node_modules/@typescript-eslint/eslint-plugin": { "version": "8.60.1", "resolved": "https://registry.npmjs.org/@typescript-eslint/eslint-plugin/-/eslint-plugin-8.60.1.tgz", @@ -5912,19 +5881,6 @@ "csstype": "^3.0.2" } }, - "node_modules/dotenv": { - "version": "17.2.3", - "resolved": "https://registry.npmjs.org/dotenv/-/dotenv-17.2.3.tgz", - "integrity": "sha512-JVUnt+DUIzu87TABbhPmNfVdBDt18BLOWjMUFJMSi/Qqg7NTYtabbvSNJGOJ7afbRuv9D/lngizHtP7QyLQ+9w==", - "dev": true, - "license": "BSD-2-Clause", - "engines": { - "node": ">=12" - }, - "funding": { - "url": "https://dotenvx.com" - } - }, "node_modules/dunder-proto": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz", @@ -13770,9 +13726,9 @@ } }, "node_modules/ws": { - "version": "8.19.0", - "resolved": "https://registry.npmjs.org/ws/-/ws-8.19.0.tgz", - "integrity": "sha512-blAT2mjOEIi0ZzruJfIhb3nps74PRWTCz1IjglWEEpQl5XS/UNama6u2/rjFkDDouqr4L67ry+1aGIALViWjDg==", + "version": "8.20.1", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.20.1.tgz", + "integrity": "sha512-It4dO0K5v//JtTXuPkfEOaI3uUN87iYPnqo/ZzqCoG3g8uhA66QUMs/SrM0YK7/NAu+r4LMh/9dq2A7k+rHs+w==", "devOptional": true, "license": "MIT", "engines": { diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index ca753e59dc5..eb6211a91d1 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -16,6 +16,8 @@ "format:check": "prettier --check .", "e2e": "playwright test --config e2e_tests/playwright.config.ts", "e2e:ui": "playwright test --ui --config e2e_tests/playwright.config.ts", + "e2e:migration": "playwright test e2e_tests/tests/migration/migratedPages.spec.ts --config e2e_tests/playwright.config.ts", + "e2e:migration:root": "playwright test --config e2e_tests/migration.serverRootPath.config.ts", "knip": "knip", "knip:fix": "knip --fix", "gen:api": "node scripts/gen-api-types.mjs" @@ -24,7 +26,6 @@ "@anthropic-ai/sdk": "0.92.0", "@headlessui/tailwindcss": "0.2.2", "@heroicons/react": "1.0.6", - "@remixicon/react": "4.9.0", "@tanstack/react-pacer": "0.2.0", "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", @@ -57,18 +58,15 @@ "@testing-library/jest-dom": "6.9.1", "@testing-library/react": "16.3.2", "@testing-library/user-event": "14.6.1", - "@types/babel__traverse": "7.28.0", "@types/lodash": "4.17.23", "@types/node": "20.19.37", "@types/react": "18.2.48", "@types/react-copy-to-clipboard": "5.0.7", "@types/react-dom": "18.3.7", "@types/react-syntax-highlighter": "15.5.13", - "@types/uuid": "10.0.0", "@vitest/coverage-v8": "3.2.4", "@vitest/ui": "3.2.4", "autoprefixer": "10.4.24", - "dotenv": "17.2.3", "eslint": "9.39.2", "eslint-config-next": "16.2.6", "eslint-config-prettier": "10.1.8", @@ -81,7 +79,6 @@ "tailwindcss": "3.4.19", "typescript": "5.9.3", "typescript-eslint": "8.60.1", - "vite": "7.3.2", "vitest": "3.2.4" }, "overrides": { @@ -90,7 +87,7 @@ "glob": "13.0.0", "minimatch": "10.2.4", "lodash": "4.18.1", - "ws": "8.19.0", + "ws": "8.20.1", "braces": "3.0.3", "axios": "1.13.6", "postcss": "8.5.13" diff --git a/ui/litellm-dashboard/scripts/gen-api-types.mjs b/ui/litellm-dashboard/scripts/gen-api-types.mjs index 66dfd2ee7b8..3c9373ec547 100644 --- a/ui/litellm-dashboard/scripts/gen-api-types.mjs +++ b/ui/litellm-dashboard/scripts/gen-api-types.mjs @@ -23,9 +23,17 @@ const specDir = mkdtempSync(join(tmpdir(), "litellm-openapi-")); const specPath = join(specDir, "openapi.json"); const python = (process.env.LITELLM_PYTHON ?? "python3").split(" "); +// The dashboard calls internal UI routes that the public /openapi.json hides via +// include_in_schema=False. Force them in so they get typed here; this mutates a +// throwaway interpreter, so the spec the proxy actually serves is unchanged. const dumpSpec = [ "import json, sys", "from litellm.proxy.proxy_server import app", + "from fastapi.routing import APIRoute", + "for route in app.routes:", + " if isinstance(route, APIRoute):", + " route.include_in_schema = True", + "app.openapi_schema = None", "with open(sys.argv[1], 'w') as f: json.dump(app.openapi(), f, sort_keys=True)", ].join("\n"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx index 02bed1adbe5..a4a4d3d0f43 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx @@ -1,10 +1,12 @@ "use client"; import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; const APIReferencePage = () => { - const proxySettings = useProxySettings(); + const { accessToken } = useAuthorized(); + const proxySettings = useProxySettings(accessToken); return ; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx deleted file mode 100644 index 90f498912a8..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx +++ /dev/null @@ -1,472 +0,0 @@ -"use client"; - -import { Layout, Menu, ConfigProvider } from "antd"; -import { - KeyOutlined, - PlayCircleOutlined, - BlockOutlined, - BarChartOutlined, - TeamOutlined, - BankOutlined, - UserOutlined, - SettingOutlined, - ApiOutlined, - AppstoreOutlined, - DatabaseOutlined, - FileTextOutlined, - LineChartOutlined, - SafetyOutlined, - ExperimentOutlined, - ToolOutlined, - TagsOutlined, - AuditOutlined, -} from "@ant-design/icons"; -// import { -// all_admin_roles, -// rolesWithWriteAccess, -// internalUserRoles, -// isAdminRole, -// } from "../utils/roles"; -// import UsageIndicator from "./usage_indicator"; -import * as React from "react"; -import { useRouter, usePathname } from "next/navigation"; -import { all_admin_roles, internalUserRoles, isAdminRole, rolesWithWriteAccess } from "@/utils/roles"; -import UsageIndicator from "@/components/UsageIndicator"; -import { serverRootPath } from "@/components/networking"; - -const { Sider } = Layout; - -// -------- Types -------- -interface SidebarProps { - accessToken: string | null; - userRole: string; - /** Fallback selection id (legacy), used if path can't be matched */ - defaultSelectedKey: string; - collapsed?: boolean; -} - -interface MenuItemCfg { - key: string; - newTab?: boolean; - page: string; // legacy id; we map this to a path below - label: string; - roles?: string[]; - children?: MenuItemCfg[]; - icon?: React.ReactNode; -} - -/** ---------- Base URL helpers ---------- */ -/** - * Normalizes NEXT_PUBLIC_BASE_URL to either "/" or "/ui/" (always with a trailing slash). - * Supported env values: "" or "ui/". - * Also considers the serverRootPath from the proxy config (e.g., "/my-custom-path"). - */ -const getBasePath = () => { - const raw = process.env.NEXT_PUBLIC_BASE_URL ?? ""; - const trimmed = raw.replace(/^\/+|\/+$/g, ""); // strip leading/trailing slashes - const uiPath = trimmed ? `/${trimmed}/` : "/"; - - // If serverRootPath is set and not "/", prepend it to the UI path - if (serverRootPath && serverRootPath !== "/") { - // Remove trailing slash from serverRootPath and ensure uiPath has no leading slash for proper joining - const cleanServerRoot = serverRootPath.replace(/\/+$/, ""); - const cleanUiPath = uiPath.replace(/^\/+/, ""); - return `${cleanServerRoot}/${cleanUiPath}`; - } - - return uiPath; -}; - -/** Map legacy `page` ids to real app routes (relative, no leading slash). */ -const routeFor = (slug: string): string => { - switch (slug) { - // top level - case "api-keys": - return "virtual-keys"; - case "llm-playground": - return "test-key"; - case "models": - return "models-and-endpoints"; - case "new_usage": - return "usage"; - case "teams": - return "teams"; - case "organizations": - return "organizations"; - case "users": - return "users"; - case "api_ref": - return "api-reference"; - case "model-hub-table": - // If you intend the newer in-dashboard page, use "model-hub". - return "model-hub"; - case "logs": - return "logs"; - case "guardrails": - return "guardrails"; - case "policies": - return "policies"; - case "chat": - return "chat"; - - // tools - case "mcp-servers": - return "tools/mcp-servers"; - case "vector-stores": - return "tools/vector-stores"; - case "byok-demo": - return "tools/byok-demo"; - - // experimental - case "caching": - return "experimental/caching"; - case "prompts": - return "experimental/prompts"; - case "budgets": - return "experimental/budgets"; - case "transform-request": - return "experimental/api-playground"; - case "tag-management": - return "experimental/tag-management"; - case "claude-code-plugins": - return "experimental/claude-code-plugins"; - case "usage": // "Old Usage" - return "experimental/old-usage"; - - // settings - case "general-settings": - return "settings/router-settings"; - case "settings": // "Logging & Alerts" - return "settings/logging-and-alerts"; - case "admin-panel": - return "settings/admin-settings"; - case "ui-theme": - return "settings/ui-theme"; - - default: - // treat as already a relative path - return slug.replace(/^\/+/, ""); - } -}; - -/** Prefix base path ("/" or "/ui/") */ -const toHref = (slugOrPath: string) => { - const base = getBasePath(); // "/" or "/ui/" - const rel = routeFor(slugOrPath).replace(/^\/+|\/+$/g, ""); - return `${base}${rel}`; -}; - -// ----- Menu config (unchanged labels/icons; same appearance) ----- -const menuItems: MenuItemCfg[] = [ - { key: "1", page: "api-keys", label: "Virtual Keys", icon: }, - { - key: "3", - page: "llm-playground", - label: "Test Key", - icon: , - roles: rolesWithWriteAccess, - }, - { - key: "2", - page: "models", - label: "Models + Endpoints", - icon: , - roles: rolesWithWriteAccess, - }, - { - key: "12", - page: "new_usage", - label: "Usage", - icon: , - roles: [...all_admin_roles, ...internalUserRoles], - }, - { key: "6", page: "teams", label: "Teams", icon: }, - { - key: "17", - page: "organizations", - label: "Organizations", - icon: , - roles: all_admin_roles, - }, - { - key: "5", - page: "users", - label: "Internal Users", - icon: , - roles: all_admin_roles, - }, - { key: "14", page: "api-reference", label: "API Reference", icon: }, - { - key: "16", - page: "model-hub-table", - label: "Model Hub", - icon: , - }, - { key: "15", page: "logs", label: "Logs", icon: }, - { - key: "11", - page: "guardrails", - label: "Guardrails", - icon: , - roles: all_admin_roles, - }, - { - key: "28", - page: "policies", - label: "Policies", - icon: , - roles: all_admin_roles, - }, - { - key: "26", - page: "tools", - label: "Tools", - icon: , - children: [ - { key: "18", page: "mcp-servers", label: "MCP Servers", icon: }, - { - key: "21", - page: "vector-stores", - label: "Vector Stores", - icon: , - roles: all_admin_roles, - }, - ], - }, - { - key: "experimental", - page: "experimental", - label: "Experimental", - icon: , - children: [ - { - key: "9", - page: "caching", - label: "Caching", - icon: , - roles: all_admin_roles, - }, - { - key: "25", - page: "prompts", - label: "Prompts", - icon: , - roles: all_admin_roles, - }, - { - key: "10", - page: "budgets", - label: "Budgets", - icon: , - roles: all_admin_roles, - }, - { - key: "20", - page: "transform-request", - label: "API Playground", - icon: , - roles: [...all_admin_roles, ...internalUserRoles], - }, - { - key: "19", - page: "tag-management", - label: "Tag Management", - icon: , - roles: all_admin_roles, - }, - { - key: "27", - page: "claude-code-plugins", - label: "Claude Code Plugins", - icon: , - roles: all_admin_roles, - }, - { key: "4", page: "usage", label: "Old Usage", icon: }, - ], - }, - { - key: "settings", - page: "settings", - label: "Settings", - icon: , - roles: all_admin_roles, - children: [ - { - key: "11", - page: "general-settings", - label: "Router Settings", - icon: , - roles: all_admin_roles, - }, - { - key: "8", - page: "settings", - label: "Logging & Alerts", - icon: , - roles: all_admin_roles, - }, - { - key: "13", - page: "admin-panel", - label: "Admin Settings", - icon: , - roles: all_admin_roles, - }, - { - key: "14", - page: "ui-theme", - label: "UI Theme", - icon: , - roles: all_admin_roles, - }, - ], - }, -]; - -const Sidebar2: React.FC = ({ accessToken, userRole, defaultSelectedKey, collapsed = false }) => { - const router = useRouter(); - const pathname = usePathname() || "/"; - - // ----- Filter by role without mutating originals ----- - const filteredMenuItems = React.useMemo(() => { - return menuItems - .filter((item) => !item.roles || item.roles.includes(userRole)) - .map((item) => ({ - ...item, - children: item.children ? item.children.filter((c) => !c.roles || c.roles.includes(userRole)) : undefined, - })); - }, [userRole]); - - // ----- Compute selected key from current path ----- - const selectedMenuKey = React.useMemo(() => { - const base = getBasePath(); - // strip base prefix and leading slash -> "virtual-keys", "tools/mcp-servers", etc. - const rel = pathname.startsWith(base) ? pathname.slice(base.length) : pathname.replace(/^\/+/, ""); - const relLower = rel.toLowerCase(); - - const matchesPath = (slug: string) => { - const route = routeFor(slug).toLowerCase(); - return relLower === route || relLower.startsWith(`${route}/`); - }; - - // search top-level - for (const item of filteredMenuItems) { - if (!item.children && matchesPath(item.page)) return item.key; - if (item.children) { - for (const child of item.children) { - if (matchesPath(child.page)) return child.key; - } - } - } - - // fallback to legacy defaultSelectedKey mapping - const fallback = filteredMenuItems.find((i) => i.page === defaultSelectedKey)?.key; - if (fallback) return fallback; - - for (const item of filteredMenuItems) { - if (item.children?.some((c) => c.page === defaultSelectedKey)) { - const child = item.children.find((c) => c.page === defaultSelectedKey)!; - return child.key; - } - } - - return "1"; - }, [pathname, filteredMenuItems, defaultSelectedKey]); - - // ----- Navigation ----- - const goTo = (slug: string, newTab?: boolean) => { - const href = toHref(slug); - if (newTab) { - window.open(href, "_blank"); - } else { - router.push(href); - } - }; - - // Wrap label in so every nav item supports right-click → "Open in new tab" - // and Ctrl/Cmd+click to open in a new tab, while preserving SPA navigation for normal clicks. - const renderNavLink = (label: string, page: string, newTab?: boolean): React.ReactNode => { - const href = toHref(page); - return ( - { - if (newTab) { - e.stopPropagation(); - return; - } - if (e.metaKey || e.ctrlKey || e.shiftKey || e.button === 1) { - e.stopPropagation(); - return; - } - e.preventDefault(); - }} - style={{ color: "inherit", textDecoration: "none" }} - > - {label} - - ); - }; - - return ( - - - - ({ - key: item.key, - icon: item.icon, - label: renderNavLink(item.label, item.page, item.newTab), - children: item.children?.map((child) => ({ - key: child.key, - icon: child.icon, - label: renderNavLink(child.label, child.page, child.newTab), - onClick: () => goTo(child.page, child.newTab), - })), - onClick: !item.children ? () => goTo(item.page, item.newTab) : undefined, - }))} - /> - - {isAdminRole(userRole) && !collapsed && } - - - ); -}; - -export default Sidebar2; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.test.ts index 13b9107bdc1..39d1b28303d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.test.ts @@ -115,8 +115,16 @@ describe("useProjects", () => { expect(global.fetch).not.toHaveBeenCalled(); }); - it("should not fetch when userRole is not an admin role", () => { + it("should fetch when userRole is an internal user role", async () => { mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "Internal User" }); + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => mockProjects }); + const { result } = renderHook(() => useProjects(), { wrapper: makeWrapper(queryClient) }); + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(global.fetch).toHaveBeenCalled(); + }); + + it("should not fetch when userRole cannot read projects", () => { + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "regular_user" }); const { result } = renderHook(() => useProjects(), { wrapper: makeWrapper(queryClient) }); expect(result.current.isFetched).toBe(false); expect(global.fetch).not.toHaveBeenCalled(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts index 7bdc8a4fe6d..c240dbb0170 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts @@ -2,7 +2,7 @@ import { useQuery } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { all_admin_roles } from "@/utils/roles"; +import { all_admin_roles, internalUserRoles } from "@/utils/roles"; // ── Types ──────────────────────────────────────────────────────────────────── @@ -42,6 +42,8 @@ export interface ProjectResponse { export const projectKeys = createQueryKeys("projects"); +const projectReaderRoles = [...all_admin_roles, ...internalUserRoles]; + // ── Fetch function ─────────────────────────────────────────────────────────── const fetchProjects = async (accessToken: string): Promise => { @@ -74,6 +76,6 @@ export const useProjects = () => { return useQuery({ queryKey: projectKeys.list({}), queryFn: async () => fetchProjects(accessToken!), - enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!), + enabled: Boolean(accessToken) && projectReaderRoles.includes(userRole!), }); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts index d4fb3073856..82cefd800f4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts @@ -1,21 +1,26 @@ -import { useState, useEffect } from "react"; import { fetchProxySettings } from "@/utils/proxyUtils"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; -export default function useProxySettings() { - const { accessToken } = useAuthorized(); - const [proxySettings, setProxySettings] = useState({ - PROXY_BASE_URL: "", - PROXY_LOGOUT_URL: "", - LITELLM_UI_API_DOC_BASE_URL: null as string | null, - }); +export const proxySettingsKeys = createQueryKeys("proxySettings"); + +export interface ProxySettings { + PROXY_BASE_URL: string; + PROXY_LOGOUT_URL: string; + LITELLM_UI_API_DOC_BASE_URL?: string | null; +} - useEffect(() => { - if (!accessToken) return; - fetchProxySettings(accessToken).then((settings) => { - if (settings) setProxySettings(settings); - }); - }, [accessToken]); +const EMPTY_PROXY_SETTINGS: ProxySettings = { + PROXY_BASE_URL: "", + PROXY_LOGOUT_URL: "", + LITELLM_UI_API_DOC_BASE_URL: null, +}; - return proxySettings; +export default function useProxySettings(accessToken: string | null): ProxySettings { + const { data } = useQuery({ + queryKey: [...proxySettingsKeys.all, accessToken], + queryFn: () => fetchProxySettings(accessToken), + enabled: Boolean(accessToken), + }); + return data ?? EMPTY_PROXY_SETTINGS; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index a611d619cc1..df5b2ab4511 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -1,92 +1,63 @@ "use client"; -import React, { Suspense, useEffect, useState } from "react"; +import React, { Suspense, useState } from "react"; import Navbar from "@/components/navbar"; +import LoadingScreen from "@/components/common_components/LoadingScreen"; import { ThemeProvider } from "@/contexts/ThemeContext"; +import { useAuth } from "@/contexts/AuthContext"; import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { useRouter, useSearchParams } from "next/navigation"; +import { useRouter, useSearchParams, usePathname } from "next/navigation"; import { DebugWarningBanner } from "@/components/DebugWarningBanner"; +import { MIGRATED_PAGES, migratedHref, legacyPageHref, legacyKeyForPathname } from "@/utils/migratedPages"; -/** ---- BASE URL HELPERS ---- */ -function normalizeBasePrefix(raw: string | undefined | null): string { - const trimmed = (raw ?? "").trim(); - if (!trimmed) return ""; - const core = trimmed.replace(/^\/+/, "").replace(/\/+$/, ""); - return core ? `/${core}/` : "/"; -} -const BASE_PREFIX = normalizeBasePrefix(process.env.NEXT_PUBLIC_BASE_URL); -function withBase(path: string): string { - const body = path.startsWith("/") ? path.slice(1) : path; - const combined = `${BASE_PREFIX}${body}`; - return combined.startsWith("/") ? combined : `/${combined}`; -} -/** -------------------------------- */ - -/** - * Pages that have been migrated to path-based routing under (dashboard)/. - * When the leftnav triggers one of these, navigate to the path route instead - * of the legacy query-param root page. - * - * Key = legacy page id used in leftnav, Value = route segment under (dashboard)/ - */ -const MIGRATED_PAGES: Record = {}; - -function LayoutContent({ children }: { children: React.ReactNode }) { +function DashboardShell({ children }: { children: React.ReactNode }) { const router = useRouter(); const searchParams = useSearchParams(); - const { accessToken } = useAuthorized(); - const [sidebarCollapsed, setSidebarCollapsed] = React.useState(false); - const [page, setPage] = useState(() => { - return searchParams.get("page") || "api-keys"; - }); + const pathname = usePathname(); + const { accessToken } = useAuth(); + const [sidebarCollapsed, setSidebarCollapsed] = useState(false); - const handleSetPage = (newPage: string) => { - // If the page has been migrated to path routing, navigate there - const migratedRoute = MIGRATED_PAGES[newPage]; - if (migratedRoute) { - router.push(withBase(migratedRoute)); - setPage(newPage); - return; - } + const page = legacyKeyForPathname(pathname) || searchParams.get("page") || "api-keys"; - // Otherwise, navigate back to the legacy root page with query params - router.push(withBase(`?page=${newPage}`)); - setPage(newPage); + const navigateToPage = (newPage: string) => { + const migratedRoute = MIGRATED_PAGES[newPage]; + router.push(migratedRoute ? migratedHref(migratedRoute) : legacyPageHref(newPage)); }; - useEffect(() => { - setPage(searchParams.get("page") || "api-keys"); - }, [searchParams]); - - const toggleSidebar = () => setSidebarCollapsed((v) => !v); - return ( - -
- {}} - accessToken={accessToken} - /> - -
-
- -
-
{children}
+
+ setSidebarCollapsed((v) => !v)} + /> + +
+
+
+
{children}
+
+ ); +} + +function LayoutContent({ children }: { children: React.ReactNode }) { + const searchParams = useSearchParams(); + const { accessToken } = useAuth(); + const isInvitationFlow = Boolean(searchParams.get("invitation_id")); + + return ( + + {isInvitationFlow ? children : {children}} ); } export default function Layout({ children }: { children: React.ReactNode }) { return ( - Loading...
}> + }> {children} ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx deleted file mode 100644 index 77496aef3e6..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ /dev/null @@ -1,26 +0,0 @@ -"use client"; - -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import useTeams from "@/app/(dashboard)/hooks/useTeams"; -import { useState } from "react"; -import ModelsAndEndpointsView from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; - -const ModelsAndEndpointsPage = () => { - const { token, premiumUser } = useAuthorized(); - const [keys, setKeys] = useState([]); - - const { teams } = useTeams(); - - return ( - {}} - premiumUser={premiumUser} - teams={teams} - /> - ); -}; - -export default ModelsAndEndpointsPage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/page.tsx deleted file mode 100644 index 6112fac3161..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/page.tsx +++ /dev/null @@ -1,34 +0,0 @@ -"use client"; - -import Organizations, { fetchOrganizations } from "@/components/organizations"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { useEffect, useState } from "react"; -import { Organization } from "@/components/networking"; -import { fetchUserModels } from "@/components/organisms/create_key_button"; - -const OrganizationsPage = () => { - const { userId: userID, accessToken, userRole, premiumUser } = useAuthorized(); - const [organizations, setOrganizations] = useState([]); - const [userModels, setUserModels] = useState([]); - - useEffect(() => { - fetchOrganizations(accessToken, setOrganizations).then(() => {}); - }, [accessToken]); - - useEffect(() => { - fetchUserModels(userID, userRole, accessToken, setUserModels).then(() => {}); - }, [userID, userRole, accessToken]); - - return ( - - ); -}; - -export default OrganizationsPage; diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx similarity index 53% rename from ui/litellm-dashboard/src/app/page.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/page.tsx index da6a0d5a76f..0854b085fae 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx @@ -1,8 +1,6 @@ "use client"; -import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView"; -import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; -import OldModelDashboard from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; +import ModelsAndEndpointsView from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; import PlaygroundPage from "@/app/(dashboard)/playground/page"; import AdminPanel from "@/components/AdminPanel"; import AgentsPanel from "@/components/agents"; @@ -10,6 +8,8 @@ import BudgetPanel from "@/components/budgets/budget_panel"; import CacheDashboard from "@/components/cache_dashboard"; import ClaudeCodePluginsPanel from "@/components/claude_code_plugins"; import { teamListCall as v2TeamListCall } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; +import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; import LoadingScreen from "@/components/common_components/LoadingScreen"; import { CostTrackingSettings } from "@/components/CostTrackingSettings"; import GeneralSettings from "@/components/general_settings"; @@ -19,7 +19,6 @@ import PoliciesPanel from "@/components/policies"; import { Team } from "@/components/key_team_helpers/key_list"; import { MCPServers } from "@/components/mcp_tools"; import ModelHubTable from "@/components/AIHub/ModelHubTable"; -import Navbar from "@/components/navbar"; import { Organization, proxyBaseUrl, getInProductNudgesCall } from "@/components/networking"; import NewUsagePage from "@/components/UsagePage/components/UsagePageView"; import OldTeams from "@/components/OldTeams"; @@ -44,7 +43,6 @@ import { MemoryView } from "@/components/MemoryView"; import WorkflowRuns from "@/components/workflow_runs"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; -import { ThemeProvider } from "@/contexts/ThemeContext"; import { useAuth } from "@/contexts/AuthContext"; import { buildLoginUrlWithReturn, @@ -54,23 +52,9 @@ import { storeReturnUrl, } from "@/utils/returnUrlUtils"; import { isAdminRole } from "@/utils/roles"; -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { MIGRATED_PAGES, migratedHref } from "@/utils/migratedPages"; import { useRouter, useSearchParams } from "next/navigation"; import { Suspense, useEffect, useMemo, useRef, useState } from "react"; -import { ConfigProvider, theme } from "antd"; - -interface ProxySettings { - PROXY_BASE_URL: string; - PROXY_LOGOUT_URL: string; - LITELLM_UI_API_DOC_BASE_URL?: string | null; -} - -/** - * Map of legacy query-param page keys → new path-based route segments. - * When a user visits ?page=, they are redirected to /ui/. - * Add entries here as pages are migrated from the if/else chain to path-based routes. - */ -const LEGACY_REDIRECTS: Record = {}; function CreateKeyPageContent() { const { authLoading, token, userID, userRole, userEmail, accessToken, premiumUser, setUserRole, setUserEmail } = @@ -80,16 +64,16 @@ function CreateKeyPageContent() { const [keys, setKeys] = useState([]); const [organizations, setOrganizations] = useState([]); const [userModels, setUserModels] = useState([]); - const [proxySettings, setProxySettings] = useState({ - PROXY_BASE_URL: "", - PROXY_LOGOUT_URL: "", - }); + const proxySettings = useProxySettings(accessToken); const router = useRouter(); const searchParams = useSearchParams()!; const [modelData, setModelData] = useState({ data: [] }); const [createClicked, setCreateClicked] = useState(false); + const { data: uiSettingsData, isLoading: uiSettingsLoading } = useUISettings(); + const nudgesDisabled = uiSettingsLoading || Boolean(uiSettingsData?.values?.disable_ui_nudges); + // Survey state - always show by default const [showSurveyPrompt, setShowSurveyPrompt] = useState(true); const [showSurveyModal, setShowSurveyModal] = useState(false); @@ -99,12 +83,6 @@ function CreateKeyPageContent() { const [showClaudeCodePrompt, setShowClaudeCodePrompt] = useState(false); const [showClaudeCodeModal, setShowClaudeCodeModal] = useState(false); - // Dark mode state - const [isDarkMode, setIsDarkMode] = useState(false); - const toggleDarkMode = () => { - setIsDarkMode(!isDarkMode); - }; - const invitation_id = searchParams.get("invitation_id"); // Parse URL query parameters for pre-filling the create key form @@ -157,32 +135,11 @@ function CreateKeyPageContent() { }; }, [searchParams, autoOpenCreate]); - // Get page from URL, default to 'api-keys' if not present - const [page, setPage] = useState(() => { - return searchParams.get("page") || "api-keys"; - }); - - // Custom setPage function that updates URL - const updatePage = (newPage: string) => { - // Update URL without full page reload - const newSearchParams = new URLSearchParams(searchParams); - newSearchParams.set("page", newPage); - - // Use Next.js router to update URL - window.history.pushState(null, "", `?${newSearchParams.toString()}`); - - setPage(newPage); - }; - - const [sidebarCollapsed, setSidebarCollapsed] = useState(false); + const page = searchParams.get("page") || "api-keys"; // Track if we've already attempted a return URL redirect to prevent race conditions const hasAttemptedReturnRedirectRef = useRef(false); - const toggleSidebar = () => { - setSidebarCollapsed(!sidebarCollapsed); - }; - const addKey = (data: any) => { setKeys((prevData) => (prevData ? [...prevData, data] : [data])); setCreateClicked(() => !createClicked); @@ -202,11 +159,10 @@ function CreateKeyPageContent() { }, [redirectToLogin]); // Redirect legacy query-param pages to their new path-based routes - const isLegacyRedirect = page in LEGACY_REDIRECTS; + const isLegacyRedirect = page in MIGRATED_PAGES; useEffect(() => { if (!authLoading && isLegacyRedirect) { - const base = (proxyBaseUrl || "") + "/ui"; - router.replace(`${base}/${LEGACY_REDIRECTS[page]}`); + router.replace(migratedHref(MIGRATED_PAGES[page])); } }, [authLoading, isLegacyRedirect, page, router]); @@ -265,6 +221,9 @@ function CreateKeyPageContent() { // Fetch in-product nudges configuration from backend useEffect(() => { + if (nudgesDisabled) { + return; + } if (accessToken && token) { (async () => { try { @@ -284,7 +243,7 @@ function CreateKeyPageContent() { } })(); } - }, [accessToken, token]); + }, [accessToken, token, nudgesDisabled]); // Auto-dismiss survey prompt after 15 seconds useEffect(() => { @@ -349,14 +308,26 @@ function CreateKeyPageContent() { } return ( - }> - - - {invitation_id ? ( + <> + {invitation_id ? ( + + ) : ( + <> + {page == "api-keys" ? ( - ) : ( -
- + ) : page == "llm-playground" ? ( + + ) : page == "users" ? ( + + ) : page == "teams" ? ( + + ) : page == "organizations" ? ( + + ) : page == "admin-panel" ? ( + + ) : page == "logging-and-alerts" ? ( + + ) : page == "budgets" ? ( + + ) : page == "guardrails" ? ( + + ) : page == "policies" ? ( + + ) : page == "agents" ? ( + + ) : page == "prompts" ? ( + + ) : page == "transform-request" ? ( + + ) : page == "router-settings" ? ( + + ) : page == "ui-theme" ? ( + + ) : page == "cost-tracking" ? ( + + ) : page == "model-hub-table" ? ( + isAdminRole(userRole) ? ( + -
-
- -
- {page == "api-keys" ? ( - - ) : page == "models" ? ( - - ) : page == "llm-playground" ? ( - - ) : page == "users" ? ( - - ) : page == "teams" ? ( - - ) : page == "organizations" ? ( - - ) : page == "admin-panel" ? ( - - ) : page == "api_ref" || page == "api-reference" ? ( - - ) : page == "logging-and-alerts" ? ( - - ) : page == "budgets" ? ( - - ) : page == "guardrails" ? ( - - ) : page == "policies" ? ( - - ) : page == "agents" ? ( - - ) : page == "prompts" ? ( - - ) : page == "transform-request" ? ( - - ) : page == "router-settings" ? ( - - ) : page == "ui-theme" ? ( - - ) : page == "cost-tracking" ? ( - - ) : page == "model-hub-table" ? ( - isAdminRole(userRole) ? ( - - ) : ( - - ) - ) : page == "caching" ? ( - - ) : page == "pass-through-settings" ? ( - - ) : page == "logs" ? ( - - ) : page == "mcp-servers" ? ( - - ) : page == "search-tools" ? ( - - ) : page == "tag-management" ? ( - - ) : page == "skills" || page == "claude-code-plugins" ? ( - - ) : page == "access-groups" ? ( - - ) : page == "projects" ? ( - - ) : page == "vector-stores" ? ( - - ) : page == "tool-policies" ? ( - - ) : page == "workflows" ? ( - - ) : page == "memory" ? ( - - ) : page == "guardrails-monitor" ? ( - - ) : page == "new_usage" ? ( - - ) : ( - - )} -
- - {/* Survey Components */} - - - - {/* Claude Code Components */} - - -
+ ) : ( + + ) + ) : page == "caching" ? ( + + ) : page == "pass-through-settings" ? ( + + ) : page == "logs" ? ( + + ) : page == "mcp-servers" ? ( + + ) : page == "search-tools" ? ( + + ) : page == "tag-management" ? ( + + ) : page == "skills" || page == "claude-code-plugins" ? ( + + ) : page == "access-groups" ? ( + + ) : page == "projects" ? ( + + ) : page == "vector-stores" ? ( + + ) : page == "tool-policies" ? ( + + ) : page == "workflows" ? ( + + ) : page == "memory" ? ( + + ) : page == "guardrails-monitor" ? ( + + ) : page == "new_usage" ? ( + + ) : ( + )} - - - + + {/* Survey Components */} + + + + {/* Claude Code Components */} + + + + )} + ); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/virtual-keys/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/virtual-keys/page.tsx deleted file mode 100644 index 226616474eb..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/virtual-keys/page.tsx +++ /dev/null @@ -1,47 +0,0 @@ -"use client"; - -import { useState } from "react"; -import useKeyList from "@/components/key_team_helpers/key_list"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import UserDashboard from "@/components/user_dashboard"; -import useTeams from "@/app/(dashboard)/hooks/useTeams"; -import { Organization } from "@/components/networking"; - -const VirtualKeysPage = () => { - const { accessToken, userRole, userId, premiumUser, userEmail } = useAuthorized(); - const { teams, setTeams } = useTeams(); - const [createClicked, setCreateClicked] = useState(false); - const [organizations, setOrganizations] = useState([]); - - const { keys, isLoading, error, pagination, refresh, setKeys } = useKeyList({ - selectedKeyAlias: null, - currentOrg: null, - accessToken: accessToken || "", - createClicked, - }); - - const addKey = (data: any) => { - setKeys((prevData) => (prevData ? [...prevData, data] : [data])); - setCreateClicked(() => !createClicked); - }; - - return ( - {}} - setUserEmail={() => {}} - setTeams={setTeams} - setKeys={setKeys} - premiumUser={premiumUser} - organizations={organizations} - addKey={addKey} - createClicked={createClicked} - /> - ); -}; - -export default VirtualKeysPage; diff --git a/ui/litellm-dashboard/src/components/AIHub/ClaudeCodeMarketplaceTab.tsx b/ui/litellm-dashboard/src/components/AIHub/ClaudeCodeMarketplaceTab.tsx deleted file mode 100644 index 762b0836921..00000000000 --- a/ui/litellm-dashboard/src/components/AIHub/ClaudeCodeMarketplaceTab.tsx +++ /dev/null @@ -1,140 +0,0 @@ -import { SearchOutlined } from "@ant-design/icons"; -import { Card, Tab, TabGroup, TabList, TabPanel, TabPanels, Text } from "@tremor/react"; -import { Input } from "antd"; -import React, { useEffect, useMemo, useState } from "react"; -import { extractCategories, filterPluginsByCategory, filterPluginsBySearch } from "../claude_code_plugins/helpers"; -import { MarketplaceResponse } from "../claude_code_plugins/types"; -import { ModelDataTable } from "../model_dashboard/table"; -import NotificationsManager from "../molecules/notifications_manager"; -import { getClaudeCodeMarketplace } from "../networking"; -import { getMarketplaceTableColumns } from "./marketplace_table_columns"; - -interface ClaudeCodeMarketplaceTabProps { - publicPage?: boolean; -} - -const ClaudeCodeMarketplaceTab: React.FC = ({ publicPage = false }) => { - const [marketplaceData, setMarketplaceData] = useState(null); - const [isLoading, setIsLoading] = useState(true); - const [searchTerm, setSearchTerm] = useState(""); - const [selectedCategoryIndex, setSelectedCategoryIndex] = useState(0); - - useEffect(() => { - fetchMarketplace(); - }, []); - - const fetchMarketplace = async () => { - setIsLoading(true); - try { - const data: MarketplaceResponse = await getClaudeCodeMarketplace(); - console.log("Claude Code marketplace:", data); - setMarketplaceData(data); - } catch (error) { - console.error("Error fetching marketplace:", error); - } finally { - setIsLoading(false); - } - }; - - const copyToClipboard = (text: string) => { - navigator.clipboard.writeText(text); - NotificationsManager.success("Copied to clipboard!"); - }; - - // Extract unique categories from plugins - const categories = useMemo(() => { - if (!marketplaceData) return ["All"]; - return extractCategories(marketplaceData.plugins); - }, [marketplaceData]); - - // Get selected category name - const selectedCategory = categories[selectedCategoryIndex] || "All"; - - // Filter plugins by search and category - const filteredPlugins = useMemo(() => { - if (!marketplaceData) return []; - - let plugins = marketplaceData.plugins; - - // Apply category filter - plugins = filterPluginsByCategory(plugins, selectedCategory); - - // Apply search filter - plugins = filterPluginsBySearch(plugins, searchTerm); - - return plugins; - }, [marketplaceData, selectedCategory, searchTerm]); - - const columns = useMemo(() => getMarketplaceTableColumns(copyToClipboard, publicPage), [publicPage]); - - if (!marketplaceData && !isLoading) { - return ( - -
- Failed to load marketplace. Please try again later. -
- - ); - } - - return ( -
- {/* Search Bar */} -
- } - value={searchTerm} - onChange={(e) => setSearchTerm(e.target.value)} - allowClear - size="large" - /> -
- - {/* Category Tabs */} - - - {categories.map((category) => { - // Count plugins in this category - const categoryPlugins = filterPluginsByCategory(marketplaceData?.plugins || [], category); - const count = filterPluginsBySearch(categoryPlugins, searchTerm).length; - - return ( - - {category} {count > 0 && `(${count})`} - - ); - })} - - - - {categories.map((category) => ( - - - {/* Plugin Table */} - - - - {/* Footer Info */} -
- - Showing {filteredPlugins.length} of {marketplaceData?.plugins.length || 0} plugin - {marketplaceData?.plugins.length !== 1 ? "s" : ""} - {searchTerm && ` matching "${searchTerm}"`} - {selectedCategory !== "All" && ` in ${selectedCategory}`} - -
- - ))} - - -
- ); -}; - -export default ClaudeCodeMarketplaceTab; diff --git a/ui/litellm-dashboard/src/components/AIHub/marketplace_table_columns.tsx b/ui/litellm-dashboard/src/components/AIHub/marketplace_table_columns.tsx deleted file mode 100644 index fa383197b3d..00000000000 --- a/ui/litellm-dashboard/src/components/AIHub/marketplace_table_columns.tsx +++ /dev/null @@ -1,172 +0,0 @@ -import { ColumnDef } from "@tanstack/react-table"; -import { Button, Badge, Text } from "@tremor/react"; -import { Tooltip } from "antd"; -import { CopyOutlined } from "@ant-design/icons"; -import { MarketplacePluginEntry } from "@/components/claude_code_plugins/types"; -import { - formatInstallCommand, - getCategoryBadgeColor, - getSourceDisplayText, -} from "@/components/claude_code_plugins/helpers"; - -export const getMarketplaceTableColumns = ( - copyToClipboard: (text: string) => void, - publicPage: boolean = false, -): ColumnDef[] => { - const allColumns: ColumnDef[] = [ - { - header: "Plugin Name", - accessorKey: "name", - enableSorting: true, - sortingFn: "alphanumeric", - cell: ({ row }) => { - const plugin = row.original; - const installCommand = formatInstallCommand(plugin); - - return ( -
-
- {plugin.name} - - copyToClipboard(installCommand)} - className="cursor-pointer text-gray-500 hover:text-blue-500 text-xs" - /> - -
- {/* Show description on mobile */} -
- {plugin.description || "No description"} -
-
- ); - }, - }, - { - header: "Description", - accessorKey: "description", - enableSorting: true, - sortingFn: "alphanumeric", - cell: ({ row }) => { - const plugin = row.original; - - return {plugin.description || "-"}; - }, - meta: { - className: "hidden md:table-cell", - }, - }, - { - header: "Version", - accessorKey: "version", - enableSorting: true, - sortingFn: "alphanumeric", - cell: ({ row }) => { - const plugin = row.original; - - return plugin.version ? ( - - v{plugin.version} - - ) : ( - - - ); - }, - meta: { - className: "hidden lg:table-cell", - }, - }, - { - header: "Category", - accessorKey: "category", - enableSorting: true, - sortingFn: "alphanumeric", - cell: ({ row }) => { - const plugin = row.original; - const badgeColor = getCategoryBadgeColor(plugin.category); - - return plugin.category ? ( - - {plugin.category} - - ) : ( - - Uncategorized - - ); - }, - meta: { - className: "hidden lg:table-cell", - }, - }, - { - header: "Source", - accessorKey: "source", - enableSorting: false, - cell: ({ row }) => { - const plugin = row.original; - const sourceText = getSourceDisplayText(plugin.source); - - return {sourceText}; - }, - meta: { - className: "hidden xl:table-cell", - }, - }, - { - header: "Keywords", - accessorKey: "keywords", - enableSorting: false, - cell: ({ row }) => { - const plugin = row.original; - const keywords = plugin.keywords?.slice(0, 3) || []; - const remaining = (plugin.keywords?.length || 0) - 3; - - return ( -
- {keywords.map((keyword, index) => ( - - {keyword} - - ))} - {remaining > 0 && ( - - +{remaining} - - )} -
- ); - }, - meta: { - className: "hidden xl:table-cell", - }, - }, - { - header: "Install Command", - id: "install_command", - enableSorting: false, - cell: ({ row }) => { - const plugin = row.original; - const installCommand = formatInstallCommand(plugin); - - return ( -
- - {installCommand} - - -
- ); - }, - }, - ]; - - return allColumns; -}; diff --git a/ui/litellm-dashboard/src/components/Projects/types.ts b/ui/litellm-dashboard/src/components/Projects/types.ts deleted file mode 100644 index 51429902dff..00000000000 --- a/ui/litellm-dashboard/src/components/Projects/types.ts +++ /dev/null @@ -1,14 +0,0 @@ -export interface Project { - id: string; - name: string; - description: string; - teamId: string; - teamAlias: string; - models: string[]; - status: "active" | "blocked"; - spend: number; - createdAt: string; - createdBy: string; - updatedAt: string; - updatedBy: string; -} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index 9ce0d908838..25865c48f9b 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -26,6 +26,7 @@ export default function UISettings() { const allowVectorStoresTeamAdminsProperty = schema?.properties?.allow_vector_stores_for_team_admins; const scopeUserSearchProperty = schema?.properties?.scope_user_search_to_org; const disableCustomApiKeysProperty = schema?.properties?.disable_custom_api_keys; + const disableUINudgesProperty = schema?.properties?.disable_ui_nudges; const values = data?.values ?? {}; const isDisabledForInternalUsers = Boolean(values.disable_model_add_for_internal_users); const isDisabledTeamAdminDeleteTeamUser = Boolean(values.disable_team_admin_delete_team_user); @@ -60,6 +61,20 @@ export default function UISettings() { ); }; + const handleToggleDisableUINudges = (checked: boolean) => { + updateSettings( + { disable_ui_nudges: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + const handleUpdatePageVisibility = (settings: { enabled_ui_pages_internal_users: string[] | null }) => { updateSettings(settings, { onSuccess: () => { @@ -451,6 +466,26 @@ export default function UISettings() { + {/* Disable in-product UI nudges */} + + + + Disable UI nudges + + {disableUINudgesProperty?.description ?? + "If true, suppresses in-product UI nudges (survey and Claude Code feedback popups) for all users."} + + + + + + {/* Page Visibility for Internal Users */} void; - accessToken: string | null; - onAgentUpdated: () => void; - isAdmin: boolean; - onAgentClick: (agentId: string) => void; -} - -const AgentTable: React.FC = ({ - agentsList, - isLoading, - onDeleteClick, - accessToken, - onAgentUpdated, - isAdmin, - onAgentClick, -}) => { - const [sorting, setSorting] = useState([{ id: "created_at", desc: true }]); - - const formatDate = (dateString?: string) => { - if (!dateString) return "-"; - const date = new Date(dateString); - return date.toLocaleString(); - }; - - const copyToClipboard = (text: string) => { - navigator.clipboard.writeText(text); - }; - - const columns: ColumnDef[] = [ - { - header: "Agent Name", - accessorKey: "agent_name", - cell: ({ row }) => { - const agent = row.original; - const name = agent.agent_name || ""; - return ( -
- - - - - { - e.stopPropagation(); - copyToClipboard(agent.agent_id); - }} - className="cursor-pointer text-gray-500 hover:text-blue-500 text-xs" - /> - -
- ); - }, - }, - { - header: "Description", - accessorKey: "agent_card_params.description", - cell: ({ row }) => { - const description = row.original.agent_card_params?.description || "No description"; - return {description}; - }, - }, - { - header: "Created At", - accessorKey: "created_at", - cell: ({ row }) => { - const agent = row.original; - return ( - - {formatDate(agent.created_at)} - - ); - }, - }, - ...(isAdmin - ? [ - { - header: "Actions", - id: "actions", - enableSorting: false, - cell: ({ row }: any) => { - const agent = row.original; - - return ( -
- -
- ); - }, - }, - ] - : []), - ]; - - const table = useReactTable({ - data: agentsList, - columns, - state: { - sorting, - }, - onSortingChange: setSorting, - getCoreRowModel: getCoreRowModel(), - getSortedRowModel: getSortedRowModel(), - enableSorting: true, - }); - - return ( -
-
- - - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - -
-
- {header.isPlaceholder ? null : flexRender(header.column.columnDef.header, header.getContext())} -
-
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
-
- - ))} - - ))} - - - {isLoading ? ( - - -
-

Loading...

-
- - - ) : agentsList && agentsList.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - - {row.getVisibleCells().map((cell) => ( - - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No agents found. Create one to get started.

-
- - - )} - -
-
-
- ); -}; - -export default AgentTable; diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/plugin_info.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/plugin_info.tsx deleted file mode 100644 index 55347025201..00000000000 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/plugin_info.tsx +++ /dev/null @@ -1,313 +0,0 @@ -import { CopyOutlined } from "@ant-design/icons"; -import { ArrowLeftIcon, ExternalLinkIcon } from "@heroicons/react/outline"; -import { Badge, Button, Card, Grid, Text, Title } from "@tremor/react"; -import { Spin, Switch, Tooltip } from "antd"; -import React, { useEffect, useState } from "react"; -import NotificationsManager from "../molecules/notifications_manager"; -import { disableClaudeCodePlugin, enableClaudeCodePlugin, getClaudeCodePluginDetails } from "../networking"; -import { - formatDateString, - formatInstallCommand, - getCategoryBadgeColor, - getSourceDisplayText, - getSourceLink, -} from "./helpers"; -import { Plugin } from "./types"; - -interface PluginInfoViewProps { - pluginId: string; - onClose: () => void; - accessToken: string | null; - isAdmin: boolean; - onPluginUpdated: () => void; -} - -const PluginInfoView: React.FC = ({ - pluginId, - onClose, - accessToken, - isAdmin, - onPluginUpdated, -}) => { - const [plugin, setPlugin] = useState(null); - const [isLoading, setIsLoading] = useState(true); - const [isToggling, setIsToggling] = useState(false); - - useEffect(() => { - fetchPluginInfo(); - }, [pluginId, accessToken]); - - const fetchPluginInfo = async () => { - if (!accessToken) return; - - setIsLoading(true); - try { - // The backend expects plugin name, not ID - // We'll need to find the plugin by ID from the list - // For now, assume pluginId is actually the plugin name - const data = await getClaudeCodePluginDetails(accessToken, pluginId as string); - setPlugin(data.plugin); - } catch (error) { - console.error("Error fetching plugin info:", error); - NotificationsManager.error("Failed to load plugin information"); - } finally { - setIsLoading(false); - } - }; - - const handleToggleEnabled = async () => { - if (!accessToken || !plugin) return; - - setIsToggling(true); - try { - if (plugin.enabled) { - await disableClaudeCodePlugin(accessToken, plugin.name); - NotificationsManager.success(`Plugin "${plugin.name}" disabled`); - } else { - await enableClaudeCodePlugin(accessToken, plugin.name); - NotificationsManager.success(`Plugin "${plugin.name}" enabled`); - } - onPluginUpdated(); - fetchPluginInfo(); - } catch (error) { - NotificationsManager.error("Failed to toggle plugin status"); - } finally { - setIsToggling(false); - } - }; - - const copyToClipboard = (text: string) => { - navigator.clipboard.writeText(text); - NotificationsManager.success("Copied to clipboard!"); - }; - - if (isLoading) { - return ( -
- -
- ); - } - - if (!plugin) { - return ( -
-

Plugin not found

- -
- ); - } - - const installCommand = formatInstallCommand(plugin); - const sourceLink = getSourceLink(plugin.source); - const categoryBadgeColor = getCategoryBadgeColor(plugin.category); - - return ( -
- {/* Header with Back Button */} -
- -

{plugin.name}

- {plugin.version && ( - - v{plugin.version} - - )} - {plugin.category && ( - - {plugin.category} - - )} - - {plugin.enabled ? "Enabled" : "Disabled"} - -
- - {/* Install Command */} - -
-
- Install Command -
{installCommand}
-
- - - -
- - - {/* Plugin Details */} - - Plugin Details - - {/* Plugin ID */} -
- Plugin ID -
- {plugin.id} - copyToClipboard(plugin.id)} - /> -
-
- - {/* Name */} -
- Name - {plugin.name} -
- - {/* Version */} -
- Version - {plugin.version || "N/A"} -
- - {/* Source */} -
- Source -
- {getSourceDisplayText(plugin.source)} - {sourceLink && ( - - - - )} -
-
- - {/* Category */} -
- Category -
- {plugin.category ? ( - - {plugin.category} - - ) : ( - Uncategorized - )} -
-
- - {/* Enabled Status */} - {isAdmin && ( -
- Status -
- - - {plugin.enabled - ? "Plugin is enabled and visible in marketplace" - : "Plugin is disabled and hidden from marketplace"} - -
-
- )} - - - - {/* Description */} - {plugin.description && ( - - Description - {plugin.description} - - )} - - {/* Keywords */} - {plugin.keywords && plugin.keywords.length > 0 && ( - - Keywords -
- {plugin.keywords.map((keyword, index) => ( - - {keyword} - - ))} -
- - )} - - {/* Author Information */} - {plugin.author && ( - - Author Information - - {plugin.author.name && ( -
- Name - {plugin.author.name} -
- )} - {plugin.author.email && ( - - )} - - - )} - - {/* Additional Links */} - {plugin.homepage && ( - - Homepage - - {plugin.homepage} - - - - )} - - {/* Timestamps */} - - Metadata - -
- Created At - {formatDateString(plugin.created_at)} -
-
- Updated At - {formatDateString(plugin.updated_at)} -
- {plugin.created_by && ( -
- Created By - {plugin.created_by} -
- )} - - -
- ); -}; - -export default PluginInfoView; diff --git a/ui/litellm-dashboard/src/components/guardrails.tsx b/ui/litellm-dashboard/src/components/guardrails.tsx index 4df3f93c275..4d1d3f65e3e 100644 --- a/ui/litellm-dashboard/src/components/guardrails.tsx +++ b/ui/litellm-dashboard/src/components/guardrails.tsx @@ -133,7 +133,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole return (
({ + createGuardrailCall: vi.fn(), + getGuardrailProviderSpecificParams: vi.fn().mockResolvedValue({}), + getGuardrailUISettings: vi.fn().mockResolvedValue({}), + modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), +})); + +const renderForm = () => { + const onClose = vi.fn(); + renderWithProviders(); + return { onClose }; +}; + +describe("AddGuardrailForm close behavior", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("does not close when the user clicks outside the modal on the mask", () => { + const { onClose } = renderForm(); + expect(screen.getByText("Create guardrail")).toBeInTheDocument(); + + const wrap = document.querySelector(".ant-modal-wrap") as HTMLElement; + expect(wrap).toBeTruthy(); + fireEvent.mouseDown(wrap); + fireEvent.mouseUp(wrap); + fireEvent.click(wrap); + + expect(onClose).not.toHaveBeenCalled(); + }); + + it("closes when the user clicks the explicit close button", () => { + const { onClose } = renderForm(); + fireEvent.click(screen.getByRole("button", { name: "✕" })); + expect(onClose).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx index 05a65d1c9e5..41a2dd7f67c 100644 --- a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx @@ -1144,6 +1144,7 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a title={null} open={visible} onCancel={handleClose} + maskClosable={false} footer={null} width={1000} closable={false} diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 43a28d94046..f4946b81d68 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -43,31 +43,9 @@ import { import NewBadge from "./common_components/NewBadge"; import type { Organization } from "./networking"; import UsageIndicator from "./UsageIndicator"; -import { serverRootPath } from "./networking"; +import { MIGRATED_PAGES, migratedHref, legacyPageHref } from "@/utils/migratedPages"; const { Sider } = Layout; -/** - * Pages migrated to path-based routing under (dashboard)/. - * Key = legacy page id, Value = route segment. - * Keep in sync with MIGRATED_PAGES in (dashboard)/layout.tsx. - */ -const MIGRATED_PAGES: Record = {}; - -/** Build an absolute href for a migrated page, respecting base URL + serverRootPath. */ -function migratedHref(routeSegment: string): string { - const raw = process.env.NEXT_PUBLIC_BASE_URL ?? ""; - const trimmed = raw.replace(/^\/+|\/+$/g, ""); - let base = trimmed ? `/${trimmed}/` : "/"; - - if (serverRootPath && serverRootPath !== "/") { - const cleanRoot = serverRootPath.replace(/\/+$/, ""); - const cleanBase = base.replace(/^\/+/, ""); - base = `${cleanRoot}/${cleanBase}`; - } - - return `${base}${routeSegment}`; -} - // Define the props type interface SidebarProps { setPage: (page: string) => void; @@ -440,18 +418,9 @@ const Sidebar: React.FC = ({ // Check if user is a team admin for any team const isTeamAdmin = useMemo(() => isUserTeamAdminForAnyTeam(teams ?? null, userId ?? ""), [teams, userId]); - // Navigate to page helper - const navigateToPage = (page: string) => { - // For migrated pages, just call setPage — the parent layout handles routing - if (MIGRATED_PAGES[page]) { - setPage(page); - return; - } - const newSearchParams = new URLSearchParams(window.location.search); - newSearchParams.set("page", page); - window.history.pushState(null, "", `?${newSearchParams.toString()}`); - setPage(page); - }; + // The parent (legacy root page or dashboard layout) owns navigation for both + // migrated and legacy pages; the sidebar only reports the selected page. + const navigateToPage = (page: string) => setPage(page); // Wrap label in so every nav item supports right-click → "Open in new tab" // and Ctrl/Cmd+click to open in a new tab, while preserving SPA navigation for normal clicks. @@ -469,15 +438,8 @@ const Sidebar: React.FC = ({ ); } - // For migrated pages, generate a path-based href for right-click "Open in new tab" const migratedRoute = MIGRATED_PAGES[page]; - const href = migratedRoute - ? migratedHref(migratedRoute) - : (() => { - const params = new URLSearchParams(window.location.search); - params.set("page", page); - return `?${params.toString()}`; - })(); + const href = migratedRoute ? migratedHref(migratedRoute) : legacyPageHref(page); return ( ({ })); // Mutable holder so individual tests can simulate "Authorize & Fetch" having -// produced a token before submit. -const oauthHook = vi.hoisted(() => ({ tokenResponse: null as Record | null })); +// produced a token before submit, and inspect the reset wiring. +const oauthHook = vi.hoisted(() => ({ + tokenResponse: null as Record | null, + reset: vi.fn(), + onTokenReceived: null as ((token: Record | null) => void) | null, +})); vi.mock("@/hooks/useMcpOAuthFlow", () => ({ - useMcpOAuthFlow: () => ({ - startOAuthFlow: vi.fn(), - status: "idle", - error: null, - tokenResponse: oauthHook.tokenResponse, - }), + useMcpOAuthFlow: (opts: { onTokenReceived: (token: Record | null) => void }) => { + oauthHook.onTokenReceived = opts.onTokenReceived; + return { + startOAuthFlow: vi.fn(), + status: "idle", + error: null, + tokenResponse: oauthHook.tokenResponse, + reset: oauthHook.reset, + }; + }, })); vi.mock("./mcp_server_cost_config", () => ({ @@ -59,7 +67,9 @@ vi.mock("./mcp_tool_configuration", () => ({ })); vi.mock("./mcp_connection_status", () => ({ - default: () =>
, + default: ({ tools }: { tools?: any[] }) => ( +
+ ), })); vi.mock("./StdioConfiguration", () => ({ @@ -121,6 +131,7 @@ describe("CreateMCPServer", () => { beforeEach(() => { vi.clearAllMocks(); oauthHook.tokenResponse = null; + oauthHook.onTokenReceived = null; }); it("should render the modal with title when visible", () => { @@ -614,6 +625,100 @@ describe("CreateMCPServer", () => { expect(defaultProps.setModalVisible).toHaveBeenCalledWith(false); }); + + it("does not leak a previous server's OAuth token into the next add-server session", async () => { + const usedToken = (token: string) => + vi.mocked(networking.testMCPToolsListRequest).mock.calls.some((call) => call[2] === token); + + const { rerender } = render(); + + await selectAntOption("Transport Type", "Streamable HTTP"); + await waitFor(() => { + expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument(); + }); + await selectAntOption("Authentication", "OAuth"); + await waitFor(() => { + expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument(); + }); + + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://server-a.example.com/mcp" } }); + }); + + // Simulate "Authorize & Fetch Token" completing for server A. + await act(async () => { + oauthHook.onTokenReceived?.({ access_token: "stale-token-A", expires_in: 3600 }); + }); + + // Precondition: the freshly fetched token drives the tool preview for server A. + await waitFor(() => { + expect(usedToken("stale-token-A")).toBe(true); + }); + + // Parent hides the modal (Cancel / successful create both flip this prop). + rerender(); + + // The OAuth flow state (source of the "Token fetched" badge) is reset on close. + expect(oauthHook.reset).toHaveBeenCalled(); + + vi.mocked(networking.testMCPToolsListRequest).mockClear(); + + // Reopen for a brand-new server and enter a different URL without re-authorizing. + rerender(); + const reopenedUrlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(reopenedUrlInput, { target: { value: "https://server-b.example.com/mcp" } }); + }); + + // The previous server's token must never be replayed for the new session. + expect(usedToken("stale-token-A")).toBe(false); + }); + + it("clears the tool list and form fields when a parent dismisses the modal", async () => { + vi.mocked(networking.testMCPToolsListRequest).mockResolvedValue({ + tools: [{ name: "tool_a" }], + error: null, + }); + const toolCount = () => screen.getByTestId("mcp-connection-status").getAttribute("data-tool-count"); + + const { rerender } = render(); + + await selectAntOption("Transport Type", "Streamable HTTP"); + await waitFor(() => { + expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument(); + }); + await selectAntOption("Authentication", "OAuth"); + await waitFor(() => { + expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument(); + }); + + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://server-a.example.com/mcp" } }); + }); + await act(async () => { + oauthHook.onTokenReceived?.({ access_token: "stale-token-A", expires_in: 3600 }); + }); + + // Precondition: a tool list is shown for server A. + await waitFor(() => { + expect(toolCount()).toBe("1"); + }); + + // Parent dismisses the modal without routing through Cancel or create. + rerender(); + + // Stale tools are cleared even though neither handler ran. + await waitFor(() => { + expect(toolCount()).toBe("0"); + }); + + // Reopening starts clean: the URL the prior server left in the Ant form store is gone. + rerender(); + const reopenedUrlInput = screen.getByPlaceholderText("https://your-mcp-server.com") as HTMLInputElement; + expect(reopenedUrlInput.value).toBe(""); + }); }); describe("when stdio transport is selected", () => { diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 6a8dc353f21..ddcc9f65d38 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -134,6 +134,7 @@ const CreateMCPServer: React.FC = ({ status: oauthStatus, error: oauthError, tokenResponse: oauthTokenResponse, + reset: resetOAuthFlow, } = useMcpOAuthFlow({ accessToken, getCredentials: () => form.getFieldValue("credentials"), @@ -188,6 +189,7 @@ const CreateMCPServer: React.FC = ({ } }, onBeforeRedirect: persistCreateUiState, + flowSource: "create", }); React.useEffect(() => { @@ -553,12 +555,19 @@ const CreateMCPServer: React.FC = ({ } }, [formValues.server_name]); - // Clear formValues when modal closes to reset child components + // Clear form, tools, and OAuth state when the modal closes so a previous server's + // authorization, credentials, or tool list never bleed into the next "Add New MCP + // Server" session, including when a parent dismisses the modal without routing + // through handleCancel or handleCreate. React.useEffect(() => { if (!isModalVisible) { + form.resetFields(); setFormValues({}); + setOauthAccessToken(null); + clearTools(); + resetOAuthFlow(); } - }, [isModalVisible]); + }, [isModalVisible, form, clearTools, resetOAuthFlow]); const isAdmin = isAdminRole(userRole); @@ -1088,7 +1097,6 @@ const CreateMCPServer: React.FC = ({
void; -}> = ({ server, isLoadingHealth, isRechecking, onRecheck }) => { - const [isHovered, setIsHovered] = useState(false); - const status = server.status || "unknown"; - const lastCheck = server.last_health_check; - const error = server.health_check_error; - - if (isLoadingHealth || isRechecking) { - return ( - - - Checking - - ); - } - - const getStatusColor = (status: string) => { - switch (status) { - case "healthy": - return "text-green-700 bg-green-50 border border-green-200"; - case "unhealthy": - return "text-red-700 bg-red-50 border border-red-200"; - default: - return "text-gray-600 bg-gray-50 border border-gray-200"; - } - }; - - const getStatusIcon = (status: string) => { - switch (status) { - case "healthy": - return "✓"; - case "unhealthy": - return "✗"; - default: - return "?"; - } - }; - - const isClickable = !!onRecheck; - - const tooltipContent = ( -
-
Health Status: {status}
- {lastCheck &&
Last Check: {new Date(lastCheck).toLocaleString()}
} - {error && ( -
-
Error:
-
{error}
-
- )} - {!lastCheck && !error &&
No health check data available
} - {isClickable &&
Click to recheck
} -
- ); - - return ( - - setIsHovered(true)} - onMouseLeave={() => setIsHovered(false)} - onClick={isClickable ? () => onRecheck(server.server_id) : undefined} - > - {isHovered && isClickable ? "↻" : getStatusIcon(status)} - {isHovered && isClickable ? "Recheck" : status.charAt(0).toUpperCase() + status.slice(1)} - - - ); -}; - -export const mcpServerColumns = ( - userRole: string, - onView: (serverId: string) => void, - onEdit: (serverId: string) => void, - onDelete: (serverId: string) => void, - isLoadingHealth?: boolean, - onByokConnect?: (server: MCPServer) => void, - onRecheckHealth?: (serverId: string) => void, - recheckingServerIds?: Set, -): ColumnDef[] => [ - { - accessorKey: "server_id", - header: "Server ID", - enableSorting: true, - cell: ({ row }) => ( - - ), - }, - { - accessorKey: "server_name", - header: "Name", - enableSorting: true, - cell: ({ row }) => { - const logoUrl = row.original.mcp_info?.logo_url; - const name = row.original.server_name; - return ( -
- {logoUrl ? ( - {`${name { - (e.target as HTMLImageElement).style.display = "none"; - }} - /> - ) : null} - {name} -
- ); - }, - }, - { - accessorKey: "alias", - header: "Alias", - enableSorting: true, - }, - { - id: "url", - header: "URL", - cell: ({ row }) => { - const url = row.original.url; - if (!url) { - return —; - } - const { maskedUrl } = getMaskedAndFullUrl(url); - return {maskedUrl}; - }, - }, - { - accessorKey: "transport", - header: "Transport", - enableSorting: true, - cell: ({ row }) => { - const transport = row.original.transport || "http"; - const specPath = row.original.spec_path; - const displayTransport = specPath && transport !== "stdio" ? "OPENAPI" : transport; - const label = displayTransport.toUpperCase(); - return ( - - {label} - - ); - }, - }, - { - accessorKey: "auth_type", - header: "Auth Type", - enableSorting: true, - cell: ({ getValue }) => { - const authType = (getValue() as string) || "none"; - return ( - - {authType} - - ); - }, - }, - { - id: "health_status", - header: "Health Status", - cell: ({ row }) => ( - - ), - }, - { - id: "mcp_access_groups", - header: "Access Groups", - cell: ({ row }) => { - const groups = row.original.mcp_access_groups; - if (Array.isArray(groups) && groups.length > 0) { - if (typeof groups[0] === "string") { - const joined = groups.join(", "); - return ( - -
- - {groups[0]} - - {groups.length > 1 && +{groups.length - 1}} -
- - ); - } - } - return —; - }, - }, - { - id: "available_on_public_internet", - header: "Network Access", - cell: ({ row }) => { - const isPublic = row.original.available_on_public_internet; - return isPublic ? ( - - - Public - - ) : ( - - - Internal - - ); - }, - }, - { - header: "Created", - accessorKey: "created_at", - enableSorting: true, - sortingFn: "datetime", - cell: ({ row }) => { - const server = row.original; - if (!server.created_at) return —; - const date = new Date(server.created_at); - return ( - - {date.toLocaleDateString()} - - ); - }, - }, - { - header: "Updated", - accessorKey: "updated_at", - enableSorting: true, - sortingFn: "datetime", - cell: ({ row }) => { - const server = row.original; - if (!server.updated_at) return —; - const date = new Date(server.updated_at); - return ( - - {date.toLocaleDateString()} - - ); - }, - }, - { - id: "byok_credential", - header: "Credential", - cell: ({ row }) => { - const server = row.original; - if (!server.is_byok) { - return —; - } - if (server.has_user_credential) { - return ( -
- - Connected - - {onByokConnect && ( - - )} -
- ); - } - return onByokConnect ? ( - - ) : null; - }, - }, - { - id: "actions", - header: "Actions", - cell: ({ row }) => ( -
- - - - - - -
- ), - }, -]; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx index 4070a5dd1af..d7c1241b044 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx @@ -8,6 +8,7 @@ import NotificationsManager from "../molecules/notifications_manager"; vi.mock("../networking", () => ({ updateMCPServer: vi.fn(), listMCPTools: vi.fn().mockResolvedValue({ tools: [], error: null }), + storeMCPOAuthUserCredential: vi.fn().mockResolvedValue({}), })); vi.mock("../molecules/notifications_manager", () => ({ @@ -17,12 +18,13 @@ vi.mock("../molecules/notifications_manager", () => ({ }, })); +const mockOauth: { tokenResponse: any } = { tokenResponse: null }; vi.mock("@/hooks/useMcpOAuthFlow", () => ({ useMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle", error: null, - tokenResponse: null, + tokenResponse: mockOauth.tokenResponse, }), })); @@ -37,12 +39,19 @@ vi.mock("./MCPPermissionManagement", () => ({ vi.mock("./mcp_tool_configuration", () => ({ default: ({ existingAllowedTools, + externalTools, + externalError, onAllowedToolsChange, onToolAllowlistInteraction, onToolNameToDisplayNameChange, onToolNameToDescriptionChange, }: any) => ( -
+