scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge branch 'litellm_internal_staging' of https://github.com/BerriAI/litellm into litellm_fix_oauth_credential_forwarding

Browse source
This commit is contained in:
mateo-berri 2026-08-22 08:45:05 -07:00
commit 092449b32f
41 changed files with 1534 additions and 2044 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/litellm_core_utils/ptu_pricing.py
View file

@ -124,6 +124,14 @@ def ptu_identity_error(
return None
PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to")
def declares_ptu(model_info: Mapping[str, object]) -> bool:
"""Whether any PTU field is set here, including one too malformed to charge."""
return any(model_info.get(field) is not None for field in PTU_MODEL_INFO_FIELDS)
def ptu_config_error(model_info: Mapping[str, object], *, model_name: str | None = None) -> str | None:
"""Why this PTU configuration cannot be honoured, else None.

6
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -27,6 +27,7 @@ from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.litellm_core_utils.ptu_pricing import (
CUSTOM_PRICING_FIELDS,
PTU_EMPTIED_PRICING_FIELDS,
PTU_MODEL_INFO_FIELDS,
PTU_ZEROED_PRICING_FIELDS,
PTU_ZEROED_TABLE_FIELDS,
SEARCH_CONTEXT_SIZES,
@ -247,7 +248,6 @@ def _raise_on_strategy_router_write_violation(
)
_PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to")
_PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"})
@ -261,7 +261,7 @@ def _explicitly_cleared_ptu_fields(model_info: ModelInfo | None) -> frozenset[st
return frozenset()
return frozenset(
field
for field in _PTU_MODEL_INFO_FIELDS
for field in PTU_MODEL_INFO_FIELDS
if field in model_info.model_fields_set and getattr(model_info, field) is None
)
@ -294,7 +294,7 @@ def _raise_if_ptu_cost_attribution_disabled(incoming_model_info: Mapping[str, ob
"""
if is_ptu_cost_attribution_enabled():
return
supplied: Final = tuple(field for field in _PTU_MODEL_INFO_FIELDS if incoming_model_info.get(field) is not None)
supplied: Final = tuple(field for field in PTU_MODEL_INFO_FIELDS if incoming_model_info.get(field) is not None)
if not supplied:
return
raise HTTPException(

17
litellm/router.py
View file

@ -66,6 +66,8 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.ptu_pricing import (
PTU_COST_ATTRIBUTION_ENV_VAR,
declares_ptu,
is_ptu_cost_attribution_enabled,
ptu_config_error,
ptu_identity_error,
@ -8234,6 +8236,21 @@ class Router:
)
duplicate_ids: Final = frozenset(model_id for model_id in declared_ids if declared_ids.count(model_id) > 1)
ptu_declared: Final = tuple(
str(entry.get("model_name"))
for entry in original_model_list
if isinstance(entry.get("model_info"), dict)
and entry["model_info"].get("db_model") is not True
and declares_ptu(entry["model_info"])
)
if ptu_declared and not is_ptu_cost_attribution_enabled():
verbose_router_logger.warning(
"PTU fields are set on config.yaml deployment(s) %s, but PTU cost attribution is disabled, so no "
"flat cost accrues and this traffic is billed per token. Set %s=True to enable it",
", ".join(ptu_declared),
PTU_COST_ATTRIBUTION_ENV_VAR,
)
for model in original_model_list:
_model_name = model.pop("model_name")
_litellm_params = model.pop("litellm_params")

1
terraform/provider/.goreleaser.yml
View file

@ -72,6 +72,7 @@ signs:
- "--detach-sign"
- "${artifact}"
release:
prerelease: auto
extra_files:
- glob: 'terraform-registry-manifest.json'
name_template: '{{ .ProjectName }}_{{ .Version }}_manifest.json'

15
terraform/provider/CHANGELOG.md
View file

@ -2,11 +2,22 @@
All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
Up to `0.4.0` the provider had its own version line, cut from the headings in
this file. It now ships at the **LiteLLM version**, on every LiteLLM release
channel, built from the same commit as the proxy (see `RELEASING.md`). The
headings below no longer drive a release; they record what changed and which
LiteLLM line first carried it. A change that breaks existing configurations
or state must be called out loudly here, because the version number can no
longer signal it.
## [Unreleased]
### Changed
- **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying
## [0.4.0] - 2026-08-06
### Fixed

16
terraform/provider/README.md
View file

@ -6,6 +6,18 @@ This Terraform provider allows you to manage LiteLLM resources through Infrastru
This directory (`terraform/provider/` in [BerriAI/litellm](https://github.com/BerriAI/litellm)) is the source of truth for the provider. [BerriAI/terraform-provider-litellm](https://github.com/BerriAI/terraform-provider-litellm) is a thin release mirror that the public Terraform Registry ingests from; do not open PRs there. Changes land here, where CI builds the provider, runs its tests, and statically audits every endpoint the provider calls against the proxy's generated OpenAPI schema (`tools/endpointaudit/`), so the provider cannot drift from the LiteLLM API silently. Releases are published by mirroring this directory into the split repo and tagging it, which triggers the goreleaser workflow there (see `RELEASING.md`)
## Versioning
The provider version **is the LiteLLM version**. Every LiteLLM release (dev, rc and stable) publishes the provider at the same version as the proxy, built from the same commit, so `1.99.0` of the provider is the one that shipped with `1.99.0` of the proxy and was audited against that proxy's API. Pin the provider to the line your proxy runs:
```hcl
version = "~> 1.99.0"
```
Pre-release versions (`1.99.0-rc.1`, `1.99.0-dev.1`) are published too; Terraform only selects one when it is pinned exactly.
Versions `0.1.0` through `0.4.0` predate this scheme and sit on their own line. They stay in the registry, but **a `~> 0.4` constraint will never pick up another release**: re-pin to the LiteLLM version to keep receiving updates.
## Features
- Manage LiteLLM model configurations
@ -32,7 +44,7 @@ terraform {
required_providers {
litellm = {
source = "BerriAI/litellm"
version = "~> 0.1.1" #HERE UPDATE VERSION ACCORDINGLY
version = "~> 1.99.0" # the LiteLLM version your proxy runs
}
}
}
@ -218,6 +230,6 @@ This project is licensed under the Apache License 2.0 - see the [LICENSE](LICENS
- Always use environment variables or secure secret management solutions to handle sensitive information like API keys and AWS credentials.
- Refer to the comprehensive documentation in the `docs/` directory for detailed usage examples and configuration options.
- Make sure to keep your provider version updated for the latest features and bug fixes.
- Keep the provider version in step with the LiteLLM version your proxy runs; see [Versioning](#versioning).
- The provider now supports AWS cross-account access with `aws_session_name` and `aws_role_name` parameters in the model resource.
- All example configurations have been consolidated into the documentation for better organization and maintenance.

139
terraform/provider/RELEASING.md
View file

@ -4,7 +4,16 @@ This document describes the release process for the LiteLLM Terraform Provider.
## Overview
Releases are automated via GitHub Actions when a version tag is pushed. The workflow builds the provider for multiple platforms, signs the artifacts with GPG, and publishes them to GitHub Releases.
The provider is released **in lockstep with LiteLLM**: every LiteLLM release (dev, rc and stable) publishes the provider at the LiteLLM version, built from the same commit as the proxy. There is no separate provider release to cut.
The flow, end to end:
1. `BerriAI/project-releaser`'s release pipeline resolves the commit to release (`main` HEAD for dev; `main` HEAD or an operator-supplied SHA for rc/stable) and passes the release approval gate
2. Its componentized terraform job rsyncs `terraform/provider/` from that commit into `BerriAI/terraform-provider-litellm`, commits, and pushes the tag `v<litellm version>` (for example `v1.99.0`, `v1.99.0-rc.1`, `v1.99.0-dev.1`), alongside the `terraform-aws-litellm` / `terraform-google-litellm` module mirrors which get the same tag
3. The tag push triggers the mirror's own `Release` workflow (goreleaser): multi-platform build, GPG-signed checksums, GitHub release. It runs unattended; project-releaser does not wait for it
4. The public Terraform Registry ingests the GitHub release as provider version `<litellm version>`
`terraform/provider/` only exists from LiteLLM ~1.95, so a stable patch cut from an older line skips the provider and publishes only the modules.
## Prerequisites
@ -68,113 +77,26 @@ Before publishing to the Terraform Registry:
**Note**: The public key fingerprint must match the key used to sign the provider releases.
## Release Steps
## What a change needs
### 1. Prepare the Release
1. **Land it in `BerriAI/litellm`.** Open a PR against `litellm_internal_staging` with the source change and a `CHANGELOG.md` entry under `[Unreleased]`. CI runs `gofmt`, `go vet`, build, tests and the endpoint-drift audit. A change that breaks existing configurations or state must say so in the changelog: the version number cannot signal it any more
2. **Wait for the next LiteLLM release.** The nightly dev release carries it within a day; it reaches a stable version on the next stable cut
3. **Verify** (optional): the version appears at https://registry.terraform.io/providers/BerriAI/litellm and https://github.com/BerriAI/terraform-provider-litellm/releases. If the tag is on the mirror but there is no release, the goreleaser run failed: https://github.com/BerriAI/terraform-provider-litellm/actions
Before creating a release:
Locally, before opening the PR:
1. **Update CHANGELOG.md**
- Move items from `[Unreleased]` section to a new version section
- Follow [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) format
- Use [Semantic Versioning](https://semver.org/spec/v2.0.0.html) for version numbers
- Include all notable changes since the last release
```bash
make test
make build
```
Example:
```markdown
## [0.1.2] - 2026-02-20
## Out-of-band publish or recovery
### Added
- New feature description
Dispatch `Build and Publish Componentized Images + Chart` in `BerriAI/project-releaser` by hand with only `publish_terraform` enabled and the `git_ref` / `tag` of the release to (re)publish. The run waits on project-releaser's release approval, then mirrors and tags exactly as the pipeline does.
### Fixed
- Bug fix description
The mirror is push-only: do not commit or tag `BerriAI/terraform-provider-litellm` directly. The publish refuses to overwrite an existing tag; a version that failed in goreleaser is recovered by re-running the mirror's `Release` workflow for that tag, not by re-tagging.
### Changed
- Changed behavior description
```
2. **Verify tests pass**
```bash
make test
```
3. **Verify the build works locally**
```bash
make build
```
4. **Land the changes in BerriAI/litellm**
Open a PR to `BerriAI/litellm` updating `terraform/provider/CHANGELOG.md` (and any source changes) and merge it
### 2. Mirror and Tag via project-releaser
The provider source lives at `terraform/provider/` in `BerriAI/litellm`; `BerriAI/terraform-provider-litellm` is a thin release mirror. Do not commit or tag the mirror directly
Normally there is nothing to do here. `BerriAI/project-releaser`'s release pipeline runs the same check on every release except `adhoc`, nightly included: it reads the topmost released heading in `terraform/provider/CHANGELOG.md`, probes the mirror for `v<version>`, and dispatches `Publish Terraform provider` only when the changelog has moved ahead of what the mirror carries. Cutting the version heading in step 1 is therefore what releases the provider, and the next release picks it up, so the wait is a day rather than a week
Dispatch by hand only for an out-of-band release, or to recover a run that failed:
1. Go to `BerriAI/project-releaser` > **Actions** > `Publish Terraform provider`
2. Click **Run workflow**:
- `git_ref`: full 40-char commit SHA from `BerriAI/litellm` to release from
- `provider_version`: the new version without the `v` prefix (e.g. `0.3.0`)
- `dry_run`: optional; validates without pushing
Automatic or manual, the run waits on the `production-release` approval in `project-releaser`, then rsyncs `terraform/provider/` into the mirror repo, commits, and pushes tag `v<provider_version>`. That approval is the only one in the flow. The tag push triggers the mirror's `Release` workflow (goreleaser), which runs unattended
**Important**:
- Tags must follow the format: `v<MAJOR>.<MINOR>.<PATCH>` (e.g., `v0.1.2`, `v1.0.0`)
- The workflow refuses to overwrite an existing tag; publish a new version instead
### 3. Monitor the Release Workflow
1. Go to: https://github.com/BerriAI/terraform-provider-litellm/actions
2. Find the "Release" workflow run for your tag
3. Monitor the progress and check for any errors
The workflow will:
- Check out the code
- Set up Go
- Import the GPG key
- Run `go mod tidy`
- Build binaries for multiple platforms (Linux, macOS, Windows, FreeBSD)
- Create archives and checksums
- Sign the checksums with GPG
- Create a GitHub release
- Upload all artifacts
### 4. Verify the Release
After the workflow completes successfully:
1. **Check the GitHub Release**
- Go to: https://github.com/BerriAI/terraform-provider-litellm/releases
- Verify the release was created with the correct version
- Confirm all artifacts are present:
- Binary archives for each platform
- SHA256SUMS file
- SHA256SUMS.sig (GPG signature)
- terraform-registry-manifest.json
2. **Verify the signature** (optional)
```bash
# Download the checksums and signature
wget https://github.com/BerriAI/terraform-provider-litellm/releases/download/v0.1.2/terraform-provider-litellm_0.1.2_SHA256SUMS
wget https://github.com/BerriAI/terraform-provider-litellm/releases/download/v0.1.2/terraform-provider-litellm_0.1.2_SHA256SUMS.sig
# Verify the signature
gpg --verify terraform-provider-litellm_0.1.2_SHA256SUMS.sig terraform-provider-litellm_0.1.2_SHA256SUMS
```
### 5. Publish to Terraform Registry (Optional)
If this provider is published to the Terraform Registry:
1. The registry should automatically detect the new release via the GitHub webhook
2. If not, you may need to manually trigger a sync on the Terraform Registry dashboard
3. Verify the new version appears at: https://registry.terraform.io/providers/BerriAI/litellm/latest
The mirror's `.github/` directory (the `Release` workflow) is the one thing the rsync preserves, so a change to the goreleaser *workflow* is a direct PR on the mirror; a change to `.goreleaser.yml` itself lands here like any other source change.
## Troubleshooting
@ -207,21 +129,15 @@ If this provider is published to the Terraform Registry:
### Tag Already Exists
**Error**: The publish workflow refuses to push because the tag already exists on the mirror
**Error**: The publish job refuses to push because the tag already exists on the mirror
**Solution**: Tags are immutable by design. Re-run the workflow with a new patch version instead of deleting or moving an existing tag
**Solution**: Tags are immutable by design and the version is the LiteLLM version, so this means the provider was already mirrored for this release. If the registry is missing the version, re-run the mirror's `Release` workflow for the existing tag rather than re-tagging
## Version Numbering
This project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html):
The provider version is the LiteLLM version, verbatim: `X.Y.Z` for a stable release, `X.Y.Z-rc.N` for a release candidate and `X.Y.Z-dev.N` for a nightly. It says which proxy the provider shipped with and was audited against; it does not follow SemVer's break-signalling, so breaking changes are announced in `CHANGELOG.md` and the registry docs instead.
- **MAJOR** version (1.0.0): Incompatible API changes
- **MINOR** version (0.1.0): New functionality in a backward-compatible manner
- **PATCH** version (0.0.1): Backward-compatible bug fixes
For pre-1.0 releases:
- Breaking changes may occur in minor versions
- Patch versions should only contain bug fixes
Versions `0.1.0` to `0.4.0` predate this and remain in the registry on their own line. A `~> 0.4` constraint never receives another release.
## Security Considerations
@ -237,5 +153,4 @@ For pre-1.0 releases:
- [Terraform Provider Publishing](https://www.terraform.io/docs/registry/providers/publishing.html)
- [HashiCorp GPG Signing Requirements](https://www.terraform.io/docs/registry/providers/publishing.html#signing-releases)
- [GitHub Actions Secrets](https://docs.github.com/en/actions/security-guides/encrypted-secrets)
- [Semantic Versioning](https://semver.org/)
- [Keep a Changelog](https://keepachangelog.com/)

8
test-quality-budget.json
View file

@ -1,18 +1,18 @@
{
"TQ001": {
"limit": 746
"limit": 744
},
"TQ002": {
"limit": 742
},
"TQ003": {
"limit": 1078
"limit": 1068
},
"TQ004": {
"limit": 557
"limit": 469
},
"TQ005": {
"limit": 2810
"limit": 2436
},
"TQ006": {
"limit": 34

1
tests/e2e/llm_translation/endpoints_client.py
View file

@ -87,6 +87,7 @@ class RichMessagesRequest(BaseModel):
max_tokens: int = 64
system: list[TextBlock]
messages: list[RichMessage]
cache: dict[str, bool] = {"no-cache": True}
class CompletionsRequest(BaseModel):

2
tests/e2e/router/reliability_support.py
View file

@ -56,7 +56,7 @@ def chat_override(
json=ReliabilityChatBody(
model=model,
messages=[ChatMessage(role="user", content=content)],
max_tokens=16,
max_tokens=64,
stream=stream,
router_settings_override=override,
),

4
tests/e2e/router/test_reliability_fallbacks_e2e.py
View file

@ -49,7 +49,7 @@ class TestReliabilityFallbacks:
resources.defer(lambda: client.proxy.delete_model(model_id))
resp = chat_override(
client.proxy, scoped_key, primary, "say hi",
client.proxy, scoped_key, primary, f"say hi {unique_marker()}",
override=RouterSettingsOverride(fallbacks=[{primary: ["gpt-5.5"]}]),
)
_assert_served_by_fallback(resp)
@ -63,7 +63,7 @@ class TestReliabilityFallbacks:
resources.defer(lambda: client.proxy.delete_model(model_id))
resp = chat_override(
client.proxy, scoped_key, primary, "say hi",
client.proxy, scoped_key, primary, f"say hi {unique_marker()}",
override=RouterSettingsOverride(fallbacks=[{primary: ["gpt-5.5"]}]),
)
_assert_served_by_fallback(resp)

36
tests/test_litellm/integrations/datadog/test_datadog_cost_management.py
View file

@ -1,4 +1,3 @@
import os
import time
from unittest.mock import AsyncMock
@ -12,34 +11,13 @@ from litellm.types.utils import StandardLoggingPayload
@pytest.fixture
def clean_env():
# Save original env
original_api_key = os.environ.get("DD_API_KEY")
original_app_key = os.environ.get("DD_APP_KEY")
original_site = os.environ.get("DD_SITE")
# Set test env
os.environ["DD_API_KEY"] = "test_api_key"
os.environ["DD_APP_KEY"] = "test_app_key"
os.environ["DD_SITE"] = "test.datadoghq.com"
yield
# Restore original env
if original_api_key:
os.environ["DD_API_KEY"] = original_api_key
else:
del os.environ["DD_API_KEY"]
if original_app_key:
os.environ["DD_APP_KEY"] = original_app_key
else:
del os.environ["DD_APP_KEY"]
if original_site:
os.environ["DD_SITE"] = original_site
else:
del os.environ["DD_SITE"]
def clean_env(monkeypatch: pytest.MonkeyPatch) -> None:
for key, value in (
("DD_API_KEY", "test_api_key"),
("DD_APP_KEY", "test_app_key"),
("DD_SITE", "test.datadoghq.com"),
):
monkeypatch.setenv(key, value)
@pytest.mark.asyncio

30
tests/test_litellm/integrations/datadog/test_datadog_metrics.py
View file

@ -1,4 +1,3 @@
import os
import time
from datetime import datetime, timedelta
from unittest.mock import AsyncMock
@ -11,25 +10,16 @@ from litellm.types.utils import StandardLoggingPayload
@pytest.fixture
def clean_env():
"""Set test env vars and restore originals after test."""
keys = ["DD_API_KEY", "DD_APP_KEY", "DD_SITE", "DD_ENV", "DD_SERVICE", "DD_VERSION"]
originals = {k: os.environ.get(k) for k in keys}
os.environ["DD_API_KEY"] = "test_api_key"
os.environ["DD_APP_KEY"] = "test_app_key"
os.environ["DD_SITE"] = "test.datadoghq.com"
os.environ["DD_ENV"] = "test-env"
os.environ["DD_SERVICE"] = "test-service"
os.environ["DD_VERSION"] = "1.0.0"
yield
for k, v in originals.items():
if v is not None:
os.environ[k] = v
elif k in os.environ:
del os.environ[k]
def clean_env(monkeypatch: pytest.MonkeyPatch) -> None:
for key, value in (
("DD_API_KEY", "test_api_key"),
("DD_APP_KEY", "test_app_key"),
("DD_SITE", "test.datadoghq.com"),
("DD_ENV", "test-env"),
("DD_SERVICE", "test-service"),
("DD_VERSION", "1.0.0"),
):
monkeypatch.setenv(key, value)
@pytest.mark.asyncio

316
tests/test_litellm/integrations/test_s3_v2.py
View file

@ -751,7 +751,7 @@ async def test_strip_base64_mixed_nested_objects():
@pytest.mark.asyncio
async def test_s3_verify_false_handling():
async def test_s3_verify_false_handling(monkeypatch: pytest.MonkeyPatch):
"""
Test that s3_verify=False is properly handled and not treated as None.
@ -763,15 +763,19 @@ async def test_s3_verify_false_handling():
import litellm
# Set up s3_callback_params with s3_verify=False
litellm.s3_callback_params = {
"s3_bucket_name": "test-bucket",
"s3_endpoint_url": "https://localhost:443",
"s3_aws_access_key_id": "minioadmin",
"s3_aws_secret_access_key": "minioadmin",
"s3_region_name": "us-east-1",
"s3_verify": False, # This should NOT be ignored
"s3_use_ssl": False, # This should also NOT be ignored
}
monkeypatch.setattr(
litellm,
"s3_callback_params",
{
"s3_bucket_name": "test-bucket",
"s3_endpoint_url": "https://localhost:443",
"s3_aws_access_key_id": "minioadmin",
"s3_aws_secret_access_key": "minioadmin",
"s3_region_name": "us-east-1",
"s3_verify": False, # This should NOT be ignored
"s3_use_ssl": False, # This should also NOT be ignored
},
)
with patch("asyncio.create_task"):
with patch(
@ -801,12 +805,9 @@ async def test_s3_verify_false_handling():
"ssl_verify": False
}, f"Expected ssl_verify=False in params, got {call_kwargs.get('params')}"
# Clean up
litellm.s3_callback_params = None
@pytest.mark.asyncio
async def test_s3_verify_none_handling():
async def test_s3_verify_none_handling(monkeypatch: pytest.MonkeyPatch):
"""
Test that s3_verify=None uses default behavior.
"""
@ -815,12 +816,16 @@ async def test_s3_verify_none_handling():
import litellm
# Set up s3_callback_params without s3_verify
litellm.s3_callback_params = {
"s3_bucket_name": "test-bucket",
"s3_aws_access_key_id": "test-key",
"s3_aws_secret_access_key": "test-secret",
"s3_region_name": "us-east-1",
}
monkeypatch.setattr(
litellm,
"s3_callback_params",
{
"s3_bucket_name": "test-bucket",
"s3_aws_access_key_id": "test-key",
"s3_aws_secret_access_key": "test-secret",
"s3_region_name": "us-east-1",
},
)
with patch("asyncio.create_task"):
with patch(
@ -846,12 +851,9 @@ async def test_s3_verify_none_handling():
assert call_kwargs["params"].get("ssl_verify") is None
# Either params is None or params={'ssl_verify': None} is acceptable
# Clean up
litellm.s3_callback_params = None
@pytest.mark.asyncio
async def test_s3_verify_false_creates_httpx_client_with_verify_false():
async def test_s3_verify_false_creates_httpx_client_with_verify_false(monkeypatch: pytest.MonkeyPatch):
"""
Test that when s3_verify=False, the actual httpx client has verify=False.
@ -862,14 +864,18 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false():
import litellm
# Set up s3_callback_params with s3_verify=False
litellm.s3_callback_params = {
"s3_bucket_name": "test-bucket",
"s3_endpoint_url": "https://localhost:443",
"s3_aws_access_key_id": "minioadmin",
"s3_aws_secret_access_key": "minioadmin",
"s3_region_name": "us-east-1",
"s3_verify": False,
}
monkeypatch.setattr(
litellm,
"s3_callback_params",
{
"s3_bucket_name": "test-bucket",
"s3_endpoint_url": "https://localhost:443",
"s3_aws_access_key_id": "minioadmin",
"s3_aws_secret_access_key": "minioadmin",
"s3_region_name": "us-east-1",
"s3_verify": False,
},
)
with patch("asyncio.create_task"):
# Create logger - this creates the httpx client
@ -888,12 +894,9 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false():
httpx_client._verify is False
), f"Expected httpx client _verify=False, got {httpx_client._verify}"
# Clean up
litellm.s3_callback_params = None
@pytest.mark.asyncio
async def test_s3_verify_false_async_client():
async def test_s3_verify_false_async_client(monkeypatch: pytest.MonkeyPatch):
"""
Test that the async httpx client respects s3_verify=False.
"""
@ -903,14 +906,18 @@ async def test_s3_verify_false_async_client():
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
# Set up s3_callback_params with s3_verify=False
litellm.s3_callback_params = {
"s3_bucket_name": "test-bucket",
"s3_endpoint_url": "https://localhost:443",
"s3_aws_access_key_id": "minioadmin",
"s3_aws_secret_access_key": "minioadmin",
"s3_region_name": "us-east-1",
"s3_verify": False,
}
monkeypatch.setattr(
litellm,
"s3_callback_params",
{
"s3_bucket_name": "test-bucket",
"s3_endpoint_url": "https://localhost:443",
"s3_aws_access_key_id": "minioadmin",
"s3_aws_secret_access_key": "minioadmin",
"s3_region_name": "us-east-1",
"s3_verify": False,
},
)
with patch("asyncio.create_task"):
logger = S3Logger()
@ -945,9 +952,6 @@ async def test_s3_verify_false_async_client():
httpx_client._verify is False
), f"Expected async httpx client _verify=False, got {httpx_client._verify}"
# Clean up
litellm.s3_callback_params = None
@pytest.mark.asyncio
async def test_strip_base64_recursive_redaction():
@ -1169,26 +1173,22 @@ def test_create_s3_batch_logging_element_flat_key_for_arn_response_id():
# --------------------------------------------------------------
# params_source / s3_callback_params_override (audit-log decoupling)
# --------------------------------------------------------------
def test_s3_callback_params_override_uses_alternate_dict():
def test_s3_callback_params_override_uses_alternate_dict(monkeypatch):
"""`s3_callback_params_override` makes the logger read its config from
the override dict instead of `litellm.s3_callback_params`."""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {"s3_bucket_name": "normal-bucket"}
try:
logger = S3Logger(
s3_callback_params_override={
"s3_bucket_name": "audit-bucket",
"s3_path": "audit-prefix",
"s3_region_name": "us-west-2",
}
)
assert logger.s3_bucket_name == "audit-bucket"
assert logger.s3_path == "audit-prefix"
assert logger.s3_region_name == "us-west-2"
finally:
litellm.s3_callback_params = original
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-bucket"})
logger = S3Logger(
s3_callback_params_override={
"s3_bucket_name": "audit-bucket",
"s3_path": "audit-prefix",
"s3_region_name": "us-west-2",
}
)
assert logger.s3_bucket_name == "audit-bucket"
assert logger.s3_path == "audit-prefix"
assert logger.s3_region_name == "us-west-2"
def test_s3_callback_params_override_does_not_mutate_inputs(monkeypatch):
@ -1198,43 +1198,31 @@ def test_s3_callback_params_override_does_not_mutate_inputs(monkeypatch):
monkeypatch.setenv("MY_AUDIT_BUCKET", "resolved-bucket")
override = {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"}
original_global = litellm.s3_callback_params
litellm.s3_callback_params = {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"}
try:
logger = S3Logger(s3_callback_params_override=override)
assert logger.s3_bucket_name == "resolved-bucket"
assert override["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET"
assert (
litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET"
)
finally:
litellm.s3_callback_params = original_global
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"})
logger = S3Logger(s3_callback_params_override=override)
assert logger.s3_bucket_name == "resolved-bucket"
assert override["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET"
assert (
litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET"
)
def test_s3_callback_params_override_none_falls_back_to_global():
def test_s3_callback_params_override_none_falls_back_to_global(monkeypatch):
"""No override → behaves exactly as today (reads `litellm.s3_callback_params`)."""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {"s3_bucket_name": "from-global"}
try:
logger = S3Logger()
assert logger.s3_bucket_name == "from-global"
finally:
litellm.s3_callback_params = original
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "from-global"})
logger = S3Logger()
assert logger.s3_bucket_name == "from-global"
def test_s3_callback_params_override_empty_dict_is_opt_in():
def test_s3_callback_params_override_empty_dict_is_opt_in(monkeypatch):
"""An empty override dict skips the global entirely (env/IAM-only config)."""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {"s3_bucket_name": "from-global"}
try:
logger = S3Logger(s3_callback_params_override={})
assert logger.s3_bucket_name is None
finally:
litellm.s3_callback_params = original
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "from-global"})
logger = S3Logger(s3_callback_params_override={})
assert logger.s3_bucket_name is None
def _expected_content_md5(payload: dict) -> str:
@ -1374,20 +1362,20 @@ async def test_async_upload_sets_server_side_encryption_header_when_configured()
assert headers["x-amz-server-side-encryption"] == "aws:kms"
def test_s3_server_side_encryption_read_from_callback_params():
def test_s3_server_side_encryption_read_from_callback_params(monkeypatch):
"""s3_server_side_encryption can be configured via s3_callback_params."""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {
"s3_bucket_name": "from-global",
"s3_server_side_encryption": "aws:kms",
}
try:
logger = S3Logger()
assert logger.s3_server_side_encryption == "aws:kms"
finally:
litellm.s3_callback_params = original
monkeypatch.setattr(
litellm,
"s3_callback_params",
{
"s3_bucket_name": "from-global",
"s3_server_side_encryption": "aws:kms",
},
)
logger = S3Logger()
assert logger.s3_server_side_encryption == "aws:kms"
@pytest.mark.asyncio
@ -1505,21 +1493,21 @@ async def test_async_upload_omits_kms_key_id_header_when_not_configured():
assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers
def test_s3_sse_kms_key_id_read_from_callback_params():
def test_s3_sse_kms_key_id_read_from_callback_params(monkeypatch):
"""s3_sse_kms_key_id can be configured via s3_callback_params."""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {
"s3_bucket_name": "from-global",
"s3_server_side_encryption": "aws:kms",
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
}
try:
logger = S3Logger()
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id")
finally:
litellm.s3_callback_params = original
monkeypatch.setattr(
litellm,
"s3_callback_params",
{
"s3_bucket_name": "from-global",
"s3_server_side_encryption": "aws:kms",
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
},
)
logger = S3Logger()
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id")
@pytest.mark.asyncio
@ -1561,83 +1549,79 @@ async def test_async_upload_infers_aws_kms_when_only_key_id_set():
)
def test_s3_sse_kms_key_id_read_from_audit_override_params():
def test_s3_sse_kms_key_id_read_from_audit_override_params(monkeypatch):
"""The audit-log override path must honor s3_sse_kms_key_id too."""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {"s3_bucket_name": "normal-logs-bucket"}
try:
logger = S3Logger(
s3_callback_params_override={
"s3_bucket_name": "audit-logs-bucket",
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/audit-key-id",
}
)
assert logger.s3_bucket_name == "audit-logs-bucket"
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/audit-key-id")
finally:
litellm.s3_callback_params = original
monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-logs-bucket"})
logger = S3Logger(
s3_callback_params_override={
"s3_bucket_name": "audit-logs-bucket",
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/audit-key-id",
}
)
assert logger.s3_bucket_name == "audit-logs-bucket"
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/audit-key-id")
def test_kms_key_id_dropped_when_algorithm_is_not_kms():
def test_kms_key_id_dropped_when_algorithm_is_not_kms(monkeypatch):
"""
AES256 plus a KMS key id is an invalid S3 combination; the key id must be
dropped at init so uploads keep working instead of silently 400ing.
"""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {
"s3_bucket_name": "from-global",
"s3_server_side_encryption": "AES256",
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
}
try:
logger = S3Logger()
assert logger.s3_server_side_encryption == "AES256"
assert logger.s3_sse_kms_key_id is None
finally:
litellm.s3_callback_params = original
monkeypatch.setattr(
litellm,
"s3_callback_params",
{
"s3_bucket_name": "from-global",
"s3_server_side_encryption": "AES256",
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
},
)
logger = S3Logger()
assert logger.s3_server_side_encryption == "AES256"
assert logger.s3_sse_kms_key_id is None
def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued():
def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(monkeypatch):
"""
A YAML boolean in s3_server_side_encryption must not crash logger init and
must not discard the valid key id; aws:kms is inferred from the key id.
"""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {
"s3_bucket_name": "from-global",
"s3_server_side_encryption": True,
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
}
try:
logger = S3Logger()
assert logger.s3_server_side_encryption == "aws:kms"
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id")
finally:
litellm.s3_callback_params = original
monkeypatch.setattr(
litellm,
"s3_callback_params",
{
"s3_bucket_name": "from-global",
"s3_server_side_encryption": True,
"s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id",
},
)
logger = S3Logger()
assert logger.s3_server_side_encryption == "aws:kms"
assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id")
def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept():
def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(monkeypatch):
"""A mistyped key id (unquoted YAML number) must not disable the valid algorithm."""
import litellm
original = litellm.s3_callback_params
litellm.s3_callback_params = {
"s3_bucket_name": "from-global",
"s3_server_side_encryption": "aws:kms",
"s3_sse_kms_key_id": 12345,
}
try:
logger = S3Logger()
assert logger.s3_server_side_encryption == "aws:kms"
assert logger.s3_sse_kms_key_id is None
finally:
litellm.s3_callback_params = original
monkeypatch.setattr(
litellm,
"s3_callback_params",
{
"s3_bucket_name": "from-global",
"s3_server_side_encryption": "aws:kms",
"s3_sse_kms_key_id": 12345,
},
)
logger = S3Logger()
assert logger.s3_server_side_encryption == "aws:kms"
assert logger.s3_sse_kms_key_id is None
_ACCESS_KEY = "AKIAIOSFODNN7EXAMPLE"

240
tests/test_litellm/interactions/test_gemini_interactions_transformation.py
View file

@ -86,29 +86,21 @@ class TestValidateEnvironment:
assert headers["X-Custom"] == "value"
assert headers["x-goog-api-key"] == "test-key"
def test_api_revision_new_schema_by_default(self, config):
def test_api_revision_new_schema_by_default(self, config, monkeypatch: pytest.MonkeyPatch):
# Default: use_legacy_interactions_schema=False → new steps schema
original = litellm.use_legacy_interactions_schema
try:
litellm.use_legacy_interactions_schema = False
headers = config.validate_environment(
headers={}, model="gemini-2.5-flash", litellm_params=None
)
assert headers["Api-Revision"] == "2026-05-20"
finally:
litellm.use_legacy_interactions_schema = original
monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False)
headers = config.validate_environment(
headers={}, model="gemini-2.5-flash", litellm_params=None
)
assert headers["Api-Revision"] == "2026-05-20"
def test_api_revision_legacy_schema_when_flag_set(self, config):
def test_api_revision_legacy_schema_when_flag_set(self, config, monkeypatch: pytest.MonkeyPatch):
# Flag on → legacy outputs schema until June 8, 2026
original = litellm.use_legacy_interactions_schema
try:
litellm.use_legacy_interactions_schema = True
headers = config.validate_environment(
headers={}, model="gemini-2.5-flash", litellm_params=None
)
assert headers["Api-Revision"] == "2026-05-07"
finally:
litellm.use_legacy_interactions_schema = original
monkeypatch.setattr(litellm, "use_legacy_interactions_schema", True)
headers = config.validate_environment(
headers={}, model="gemini-2.5-flash", litellm_params=None
)
assert headers["Api-Revision"] == "2026-05-07"
class TestGetCompleteUrl:
@ -561,23 +553,19 @@ class TestInteractionOperationUrls:
class TestTransformRequestSchemaCoalescing:
"""Test new-schema request coalescing (Api-Revision: 2026-05-20)."""
def test_response_mime_type_folded_into_response_format(self, config):
original = litellm.use_legacy_interactions_schema
try:
litellm.use_legacy_interactions_schema = False
body = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="summarise",
optional_params={
"response_mime_type": "application/json",
"response_format": {"type": "object", "properties": {}},
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
finally:
litellm.use_legacy_interactions_schema = original
def test_response_mime_type_folded_into_response_format(self, config, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False)
body = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="summarise",
optional_params={
"response_mime_type": "application/json",
"response_format": {"type": "object", "properties": {}},
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
# response_mime_type must not appear as a top-level body key
assert "response_mime_type" not in body
@ -586,25 +574,21 @@ class TestTransformRequestSchemaCoalescing:
assert rf["mime_type"] == "application/json"
assert "schema" in rf
def test_image_config_moved_to_response_format(self, config):
original = litellm.use_legacy_interactions_schema
try:
litellm.use_legacy_interactions_schema = False
body = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="draw a sunset",
optional_params={
"generation_config": {
"temperature": 0.7,
"image_config": {"aspect_ratio": "1:1", "image_size": "1K"},
}
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
finally:
litellm.use_legacy_interactions_schema = original
def test_image_config_moved_to_response_format(self, config, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False)
body = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="draw a sunset",
optional_params={
"generation_config": {
"temperature": 0.7,
"image_config": {"aspect_ratio": "1:1", "image_size": "1K"},
}
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
# image_config removed from generation_config
assert "image_config" not in body.get("generation_config", {})
@ -613,95 +597,85 @@ class TestTransformRequestSchemaCoalescing:
assert rf["type"] == "image"
assert rf["aspect_ratio"] == "1:1"
def test_response_mime_type_skipped_when_response_format_is_list(self, config):
def test_response_mime_type_skipped_when_response_format_is_list(self, config, monkeypatch: pytest.MonkeyPatch):
"""Lists are already polymorphic; do not wrap them into schema."""
original = litellm.use_legacy_interactions_schema
try:
litellm.use_legacy_interactions_schema = False
rf_list = [
{"type": "text", "mime_type": "application/json"},
{"type": "image", "aspect_ratio": "1:1"},
]
body = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="multimodal",
optional_params={
"response_format": rf_list,
"response_mime_type": "application/json",
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
finally:
litellm.use_legacy_interactions_schema = original
monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False)
rf_list = [
{"type": "text", "mime_type": "application/json"},
{"type": "image", "aspect_ratio": "1:1"},
]
body = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="multimodal",
optional_params={
"response_format": rf_list,
"response_mime_type": "application/json",
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert body["response_format"] == rf_list
assert "response_mime_type" not in body
def test_image_config_appended_to_response_format_list_without_mutating_input(
self, config
self,
config,
monkeypatch: pytest.MonkeyPatch,
):
"""When response_format is already a list, image_config must not mutate optional_params."""
original = litellm.use_legacy_interactions_schema
try:
litellm.use_legacy_interactions_schema = False
text_rf = {"type": "text", "mime_type": "application/json"}
optional_params = {
"response_format": [text_rf],
"generation_config": {
"image_config": {"aspect_ratio": "16:9", "image_size": "2K"},
},
}
original_rf = optional_params["response_format"]
monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False)
text_rf = {"type": "text", "mime_type": "application/json"}
optional_params = {
"response_format": [text_rf],
"generation_config": {
"image_config": {"aspect_ratio": "16:9", "image_size": "2K"},
},
}
original_rf = optional_params["response_format"]
body = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="draw and summarise",
optional_params=optional_params,
litellm_params=GenericLiteLLMParams(),
headers={},
)
body = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="draw and summarise",
optional_params=optional_params,
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert optional_params["response_format"] is original_rf
assert len(optional_params["response_format"]) == 1
assert body["response_format"] == [
text_rf,
{"type": "image", "aspect_ratio": "16:9", "image_size": "2K"},
]
assert optional_params["response_format"] is original_rf
assert len(optional_params["response_format"]) == 1
assert body["response_format"] == [
text_rf,
{"type": "image", "aspect_ratio": "16:9", "image_size": "2K"},
]
# Retry must not append a second image entry into the caller's list.
body_retry = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="draw and summarise",
optional_params=optional_params,
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert len(optional_params["response_format"]) == 1
assert body_retry["response_format"] == body["response_format"]
finally:
litellm.use_legacy_interactions_schema = original
# Retry must not append a second image entry into the caller's list.
body_retry = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="draw and summarise",
optional_params=optional_params,
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert len(optional_params["response_format"]) == 1
assert body_retry["response_format"] == body["response_format"]
def test_legacy_schema_passes_fields_unchanged(self, config):
original = litellm.use_legacy_interactions_schema
try:
litellm.use_legacy_interactions_schema = True
body = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="hello",
optional_params={
"response_mime_type": "application/json",
"generation_config": {"image_config": {"aspect_ratio": "16:9"}},
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
finally:
litellm.use_legacy_interactions_schema = original
def test_legacy_schema_passes_fields_unchanged(self, config, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "use_legacy_interactions_schema", True)
body = config.transform_request(
model="gemini/gemini-2.5-flash",
agent=None,
input="hello",
optional_params={
"response_mime_type": "application/json",
"generation_config": {"image_config": {"aspect_ratio": "16:9"}},
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert body["response_mime_type"] == "application/json"
assert body["generation_config"]["image_config"]["aspect_ratio"] == "16:9"

209
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -1,6 +1,4 @@
import json
import os
import sys
import pytest
from fastapi.testclient import TestClient
@ -28,10 +26,6 @@ from litellm.types.utils import (
StandardBuiltInToolsParams,
)
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm.litellm_core_utils.llm_cost_calc.utils import (
PromptTokensDetailsResult,
TokenTypeCostBreakdown,
@ -44,13 +38,17 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
from litellm.types.utils import CacheCreationTokenDetails, Usage
def test_reasoning_tokens_no_price_set():
@pytest.fixture
def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
def test_reasoning_tokens_no_price_set(_local_model_cost_map):
# Use o1 - o1-mini was deprecated/renamed; o1 has same reasoning-token semantics
# (no separate output_cost_per_reasoning_token, so all completion tokens use output_cost_per_token)
model = "o1"
custom_llm_provider = "openai"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model_cost_map = litellm.model_cost[model]
usage = Usage(
completion_tokens=1578,
@ -87,11 +85,9 @@ def test_reasoning_tokens_no_price_set():
)
def test_reasoning_tokens_gemini():
def test_reasoning_tokens_gemini(_local_model_cost_map):
model = "gemini-2.5-flash"
custom_llm_provider = "gemini"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
usage = Usage(
completion_tokens=1578,
@ -132,12 +128,10 @@ def test_reasoning_tokens_gemini():
)
def test_reasoning_tokens_gemini_3_1_flash_lite():
def test_reasoning_tokens_gemini_3_1_flash_lite(_local_model_cost_map):
"""Test cost calculation for gemini-3.1-flash-lite-preview with reasoning tokens"""
model = "gemini-3.1-flash-lite-preview"
custom_llm_provider = "gemini"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
usage = Usage(
completion_tokens=1000,
@ -270,11 +264,9 @@ def test_image_tokens_fallback_to_base_cost():
assert round(completion_cost, 12) == round(expected_completion_cost, 12)
def test_video_output_tokens_gemini_omni_flash_preview():
def test_video_output_tokens_gemini_omni_flash_preview(_local_model_cost_map):
"""Video output tokens are billed at output_cost_per_video_token, not the text rate and not zero."""
model = "gemini-omni-flash-preview"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
text_tokens = 100
video_tokens = 46336
@ -310,11 +302,9 @@ def test_video_output_tokens_gemini_omni_flash_preview():
)
def test_video_input_tokens_gemini_omni_flash_preview():
def test_video_input_tokens_gemini_omni_flash_preview(_local_model_cost_map):
"""Video input tokens are billed at the standard input rate instead of being dropped."""
model = "gemini-omni-flash-preview"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
usage = Usage(
completion_tokens=10,
@ -369,12 +359,10 @@ def test_video_tokens_fallback_to_base_cost():
assert round(completion_cost, 12) == round((600 + 1120) * 2e-6, 12)
def test_generic_cost_per_token_above_200k_tokens():
def test_generic_cost_per_token_above_200k_tokens(_local_model_cost_map):
# gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing
model = "gemini-2.5-pro"
custom_llm_provider = "vertex_ai"
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 = 220 * 1e6
@ -420,12 +408,10 @@ def test_get_token_base_cost_picks_highest_crossed_tier():
assert prompt_base_cost == 9e-6
def test_generic_cost_per_token_gpt54_above_272k_tokens():
def test_generic_cost_per_token_gpt54_above_272k_tokens(_local_model_cost_map):
"""GPT-5.4/5.4-pro: prompts >272K input tokens priced at 2x input, 1.5x output."""
model = "gpt-5.4"
custom_llm_provider = "openai"
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 = 273000 # Above 272K threshold
@ -450,12 +436,10 @@ 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():
def test_generic_cost_per_token_minimax_m3_above_512k_tokens(_local_model_cost_map):
"""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
@ -493,10 +477,8 @@ def test_generic_cost_per_token_minimax_m3_above_512k_tokens():
"bedrock_mantle/openai.gpt-5.6-luna",
],
)
def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(model):
def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(_local_model_cost_map, model):
"""Bedrock GPT-5.6 supports a 1M context window, billed at the long-context rates above 272K."""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model_cost_map = litellm.model_cost[model]
assert model_cost_map["max_input_tokens"] == 1000000
@ -827,12 +809,10 @@ def test_generic_cost_per_token_tiered_pricing_bills_reasoning_at_tier_rate():
litellm.model_cost.pop(model, None)
def test_generic_cost_per_token_gpt55():
def test_generic_cost_per_token_gpt55(_local_model_cost_map):
"""gpt-5.5: base pricing — $5/1M input, $30/1M output, $0.50/1M cached input."""
model = "gpt-5.5"
custom_llm_provider = "openai"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model_cost_map = litellm.model_cost[model]
@ -867,12 +847,10 @@ def test_generic_cost_per_token_gpt55():
)
def test_generic_cost_per_token_gpt55_pro():
def test_generic_cost_per_token_gpt55_pro(_local_model_cost_map):
"""gpt-5.5-pro: responses-only model — $30/1M input, $180/1M output, $3/1M cached input."""
model = "gpt-5.5-pro"
custom_llm_provider = "openai"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model_cost_map = litellm.model_cost[model]
@ -919,7 +897,7 @@ def test_generic_cost_per_token_gpt55_pro():
("gpt-5.6-luna", 2e-7, 1.2e-6, 2e-8, 2.5e-7),
],
)
def test_generic_cost_per_token_gpt56(
def test_generic_cost_per_token_gpt56(_local_model_cost_map,
model, input_cost, output_cost, cache_read_cost, cache_write_cost
):
"""gpt-5.6 (sol/terra/luna): base pricing + new cache-write cost.
@ -927,8 +905,6 @@ def test_generic_cost_per_token_gpt56(
Cache writes are billed at 1.25x the uncached input rate for this family.
"""
custom_llm_provider = "openai"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model_cost_map = litellm.model_cost[model]
@ -989,7 +965,7 @@ def test_gpt_5_6_alias_prices_match_sol(local_model_cost_map):
("gpt-5.6-luna", 2e-7, 9e-7),
],
)
def test_generic_cost_per_token_gpt56_flex_above_272k(
def test_generic_cost_per_token_gpt56_flex_above_272k(_local_model_cost_map,
model, flex_long_input_cost, flex_long_output_cost
):
"""A >272K flex request bills the flex long-context rate, not the standard one.
@ -998,8 +974,6 @@ def test_generic_cost_per_token_gpt56_flex_above_272k(
``*_above_272k_tokens_flex`` keys these requests silently fell back to the
standard long-context price, billing 2x what OpenAI charges.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
prompt_tokens = 300000
completion_tokens = 1000
@ -1038,11 +1012,9 @@ def test_generic_cost_per_token_gpt56_flex_above_272k(
("flex", 300000, 2e-6, 2.5e-6, 2e-7),
],
)
def test_generic_cost_per_token_gpt56_terra_cache_costs_by_tier_and_context(
def test_generic_cost_per_token_gpt56_terra_cache_costs_by_tier_and_context(_local_model_cost_map,
service_tier, prompt_tokens, input_rate, cache_write_rate, cache_read_rate
):
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
cached_tokens = 50000
cache_write_tokens = 40000
@ -1130,7 +1102,7 @@ def test_generic_cost_per_token_gpt56_cyber(
("azure/eu/gpt-5.6-luna", 2.2e-7, 1.32e-6, 2.2e-8),
],
)
def test_generic_cost_per_token_azure_gpt56(
def test_generic_cost_per_token_azure_gpt56(_local_model_cost_map,
model, input_cost, output_cost, cache_read_cost
):
"""Azure gpt-5.6 (global + us/eu regional): Azure prices this family on its own
@ -1138,8 +1110,6 @@ def test_generic_cost_per_token_azure_gpt56(
promotional cut OpenAI applied to gpt-5.6-sol, so these rates deliberately sit
above the openai ones and must not be lowered to match them.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model_cost_map = litellm.model_cost[model]
assert model_cost_map["litellm_provider"] == "azure"
@ -1180,7 +1150,7 @@ def test_generic_cost_per_token_azure_gpt56(
("gpt-5.5-pro-2026-04-23", False, True, False),
],
)
def test_gpt55_reasoning_effort_flags_match_live_openai_api(
def test_gpt55_reasoning_effort_flags_match_live_openai_api(_local_model_cost_map,
model, expected_none, expected_xhigh, expected_minimal
):
"""Pin reasoning_effort capability flags to OpenAI's actual API contract.
@ -1189,8 +1159,6 @@ def test_gpt55_reasoning_effort_flags_match_live_openai_api(
``Unsupported value: 'reasoning_effort' does not support 'minimal' with
this model``. gpt-5.5-pro additionally rejects 'none' and 'low'.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
m = litellm.model_cost[model]
assert (
@ -1211,7 +1179,7 @@ def test_gpt55_reasoning_effort_flags_match_live_openai_api(
("gpt-5.5-pro", "gpt-5.5-pro-2026-04-23"),
],
)
def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(
def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(_local_model_cost_map,
base_model, dated_model
):
"""Dated snapshots must carry the same reasoning_effort capability flags as
@ -1223,8 +1191,6 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(
behavior between ``gpt-5.5`` and ``gpt-5.5-2026-04-23``. Pinning to a
dated variant must never lose capabilities relative to the base alias.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
base = litellm.model_cost[base_model]
dated = litellm.model_cost[dated_model]
@ -1251,7 +1217,7 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(
("azure/gpt-5.5-pro-2026-04-23", "responses", 3e-5, 1.8e-4, 3e-6),
],
)
def test_azure_gpt55_entries_present_with_correct_pricing(
def test_azure_gpt55_entries_present_with_correct_pricing(_local_model_cost_map,
model, expected_mode, expected_input, expected_output, expected_cache_read
):
"""Day-0 Azure entries for GPT-5.5 mirror the OpenAI pricing structure.
@ -1260,8 +1226,6 @@ def test_azure_gpt55_entries_present_with_correct_pricing(
on 2026-04-24): $5/$30 input/output per 1M for chat, $30/$180 for pro.
Cache discount is 10% of input.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
m = litellm.model_cost[model]
assert m["litellm_provider"] == "azure"
@ -1286,12 +1250,10 @@ def test_azure_gpt55_entries_present_with_correct_pricing(
("azure/gpt-5.5-pro", False, False, True),
],
)
def test_azure_gpt55_reasoning_effort_flags_match_live_openai_api(
def test_azure_gpt55_reasoning_effort_flags_match_live_openai_api(_local_model_cost_map,
model, expected_none, expected_minimal, expected_xhigh
):
"""Azure entries pin reasoning_effort flags to OpenAI's actual API contract."""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
m = litellm.model_cost[model]
assert m.get("supports_none_reasoning_effort") is expected_none
@ -1671,11 +1633,9 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details():
assert round(result, 6) == round(expected, 6)
def test_service_tier_flex_pricing():
def test_service_tier_flex_pricing(_local_model_cost_map):
"""Test that flex service tier uses correct pricing (approximately 50% of standard)."""
# Set up environment for local model cost map
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
# Test with gpt-5-nano which has flex pricing
model = "gpt-5-nano"
@ -1728,11 +1688,9 @@ def test_service_tier_flex_pricing():
), f"Flex total cost mismatch: {flex_total} vs {expected_flex_total}"
def test_service_tier_default_pricing():
def test_service_tier_default_pricing(_local_model_cost_map):
"""Test that when no service tier is provided, standard pricing is used."""
# Set up environment for local model cost map
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
# Test with gpt-5-nano
model = "gpt-5-nano"
@ -1779,11 +1737,9 @@ def test_service_tier_default_pricing():
), f"Standard completion cost mismatch: {default_cost[1]} vs {expected_standard_completion}"
def test_service_tier_fallback_pricing():
def test_service_tier_fallback_pricing(_local_model_cost_map):
"""Test that when service tier is provided but model doesn't have those keys, it falls back to standard pricing."""
# Set up environment for local model cost map
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
# Test with gpt-4 which doesn't have flex pricing keys
model = "gpt-4"
@ -1891,15 +1847,13 @@ def test_service_tier_ultrafast_pricing():
assert completion_cost == pytest.approx(400 * 3e-04)
def test_service_tier_ultrafast_fallback_pricing():
def test_service_tier_ultrafast_fallback_pricing(_local_model_cost_map):
"""Without *_ultrafast keys an ultrafast request bills the standard rate, not zero.
Guards the suffix fallback in _get_cost_per_unit: "_fast" is a substring of
"_ultrafast", so a shortest-first suffix match would strip the wrong suffix
and price the request at 0.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
@ -1929,7 +1883,7 @@ def test_service_tier_ultrafast_fallback_pricing():
"gemini-3.1-flash-lite-image",
],
)
def test_gemini_image_generation_cost_with_zero_text_tokens(model: str):
def test_gemini_image_generation_cost_with_zero_text_tokens(_local_model_cost_map, model: str):
"""
Test that image_tokens are correctly costed when text_tokens=0.
@ -1939,8 +1893,6 @@ def test_gemini_image_generation_cost_with_zero_text_tokens(model: str):
https://github.com/BerriAI/litellm/issues/17410
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
custom_llm_provider = "vertex_ai"
@ -1995,13 +1947,11 @@ def test_gemini_image_generation_cost_with_zero_text_tokens(model: str):
), f"Expected completion cost ${expected_completion_cost:.6f}, got ${completion_cost:.6f}"
def test_vertex_image_generation_cost_prefers_token_usage_metadata():
def test_vertex_image_generation_cost_prefers_token_usage_metadata(_local_model_cost_map):
"""
When usage metadata exists on image responses, Vertex image generation cost
should be calculated from token pricing, not flat output_cost_per_image.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gemini-3.1-flash-image-preview"
model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai")
@ -2040,13 +1990,11 @@ def test_vertex_image_generation_cost_prefers_token_usage_metadata():
assert cost != len(image_response.data) * model_info["output_cost_per_image"]
def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing():
def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing(_local_model_cost_map):
"""
Without usage metadata, Vertex image generation cost should fall back to
output_cost_per_image * number_of_images.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gemini-3.1-flash-image-preview"
model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai")
@ -2064,13 +2012,11 @@ def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing():
assert round(cost, 10) == round(expected_cost, 10)
def test_gemini_image_generation_cost_prefers_token_usage_metadata():
def test_gemini_image_generation_cost_prefers_token_usage_metadata(_local_model_cost_map):
"""
When usage metadata exists on image responses, Gemini image generation cost
should be calculated from token pricing, not flat output_cost_per_image.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gemini/gemini-3-pro-image-preview"
model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini")
@ -2109,13 +2055,11 @@ def test_gemini_image_generation_cost_prefers_token_usage_metadata():
assert cost != len(image_response.data) * model_info["output_cost_per_image"]
def test_gemini_image_generation_cost_falls_back_to_flat_image_pricing():
def test_gemini_image_generation_cost_falls_back_to_flat_image_pricing(_local_model_cost_map):
"""
Without usage metadata, Gemini image generation cost should fall back to
output_cost_per_image * number_of_images.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gemini/gemini-3-pro-image-preview"
model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini")
@ -2212,7 +2156,7 @@ def test_reasoning_tokens_without_text_tokens_gpt5_nano():
), "Bug detected: Cost calculation is using only reasoning_tokens instead of all completion_tokens!"
def test_image_count_prevents_text_tokens_fallback():
def test_image_count_prevents_text_tokens_fallback(_local_model_cost_map):
"""
Test that the text_tokens fallback in generic_cost_per_token does not
override text_tokens=0 when image_count > 0.
@ -2221,8 +2165,6 @@ def test_image_count_prevents_text_tokens_fallback():
When image_count > 0, text_tokens=0 is intentional (image-only request),
not "text_tokens not set by provider."
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
# Simulate Nova image-only embedding: prompt_tokens estimated from
# embedding dimensions (768 for 3072-dim), image_count=1
@ -2256,20 +2198,6 @@ def test_image_count_prevents_text_tokens_fallback():
# ---------------------------------------------------------------------------
@pytest.fixture
def _local_model_cost_map():
prev_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP")
prev_model_cost = litellm.model_cost
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
try:
yield
finally:
litellm.model_cost = prev_model_cost
if prev_env is None:
os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None)
else:
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = prev_env
@pytest.mark.parametrize("model", ["gpt-5.4", "gpt-realtime-2.1", "gpt-realtime-2.1-mini"])
@ -2603,7 +2531,7 @@ def test_threshold_keys_exclude_service_tier_variants():
("cerebras/qwen-3-32b", "cerebras", 250, 0),
],
)
def test_token_type_cost_breakdown_is_provider_agnostic(
def test_token_type_cost_breakdown_is_provider_agnostic(_local_model_cost_map,
model, custom_llm_provider, reasoning_tokens, cached_tokens
):
"""
@ -2615,8 +2543,6 @@ def test_token_type_cost_breakdown_is_provider_agnostic(
there - not the top-level cache_read_input_tokens attribute the old breakdown code
relied on - is what makes Vertex/OpenAI/Azure cache costs show up at all.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
usage = Usage(
prompt_tokens=1000,
@ -2647,10 +2573,8 @@ def test_token_type_cost_breakdown_is_provider_agnostic(
assert breakdown.cache_read_cost == pytest.approx(cached_tokens * cache_read_rate)
def test_token_type_cost_breakdown_matches_real_gemini_numbers():
def test_token_type_cost_breakdown_matches_real_gemini_numbers(_local_model_cost_map):
"""Hard-coded against the exact gemini-2.5-flash response that exposed the gap."""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
usage = Usage(
prompt_tokens=209,
@ -2673,9 +2597,7 @@ def test_token_type_cost_breakdown_matches_real_gemini_numbers():
assert breakdown.cache_creation_cost == 0.0
def test_token_type_cost_breakdown_xai_at_exactly_200k_uses_higher_tier_rates():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_token_type_cost_breakdown_xai_at_exactly_200k_uses_higher_tier_rates(_local_model_cost_map):
usage = Usage(
prompt_tokens=200_000,
@ -2697,9 +2619,7 @@ def test_token_type_cost_breakdown_xai_at_exactly_200k_uses_higher_tier_rates():
assert breakdown.cache_read_cost == pytest.approx(50_000 * 4e-07)
def test_token_type_cost_breakdown_xai_just_below_200k_uses_base_tier_rates():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_token_type_cost_breakdown_xai_just_below_200k_uses_base_tier_rates(_local_model_cost_map):
usage = Usage(
prompt_tokens=199_999,
@ -2721,14 +2641,12 @@ def test_token_type_cost_breakdown_xai_just_below_200k_uses_base_tier_rates():
assert breakdown.cache_read_cost == pytest.approx(50_000 * 2e-07)
def test_token_type_cost_breakdown_includes_cache_creation_from_top_level_usage():
def test_token_type_cost_breakdown_includes_cache_creation_from_top_level_usage(_local_model_cost_map):
"""
Bedrock/Anthropic report cache tokens as top-level usage fields; the Usage
constructor maps them onto prompt_tokens_details, so the breakdown must still
pick up both cache-read and cache-creation costs.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "anthropic.claude-3-5-haiku-20241022-v1:0"
usage = Usage(
@ -2752,14 +2670,12 @@ def test_token_type_cost_breakdown_includes_cache_creation_from_top_level_usage(
)
def test_token_type_cost_breakdown_reads_cache_write_tokens():
def test_token_type_cost_breakdown_reads_cache_write_tokens(_local_model_cost_map):
"""
Some OpenAI-compatible providers (e.g. kimi-k2) report cache-write tokens under
`cache_write_tokens` rather than `cache_creation_tokens`. The breakdown must read
it the same way the total-cost normalization does, so the two agree.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "anthropic.claude-3-5-haiku-20241022-v1:0"
usage = Usage(
@ -2780,7 +2696,7 @@ def test_token_type_cost_breakdown_reads_cache_write_tokens():
)
def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6():
def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6(_local_model_cost_map):
"""
Regression: OpenAI gpt-5.6 reports cache-write tokens under
prompt_tokens_details.cache_write_tokens (not the Anthropic cache_creation_tokens
@ -2788,8 +2704,6 @@ def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6():
input rate. Customer report: cache creation tokens were never counted for the
GPT-5.6 series, so cost was undercounted on cache-write requests.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gpt-5.6"
usage = Usage(
@ -2811,14 +2725,12 @@ def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6():
assert prompt_cost > 1000 * info["input_cost_per_token"]
def test_generic_cost_per_token_backs_out_cache_write_tokens_from_text_tokens():
def test_generic_cost_per_token_backs_out_cache_write_tokens_from_text_tokens(_local_model_cost_map):
"""
Regression for #34801: when a provider reports text_tokens covering the whole
prompt alongside cache-write tokens (and no cache reads), the cache-write tokens
must be backed out of the text total instead of being billed twice.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gpt-5.6"
usage = Usage(
@ -2837,15 +2749,13 @@ def test_generic_cost_per_token_backs_out_cache_write_tokens_from_text_tokens():
assert prompt_cost == pytest.approx(expected_prompt)
def test_token_type_cost_breakdown_reconciles_with_generic_total():
def test_token_type_cost_breakdown_reconciles_with_generic_total(_local_model_cost_map):
"""
Both-ways check: the reasoning subset must sum with the remaining (text) output
cost to exactly the completion total, and the cache-read subset with the remaining
input cost to exactly the prompt total, as computed by generic_cost_per_token.
A mismatch here would mean the breakdown misrepresents what was actually billed.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gemini-2.5-flash"
custom_llm_provider = "vertex_ai"
@ -2878,9 +2788,7 @@ def test_token_type_cost_breakdown_reconciles_with_generic_total():
assert text_input_cost + breakdown.cache_read_cost == pytest.approx(prompt_cost)
def test_token_type_cost_breakdown_zero_without_special_tokens():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_token_type_cost_breakdown_zero_without_special_tokens(_local_model_cost_map):
usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
breakdown = get_token_type_cost_breakdown(
@ -2917,7 +2825,7 @@ def test_token_type_cost_breakdown_zero_without_special_tokens():
),
],
)
def test_token_type_cost_breakdown_openai_responses_api_cache_write_read(
def test_token_type_cost_breakdown_openai_responses_api_cache_write_read(_local_model_cost_map,
raw_usage, expect_read, expect_write
):
"""Regression for #34309: OpenAI Responses API reports cache tokens under
@ -2926,8 +2834,6 @@ def test_token_type_cost_breakdown_openai_responses_api_cache_write_read(
cache_read_cost / cache_creation_cost from the transformed usage."""
from litellm.responses.utils import ResponseAPILoggingUtils
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gpt-5.6"
usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_usage)
@ -2968,15 +2874,13 @@ def test_token_type_cost_breakdown_handles_unknown_model_gracefully():
)
def test_token_type_cost_breakdown_applies_regional_uplift():
def test_token_type_cost_breakdown_applies_regional_uplift(_local_model_cost_map):
"""
Regional OpenAI hosts (eu./us.) apply a flat uplift to every token cost. The
per-type breakdown must apply the same uplift via data_residency so it stays
reconciled with the uplifted input_cost/output_cost totals, instead of being
logged at the base rate.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gpt-5.4"
custom_llm_provider = "openai"
@ -3024,15 +2928,13 @@ def test_token_type_cost_breakdown_applies_regional_uplift():
assert text_input_cost + eu.cache_read_cost == pytest.approx(prompt_cost)
def test_token_type_cost_breakdown_applies_vertex_regional_uplift():
def test_token_type_cost_breakdown_applies_vertex_regional_uplift(_local_model_cost_map):
"""
Non-global Vertex endpoints apply a flat 1.1x uplift to every token cost. The
per-type breakdown must apply the same uplift via vertex_location so it stays
reconciled with the uplifted input_cost/output_cost totals, instead of being
logged at the global rate.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "claude-haiku-4-5@20251001"
custom_llm_provider = "vertex_ai"
@ -3075,7 +2977,7 @@ def test_token_type_cost_breakdown_applies_vertex_regional_uplift():
assert text_input_cost + regional.cache_read_cost == pytest.approx(prompt_cost)
def test_token_type_cost_breakdown_applies_anthropic_geo_multiplier(monkeypatch):
def test_token_type_cost_breakdown_applies_anthropic_geo_multiplier(_local_model_cost_map, monkeypatch):
"""
Anthropic's regional (geo) uplift lives in provider_specific_entry and is
applied to every token type in the totals, so the per-type breakdown must
@ -3088,7 +2990,6 @@ def test_token_type_cost_breakdown_applies_anthropic_geo_multiplier(monkeypatch)
)
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "claude-test-geo-breakdown-model"
litellm.register_model(
@ -3209,9 +3110,7 @@ GEMINI_DAY0_LAUNCH_PRICING = [
@pytest.mark.parametrize("model,input_cost,output_cost,cache_read_cost", GEMINI_DAY0_LAUNCH_PRICING)
def test_gemini_36_flash_and_35_flash_lite_launch_pricing(model, input_cost, output_cost, cache_read_cost):
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_gemini_36_flash_and_35_flash_lite_launch_pricing(_local_model_cost_map, model, input_cost, output_cost, cache_read_cost):
model_cost_map = litellm.model_cost[model]
assert model_cost_map["input_cost_per_token"] == input_cost
@ -3224,9 +3123,7 @@ def test_gemini_36_flash_and_35_flash_lite_launch_pricing(model, input_cost, out
assert model_cost_map["max_input_tokens"] == 1048576
def test_generic_cost_per_token_gemini_36_flash():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_generic_cost_per_token_gemini_36_flash(_local_model_cost_map):
usage = Usage(
prompt_tokens=1000,
@ -3292,9 +3189,7 @@ def test_gemini_36_flash_batch_introductory_pricing(model, _local_model_cost_map
assert model_cost_map["output_cost_per_token_batches"] == 1.875e-06
def test_generic_cost_per_token_gemini_35_flash_lite():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_generic_cost_per_token_gemini_35_flash_lite(_local_model_cost_map):
usage = Usage(
prompt_tokens=1000,

134
tests/test_litellm/litellm_core_utils/test_realtime_streaming.py
View file

@ -1,6 +1,4 @@
import json
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -8,7 +6,6 @@ from websockets.exceptions import ConnectionClosed
import litellm
sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.realtime_streaming import (
@ -1326,7 +1323,7 @@ async def test_log_messages_includes_tools_in_model_call_details():
@pytest.mark.asyncio
async def test_realtime_guardrail_blocks_prompt_injection():
async def test_realtime_guardrail_blocks_prompt_injection(monkeypatch: pytest.MonkeyPatch):
"""
Test that when a transcription event containing prompt injection arrives from the
backend, a registered guardrail blocks it — sending a warning to the client
@ -1350,7 +1347,7 @@ async def test_realtime_guardrail_blocks_prompt_injection():
event_hook=GuardrailEventHooks.realtime_input_transcription,
default_on=True,
)
litellm.callbacks = [guardrail]
monkeypatch.setattr(litellm, "callbacks", [guardrail])
# --- client websocket mock ---
client_ws = MagicMock()
@ -1405,11 +1402,10 @@ async def test_realtime_guardrail_blocks_prompt_injection():
f"Expected guardrail_violation error type, got: {error_events[0]}"
)
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
async def test_realtime_guardrail_allows_clean_transcript():
async def test_realtime_guardrail_allows_clean_transcript(monkeypatch: pytest.MonkeyPatch):
"""
Test that a clean transcript passes through the guardrail and triggers
response.create to the backend.
@ -1430,7 +1426,7 @@ async def test_realtime_guardrail_allows_clean_transcript():
event_hook=GuardrailEventHooks.realtime_input_transcription,
default_on=True,
)
litellm.callbacks = [guardrail]
monkeypatch.setattr(litellm, "callbacks", [guardrail])
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
@ -1463,11 +1459,10 @@ async def test_realtime_guardrail_allows_clean_transcript():
response_creates = [e for e in sent_to_backend if e.get("type") == "response.create"]
assert len(response_creates) == 1, f"Clean transcript should trigger response.create, got: {sent_to_backend}"
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
async def test_realtime_text_input_guardrail_blocks_and_returns_error():
async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch):
"""
Test that when conversation.item.create arrives with text that triggers a guardrail,
the proxy blocks it (doesn't forward to backend) and returns an error event directly
@ -1495,7 +1490,7 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error():
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
litellm.callbacks = [guardrail]
monkeypatch.setattr(litellm, "callbacks", [guardrail])
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
@ -1558,11 +1553,10 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error():
]
assert len(original_items) == 0, f"Blocked item should not be forwarded to backend, got: {original_items}"
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
async def test_realtime_function_call_output_guardrail_blocks_and_returns_error():
async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch):
"""
Test that a client-supplied function_call_output whose content triggers a
guardrail is blocked: it is not forwarded to the backend, and an error
@ -1590,7 +1584,7 @@ async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
litellm.callbacks = [guardrail]
monkeypatch.setattr(litellm, "callbacks", [guardrail])
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
@ -1648,11 +1642,10 @@ async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(
assert sanitized_item["call_id"] == "call_123"
assert "test@example.com" not in sanitized_item["output"]
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
async def test_realtime_function_call_output_guardrail_allows_clean_output():
async def test_realtime_function_call_output_guardrail_allows_clean_output(monkeypatch: pytest.MonkeyPatch):
"""
Test that a clean function_call_output passes through and reaches the backend
when guardrails are configured.
@ -1670,7 +1663,7 @@ async def test_realtime_function_call_output_guardrail_allows_clean_output():
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
litellm.callbacks = [guardrail]
monkeypatch.setattr(litellm, "callbacks", [guardrail])
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
@ -1714,11 +1707,10 @@ async def test_realtime_function_call_output_guardrail_allows_clean_output():
]
assert len(forwarded) == 1, f"Clean function_call_output should be forwarded, got: {forwarded}"
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
async def test_realtime_text_input_guardrail_uses_pre_call_mode():
async def test_realtime_text_input_guardrail_uses_pre_call_mode(monkeypatch: pytest.MonkeyPatch):
"""
Test that _has_realtime_guardrails returns True for a guardrail configured with
pre_call mode (not just realtime_input_transcription).
@ -1736,7 +1728,7 @@ async def test_realtime_text_input_guardrail_uses_pre_call_mode():
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
litellm.callbacks = [guardrail]
monkeypatch.setattr(litellm, "callbacks", [guardrail])
client_ws = MagicMock()
backend_ws = MagicMock()
@ -1751,11 +1743,10 @@ async def test_realtime_text_input_guardrail_uses_pre_call_mode():
"pre_call-only guardrail must not disable server_vad auto-response"
)
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
async def test_realtime_session_created_injects_session_update_for_audio_guardrail():
async def test_realtime_session_created_injects_session_update_for_audio_guardrail(monkeypatch: pytest.MonkeyPatch):
"""
Test that when an audio transcription guardrail is configured, a session.created
event from the backend triggers a session.update injection (create_response: false)
@ -1775,7 +1766,7 @@ async def test_realtime_session_created_injects_session_update_for_audio_guardra
event_hook=GuardrailEventHooks.realtime_input_transcription,
default_on=True,
)
litellm.callbacks = [guardrail]
monkeypatch.setattr(litellm, "callbacks", [guardrail])
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
@ -1809,11 +1800,12 @@ async def test_realtime_session_created_injects_session_update_for_audio_guardra
"GA session.update must nest turn_detection under audio.input"
)
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
async def test_realtime_session_created_does_not_inject_session_update_for_pre_call_only():
async def test_realtime_session_created_does_not_inject_session_update_for_pre_call_only(
monkeypatch: pytest.MonkeyPatch,
):
"""
pre_call-only guardrails must not inject create_response:false on realtime
sessions — that breaks server_vad for audio-only voice agents (e.g. Model Armor).
@ -1831,7 +1823,7 @@ async def test_realtime_session_created_does_not_inject_session_update_for_pre_c
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
litellm.callbacks = [guardrail]
monkeypatch.setattr(litellm, "callbacks", [guardrail])
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
@ -1853,11 +1845,10 @@ async def test_realtime_session_created_does_not_inject_session_update_for_pre_c
session_updates = [e for e in sent_to_backend if e.get("type") == "session.update"]
assert len(session_updates) == 0, f"pre_call-only guardrail must not inject session.update, got: {sent_to_backend}"
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad():
async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monkeypatch: pytest.MonkeyPatch):
"""Model Armor-style pre_call + post_call must not gate audio VAD."""
import litellm
from litellm.integrations.custom_guardrail import CustomGuardrail
@ -1867,18 +1858,22 @@ async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad():
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return inputs
litellm.callbacks = [
ModelArmorStyleGuardrail(
guardrail_name="model_armor_all_pre_call",
event_hook=GuardrailEventHooks.pre_call,
default_on=False,
),
ModelArmorStyleGuardrail(
guardrail_name="model_armor_all_post_call",
event_hook=GuardrailEventHooks.post_call,
default_on=False,
),
]
monkeypatch.setattr(
litellm,
"callbacks",
[
ModelArmorStyleGuardrail(
guardrail_name="model_armor_all_pre_call",
event_hook=GuardrailEventHooks.pre_call,
default_on=False,
),
ModelArmorStyleGuardrail(
guardrail_name="model_armor_all_post_call",
event_hook=GuardrailEventHooks.post_call,
default_on=False,
),
],
)
client_ws = MagicMock()
backend_ws = MagicMock()
@ -1900,11 +1895,10 @@ async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad():
assert streaming._has_realtime_guardrails() is True
assert streaming._has_audio_transcription_guardrails() is False
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
async def test_end_session_after_n_fails_closes_connection():
async def test_end_session_after_n_fails_closes_connection(monkeypatch: pytest.MonkeyPatch):
"""
Test that end_session_after_n_fails=2 closes the backend websocket after
the second guardrail violation in a session.
@ -1923,7 +1917,7 @@ async def test_end_session_after_n_fails_closes_connection():
default_on=True,
end_session_after_n_fails=2,
)
litellm.callbacks = [guardrail]
monkeypatch.setattr(litellm, "callbacks", [guardrail])
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
@ -1948,11 +1942,10 @@ async def test_end_session_after_n_fails_closes_connection():
assert backend_ws.close.called, "Expected backend_ws.close() to be called after 2 violations"
assert streaming._violation_count == 2
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
async def test_on_violation_end_session_closes_on_first_fail():
async def test_on_violation_end_session_closes_on_first_fail(monkeypatch: pytest.MonkeyPatch):
"""
Test that on_violation='end_session' closes the session immediately on the
first violation, regardless of end_session_after_n_fails.
@ -1971,7 +1964,7 @@ async def test_on_violation_end_session_closes_on_first_fail():
default_on=True,
on_violation="end_session",
)
litellm.callbacks = [guardrail]
monkeypatch.setattr(litellm, "callbacks", [guardrail])
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
@ -1995,7 +1988,6 @@ async def test_on_violation_end_session_closes_on_first_fail():
assert backend_ws.close.called, "Expected session to close immediately with on_violation=end_session"
assert streaming._violation_count == 1
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
@ -2898,53 +2890,47 @@ def _transcription_guardrail():
)
def test_setup_folds_in_auto_response_disable_when_transcription_guardrail_active():
def test_setup_folds_in_auto_response_disable_when_transcription_guardrail_active(monkeypatch: pytest.MonkeyPatch):
"""Gemini rejects a second setup, so a transcription guardrail's auto-response
disable must be folded into the one-and-only setup; otherwise the model
auto-responds and the guardrail is bypassed."""
import litellm
litellm.callbacks = [_transcription_guardrail()]
try:
streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock())
setup = json.dumps(
{
"setup": {
"model": "models/gemini-3.1-flash-live-preview",
"generationConfig": {"responseModalities": ["AUDIO"]},
"inputAudioTranscription": {},
}
monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()])
streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock())
setup = json.dumps(
{
"setup": {
"model": "models/gemini-3.1-flash-live-preview",
"generationConfig": {"responseModalities": ["AUDIO"]},
"inputAudioTranscription": {},
}
)
out = json.loads(streaming._maybe_inject_guardrail_auto_response_disable(setup))
aad = out["setup"]["realtimeInputConfig"]["automaticActivityDetection"]
assert aad["disabled"] is True
finally:
litellm.callbacks = []
}
)
out = json.loads(streaming._maybe_inject_guardrail_auto_response_disable(setup))
aad = out["setup"]["realtimeInputConfig"]["automaticActivityDetection"]
assert aad["disabled"] is True
def test_setup_unchanged_without_transcription_guardrail():
def test_setup_unchanged_without_transcription_guardrail(monkeypatch: pytest.MonkeyPatch):
import litellm
litellm.callbacks = []
monkeypatch.setattr(litellm, "callbacks", [])
streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock())
setup = json.dumps({"setup": {"model": "x", "generationConfig": {"responseModalities": ["AUDIO"]}}})
out = streaming._maybe_inject_guardrail_auto_response_disable(setup)
assert json.loads(out) == json.loads(setup)
def test_non_bidi_setup_left_untouched_for_followup_capable_providers():
def test_non_bidi_setup_left_untouched_for_followup_capable_providers(monkeypatch: pytest.MonkeyPatch):
"""OpenAI realtime accepts a follow-up session.update, so a non-bidi message
(no top-level 'setup' key) must be left untouched even with a guardrail on."""
import litellm
litellm.callbacks = [_transcription_guardrail()]
try:
streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock())
msg = json.dumps({"type": "session.update", "session": {"instructions": "hi"}})
assert streaming._maybe_inject_guardrail_auto_response_disable(msg) == msg
finally:
litellm.callbacks = []
monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()])
streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock())
msg = json.dumps({"type": "session.update", "session": {"instructions": "hi"}})
assert streaming._maybe_inject_guardrail_auto_response_disable(msg) == msg
@pytest.mark.asyncio

1
tests/test_litellm/llms/azure/test_azure_common_utils.py
View file

@ -812,7 +812,6 @@ async def test_azure_client_reuse(function_name, is_async, args):
"""
Test that multiple Azure API calls reuse the same Azure OpenAI client
"""
litellm.set_verbose = True
# Determine which client class to mock based on whether the test is async
client_path = (

2
tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py
View file

@ -153,7 +153,6 @@ class TestBedrockAsyncInvokeEmbedding:
def test_async_invoke_twelvelabs_embedding_with_mock(self):
"""Test async invoke embedding with mocked HTTP calls."""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
model = "bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0"
@ -193,7 +192,6 @@ class TestBedrockAsyncInvokeEmbedding:
@pytest.mark.asyncio
async def test_async_invoke_twelvelabs_embedding_async_with_mock(self):
"""Test async invoke embedding with async calls."""
litellm.set_verbose = True
client = AsyncHTTPHandler()
test_api_key = "test-bearer-token-12345"
model = "bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0"

15
tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py
View file

@ -50,7 +50,6 @@ test_image_base64 = "data:image/png,test_image_base64_data"
)
def test_bedrock_embedding_with_api_key_bearer_token(model, input_type, embed_response):
"""Test embedding functionality with bearer token authentication"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
@ -98,7 +97,6 @@ def test_bedrock_embedding_with_env_variable_bearer_token(
model, input_type, embed_response
):
"""Test embedding functionality with bearer token from environment variable"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "env-bearer-token-12345"
@ -130,7 +128,6 @@ def test_bedrock_embedding_with_env_variable_bearer_token(
@pytest.mark.asyncio
async def test_async_bedrock_embedding_with_bearer_token():
"""Test async embedding functionality with bearer token authentication"""
litellm.set_verbose = True
client = AsyncHTTPHandler()
test_api_key = "async-bearer-token-12345"
model = "bedrock/amazon.titan-embed-text-v1"
@ -160,7 +157,6 @@ async def test_async_bedrock_embedding_with_bearer_token():
def test_bedrock_embedding_with_sigv4():
"""Test embedding falls back to SigV4 auth when no bearer token is provided"""
litellm.set_verbose = True
model = "bedrock/amazon.titan-embed-text-v1"
with patch(
@ -182,7 +178,6 @@ def test_bedrock_embedding_with_sigv4():
def test_bedrock_titan_v2_encoding_format_float():
"""Test amazon.titan-embed-text-v2:0 with encoding_format=float parameter"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
model = "bedrock/amazon.titan-embed-text-v2:0"
@ -220,7 +215,6 @@ def test_bedrock_titan_v2_encoding_format_float():
def test_bedrock_titan_v2_encoding_format_base64():
"""Test amazon.titan-embed-text-v2:0 with encoding_format=base64 parameter (maps to binary)"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
model = "bedrock/amazon.titan-embed-text-v2:0"
@ -260,7 +254,6 @@ def test_bedrock_titan_v2_encoding_format_base64():
def test_twelvelabs_input_type_parameter_mapping():
"""Test that input_type parameter is correctly mapped to inputType for TwelveLabs models"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
model = "bedrock/twelvelabs.marengo-embed-2-7-v1:0"
@ -300,7 +293,6 @@ def test_twelvelabs_input_type_parameter_mapping():
def test_twelvelabs_input_type_parameter_mapping_async_invoke():
"""Test that input_type parameter is correctly mapped to inputType for TwelveLabs async invoke models"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
model = "bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0"
@ -343,7 +335,6 @@ def test_twelvelabs_input_type_parameter_mapping_async_invoke():
def test_twelvelabs_missing_input_type_error():
"""Test that missing input_type parameter defaults to 'text' for TwelveLabs models"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
@ -422,7 +413,6 @@ def test_bedrock_embedding_header_forwarding(model, embed_response):
Relevant Issue: https://github.com/BerriAI/litellm/pull/16042
"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
@ -489,7 +479,6 @@ def test_bedrock_embedding_extra_headers_and_headers_merge():
This ensures that headers from kwargs (forwarded by proxy) and extra_headers
(passed explicitly) are both included in the final headers sent to the provider.
"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
model = "bedrock/amazon.titan-embed-text-v1"
@ -557,7 +546,6 @@ def test_bedrock_cohere_v4_embedding_response_parsing():
Test parsing of Bedrock Cohere v4 embedding response which returns a dictionary of embeddings
keyed by type (e.g. 'float', 'int8') instead of a direct list.
"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
model = "bedrock/cohere.embed-v4:0"
@ -617,7 +605,6 @@ def test_bedrock_embedding_custom_headers_with_iam_role_and_custom_api_base():
Relevant Issue: Custom headers not forwarded with IAM roles + custom api_base
"""
litellm.set_verbose = True
client = HTTPHandler()
# Simulate IAM role credentials with session token
@ -734,7 +721,6 @@ async def test_bedrock_embedding_custom_headers_with_iam_role_and_custom_api_bas
This is the async version of the test above, verifying the fix works for both
sync and async embedding calls.
"""
litellm.set_verbose = True
client = AsyncHTTPHandler()
# Simulate IAM role credentials with session token
@ -977,7 +963,6 @@ def test_bedrock_cohere_embedding_types_wrapped_as_list(
Malformed input request: #/embedding_types: expected type: JSONArray, found: String
when `encoding_format` is passed as a string.
"""
litellm.set_verbose = True
client = HTTPHandler()
model = "bedrock/cohere.embed-multilingual-v3"

4
tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py
View file

@ -18,7 +18,6 @@ mock_image_response = {"images": ["base64_encoded_image_data"], "error": None}
class TestBedrockImageGeneration:
def test_image_generation_with_api_key_bearer_token(self):
"""Test image generation with bearer token authentication"""
litellm.set_verbose = True
test_api_key = "test-bearer-token-12345"
model = "bedrock/stability.sd3-large-v1:0"
prompt = "A cute baby sea otter"
@ -53,7 +52,6 @@ class TestBedrockImageGeneration:
def test_image_generation_with_env_variable_bearer_token(self, monkeypatch):
"""Test image generation with bearer token from environment variable"""
litellm.set_verbose = True
test_api_key = "env-bearer-token-12345"
model = "bedrock/stability.sd3-large-v1:0"
prompt = "A cute baby sea otter"
@ -90,7 +88,6 @@ class TestBedrockImageGeneration:
@pytest.mark.asyncio
async def test_async_image_generation_with_bearer_token(self):
"""Test async image generation with bearer token authentication"""
litellm.set_verbose = True
test_api_key = "async-bearer-token-12345"
model = "bedrock/stability.sd3-large-v1:0"
prompt = "A cute baby sea otter"
@ -125,7 +122,6 @@ class TestBedrockImageGeneration:
def test_image_generation_with_sigv4(self):
"""Test image generation falls back to SigV4 auth when no bearer token is provided"""
litellm.set_verbose = True
model = "bedrock/stability.sd3-large-v1:0"
prompt = "A cute baby sea otter"

3
tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py
View file

@ -66,7 +66,6 @@ def test_bedrock_rerank_header_forwarding_sync(model):
This test verifies the fix for the issue where headers configured via
forward_client_headers_to_llm_api were not being passed to Bedrock rerank provider.
"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
@ -160,7 +159,6 @@ async def test_bedrock_rerank_header_forwarding_async(model):
This test verifies the fix for the issue where headers configured via
forward_client_headers_to_llm_api were not being passed to Bedrock rerank provider.
"""
litellm.set_verbose = True
client = AsyncHTTPHandler()
test_api_key = "test-bearer-token-12345"
@ -332,7 +330,6 @@ def test_bedrock_rerank_extra_headers_and_headers_merge():
This ensures that headers from kwargs (forwarded by proxy) and extra_headers
(passed explicitly) are both included in the final headers sent to the provider.
"""
litellm.set_verbose = True
client = HTTPHandler()
test_api_key = "test-bearer-token-12345"
model = "bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0"

258
tests/test_litellm/llms/custom_httpx/test_http_handler.py
View file

@ -131,79 +131,62 @@ def test_sync_post_streaming_status_error_should_not_wait_forever_for_body(
@pytest.mark.asyncio
async def test_ssl_security_level(monkeypatch):
# Ensure aiohttp transport is enabled for this test
original_disable = litellm.disable_aiohttp_transport
litellm.disable_aiohttp_transport = False
monkeypatch.setattr(litellm, "disable_aiohttp_transport", False)
try:
with patch.dict(os.environ, clear=True):
# Set environment variable for SSL security level
monkeypatch.setenv("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1")
with patch.dict(os.environ, clear=True):
# Set environment variable for SSL security level
monkeypatch.setenv("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1")
# Create async client with SSL verification disabled to isolate SSL context testing
client = AsyncHTTPHandler()
# Create async client with SSL verification disabled to isolate SSL context testing
client = AsyncHTTPHandler()
try:
# Get the transport (should be LiteLLMAiohttpTransport)
transport = client.client._transport
assert isinstance(transport, LiteLLMAiohttpTransport)
try:
# Get the transport (should be LiteLLMAiohttpTransport)
transport = client.client._transport
assert isinstance(transport, LiteLLMAiohttpTransport)
# Get the aiohttp ClientSession
client_session = transport._get_valid_client_session()
# Get the aiohttp ClientSession
client_session = transport._get_valid_client_session()
# Get the connector from the session
connector = client_session.connector
assert isinstance(connector, TCPConnector)
# Get the connector from the session
connector = client_session.connector
assert isinstance(connector, TCPConnector)
# Get the SSL context from the connector
ssl_context = connector._ssl
# Get the SSL context from the connector
ssl_context = connector._ssl
# Verify that the SSL context exists and has the correct cipher string
assert isinstance(ssl_context, ssl.SSLContext)
finally:
await client.close()
finally:
# Restore original setting
litellm.disable_aiohttp_transport = original_disable
# Verify that the SSL context exists and has the correct cipher string
assert isinstance(ssl_context, ssl.SSLContext)
finally:
await client.close()
@pytest.mark.asyncio
async def test_force_ipv4_transport():
async def test_force_ipv4_transport(monkeypatch: pytest.MonkeyPatch):
"""Test transport creation with force_ipv4 enabled"""
original_force_ipv4 = litellm.force_ipv4
original_disable = litellm.disable_aiohttp_transport
litellm.force_ipv4 = True
litellm.disable_aiohttp_transport = True
monkeypatch.setattr(litellm, "force_ipv4", True)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
try:
transport = AsyncHTTPHandler._create_async_transport()
transport = AsyncHTTPHandler._create_async_transport()
# Should get an AsyncHTTPTransport (no real HTTP call — avoids CI hangs)
assert isinstance(transport, httpx.AsyncHTTPTransport)
finally:
litellm.force_ipv4 = original_force_ipv4
litellm.disable_aiohttp_transport = original_disable
# Should get an AsyncHTTPTransport (no real HTTP call — avoids CI hangs)
assert isinstance(transport, httpx.AsyncHTTPTransport)
@pytest.mark.asyncio
async def test_aiohttp_disabled_transport():
async def test_aiohttp_disabled_transport(monkeypatch: pytest.MonkeyPatch):
"""Test transport creation with aiohttp disabled"""
original_disable = litellm.disable_aiohttp_transport
original_force_ipv4 = litellm.force_ipv4
litellm.disable_aiohttp_transport = True
litellm.force_ipv4 = False
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "force_ipv4", False)
try:
transport = AsyncHTTPHandler._create_async_transport()
transport = AsyncHTTPHandler._create_async_transport()
# Should get None when both aiohttp is disabled and force_ipv4 is False
assert transport is None
finally:
litellm.disable_aiohttp_transport = original_disable
litellm.force_ipv4 = original_force_ipv4
# Should get None when both aiohttp is disabled and force_ipv4 is False
assert transport is None
@pytest.mark.asyncio
async def test_ssl_verification_with_aiohttp_transport():
async def test_ssl_verification_with_aiohttp_transport(monkeypatch: pytest.MonkeyPatch):
"""
Test aiohttp respects ssl_verify=False
@ -213,38 +196,33 @@ async def test_ssl_verification_with_aiohttp_transport():
import aiohttp
# Ensure aiohttp transport is enabled for this test
original_disable = litellm.disable_aiohttp_transport
litellm.disable_aiohttp_transport = False
monkeypatch.setattr(litellm, "disable_aiohttp_transport", False)
litellm_async_client = AsyncHTTPHandler(ssl_verify=False)
try:
litellm_async_client = AsyncHTTPHandler(ssl_verify=False)
transport = litellm_async_client.client._transport
assert isinstance(transport, LiteLLMAiohttpTransport)
transport_connector = transport._get_valid_client_session().connector
assert isinstance(transport_connector, TCPConnector)
aiohttp_session = aiohttp.ClientSession(
connector=aiohttp.TCPConnector(ssl=False)
)
try:
transport = litellm_async_client.client._transport
assert isinstance(transport, LiteLLMAiohttpTransport)
transport_connector = transport._get_valid_client_session().connector
assert isinstance(transport_connector, TCPConnector)
aiohttp_connector = aiohttp_session.connector
assert isinstance(aiohttp_connector, aiohttp.TCPConnector)
aiohttp_session = aiohttp.ClientSession(
connector=aiohttp.TCPConnector(ssl=False)
)
try:
aiohttp_connector = aiohttp_session.connector
assert isinstance(aiohttp_connector, aiohttp.TCPConnector)
# assert both litellm transport and aiohttp session have ssl_verify=False
assert transport_connector._ssl == aiohttp_connector._ssl
finally:
await aiohttp_session.close()
# assert both litellm transport and aiohttp session have ssl_verify=False
assert transport_connector._ssl == aiohttp_connector._ssl
finally:
await litellm_async_client.close()
await aiohttp_session.close()
finally:
# Restore original setting
litellm.disable_aiohttp_transport = original_disable
await litellm_async_client.close()
@pytest.mark.asyncio
async def test_ssl_verification_with_shared_session():
async def test_ssl_verification_with_shared_session(monkeypatch: pytest.MonkeyPatch):
"""
Test that ssl_verify=False is respected even with shared sessions.
@ -257,67 +235,55 @@ async def test_ssl_verification_with_shared_session():
import aiohttp
# Ensure aiohttp transport is enabled for this test
original_disable = litellm.disable_aiohttp_transport
litellm.disable_aiohttp_transport = False
monkeypatch.setattr(litellm, "disable_aiohttp_transport", False)
shared_session = aiohttp.ClientSession()
try:
# Create a shared session (simulating what happens in production)
shared_session = aiohttp.ClientSession()
# Create transport with shared session and ssl_verify=False
transport = AsyncHTTPHandler._create_aiohttp_transport(
ssl_verify=False,
shared_session=shared_session,
)
try:
# Create transport with shared session and ssl_verify=False
transport = AsyncHTTPHandler._create_aiohttp_transport(
ssl_verify=False,
shared_session=shared_session,
)
# Verify the transport uses the shared session
assert transport.client is shared_session
# Verify the transport uses the shared session
assert transport.client is shared_session
# Verify the SSL setting is stored in the transport for per-request use
assert transport._ssl_verify is False
finally:
await shared_session.close()
# Verify the SSL setting is stored in the transport for per-request use
assert transport._ssl_verify is False
finally:
# Restore original setting
litellm.disable_aiohttp_transport = original_disable
await shared_session.close()
@pytest.mark.asyncio
async def test_ssl_context_with_shared_session():
async def test_ssl_context_with_shared_session(monkeypatch: pytest.MonkeyPatch):
"""
Test that ssl_context is respected even with shared sessions.
"""
import aiohttp
# Ensure aiohttp transport is enabled for this test
original_disable = litellm.disable_aiohttp_transport
litellm.disable_aiohttp_transport = False
monkeypatch.setattr(litellm, "disable_aiohttp_transport", False)
custom_ssl_context = ssl.create_default_context()
# Create a shared session
shared_session = aiohttp.ClientSession()
try:
# Create a custom SSL context
custom_ssl_context = ssl.create_default_context()
# Create transport with shared session and custom ssl_context
transport = AsyncHTTPHandler._create_aiohttp_transport(
ssl_context=custom_ssl_context,
shared_session=shared_session,
)
# Create a shared session
shared_session = aiohttp.ClientSession()
# Verify the transport uses the shared session
assert transport.client is shared_session
try:
# Create transport with shared session and custom ssl_context
transport = AsyncHTTPHandler._create_aiohttp_transport(
ssl_context=custom_ssl_context,
shared_session=shared_session,
)
# Verify the transport uses the shared session
assert transport.client is shared_session
# Verify the SSL context is stored in the transport for per-request use
assert transport._ssl_verify is custom_ssl_context
finally:
await shared_session.close()
# Verify the SSL context is stored in the transport for per-request use
assert transport._ssl_verify is custom_ssl_context
finally:
# Restore original setting
litellm.disable_aiohttp_transport = original_disable
await shared_session.close()
def test_get_ssl_configuration():
@ -563,26 +529,22 @@ def test_ssl_ecdh_curve(
if env_curve:
monkeypatch.setenv("SSL_ECDH_CURVE", env_curve)
original_value = litellm.ssl_ecdh_curve
try:
litellm.ssl_ecdh_curve = litellm_curve
monkeypatch.setattr(litellm, "ssl_ecdh_curve", litellm_curve)
# Create a real SSL context and patch set_ecdh_curve on it
# We need a real SSLContext instance (not a MagicMock) because _create_ssl_context
# calls methods like set_ciphers() and minimum_version that require a real context.
# We patch set_ecdh_curve specifically to verify it's called with the correct curve.
real_ssl_context = ssl.create_default_context()
with patch("ssl.create_default_context", return_value=real_ssl_context):
with patch.object(real_ssl_context, "set_ecdh_curve") as mock_set_curve:
ssl_context = get_ssl_configuration()
# Create a real SSL context and patch set_ecdh_curve on it
# We need a real SSLContext instance (not a MagicMock) because _create_ssl_context
# calls methods like set_ciphers() and minimum_version that require a real context.
# We patch set_ecdh_curve specifically to verify it's called with the correct curve.
real_ssl_context = ssl.create_default_context()
with patch("ssl.create_default_context", return_value=real_ssl_context):
with patch.object(real_ssl_context, "set_ecdh_curve") as mock_set_curve:
ssl_context = get_ssl_configuration()
if should_call:
mock_set_curve.assert_called_once_with(expected_curve)
else:
mock_set_curve.assert_not_called()
assert isinstance(ssl_context, ssl.SSLContext)
finally:
litellm.ssl_ecdh_curve = original_value
if should_call:
mock_set_curve.assert_called_once_with(expected_curve)
else:
mock_set_curve.assert_not_called()
assert isinstance(ssl_context, ssl.SSLContext)
def test_default_user_agent_is_litellm_version(monkeypatch):
@ -753,46 +715,38 @@ class TestDefaultCachedClientTimeoutHonorsRequestTimeout:
no per-model timeout (e.g. Bedrock) hung for 600s.
"""
@pytest.fixture
def restore_request_timeout(self):
original_value = litellm.request_timeout
original_flag = litellm.request_timeout_explicitly_set
try:
yield
finally:
litellm.request_timeout = original_value
litellm.request_timeout_explicitly_set = original_flag
def test_default_when_request_timeout_unset(self, restore_request_timeout):
def test_default_when_request_timeout_unset(self, monkeypatch: pytest.MonkeyPatch):
from litellm.llms.custom_httpx.http_handler import (
_DEFAULT_TIMEOUT,
_default_cached_client_timeout,
)
litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS
litellm.request_timeout_explicitly_set = False
monkeypatch.setattr(
litellm, "request_timeout", litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS
)
monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False)
assert _default_cached_client_timeout() is _DEFAULT_TIMEOUT
def test_uses_explicit_request_timeout(self, restore_request_timeout):
def test_uses_explicit_request_timeout(self, monkeypatch: pytest.MonkeyPatch):
from litellm.llms.custom_httpx.http_handler import (
_default_cached_client_timeout,
)
litellm.request_timeout = 300
litellm.request_timeout_explicitly_set = True
monkeypatch.setattr(litellm, "request_timeout", 300)
monkeypatch.setattr(litellm, "request_timeout_explicitly_set", True)
resolved = _default_cached_client_timeout()
assert resolved.read == 300.0
assert resolved.connect == 5.0
def test_cached_async_client_built_with_explicit_request_timeout(
self, restore_request_timeout
self, monkeypatch: pytest.MonkeyPatch
):
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.utils import LlmProviders
litellm.request_timeout = 300
litellm.request_timeout_explicitly_set = True
monkeypatch.setattr(litellm, "request_timeout", 300)
monkeypatch.setattr(litellm, "request_timeout_explicitly_set", True)
litellm.in_memory_llm_clients_cache = LLMClientCache()
client = get_async_httpx_client(llm_provider=LlmProviders.BEDROCK)
assert client.timeout.read == 300.0

1
tests/test_litellm/llms/openai/test_openai_common_utils.py
View file

@ -86,7 +86,6 @@ async def test_openai_client_reuse(function_name, is_async, args):
"""
Test that multiple API calls reuse the same OpenAI client
"""
litellm.set_verbose = True
# Determine which client class to mock based on whether the test is async
client_path = (

6
tests/test_litellm/llms/vertex_ai/test_vertex.py
View file

@ -33,7 +33,6 @@ def test_completion_pydantic_obj_2():
from litellm.llms.custom_httpx.http_handler import HTTPHandler
litellm.set_verbose = True
class CalendarEvent(BaseModel):
name: str
@ -259,7 +258,6 @@ def test_vertex_tool_type_field_removal():
def test_function_calling_with_gemini():
from litellm.llms.custom_httpx.http_handler import HTTPHandler
litellm.set_verbose = True
client = HTTPHandler()
with patch.object(client, "post", new=MagicMock()) as mock_post:
try:
@ -310,7 +308,6 @@ def test_function_calling_with_gemini():
def test_multiple_function_call():
litellm.set_verbose = True
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
@ -420,7 +417,6 @@ def test_multiple_function_call():
def test_multiple_function_call_changed_text_pos():
litellm.set_verbose = True
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
@ -528,7 +524,6 @@ def test_multiple_function_call_changed_text_pos():
def test_function_calling_with_gemini_multiple_results():
litellm.set_verbose = True
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
@ -1103,7 +1098,6 @@ def test_logprobs_unit_test():
def test_logprobs():
litellm.set_verbose = True
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()

1
tests/test_litellm/llms/volcengine/test_volcengine_embedding.py
View file

@ -31,7 +31,6 @@ class TestVolcEngineEmbedding(BaseLLMEmbeddingTest):
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_basic_embedding(self, sync_mode):
"""Test basic embedding functionality with realistic response"""
litellm.set_verbose = True
embedding_call_args = self.get_base_embedding_call_args()
# Mock the embedding functions to avoid actual API calls

40
tests/test_litellm/llms/zai/test_zai_provider.py
View file

@ -13,6 +13,12 @@ from litellm import completion
from litellm.cost_calculator import cost_per_token
@pytest.fixture
def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
@pytest.fixture
def zai_response():
"""Mock response from Z.AI API"""
@ -51,12 +57,8 @@ def test_zai_in_provider_lists():
assert "zai" in litellm.provider_list
def test_zai_models_in_model_cost(monkeypatch):
def test_zai_models_in_model_cost(local_model_cost_map):
"""Test that ZAI models are in the model cost map"""
import os
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
zai_models = [
"zai/glm-4.7",
@ -75,12 +77,8 @@ def test_zai_models_in_model_cost(monkeypatch):
assert litellm.model_cost[model]["litellm_provider"] == "zai"
def test_zai_glm46_cost_calculation(monkeypatch):
def test_zai_glm46_cost_calculation(local_model_cost_map):
"""Test the cost calculation for glm-4.6"""
import os
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
key = "zai/glm-4.6"
info = litellm.model_cost[key]
@ -96,12 +94,8 @@ def test_zai_glm46_cost_calculation(monkeypatch):
assert math.isclose(completion_cost, 2.2, rel_tol=1e-6)
def test_zai_flash_model_is_free(monkeypatch):
def test_zai_flash_model_is_free(local_model_cost_map):
"""Test that glm-4.5-flash has zero cost"""
import os
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
key = "zai/glm-4.5-flash"
info = litellm.model_cost[key]
@ -110,12 +104,8 @@ def test_zai_flash_model_is_free(monkeypatch):
assert info["output_cost_per_token"] == 0
def test_glm47_supports_reasoning(monkeypatch):
def test_glm47_supports_reasoning(local_model_cost_map):
"""Test that GLM-4.7 supports reasoning"""
import os
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
key = "zai/glm-4.7"
assert key in litellm.model_cost, f"Model {key} not found in model_cost"
@ -124,12 +114,8 @@ def test_glm47_supports_reasoning(monkeypatch):
assert info["supports_reasoning"] is True
def test_glm47_cost_calculation(monkeypatch):
def test_glm47_cost_calculation(local_model_cost_map):
"""Test cost calculation for GLM-4.7"""
import os
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
prompt_cost, completion_cost = cost_per_token(
model="zai/glm-4.7",
@ -146,7 +132,7 @@ def test_glm47_cost_calculation(monkeypatch):
async def test_zai_completion_call(respx_mock, zai_response, monkeypatch):
"""Test completion call with zai provider using mocked response"""
monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
litellm.disable_aiohttp_transport = True
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(
json=zai_response
@ -172,7 +158,7 @@ async def test_zai_completion_call(respx_mock, zai_response, monkeypatch):
def test_zai_sync_completion(respx_mock, zai_response, monkeypatch):
"""Test synchronous completion call"""
monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
litellm.disable_aiohttp_transport = True
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(
json=zai_response

2
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py
View file

@ -26,7 +26,6 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
def test_cato_guard_config():
litellm.set_verbose = True
litellm.guardrail_name_config_map = {}
init_guardrails_v2(
@ -47,7 +46,6 @@ def test_cato_guard_config():
def test_cato_guard_config_no_api_key(monkeypatch):
monkeypatch.delenv("CATO_API_KEY", raising=False)
litellm.set_verbose = True
litellm.guardrail_name_config_map = {}
with pytest.raises(CatoNetworksGuardrailMissingSecrets, match="Couldn't get Cato Networks api key"):
init_guardrails_v2(

1
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py
View file

@ -44,7 +44,6 @@ from tests.test_litellm.proxy.guardrails.guardrail_hooks._cisco_ai_defense_test_
def test_cisco_ai_defense_config_via_init_v2_chat(monkeypatch):
monkeypatch.setenv("CISCO_AI_DEFENSE_API_KEY", "test-key")
litellm.set_verbose = True
litellm.guardrail_name_config_map = {}
init_guardrails_v2(

12
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py
View file

@ -1,10 +1,8 @@
import os
import sys
import pytest
from unittest.mock import patch, MagicMock, AsyncMock
from httpx import Response, Request
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm.proxy.guardrails.guardrail_hooks.deepkeep.deepkeep import (
@ -17,10 +15,9 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
from litellm.exceptions import GuardrailRaisedException
def test_deepkeep_guard_config(monkeypatch):
def test_deepkeep_guard_config(monkeypatch: pytest.MonkeyPatch):
"""Test DeepKeep guard configuration with init_guardrails_v2."""
litellm.set_verbose = True
litellm.guardrail_name_config_map = {}
monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
monkeypatch.setenv("DEEPKEEP_API_KEY", "test-key")
monkeypatch.setenv("DEEPKEEP_API_BASE", "https://test.deepkeep.ai")
@ -42,9 +39,6 @@ def test_deepkeep_guard_config(monkeypatch):
)
# Clean up
del os.environ["DEEPKEEP_API_KEY"]
del os.environ["DEEPKEEP_API_BASE"]
del os.environ["DEEPKEEP_FIREWALL_ID"]
class TestDeepKeepGuardrail:
@ -108,7 +102,7 @@ class TestDeepKeepGuardrail:
== "https://test.deepkeep.ai/v3/openai/beta/litellm_basic_guardrail_api"
)
def test_initialization_with_env_vars(self, monkeypatch):
def test_initialization_with_env_vars(self, monkeypatch: pytest.MonkeyPatch):
"""should initialize successfully using environment variables."""
monkeypatch.setenv("DEEPKEEP_API_KEY", "env-key")
monkeypatch.setenv("DEEPKEEP_API_BASE", "https://env.deepkeep.ai")

62
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py
View file

@ -1,5 +1,4 @@
import os
import sys
import uuid
from typing import List, cast
from unittest.mock import AsyncMock, MagicMock, patch
@ -8,7 +7,6 @@ import pytest
from fastapi import HTTPException
from httpx import Request, Response
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm import ModelResponse
@ -26,10 +24,9 @@ from litellm.types.utils import (
)
def test_hiddenlayer_config_saas(monkeypatch):
def test_hiddenlayer_config_saas(monkeypatch: pytest.MonkeyPatch):
"""Test Hiddenlayer SaaS configuration with init_guardrails_v2."""
litellm.set_verbose = True
litellm.guardrail_name_config_map = {}
monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
# Set environment variables for testing
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -50,8 +47,6 @@ def test_hiddenlayer_config_saas(monkeypatch):
)
# Clean up
if "HIDDENLAYER_API_BASE" in os.environ:
del os.environ["HIDDENLAYER_API_BASE"]
class TestHiddenlayerGuardrail:
@ -71,7 +66,7 @@ class TestHiddenlayerGuardrail:
if key in os.environ:
del os.environ[key]
def test_initialization(self, monkeypatch):
def test_initialization(self, monkeypatch: pytest.MonkeyPatch):
"""Test successful initialization with default values."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -84,17 +79,16 @@ class TestHiddenlayerGuardrail:
assert guardrail.guardrail_name == "hiddenlayer"
assert guardrail.event_hook == "pre_call"
def test_initialization_fails_when_api_key_missing(self):
def test_initialization_fails_when_api_key_missing(self, monkeypatch: pytest.MonkeyPatch):
"""Test that initialization fails when API key is not set."""
# Ensure API key is not set
if "HIDDENLAYER_CLIENT_SECRET" in os.environ:
del os.environ["HIDDENLAYER_CLIENT_SECRET"]
monkeypatch.delenv("HIDDENLAYER_CLIENT_SECRET", raising=False)
with pytest.raises(RuntimeError):
HiddenlayerGuardrail(guardrail_name="hiddenlayer", event_hook="pre_call")
@pytest.mark.asyncio
async def test_apply_guardrail_request_no_violations(self, monkeypatch):
async def test_apply_guardrail_request_no_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for request with no violations detected."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -151,7 +145,7 @@ class TestHiddenlayerGuardrail:
assert call_args.args[0] == f"{guardrail.api_base}/detection/v1/interactions"
@pytest.mark.asyncio
async def test_apply_guardrail_request_with_violations(self, monkeypatch):
async def test_apply_guardrail_request_with_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for request with violations detected."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -209,7 +203,7 @@ class TestHiddenlayerGuardrail:
assert "Blocked by Hiddenlayer" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_apply_guardrail_response_no_violations(self, monkeypatch):
async def test_apply_guardrail_response_no_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for response with no violations detected."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -279,7 +273,7 @@ class TestHiddenlayerGuardrail:
mock_post.assert_called_once()
@pytest.mark.asyncio
async def test_apply_guardrail_response_with_violations(self, monkeypatch):
async def test_apply_guardrail_response_with_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for response with violations detected."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -348,7 +342,7 @@ class TestHiddenlayerGuardrail:
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_apply_guardrail_api_error_handling(self, monkeypatch):
async def test_apply_guardrail_api_error_handling(self, monkeypatch: pytest.MonkeyPatch):
"""Test handling of API errors in apply_guardrail."""
# Set required API key
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -391,7 +385,7 @@ class TestHiddenlayerGuardrail:
assert result == inputs
@pytest.mark.asyncio
async def test_validate_with_call_hiddenlayer_method(self, monkeypatch):
async def test_validate_with_call_hiddenlayer_method(self, monkeypatch: pytest.MonkeyPatch):
"""Test the _validate_with_guard_server internal method."""
# Set required API key
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -433,7 +427,7 @@ class TestHiddenlayerGuardrail:
)
@pytest.mark.asyncio
async def test_apply_guardrail_request_with_image(self, monkeypatch):
async def test_apply_guardrail_request_with_image(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail sends multimodal content (image) to HiddenLayer v1."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -498,7 +492,7 @@ class TestHiddenlayerGuardrail:
assert result is not None
@pytest.mark.asyncio
async def test_apply_guardrail_redact_with_image_content(self, monkeypatch):
async def test_apply_guardrail_redact_with_image_content(self, monkeypatch: pytest.MonkeyPatch):
"""Test that REDACT action with multimodal content extracts text properly into inputs['texts']."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -570,10 +564,9 @@ class TestHiddenlayerGuardrail:
assert config_model.__name__ == "HiddenlayerGuardrailConfigModel"
def test_hiddenlayer_config_v2(monkeypatch):
def test_hiddenlayer_config_v2(monkeypatch: pytest.MonkeyPatch):
"""Test HiddenLayer V2 configuration with init_guardrails_v2."""
litellm.set_verbose = True
litellm.guardrail_name_config_map = {}
monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -593,8 +586,6 @@ def test_hiddenlayer_config_v2(monkeypatch):
config_file_path="",
)
if "HIDDENLAYER_API_BASE" in os.environ:
del os.environ["HIDDENLAYER_API_BASE"]
class TestHiddenlayerGuardrailV2:
@ -612,7 +603,7 @@ class TestHiddenlayerGuardrailV2:
if key in os.environ:
del os.environ[key]
def test_initialization(self, monkeypatch):
def test_initialization(self, monkeypatch: pytest.MonkeyPatch):
"""Test successful initialization with default values."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -624,16 +615,15 @@ class TestHiddenlayerGuardrailV2:
assert guardrail.guardrail_name == "hiddenlayer"
assert guardrail.event_hook == "pre_call"
def test_initialization_fails_when_api_key_missing(self):
def test_initialization_fails_when_api_key_missing(self, monkeypatch: pytest.MonkeyPatch):
"""Test that initialization fails when API key is not set for SaaS."""
if "HIDDENLAYER_CLIENT_SECRET" in os.environ:
del os.environ["HIDDENLAYER_CLIENT_SECRET"]
monkeypatch.delenv("HIDDENLAYER_CLIENT_SECRET", raising=False)
with pytest.raises(RuntimeError):
HiddenlayerGuardrailV2(guardrail_name="hiddenlayer", event_hook="pre_call")
@pytest.mark.asyncio
async def test_apply_guardrail_request_no_violations(self, monkeypatch):
async def test_apply_guardrail_request_no_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for request with no violations detected."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -691,7 +681,7 @@ class TestHiddenlayerGuardrailV2:
assert "detection/v2/request-evaluations" in call_args.args[0]
@pytest.mark.asyncio
async def test_apply_guardrail_request_with_violations(self, monkeypatch):
async def test_apply_guardrail_request_with_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for request with violations detected (block via header)."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -751,7 +741,7 @@ class TestHiddenlayerGuardrailV2:
assert "Blocked by Hiddenlayer" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_apply_guardrail_response_no_violations(self, monkeypatch):
async def test_apply_guardrail_response_no_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for response with no violations detected."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -816,7 +806,7 @@ class TestHiddenlayerGuardrailV2:
assert "detection/v2/response-evaluations" in call_args.args[0]
@pytest.mark.asyncio
async def test_apply_guardrail_response_with_violations(self, monkeypatch):
async def test_apply_guardrail_response_with_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for response with violations detected (block via header)."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -863,7 +853,7 @@ class TestHiddenlayerGuardrailV2:
assert "Blocked by Hiddenlayer" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_apply_guardrail_response_with_tool_calls(self, monkeypatch):
async def test_apply_guardrail_response_with_tool_calls(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for response containing tool calls."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -924,7 +914,7 @@ class TestHiddenlayerGuardrailV2:
assert "detection/v2/response-evaluations" in call_args.args[0]
@pytest.mark.asyncio
async def test_call_hiddenlayer_uses_correct_endpoints(self, monkeypatch):
async def test_call_hiddenlayer_uses_correct_endpoints(self, monkeypatch: pytest.MonkeyPatch):
"""Test that _call_hiddenlayer uses the v2 request/response evaluation endpoints."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -959,7 +949,7 @@ class TestHiddenlayerGuardrailV2:
assert "detection/v2/response-evaluations" in mock_post.call_args.args[0]
@pytest.mark.asyncio
async def test_apply_guardrail_request_with_image(self, monkeypatch):
async def test_apply_guardrail_request_with_image(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail sends multimodal content (image) to HiddenLayer v2."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
@ -1030,7 +1020,7 @@ class TestHiddenlayerGuardrailV2:
assert texts == ["how much is on this receipt?"]
@pytest.mark.asyncio
async def test_apply_guardrail_request_with_image_multimodal_response(self, monkeypatch):
async def test_apply_guardrail_request_with_image_multimodal_response(self, monkeypatch: pytest.MonkeyPatch):
"""Test that new_texts extraction handles multimodal content (list) returned by HiddenLayer v2."""
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")

137
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py
View file

@ -1,5 +1,3 @@
import os
import sys
import uuid
from unittest.mock import AsyncMock, MagicMock, patch
@ -8,8 +6,6 @@ import pytest
from fastapi import HTTPException
from httpx import Request, Response
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm import ModelResponse
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -18,12 +14,11 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message
def test_onyx_guard_config(monkeypatch):
def test_onyx_guard_config(monkeypatch: pytest.MonkeyPatch):
"""Test Onyx guard configuration with init_guardrails_v2."""
litellm.set_verbose = True
litellm.guardrail_name_config_map = {}
monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
monkeypatch.setattr(litellm, "callbacks", [])
# Set environment variables for testing
monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security")
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
@ -41,16 +36,15 @@ def test_onyx_guard_config(monkeypatch):
config_file_path="",
)
# Clean up
if "ONYX_API_BASE" in os.environ:
del os.environ["ONYX_API_BASE"]
if "ONYX_API_KEY" in os.environ:
del os.environ["ONYX_API_KEY"]
registered = [c for c in litellm.callbacks if isinstance(c, OnyxGuardrail)]
assert len(registered) == 1
assert registered[0].guardrail_name == "onyx-guard"
assert registered[0].default_on is True
assert registered[0].event_hook == "pre_call"
def test_onyx_guard_with_custom_timeout_from_kwargs(monkeypatch):
def test_onyx_guard_with_custom_timeout_from_kwargs(monkeypatch: pytest.MonkeyPatch):
"""Test Onyx guard instantiation with custom timeout passed via kwargs."""
# Set environment variables for testing
monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security")
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
@ -74,20 +68,13 @@ def test_onyx_guard_with_custom_timeout_from_kwargs(monkeypatch):
assert timeout_param.read == 45.0
assert timeout_param.connect == 5.0
# Clean up
if "ONYX_API_BASE" in os.environ:
del os.environ["ONYX_API_BASE"]
if "ONYX_API_KEY" in os.environ:
del os.environ["ONYX_API_KEY"]
def test_onyx_guard_with_timeout_none_uses_env_var(monkeypatch):
def test_onyx_guard_with_timeout_none_uses_env_var(monkeypatch: pytest.MonkeyPatch):
"""Test Onyx guard with timeout=None uses ONYX_TIMEOUT env var.
When timeout=None is passed (as it would be from config model with default None),
the ONYX_TIMEOUT environment variable should be used.
"""
# Set environment variables for testing
monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security")
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
monkeypatch.setenv("ONYX_TIMEOUT", "60")
@ -112,23 +99,13 @@ def test_onyx_guard_with_timeout_none_uses_env_var(monkeypatch):
assert timeout_param.read == 60.0
assert timeout_param.connect == 5.0
# Clean up
if "ONYX_API_BASE" in os.environ:
del os.environ["ONYX_API_BASE"]
if "ONYX_API_KEY" in os.environ:
del os.environ["ONYX_API_KEY"]
if "ONYX_TIMEOUT" in os.environ:
del os.environ["ONYX_TIMEOUT"]
def test_onyx_guard_with_timeout_none_defaults_to_10(monkeypatch):
def test_onyx_guard_with_timeout_none_defaults_to_10(monkeypatch: pytest.MonkeyPatch):
"""Test Onyx guard with timeout=None and no env var defaults to 10 seconds."""
# Set environment variables for testing
monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security")
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
# Ensure ONYX_TIMEOUT is not set
if "ONYX_TIMEOUT" in os.environ:
del os.environ["ONYX_TIMEOUT"]
monkeypatch.delenv("ONYX_TIMEOUT", raising=False)
with patch(
"litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client"
@ -150,33 +127,17 @@ def test_onyx_guard_with_timeout_none_defaults_to_10(monkeypatch):
assert timeout_param.read == 10.0
assert timeout_param.connect == 5.0
# Clean up
if "ONYX_API_BASE" in os.environ:
del os.environ["ONYX_API_BASE"]
if "ONYX_API_KEY" in os.environ:
del os.environ["ONYX_API_KEY"]
class TestOnyxGuardrail:
"""Test suite for Onyx Security Guardrail integration."""
def setup_method(self):
"""Setup test environment."""
# Clean up any existing environment variables
for key in ["ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"]:
if key in os.environ:
del os.environ[key]
@pytest.fixture(autouse=True)
def clear_onyx_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
for key in ("ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"):
monkeypatch.delenv(key, raising=False)
def teardown_method(self):
"""Clean up test environment."""
# Clean up any environment variables set during tests
for key in ["ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"]:
if key in os.environ:
del os.environ[key]
def test_initialization_with_defaults(self, monkeypatch):
def test_initialization_with_defaults(self, monkeypatch: pytest.MonkeyPatch):
"""Test successful initialization with default values."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
guardrail = OnyxGuardrail(
@ -189,7 +150,7 @@ class TestOnyxGuardrail:
assert guardrail.guardrail_name == "test-guard"
assert guardrail.event_hook == "pre_call"
def test_initialization_with_env_vars(self, monkeypatch):
def test_initialization_with_env_vars(self, monkeypatch: pytest.MonkeyPatch):
"""Test initialization with environment variables."""
monkeypatch.setenv("ONYX_API_BASE", "https://custom.onyx.security")
monkeypatch.setenv("ONYX_API_KEY", "custom-api-key")
@ -202,18 +163,17 @@ class TestOnyxGuardrail:
assert guardrail.api_key == "custom-api-key"
assert guardrail.event_hook == "post_call"
def test_initialization_fails_when_api_key_missing(self):
def test_initialization_fails_when_api_key_missing(self, monkeypatch: pytest.MonkeyPatch):
"""Test that initialization fails when API key is not set."""
# Ensure API key is not set
if "ONYX_API_KEY" in os.environ:
del os.environ["ONYX_API_KEY"]
monkeypatch.delenv("ONYX_API_KEY", raising=False)
with pytest.raises(
ValueError, match="ONYX_API_KEY environment variable is not set"
):
OnyxGuardrail(guardrail_name="test-guard", event_hook="pre_call")
def test_initialization_with_default_timeout(self, monkeypatch):
def test_initialization_with_default_timeout(self, monkeypatch: pytest.MonkeyPatch):
"""Test that default timeout is 10.0 seconds."""
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
@ -232,7 +192,7 @@ class TestOnyxGuardrail:
assert timeout_param.read == 10.0
assert timeout_param.connect == 5.0
def test_initialization_with_custom_timeout_parameter(self, monkeypatch):
def test_initialization_with_custom_timeout_parameter(self, monkeypatch: pytest.MonkeyPatch):
"""Test initialization with custom timeout parameter."""
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
@ -254,7 +214,7 @@ class TestOnyxGuardrail:
assert timeout_param.read == 30.0
assert timeout_param.connect == 5.0
def test_initialization_with_timeout_from_env_var(self, monkeypatch):
def test_initialization_with_timeout_from_env_var(self, monkeypatch: pytest.MonkeyPatch):
"""Test initialization with timeout from ONYX_TIMEOUT environment variable.
Note: The env var is only used when timeout=None is explicitly passed,
@ -282,7 +242,7 @@ class TestOnyxGuardrail:
assert timeout_param.read == 25.0
assert timeout_param.connect == 5.0
def test_initialization_timeout_parameter_overrides_env_var(self, monkeypatch):
def test_initialization_timeout_parameter_overrides_env_var(self, monkeypatch: pytest.MonkeyPatch):
"""Test that timeout parameter overrides ONYX_TIMEOUT environment variable."""
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
monkeypatch.setenv("ONYX_TIMEOUT", "25")
@ -306,9 +266,8 @@ class TestOnyxGuardrail:
assert timeout_param.connect == 5.0
@pytest.mark.asyncio
async def test_apply_guardrail_request_no_violations(self, monkeypatch):
async def test_apply_guardrail_request_no_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for request with no violations detected."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
# Setup guardrail
@ -372,9 +331,8 @@ class TestOnyxGuardrail:
assert call_args.kwargs["json"]["conversation_id"] == "test-call-id"
@pytest.mark.asyncio
async def test_apply_guardrail_request_with_violations(self, monkeypatch):
async def test_apply_guardrail_request_with_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for request with violations detected."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
# Setup guardrail
@ -423,9 +381,8 @@ class TestOnyxGuardrail:
assert "prompt_injection" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_apply_guardrail_response_no_violations(self, monkeypatch):
async def test_apply_guardrail_response_no_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for response with no violations detected."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
# Setup guardrail
@ -497,9 +454,8 @@ class TestOnyxGuardrail:
assert call_args.kwargs["json"]["conversation_id"] == "test-call-id-2"
@pytest.mark.asyncio
async def test_apply_guardrail_response_with_violations(self, monkeypatch):
async def test_apply_guardrail_response_with_violations(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail for response with violations detected."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
# Setup guardrail
@ -558,9 +514,8 @@ class TestOnyxGuardrail:
assert "illegal_instructions" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_apply_guardrail_api_error_handling(self, monkeypatch):
async def test_apply_guardrail_api_error_handling(self, monkeypatch: pytest.MonkeyPatch):
"""Test handling of API errors in apply_guardrail."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
guardrail = OnyxGuardrail(
@ -591,9 +546,8 @@ class TestOnyxGuardrail:
assert result == inputs
@pytest.mark.asyncio
async def test_apply_guardrail_timeout_error_handling(self, monkeypatch):
async def test_apply_guardrail_timeout_error_handling(self, monkeypatch: pytest.MonkeyPatch):
"""Test handling of timeout errors in apply_guardrail (graceful degradation)."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
guardrail = OnyxGuardrail(
@ -629,9 +583,8 @@ class TestOnyxGuardrail:
assert result == inputs
@pytest.mark.asyncio
async def test_apply_guardrail_read_timeout_error_handling(self, monkeypatch):
async def test_apply_guardrail_read_timeout_error_handling(self, monkeypatch: pytest.MonkeyPatch):
"""Test handling of read timeout errors in apply_guardrail."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
guardrail = OnyxGuardrail(
@ -667,9 +620,8 @@ class TestOnyxGuardrail:
assert result == inputs
@pytest.mark.asyncio
async def test_apply_guardrail_connect_timeout_error_handling(self, monkeypatch):
async def test_apply_guardrail_connect_timeout_error_handling(self, monkeypatch: pytest.MonkeyPatch):
"""Test handling of connect timeout errors in apply_guardrail."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
guardrail = OnyxGuardrail(
@ -705,9 +657,8 @@ class TestOnyxGuardrail:
assert result == inputs
@pytest.mark.asyncio
async def test_apply_guardrail_no_logging_obj(self, monkeypatch):
async def test_apply_guardrail_no_logging_obj(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail without logging object (uses UUID)."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
guardrail = OnyxGuardrail(
@ -747,9 +698,8 @@ class TestOnyxGuardrail:
assert call_args.kwargs["json"]["conversation_id"] == "test-uuid"
@pytest.mark.asyncio
async def test_validate_with_guard_server_method(self, monkeypatch):
async def test_validate_with_guard_server_method(self, monkeypatch: pytest.MonkeyPatch):
"""Test the _validate_with_guard_server internal method."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
guardrail = OnyxGuardrail(
@ -788,9 +738,8 @@ class TestOnyxGuardrail:
)
@pytest.mark.asyncio
async def test_validate_with_guard_server_blocked(self, monkeypatch):
async def test_validate_with_guard_server_blocked(self, monkeypatch: pytest.MonkeyPatch):
"""Test _validate_with_guard_server when request is blocked."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
guardrail = OnyxGuardrail(
@ -825,9 +774,8 @@ class TestOnyxGuardrail:
assert config_model.__name__ == "OnyxGuardrailConfigModel"
@pytest.mark.asyncio
async def test_apply_guardrail_with_modelresponse(self, monkeypatch):
async def test_apply_guardrail_with_modelresponse(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail with ModelResponse object for response type."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
guardrail = OnyxGuardrail(
@ -880,9 +828,8 @@ class TestOnyxGuardrail:
assert "payload" in call_args.kwargs["json"]
@pytest.mark.asyncio
async def test_apply_guardrail_response_error_handling(self, monkeypatch):
async def test_apply_guardrail_response_error_handling(self, monkeypatch: pytest.MonkeyPatch):
"""Test error handling when processing response data."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
guardrail = OnyxGuardrail(
@ -925,7 +872,7 @@ class TestOnyxIntegration:
"""Test integration scenarios."""
@pytest.mark.asyncio
async def test_full_guardrail_flow(self, monkeypatch):
async def test_full_guardrail_flow(self, monkeypatch: pytest.MonkeyPatch):
"""Test full guardrail flow with multiple hooks."""
# Set environment variables
monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security")
@ -966,16 +913,10 @@ class TestOnyxIntegration:
)
assert len(custom_loggers) >= 3
# Clean up
if "ONYX_API_BASE" in os.environ:
del os.environ["ONYX_API_BASE"]
if "ONYX_API_KEY" in os.environ:
del os.environ["ONYX_API_KEY"]
@pytest.mark.asyncio
async def test_apply_guardrail_empty_request_data(self, monkeypatch):
async def test_apply_guardrail_empty_request_data(self, monkeypatch: pytest.MonkeyPatch):
"""Test apply_guardrail with empty request data."""
# Set required API key
monkeypatch.setenv("ONYX_API_KEY", "test-api-key")
guardrail = OnyxGuardrail(

14
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py
View file

@ -1,11 +1,9 @@
import os
import sys
import pytest
from fastapi import HTTPException
from httpx import ConnectError, Request, Response
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm import DualCache
@ -93,23 +91,23 @@ class TestRepelloAIInitialization:
with pytest.raises(ValueError, match="asset_id"):
RepelloAIGuardrail(api_key="test-api-key", guardrail_name="t")
def test_api_key_from_env(self, monkeypatch):
def test_api_key_from_env(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("REPELLOAI_API_KEY", "env-key")
guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t")
assert guardrail.repelloai_api_key == "env-key"
def test_api_key_from_argus_env(self, monkeypatch):
def test_api_key_from_argus_env(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("ARGUS_API_KEY", "argus-key")
guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t")
assert guardrail.repelloai_api_key == "argus-key"
def test_argus_env_preferred_over_legacy(self, monkeypatch):
def test_argus_env_preferred_over_legacy(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("ARGUS_API_KEY", "argus-key")
monkeypatch.setenv("REPELLOAI_API_KEY", "legacy-key")
guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t")
assert guardrail.repelloai_api_key == "argus-key"
def test_explicit_api_key_preferred_over_env(self, monkeypatch):
def test_explicit_api_key_preferred_over_env(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("ARGUS_API_KEY", "argus-key")
guardrail = RepelloAIGuardrail(
api_key="explicit-key", asset_id="asset-123", guardrail_name="t"
@ -145,9 +143,9 @@ class TestRepelloAIInitialization:
assert guardrail.api_base == DEFAULT_REPELLOAI_API_BASE
assert guardrail.unreachable_fallback == "fail_closed"
def test_init_guardrails_v2_wiring(self, monkeypatch):
def test_init_guardrails_v2_wiring(self, monkeypatch: pytest.MonkeyPatch):
"""The guardrail registers and constructs via the config.yaml path."""
litellm.guardrail_name_config_map = {}
monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
monkeypatch.setenv("REPELLOAI_API_KEY", "test-key")
init_guardrails_v2(
all_guardrails=[

1
tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py
View file

@ -65,7 +65,6 @@ def setup_and_teardown():
asyncio.set_event_loop(loop)
# Set up litellm state
litellm.set_verbose = True
litellm.guardrail_name_config_map = {}
yield

91
tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py
View file

@ -1,5 +1,3 @@
import os
import sys
from fastapi.exceptions import HTTPException
from unittest.mock import patch, AsyncMock
from httpx import Response, Request
@ -12,19 +10,15 @@ from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security im
PromptSecurityGuardrail,
)
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
def test_prompt_security_guard_config(monkeypatch):
def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch):
"""Test guardrail initialization with proper configuration"""
litellm.set_verbose = True
litellm.guardrail_name_config_map = {}
monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
monkeypatch.setattr(litellm, "callbacks", [])
# Set environment variables for testing
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
@ -42,21 +36,19 @@ def test_prompt_security_guard_config(monkeypatch):
config_file_path="",
)
# Clean up
del os.environ["PROMPT_SECURITY_API_KEY"]
del os.environ["PROMPT_SECURITY_API_BASE"]
registered = [c for c in litellm.callbacks if isinstance(c, PromptSecurityGuardrail)]
assert len(registered) == 1
assert registered[0].guardrail_name == "prompt_security"
assert registered[0].default_on is True
assert registered[0].event_hook == "during_call"
def test_prompt_security_guard_config_no_api_key():
def test_prompt_security_guard_config_no_api_key(monkeypatch: pytest.MonkeyPatch):
"""Test that initialization fails when API key is missing"""
litellm.set_verbose = True
litellm.guardrail_name_config_map = {}
monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
# Ensure API key is not in environment
if "PROMPT_SECURITY_API_KEY" in os.environ:
del os.environ["PROMPT_SECURITY_API_KEY"]
if "PROMPT_SECURITY_API_BASE" in os.environ:
del os.environ["PROMPT_SECURITY_API_BASE"]
monkeypatch.delenv("PROMPT_SECURITY_API_KEY", raising=False)
monkeypatch.delenv("PROMPT_SECURITY_API_BASE", raising=False)
with pytest.raises(
PromptSecurityGuardrailMissingSecrets,
@ -78,7 +70,7 @@ def test_prompt_security_guard_config_no_api_key():
@pytest.mark.asyncio
async def test_apply_guardrail_block_request(monkeypatch):
async def test_apply_guardrail_block_request(monkeypatch: pytest.MonkeyPatch):
"""Test that apply_guardrail blocks malicious prompts"""
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
@ -126,13 +118,9 @@ async def test_apply_guardrail_block_request(monkeypatch):
assert "prompt_injection" in str(excinfo.value.detail)
assert "jailbreak" in str(excinfo.value.detail)
# Clean up
del os.environ["PROMPT_SECURITY_API_KEY"]
del os.environ["PROMPT_SECURITY_API_BASE"]
@pytest.mark.asyncio
async def test_apply_guardrail_modify_request(monkeypatch):
async def test_apply_guardrail_modify_request(monkeypatch: pytest.MonkeyPatch):
"""Test that apply_guardrail modifies prompts when needed"""
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
@ -177,13 +165,9 @@ async def test_apply_guardrail_modify_request(monkeypatch):
assert result["texts"] == ["User prompt with PII: SSN [REDACTED]"]
# Clean up
del os.environ["PROMPT_SECURITY_API_KEY"]
del os.environ["PROMPT_SECURITY_API_BASE"]
@pytest.mark.asyncio
async def test_apply_guardrail_allow_request(monkeypatch):
async def test_apply_guardrail_allow_request(monkeypatch: pytest.MonkeyPatch):
"""Test that apply_guardrail allows safe prompts"""
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
@ -220,13 +204,9 @@ async def test_apply_guardrail_allow_request(monkeypatch):
assert result == inputs
# Clean up
del os.environ["PROMPT_SECURITY_API_KEY"]
del os.environ["PROMPT_SECURITY_API_BASE"]
@pytest.mark.asyncio
async def test_apply_guardrail_block_response(monkeypatch):
async def test_apply_guardrail_block_response(monkeypatch: pytest.MonkeyPatch):
"""Test that apply_guardrail blocks malicious responses"""
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
@ -267,13 +247,9 @@ async def test_apply_guardrail_block_response(monkeypatch):
assert "Blocked by Prompt Security" in str(excinfo.value.detail)
assert "pii_exposure" in str(excinfo.value.detail)
# Clean up
del os.environ["PROMPT_SECURITY_API_KEY"]
del os.environ["PROMPT_SECURITY_API_BASE"]
@pytest.mark.asyncio
async def test_apply_guardrail_modify_response(monkeypatch):
async def test_apply_guardrail_modify_response(monkeypatch: pytest.MonkeyPatch):
"""Test that apply_guardrail modifies responses when needed"""
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
@ -311,13 +287,9 @@ async def test_apply_guardrail_modify_response(monkeypatch):
assert result["texts"] == ["Your SSN is [REDACTED]"]
# Clean up
del os.environ["PROMPT_SECURITY_API_KEY"]
del os.environ["PROMPT_SECURITY_API_BASE"]
@pytest.mark.asyncio
async def test_file_sanitization(monkeypatch):
async def test_file_sanitization(monkeypatch: pytest.MonkeyPatch):
"""Test file sanitization for images"""
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
@ -401,13 +373,9 @@ async def test_file_sanitization(monkeypatch):
# Should complete without errors and return the data
assert result is not None
# Clean up
del os.environ["PROMPT_SECURITY_API_KEY"]
del os.environ["PROMPT_SECURITY_API_BASE"]
@pytest.mark.asyncio
async def test_file_sanitization_block(monkeypatch):
async def test_file_sanitization_block(monkeypatch: pytest.MonkeyPatch):
"""Test that file sanitization blocks malicious files"""
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
@ -485,13 +453,9 @@ async def test_file_sanitization_block(monkeypatch):
assert "File blocked by Prompt Security" in str(excinfo.value.detail)
assert "malware_detected" in str(excinfo.value.detail)
# Clean up
del os.environ["PROMPT_SECURITY_API_KEY"]
del os.environ["PROMPT_SECURITY_API_BASE"]
@pytest.mark.asyncio
async def test_user_api_key_alias_forwarding(monkeypatch):
async def test_user_api_key_alias_forwarding(monkeypatch: pytest.MonkeyPatch):
"""Test that user API key alias is properly sent via headers and payload"""
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
@ -530,12 +494,9 @@ async def test_user_api_key_alias_forwarding(monkeypatch):
payload = call_kwargs["json"]
assert payload["user"] == "vk-alias"
del os.environ["PROMPT_SECURITY_API_KEY"]
del os.environ["PROMPT_SECURITY_API_BASE"]
@pytest.mark.asyncio
async def test_role_filtering(monkeypatch):
async def test_role_filtering(monkeypatch: pytest.MonkeyPatch):
"""Test that tool/function messages are filtered out by default"""
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
@ -594,13 +555,9 @@ async def test_role_filtering(monkeypatch):
assert len(sent_messages) == 3
assert all(msg["role"] in ["system", "user", "assistant"] for msg in sent_messages)
# Clean up
del os.environ["PROMPT_SECURITY_API_KEY"]
del os.environ["PROMPT_SECURITY_API_BASE"]
@pytest.mark.asyncio
async def test_check_tool_results_enabled(monkeypatch):
async def test_check_tool_results_enabled(monkeypatch: pytest.MonkeyPatch):
"""Test with check_tool_results=True: transforms tool/function to 'other' role"""
monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
@ -680,7 +637,3 @@ async def test_check_tool_results_enabled(monkeypatch):
assert "indirect_prompt_injection" in str(excinfo.value.detail)
# Clean up
del os.environ["PROMPT_SECURITY_API_KEY"]
del os.environ["PROMPT_SECURITY_API_BASE"]
del os.environ["PROMPT_SECURITY_CHECK_TOOL_RESULTS"]

828
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

File diff suppressed because it is too large Load diff

570
tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py
View file

@ -165,7 +165,7 @@ class ContentCheckGuardrail(CustomGuardrail):
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_escalation_step1_fails_step2_blocks():
async def test_escalation_step1_fails_step2_blocks(monkeypatch):
"""
Pipeline: simple-filter (on_fail: next) -> advanced-filter (on_fail: block)
Input: request that fails simple-filter
@ -182,36 +182,32 @@ async def test_escalation_step1_fails_step2_blocks():
],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [simple_guard, advanced_guard]
monkeypatch.setattr(litellm, "callbacks", [simple_guard, advanced_guard])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "bad content"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "bad content"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
assert simple_guard.calls == 1
assert advanced_guard.calls == 1
assert result.terminal_action == "block"
assert len(result.step_results) == 2
assert result.step_results[0].guardrail_name == "simple-filter"
assert result.step_results[0].outcome == "fail"
assert result.step_results[0].action_taken == "next"
assert result.step_results[1].guardrail_name == "advanced-filter"
assert result.step_results[1].outcome == "fail"
assert result.step_results[1].action_taken == "block"
finally:
litellm.callbacks = original_callbacks
assert simple_guard.calls == 1
assert advanced_guard.calls == 1
assert result.terminal_action == "block"
assert len(result.step_results) == 2
assert result.step_results[0].guardrail_name == "simple-filter"
assert result.step_results[0].outcome == "fail"
assert result.step_results[0].action_taken == "next"
assert result.step_results[1].guardrail_name == "advanced-filter"
assert result.step_results[1].outcome == "fail"
assert result.step_results[1].action_taken == "block"
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_block_carries_original_guardrail_exception():
async def test_block_carries_original_guardrail_exception(monkeypatch):
"""A blocking step must expose the guardrail's own raised exception on the
result so the caller can re-raise it verbatim, giving the policy path the
same response/trace as a direct guardrail attachment."""
@ -219,67 +215,52 @@ async def test_block_carries_original_guardrail_exception():
pipeline = GuardrailPipeline(
mode="pre_call",
steps=[
PipelineStep(
guardrail="moderation-filter", on_fail="block", on_pass="allow"
)
],
steps=[PipelineStep(guardrail="moderation-filter", on_fail="block", on_pass="allow")],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [guard]
monkeypatch.setattr(litellm, "callbacks", [guard])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "bad content"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "bad content"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
assert result.terminal_action == "block"
assert isinstance(result.original_exception, HTTPException)
assert result.original_exception.status_code == 400
assert result.original_exception.detail == "Content policy violation"
finally:
litellm.callbacks = original_callbacks
assert result.terminal_action == "block"
assert isinstance(result.original_exception, HTTPException)
assert result.original_exception.status_code == 400
assert result.original_exception.detail == "Content policy violation"
@pytest.mark.asyncio
async def test_unsupported_mode_yields_error_outcome_without_exception():
async def test_unsupported_mode_yields_error_outcome_without_exception(monkeypatch):
"""An unexpected hook mode must surface as an error outcome (carrying no
original exception), not crash or run the guardrail."""
guard = AlwaysPassGuardrail(guardrail_name="filter")
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [guard]
monkeypatch.setattr(litellm, "callbacks", [guard])
try:
result = await PipelineExecutor.execute_steps(
steps=[PipelineStep(guardrail="filter", on_error="block", on_fail="block")],
mode="during_call",
data={"messages": [{"role": "user", "content": "hi"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
result = await PipelineExecutor.execute_steps(
steps=[PipelineStep(guardrail="filter", on_error="block", on_fail="block")],
mode="during_call",
data={"messages": [{"role": "user", "content": "hi"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
assert guard.calls == 0
assert result.terminal_action == "block"
assert result.step_results[0].outcome == "error"
assert (
"Unsupported pipeline mode: during_call"
in result.step_results[0].error_detail
)
assert result.original_exception is None
finally:
litellm.callbacks = original_callbacks
assert guard.calls == 0
assert result.terminal_action == "block"
assert result.step_results[0].outcome == "error"
assert "Unsupported pipeline mode: during_call" in result.step_results[0].error_detail
assert result.original_exception is None
@pytest.mark.asyncio
async def test_passthrough_guardrail_failure_can_pipeline_block():
async def test_passthrough_guardrail_failure_can_pipeline_block(monkeypatch):
"""
Pipeline: passthrough guardrail (on_fail: block)
Expected: passthrough ModifyResponseException is treated as policy fail,
@ -298,35 +279,31 @@ async def test_passthrough_guardrail_failure_can_pipeline_block():
],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [passthrough_guard]
monkeypatch.setattr(litellm, "callbacks", [passthrough_guard])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={
"model": "fake-model",
"messages": [{"role": "user", "content": "bad content"}],
},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={
"model": "fake-model",
"messages": [{"role": "user", "content": "bad content"}],
},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
assert passthrough_guard.calls == 1
assert result.terminal_action == "block"
assert len(result.step_results) == 1
assert result.step_results[0].guardrail_name == "passthrough-filter"
assert result.step_results[0].outcome == "fail"
assert result.step_results[0].action_taken == "block"
assert result.error_message == "Content policy violation"
finally:
litellm.callbacks = original_callbacks
assert passthrough_guard.calls == 1
assert result.terminal_action == "block"
assert len(result.step_results) == 1
assert result.step_results[0].guardrail_name == "passthrough-filter"
assert result.step_results[0].outcome == "fail"
assert result.step_results[0].action_taken == "block"
assert result.error_message == "Content policy violation"
@pytest.mark.asyncio
async def test_custom_code_guardrail_failure_can_pipeline_block():
async def test_custom_code_guardrail_failure_can_pipeline_block(monkeypatch):
"""
Pipeline: custom code guardrail (on_fail: block)
Expected: custom code keeps its standalone passthrough block behavior, and
@ -334,10 +311,7 @@ async def test_custom_code_guardrail_failure_can_pipeline_block():
"""
custom_guard = CustomCodeGuardrail(
guardrail_name="custom-code-filter",
custom_code=(
"def apply_guardrail(inputs, request_data, input_type):\n"
' return block("SSN detected")\n'
),
custom_code=('def apply_guardrail(inputs, request_data, input_type):\n return block("SSN detected")\n'),
)
pipeline = GuardrailPipeline(
@ -351,35 +325,31 @@ async def test_custom_code_guardrail_failure_can_pipeline_block():
],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [custom_guard]
monkeypatch.setattr(litellm, "callbacks", [custom_guard])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={
"model": "fake-model",
"messages": [{"role": "user", "content": "123-45-6789"}],
},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={
"model": "fake-model",
"messages": [{"role": "user", "content": "123-45-6789"}],
},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
assert result.terminal_action == "block"
assert len(result.step_results) == 1
assert result.step_results[0].guardrail_name == "custom-code-filter"
assert result.step_results[0].outcome == "fail"
assert result.step_results[0].action_taken == "block"
assert result.error_message == "SSN detected"
finally:
litellm.callbacks = original_callbacks
assert result.terminal_action == "block"
assert len(result.step_results) == 1
assert result.step_results[0].guardrail_name == "custom-code-filter"
assert result.step_results[0].outcome == "fail"
assert result.step_results[0].action_taken == "block"
assert result.error_message == "SSN detected"
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_early_allow_step1_passes_step2_skipped():
async def test_early_allow_step1_passes_step2_skipped(monkeypatch):
"""
Pipeline: simple-filter (on_pass: allow) -> advanced-filter
Input: clean request that passes simple-filter
@ -396,32 +366,28 @@ async def test_early_allow_step1_passes_step2_skipped():
],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [simple_guard, advanced_guard]
monkeypatch.setattr(litellm, "callbacks", [simple_guard, advanced_guard])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "clean content"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "clean content"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
assert simple_guard.calls == 1
assert advanced_guard.calls == 0
assert result.terminal_action == "allow"
assert len(result.step_results) == 1
assert result.step_results[0].outcome == "pass"
assert result.step_results[0].action_taken == "allow"
finally:
litellm.callbacks = original_callbacks
assert simple_guard.calls == 1
assert advanced_guard.calls == 0
assert result.terminal_action == "allow"
assert len(result.step_results) == 1
assert result.step_results[0].outcome == "pass"
assert result.step_results[0].action_taken == "allow"
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_escalation_step1_fails_step2_passes():
async def test_escalation_step1_fails_step2_passes(monkeypatch):
"""
Pipeline: simple-filter (on_fail: next) -> advanced-filter (on_pass: allow)
Input: request that fails simple but passes advanced
@ -438,34 +404,30 @@ async def test_escalation_step1_fails_step2_passes():
],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [simple_guard, advanced_guard]
monkeypatch.setattr(litellm, "callbacks", [simple_guard, advanced_guard])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "borderline content"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "borderline content"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
assert simple_guard.calls == 1
assert advanced_guard.calls == 1
assert result.terminal_action == "allow"
assert len(result.step_results) == 2
assert result.step_results[0].outcome == "fail"
assert result.step_results[0].action_taken == "next"
assert result.step_results[1].outcome == "pass"
assert result.step_results[1].action_taken == "allow"
finally:
litellm.callbacks = original_callbacks
assert simple_guard.calls == 1
assert advanced_guard.calls == 1
assert result.terminal_action == "allow"
assert len(result.step_results) == 2
assert result.step_results[0].outcome == "fail"
assert result.step_results[0].action_taken == "next"
assert result.step_results[1].outcome == "pass"
assert result.step_results[1].action_taken == "allow"
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_data_forwarding_pii_masking():
async def test_data_forwarding_pii_masking(monkeypatch):
"""
Pipeline: pii-masker (pass_data: true, on_pass: next) -> content-check (on_pass: allow)
Input: "Hello John Smith"
@ -487,31 +449,27 @@ async def test_data_forwarding_pii_masking():
],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [pii_guard, content_guard]
monkeypatch.setattr(litellm, "callbacks", [pii_guard, content_guard])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "Hello John Smith"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="pii-then-safety",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "Hello John Smith"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="pii-then-safety",
)
assert pii_guard.calls == 1
assert content_guard.calls == 1
assert content_guard.received_messages[0]["content"] == "Hello [REDACTED]"
assert result.terminal_action == "allow"
assert result.modified_data is not None
assert result.modified_data["messages"][0]["content"] == "Hello [REDACTED]"
finally:
litellm.callbacks = original_callbacks
assert pii_guard.calls == 1
assert content_guard.calls == 1
assert content_guard.received_messages[0]["content"] == "Hello [REDACTED]"
assert result.terminal_action == "allow"
assert result.modified_data is not None
assert result.modified_data["messages"][0]["content"] == "Hello [REDACTED]"
@pytest.mark.asyncio
async def test_guardrail_not_found_uses_on_fail():
async def test_guardrail_not_found_uses_on_fail(monkeypatch):
"""
If a guardrail is not found, treat as error and use on_fail action.
"""
@ -526,29 +484,25 @@ async def test_guardrail_not_found_uses_on_fail():
],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = []
monkeypatch.setattr(litellm, "callbacks", [])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "test"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test-policy",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "test"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test-policy",
)
assert result.terminal_action == "block"
assert result.step_results[0].outcome == "error"
assert "not found" in result.step_results[0].error_detail
finally:
litellm.callbacks = original_callbacks
assert result.terminal_action == "block"
assert result.step_results[0].outcome == "error"
assert "not found" in result.step_results[0].error_detail
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_on_error_next_fallback_on_api_outage_on_fail_blocks_content():
async def test_on_error_next_fallback_on_api_outage_on_fail_blocks_content(monkeypatch):
"""
Policy intervention (400) uses on_fail; technical error (503) uses on_error.
@ -574,32 +528,28 @@ async def test_on_error_next_fallback_on_api_outage_on_fail_blocks_content():
],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [primary, fallback]
monkeypatch.setattr(litellm, "callbacks", [primary, fallback])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "any"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="mod-fallback",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "any"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="mod-fallback",
)
assert primary.calls == 1
assert fallback.calls == 1
assert result.terminal_action == "allow"
assert result.step_results[0].outcome == "error"
assert result.step_results[0].action_taken == "next"
assert result.step_results[1].outcome == "pass"
finally:
litellm.callbacks = original_callbacks
assert primary.calls == 1
assert fallback.calls == 1
assert result.terminal_action == "allow"
assert result.step_results[0].outcome == "error"
assert result.step_results[0].action_taken == "next"
assert result.step_results[1].outcome == "pass"
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_on_fail_next_on_content_on_error_block_stops_api_fallback():
async def test_on_fail_next_on_content_on_error_block_stops_api_fallback(monkeypatch):
"""
Content policy fail (400) uses on_fail: next; API error uses on_error: block (no second step).
"""
@ -625,48 +575,40 @@ async def test_on_fail_next_on_content_on_error_block_stops_api_fallback():
],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [primary_content, fallback]
monkeypatch.setattr(litellm, "callbacks", [primary_content, fallback])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline_content.steps,
mode=pipeline_content.mode,
data={"messages": [{"role": "user", "content": "bad"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test",
)
assert result.terminal_action == "allow"
assert primary_content.calls == 1
assert fallback.calls == 1
finally:
litellm.callbacks = original_callbacks
result = await PipelineExecutor.execute_steps(
steps=pipeline_content.steps,
mode=pipeline_content.mode,
data={"messages": [{"role": "user", "content": "bad"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test",
)
assert result.terminal_action == "allow"
assert primary_content.calls == 1
assert fallback.calls == 1
# API outage: on_error block -> do not run fallback
fallback.calls = 0
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [primary_api, fallback]
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline_content.steps,
mode=pipeline_content.mode,
data={"messages": [{"role": "user", "content": "ok"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test",
)
assert result.terminal_action == "block"
assert primary_api.calls == 1
assert fallback.calls == 0
assert result.step_results[0].outcome == "error"
assert result.step_results[0].action_taken == "block"
finally:
litellm.callbacks = original_callbacks
monkeypatch.setattr(litellm, "callbacks", [primary_api, fallback])
result = await PipelineExecutor.execute_steps(
steps=pipeline_content.steps,
mode=pipeline_content.mode,
data={"messages": [{"role": "user", "content": "ok"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test",
)
assert result.terminal_action == "block"
assert primary_api.calls == 1
assert fallback.calls == 0
assert result.step_results[0].outcome == "error"
assert result.step_results[0].action_taken == "block"
@pytest.mark.asyncio
async def test_guardrail_not_found_with_next_continues():
async def test_guardrail_not_found_with_next_continues(monkeypatch):
"""
If a guardrail is not found and on_fail is 'next', continue to next step.
"""
@ -688,32 +630,28 @@ async def test_guardrail_not_found_with_next_continues():
],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [pass_guard]
monkeypatch.setattr(litellm, "callbacks", [pass_guard])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "test"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test-policy",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "test"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test-policy",
)
assert result.terminal_action == "allow"
assert len(result.step_results) == 2
assert result.step_results[0].outcome == "error"
assert result.step_results[0].action_taken == "next"
assert result.step_results[1].outcome == "pass"
assert pass_guard.calls == 1
finally:
litellm.callbacks = original_callbacks
assert result.terminal_action == "allow"
assert len(result.step_results) == 2
assert result.step_results[0].outcome == "error"
assert result.step_results[0].action_taken == "next"
assert result.step_results[1].outcome == "pass"
assert pass_guard.calls == 1
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_single_step_pipeline_block():
async def test_single_step_pipeline_block(monkeypatch):
"""Single step pipeline that blocks."""
guard = AlwaysFailGuardrail(guardrail_name="blocker")
@ -722,27 +660,23 @@ async def test_single_step_pipeline_block():
steps=[PipelineStep(guardrail="blocker", on_fail="block")],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [guard]
monkeypatch.setattr(litellm, "callbacks", [guard])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "test"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "test"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test",
)
assert result.terminal_action == "block"
assert guard.calls == 1
finally:
litellm.callbacks = original_callbacks
assert result.terminal_action == "block"
assert guard.calls == 1
@pytest.mark.asyncio
async def test_single_step_pipeline_allow():
async def test_single_step_pipeline_allow(monkeypatch):
"""Single step pipeline that allows."""
guard = AlwaysPassGuardrail(guardrail_name="passer")
@ -751,27 +685,23 @@ async def test_single_step_pipeline_allow():
steps=[PipelineStep(guardrail="passer", on_pass="allow")],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [guard]
monkeypatch.setattr(litellm, "callbacks", [guard])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "test"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "test"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test",
)
assert result.terminal_action == "allow"
assert guard.calls == 1
finally:
litellm.callbacks = original_callbacks
assert result.terminal_action == "allow"
assert guard.calls == 1
@pytest.mark.asyncio
async def test_step_results_include_duration():
async def test_step_results_include_duration(monkeypatch):
"""Step results should include timing information."""
guard = AlwaysPassGuardrail(guardrail_name="timed")
@ -780,23 +710,19 @@ async def test_step_results_include_duration():
steps=[PipelineStep(guardrail="timed")],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [guard]
monkeypatch.setattr(litellm, "callbacks", [guard])
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "test"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test",
)
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "test"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="test",
)
assert result.step_results[0].duration_seconds is not None
assert result.step_results[0].duration_seconds >= 0
finally:
litellm.callbacks = original_callbacks
assert result.step_results[0].duration_seconds is not None
assert result.step_results[0].duration_seconds >= 0
class _PolicyOptOutGuardrail(CustomGuardrail):

1
tests/test_litellm/responses/test_text_format_conversion.py
View file

@ -158,7 +158,6 @@ class TestTextFormatConversion:
new=mock_handler,
):
litellm._turn_on_debug()
litellm.set_verbose = True
# Call aresponses with text_format parameter
response = await litellm.aresponses(

218
tests/test_litellm/test_cost_calculator.py
View file

@ -1,12 +1,6 @@
import os
import sys
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
from pydantic import BaseModel
@ -24,6 +18,12 @@ from litellm.types.utils import ModelInfo, ModelResponse, PromptTokensDetailsWra
from litellm.utils import TranscriptionResponse
@pytest.fixture
def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
def test_cost_per_token_duplicate_openai_prefix_matches_model_cost(monkeypatch):
"""
Router/proxy configs may use deployment ids like openai/openai/<model>. Cost lookup must
@ -93,14 +93,12 @@ def test_cost_per_token_non_string_model_does_not_hang():
assert result.get("status") in ("returned", "raised")
def test_completion_cost_uses_response_model_for_dynamic_routing():
def test_completion_cost_uses_response_model_for_dynamic_routing(_local_model_cost_map):
"""
Test that completion_cost uses the model from the response object
when the input model (e.g., azure-model-router) is not in model_cost.
This supports Azure Model Router and similar dynamic routing scenarios.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
# Simulate Azure Model Router: input is generic router, response has actual model
response = ModelResponse(
@ -139,9 +137,7 @@ def test_cost_calculator_with_response_cost_in_additional_headers():
assert result == 1000
def test_baseten_model_api_pricing_entries():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_baseten_model_api_pricing_entries(_local_model_cost_map):
expected_pricing = {
"baseten/nvidia/Nemotron-120B-A12B": (3e-07, 7.5e-07),
@ -165,9 +161,7 @@ def test_baseten_model_api_pricing_entries():
assert model_info["output_cost_per_token"] == output_cost
def test_wandb_model_api_pricing_entries():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_wandb_model_api_pricing_entries(_local_model_cost_map):
expected_pricing = {
"wandb/moonshotai/Kimi-K2.5": (6e-07, 3e-06),
@ -182,9 +176,7 @@ def test_wandb_model_api_pricing_entries():
assert model_info["output_cost_per_token"] == output_cost
def test_openrouter_qwen36_plus_model_info():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_openrouter_qwen36_plus_model_info(_local_model_cost_map):
model_info = litellm.model_cost.get("openrouter/qwen/qwen3.6-plus")
@ -208,9 +200,7 @@ def test_openrouter_qwen36_plus_model_info():
"github_copilot/mai-code-1-flash-internal",
],
)
def test_github_copilot_mai_code_1_flash_pricing(model):
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_github_copilot_mai_code_1_flash_pricing(_local_model_cost_map, model):
model_info = litellm.model_cost.get(model)
@ -238,9 +228,7 @@ def test_github_copilot_mai_code_1_flash_pricing(model):
assert completion_usd == pytest.approx(500 * 4.5e-06)
def test_cost_calculator_with_usage(monkeypatch):
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch):
usage = Usage(
prompt_tokens=120,
@ -320,11 +308,9 @@ def test_cost_calculator_with_usage(monkeypatch):
assert result == expected_cost, f"Got {result}, Expected {expected_cost}"
def test_transcription_cost_uses_token_pricing():
def test_transcription_cost_uses_token_pricing(_local_model_cost_map):
from litellm import completion_cost
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
usage = Usage(
prompt_tokens=14,
@ -348,11 +334,9 @@ def test_transcription_cost_uses_token_pricing():
assert pytest.approx(cost, rel=1e-6) == expected_cost
def test_transcription_cost_falls_back_to_duration():
def test_transcription_cost_falls_back_to_duration(_local_model_cost_map):
from litellm import completion_cost
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
response = TranscriptionResponse(text="demo text")
response.duration = 10.0
@ -368,14 +352,12 @@ def test_transcription_cost_falls_back_to_duration():
assert pytest.approx(cost, rel=1e-6) == expected_cost
def test_vertex_chirp_3_transcription_cost_from_duration():
def test_vertex_chirp_3_transcription_cost_from_duration(_local_model_cost_map):
"""Regression: the chirp_3 cost map entry shipped with output_cost_per_second 0.0,
and cost_per_second prefers output_cost_per_second whenever it is not None, so
every transcription priced to $0.00 instead of using input_cost_per_second."""
from litellm import completion_cost
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
response = TranscriptionResponse(text="demo text")
response.duration = 18.0
@ -1127,9 +1109,7 @@ def test_tiered_pricing_only_deployment_completion_cost_is_nonzero():
assert cost > 0
def test_azure_realtime_cost_calculator():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_azure_realtime_cost_calculator(_local_model_cost_map):
cost = handle_realtime_stream_cost_calculation(
results=[
@ -1152,7 +1132,7 @@ def test_azure_realtime_cost_calculator():
assert cost > 0
def test_azure_audio_output_cost_calculation():
def test_azure_audio_output_cost_calculation(_local_model_cost_map):
"""
Test that Azure audio models correctly calculate costs for audio output tokens.
@ -1162,8 +1142,6 @@ def test_azure_audio_output_cost_calculation():
"""
from litellm.types.utils import Choices, CompletionTokensDetailsWrapper, Message
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
# Scenario from issue #19764:
# Input: 17 text tokens, 0 audio tokens
@ -1672,7 +1650,7 @@ def test_gemini_25_explicit_caching_cost_direct_usage():
assert expected_actual_cost == total_cost
def test_azure_ai_cache_cost_calculation():
def test_azure_ai_cache_cost_calculation(_local_model_cost_map):
"""
Test that azure_ai provider correctly calculates cache costs using generic_cost_per_token.
@ -1683,8 +1661,6 @@ def test_azure_ai_cache_cost_calculation():
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
# Register a custom azure_ai model with cache pricing
test_model_id = "test-azure-ai-claude-model"
@ -1817,15 +1793,13 @@ def test_vertex_uplift_composes_with_above_128k_pricing(monkeypatch):
assert regional_completion == pytest.approx(global_completion * 1.10, rel=1e-9)
def test_cost_discount_vertex_ai():
def test_cost_discount_vertex_ai(monkeypatch):
"""
Test that cost discount is applied correctly for Vertex AI provider
"""
from litellm import completion_cost
from litellm.types.utils import Usage
# Save original config
original_discount_config = litellm.cost_discount_config.copy()
# Create mock response (use a model that exists in model_prices_and_context_window.json)
response = ModelResponse(
@ -1838,7 +1812,7 @@ def test_cost_discount_vertex_ai():
)
# Calculate cost without discount
litellm.cost_discount_config = {}
monkeypatch.setattr(litellm, "cost_discount_config", {})
cost_without_discount = completion_cost(
completion_response=response,
model="vertex_ai/gemini-3-pro-preview",
@ -1846,7 +1820,7 @@ def test_cost_discount_vertex_ai():
)
# Set 5% discount for vertex_ai
litellm.cost_discount_config = {"vertex_ai": 0.05}
monkeypatch.setattr(litellm, "cost_discount_config", {"vertex_ai": 0.05})
# Calculate cost with discount
cost_with_discount = completion_cost(
@ -1855,8 +1829,6 @@ def test_cost_discount_vertex_ai():
custom_llm_provider="vertex_ai",
)
# Restore original config
litellm.cost_discount_config = original_discount_config
# Verify discount is applied (5% off means 95% of original cost)
expected_cost = cost_without_discount * 0.95
@ -1868,15 +1840,13 @@ def test_cost_discount_vertex_ai():
print(f" - Savings: ${cost_without_discount - cost_with_discount:.6f}")
def test_cost_discount_not_applied_to_other_providers():
def test_cost_discount_not_applied_to_other_providers(monkeypatch):
"""
Test that cost discount only applies to configured providers
"""
from litellm import completion_cost
from litellm.types.utils import Usage
# Save original config
original_discount_config = litellm.cost_discount_config.copy()
# Create mock response for OpenAI
response = ModelResponse(
@ -1889,7 +1859,7 @@ def test_cost_discount_not_applied_to_other_providers():
)
# Set discount only for vertex_ai (not openai)
litellm.cost_discount_config = {"vertex_ai": 0.05}
monkeypatch.setattr(litellm, "cost_discount_config", {"vertex_ai": 0.05})
# Calculate cost for OpenAI - should NOT have discount applied
cost_with_selective_discount = completion_cost(
@ -1899,15 +1869,13 @@ def test_cost_discount_not_applied_to_other_providers():
)
# Clear discount config
litellm.cost_discount_config = {}
monkeypatch.setattr(litellm, "cost_discount_config", {})
cost_without_discount = completion_cost(
completion_response=response,
model="gpt-4",
custom_llm_provider="openai",
)
# Restore original config
litellm.cost_discount_config = original_discount_config
# Costs should be the same (no discount applied to OpenAI)
assert cost_with_selective_discount == cost_without_discount
@ -1917,15 +1885,13 @@ def test_cost_discount_not_applied_to_other_providers():
print(f" - Cost remains unchanged: ${cost_with_selective_discount:.6f}")
def test_cost_margin_percentage():
def test_cost_margin_percentage(monkeypatch):
"""
Test that percentage-based cost margin is applied correctly
"""
from litellm import completion_cost
from litellm.types.utils import Usage
# Save original config
original_margin_config = litellm.cost_margin_config.copy()
# Create mock response
response = ModelResponse(
@ -1938,7 +1904,7 @@ def test_cost_margin_percentage():
)
# Calculate cost without margin
litellm.cost_margin_config = {}
monkeypatch.setattr(litellm, "cost_margin_config", {})
cost_without_margin = completion_cost(
completion_response=response,
model="gpt-4",
@ -1946,7 +1912,7 @@ def test_cost_margin_percentage():
)
# Set 10% margin for openai
litellm.cost_margin_config = {"openai": 0.10}
monkeypatch.setattr(litellm, "cost_margin_config", {"openai": 0.10})
# Calculate cost with margin
cost_with_margin = completion_cost(
@ -1955,8 +1921,6 @@ def test_cost_margin_percentage():
custom_llm_provider="openai",
)
# Restore original config
litellm.cost_margin_config = original_margin_config
# Verify margin is applied (10% margin means 110% of original cost)
expected_cost = cost_without_margin * 1.10
@ -1968,15 +1932,13 @@ def test_cost_margin_percentage():
print(f" - Margin added: ${cost_with_margin - cost_without_margin:.6f}")
def test_cost_margin_fixed_amount():
def test_cost_margin_fixed_amount(monkeypatch):
"""
Test that fixed amount cost margin is applied correctly
"""
from litellm import completion_cost
from litellm.types.utils import Usage
# Save original config
original_margin_config = litellm.cost_margin_config.copy()
# Create mock response
response = ModelResponse(
@ -1989,7 +1951,7 @@ def test_cost_margin_fixed_amount():
)
# Calculate cost without margin
litellm.cost_margin_config = {}
monkeypatch.setattr(litellm, "cost_margin_config", {})
cost_without_margin = completion_cost(
completion_response=response,
model="gpt-4",
@ -1997,7 +1959,7 @@ def test_cost_margin_fixed_amount():
)
# Set $0.001 fixed margin for openai
litellm.cost_margin_config = {"openai": {"fixed_amount": 0.001}}
monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": 0.001}})
# Calculate cost with margin
cost_with_margin = completion_cost(
@ -2006,8 +1968,6 @@ def test_cost_margin_fixed_amount():
custom_llm_provider="openai",
)
# Restore original config
litellm.cost_margin_config = original_margin_config
# Verify fixed margin is applied
expected_cost = cost_without_margin + 0.001
@ -2019,15 +1979,13 @@ def test_cost_margin_fixed_amount():
print(f" - Margin added: ${cost_with_margin - cost_without_margin:.6f}")
def test_cost_margin_combined():
def test_cost_margin_combined(monkeypatch):
"""
Test that combined percentage and fixed amount margin is applied correctly
"""
from litellm import completion_cost
from litellm.types.utils import Usage
# Save original config
original_margin_config = litellm.cost_margin_config.copy()
# Create mock response
response = ModelResponse(
@ -2040,7 +1998,7 @@ def test_cost_margin_combined():
)
# Calculate cost without margin
litellm.cost_margin_config = {}
monkeypatch.setattr(litellm, "cost_margin_config", {})
cost_without_margin = completion_cost(
completion_response=response,
model="gpt-4",
@ -2048,9 +2006,9 @@ def test_cost_margin_combined():
)
# Set 8% margin + $0.0005 fixed for openai
litellm.cost_margin_config = {
monkeypatch.setattr(litellm, "cost_margin_config", {
"openai": {"percentage": 0.08, "fixed_amount": 0.0005}
}
})
# Calculate cost with margin
cost_with_margin = completion_cost(
@ -2059,8 +2017,6 @@ def test_cost_margin_combined():
custom_llm_provider="openai",
)
# Restore original config
litellm.cost_margin_config = original_margin_config
# Verify combined margin is applied
expected_cost = cost_without_margin * 1.08 + 0.0005
@ -2072,15 +2028,13 @@ def test_cost_margin_combined():
print(f" - Margin added: ${cost_with_margin - cost_without_margin:.6f}")
def test_cost_margin_global():
def test_cost_margin_global(monkeypatch):
"""
Test that global margin is applied when no provider-specific margin is configured
"""
from litellm import completion_cost
from litellm.types.utils import Usage
# Save original config
original_margin_config = litellm.cost_margin_config.copy()
# Create mock response
response = ModelResponse(
@ -2093,7 +2047,7 @@ def test_cost_margin_global():
)
# Calculate cost without margin
litellm.cost_margin_config = {}
monkeypatch.setattr(litellm, "cost_margin_config", {})
cost_without_margin = completion_cost(
completion_response=response,
model="gpt-4",
@ -2101,7 +2055,7 @@ def test_cost_margin_global():
)
# Set 5% global margin (no provider-specific margin)
litellm.cost_margin_config = {"global": 0.05}
monkeypatch.setattr(litellm, "cost_margin_config", {"global": 0.05})
# Calculate cost with global margin
cost_with_global_margin = completion_cost(
@ -2110,8 +2064,6 @@ def test_cost_margin_global():
custom_llm_provider="openai",
)
# Restore original config
litellm.cost_margin_config = original_margin_config
# Verify global margin is applied
expected_cost = cost_without_margin * 1.05
@ -2123,15 +2075,13 @@ def test_cost_margin_global():
print(f" - Margin added: ${cost_with_global_margin - cost_without_margin:.6f}")
def test_cost_margin_provider_overrides_global():
def test_cost_margin_provider_overrides_global(monkeypatch):
"""
Test that provider-specific margin overrides global margin
"""
from litellm import completion_cost
from litellm.types.utils import Usage
# Save original config
original_margin_config = litellm.cost_margin_config.copy()
# Create mock response
response = ModelResponse(
@ -2144,7 +2094,7 @@ def test_cost_margin_provider_overrides_global():
)
# Calculate cost without margin
litellm.cost_margin_config = {}
monkeypatch.setattr(litellm, "cost_margin_config", {})
cost_without_margin = completion_cost(
completion_response=response,
model="gpt-4",
@ -2152,7 +2102,7 @@ def test_cost_margin_provider_overrides_global():
)
# Set 5% global margin and 10% provider-specific margin
litellm.cost_margin_config = {"global": 0.05, "openai": 0.10}
monkeypatch.setattr(litellm, "cost_margin_config", {"global": 0.05, "openai": 0.10})
# Calculate cost - should use provider-specific margin (10%), not global (5%)
cost_with_provider_margin = completion_cost(
@ -2161,8 +2111,6 @@ def test_cost_margin_provider_overrides_global():
custom_llm_provider="openai",
)
# Restore original config
litellm.cost_margin_config = original_margin_config
# Verify provider-specific margin is used (not global)
expected_cost = cost_without_margin * 1.10 # 10% from provider, not 5% from global
@ -2176,16 +2124,13 @@ def test_cost_margin_provider_overrides_global():
print(f" - Margin added: ${cost_with_provider_margin - cost_without_margin:.6f}")
def test_cost_margin_with_discount():
def test_cost_margin_with_discount(monkeypatch):
"""
Test that margin is applied after discount (independent calculation)
"""
from litellm import completion_cost
from litellm.types.utils import Usage
# Save original configs
original_margin_config = litellm.cost_margin_config.copy()
original_discount_config = litellm.cost_discount_config.copy()
# Create mock response
response = ModelResponse(
@ -2198,8 +2143,8 @@ def test_cost_margin_with_discount():
)
# Calculate base cost
litellm.cost_margin_config = {}
litellm.cost_discount_config = {}
monkeypatch.setattr(litellm, "cost_margin_config", {})
monkeypatch.setattr(litellm, "cost_discount_config", {})
base_cost = completion_cost(
completion_response=response,
model="gpt-4",
@ -2207,8 +2152,8 @@ def test_cost_margin_with_discount():
)
# Set 5% discount and 10% margin
litellm.cost_discount_config = {"openai": 0.05}
litellm.cost_margin_config = {"openai": 0.10}
monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.05})
monkeypatch.setattr(litellm, "cost_margin_config", {"openai": 0.10})
# Calculate cost with both discount and margin
cost_with_both = completion_cost(
@ -2217,9 +2162,6 @@ def test_cost_margin_with_discount():
custom_llm_provider="openai",
)
# Restore original configs
litellm.cost_margin_config = original_margin_config
litellm.cost_discount_config = original_discount_config
# Verify: discount applied first, then margin
# Base cost -> discount: base * 0.95 -> margin: (base * 0.95) * 1.10
@ -2286,12 +2228,10 @@ def test_azure_image_generation_cost_calculator():
assert cost > 0.079
def test_completion_cost_extracts_service_tier_from_response():
def test_completion_cost_extracts_service_tier_from_response(_local_model_cost_map):
"""Test that completion_cost extracts service_tier from completion_response object."""
from litellm import completion_cost
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
# Test with gpt-5-nano which has flex pricing
model = "gpt-5-nano"
@ -2338,12 +2278,10 @@ def test_completion_cost_extracts_service_tier_from_response():
), f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}"
def test_completion_cost_extracts_service_tier_from_usage():
def test_completion_cost_extracts_service_tier_from_usage(_local_model_cost_map):
"""Test that completion_cost extracts service_tier from usage object."""
from litellm import completion_cost
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
# Test with gpt-5-nano which has flex pricing
model = "gpt-5-nano"
@ -2397,12 +2335,10 @@ def test_completion_cost_extracts_service_tier_from_usage():
), f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}"
def test_completion_cost_service_tier_priority():
def test_completion_cost_service_tier_priority(_local_model_cost_map):
"""Test that service_tier extraction follows priority: optional_params > completion_response > usage."""
from litellm import completion_cost
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
# Test with gpt-5-nano which has flex pricing
model = "gpt-5-nano"
@ -2457,12 +2393,10 @@ def test_completion_cost_service_tier_priority():
), "Costs from params and usage should be similar (both flex)"
def test_completion_cost_service_tier_for_bedrock():
def test_completion_cost_service_tier_for_bedrock(_local_model_cost_map):
"""Test that Bedrock cost calculation applies service_tier-specific pricing."""
from litellm import completion_cost
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "bedrock/us-east-1/test-bedrock-service-tier-cost-model"
litellm.register_model(
@ -2507,7 +2441,7 @@ def test_completion_cost_service_tier_for_bedrock():
assert priority_cost > default_cost > flex_cost > 0
def test_completion_cost_service_tier_for_anthropic():
def test_completion_cost_service_tier_for_anthropic(_local_model_cost_map):
"""
Anthropic priority-tier requests must be priced at the priority rate.
@ -2519,8 +2453,6 @@ def test_completion_cost_service_tier_for_anthropic():
from litellm import completion_cost
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "claude-test-service-tier-cost-model"
litellm.register_model(
@ -2561,7 +2493,7 @@ def test_completion_cost_service_tier_for_anthropic():
assert priority_cost == pytest.approx(2 * standard_cost)
def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate():
def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate(_local_model_cost_map):
"""
Proxy billing path regression for LIT-3771.
@ -2574,8 +2506,6 @@ def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate():
from litellm import completion_cost
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "claude-test-auto-tier-cost-model"
litellm.register_model(
@ -2613,7 +2543,7 @@ def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate():
assert cost == pytest.approx(expected_priority)
def test_completion_cost_non_string_service_tier_defers_to_served_tier():
def test_completion_cost_non_string_service_tier_defers_to_served_tier(_local_model_cost_map):
"""
Regression: a non-string request-level ``service_tier`` (reachable via
``allowed_openai_params``/``drop_params``) must not crash cost tracking.
@ -2627,8 +2557,6 @@ def test_completion_cost_non_string_service_tier_defers_to_served_tier():
from litellm import completion_cost
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "claude-test-non-string-tier-cost-model"
litellm.register_model(
@ -2665,7 +2593,7 @@ def test_completion_cost_non_string_service_tier_defers_to_served_tier():
assert cost == pytest.approx(expected_priority)
def test_completion_cost_non_string_response_service_tier_defers_to_served_tier():
def test_completion_cost_non_string_response_service_tier_defers_to_served_tier(_local_model_cost_map):
"""
Regression: a non-string ``service_tier`` on the response object must not
crash cost tracking.
@ -2679,8 +2607,6 @@ def test_completion_cost_non_string_response_service_tier_defers_to_served_tier(
from litellm import completion_cost
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "claude-test-response-non-string-tier-cost-model"
litellm.register_model(
@ -2718,7 +2644,7 @@ def test_completion_cost_non_string_response_service_tier_defers_to_served_tier(
assert cost == pytest.approx(expected_priority)
def test_completion_cost_non_string_usage_service_tier_prices_standard():
def test_completion_cost_non_string_usage_service_tier_prices_standard(_local_model_cost_map):
"""
Regression: a non-string ``service_tier`` on the usage object must not crash
cost tracking.
@ -2729,8 +2655,6 @@ def test_completion_cost_non_string_usage_service_tier_prices_standard():
"""
from litellm import completion_cost
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "claude-test-usage-non-string-tier-cost-model"
litellm.register_model(
@ -2764,7 +2688,7 @@ def test_completion_cost_non_string_usage_service_tier_prices_standard():
assert cost == pytest.approx(expected_standard)
def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier():
def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier(_local_model_cost_map):
"""
Regression for the cache/tier interaction in the Anthropic geo/speed path.
@ -2780,8 +2704,6 @@ def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier():
)
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "claude-test-priority-cache-fast-model"
litellm.register_model(
@ -2837,7 +2759,7 @@ def _register_anthropic_geo_cache_model(model: str) -> None:
)
def test_anthropic_geo_multiplier_applies_to_cache_tokens(monkeypatch):
def test_anthropic_geo_multiplier_applies_to_cache_tokens(_local_model_cost_map, monkeypatch):
"""
Regression: the regional (geo) uplift must scale cache read and cache write
cost too, not just non-cache input and output.
@ -2853,7 +2775,6 @@ def test_anthropic_geo_multiplier_applies_to_cache_tokens(monkeypatch):
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "claude-test-geo-cache-model"
_register_anthropic_geo_cache_model(model)
@ -2882,7 +2803,7 @@ def test_anthropic_geo_multiplier_applies_to_cache_tokens(monkeypatch):
assert geo_completion_cost == pytest.approx(base_completion_cost * 1.1)
def test_anthropic_geo_and_fast_multipliers_compose(monkeypatch):
def test_anthropic_geo_and_fast_multipliers_compose(_local_model_cost_map, monkeypatch):
"""
The ``fast`` speed multiplier stays cache-exclusive (the old explicit
``fast/`` entries kept base cache rates) while the geo multiplier scales the
@ -2895,7 +2816,6 @@ def test_anthropic_geo_and_fast_multipliers_compose(monkeypatch):
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "claude-test-geo-fast-cache-model"
_register_anthropic_geo_cache_model(model)
@ -3100,7 +3020,7 @@ def test_gemini_implicit_caching_cost_calculation():
)
def test_additional_costs_only_for_azure_ai():
def test_additional_costs_only_for_azure_ai(_local_model_cost_map):
"""
Test that _get_additional_costs is only called for azure_ai provider.
@ -3111,8 +3031,6 @@ def test_additional_costs_only_for_azure_ai():
"""
from litellm.cost_calculator import _get_additional_costs
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
# Non-azure_ai providers should return None
result = _get_additional_costs(
@ -3140,7 +3058,7 @@ def test_additional_costs_only_for_azure_ai():
assert result is None, "Vertex AI should have no additional costs"
def test_openrouter_gemini_3_1_flash_lite_preview_pricing():
def test_openrouter_gemini_3_1_flash_lite_preview_pricing(_local_model_cost_map):
"""
Test that openrouter/google/gemini-3.1-flash-lite-preview has a pricing entry.
@ -3150,8 +3068,6 @@ def test_openrouter_gemini_3_1_flash_lite_preview_pricing():
model_prices_and_context_window.json when other Gemini 3.x variants were present.
This caused ValueError: This model isn't mapped yet during router pre-call checks.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model_name = "openrouter/google/gemini-3.1-flash-lite-preview"
model_info = litellm.model_cost.get(model_name)
@ -3164,9 +3080,7 @@ def test_openrouter_gemini_3_1_flash_lite_preview_pricing():
assert model_info["max_output_tokens"] == 65536
def test_gemini_3_1_flash_lite_pricing():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_gemini_3_1_flash_lite_pricing(_local_model_cost_map):
for model_name in (
"gemini-3.1-flash-lite",
@ -3489,7 +3403,7 @@ def test_custom_pricing_without_cache_keys_preserves_legacy_behavior():
assert cost == pytest.approx(expected)
def test_openrouter_gemini_3_1_flash_lite_stable_pricing():
def test_openrouter_gemini_3_1_flash_lite_stable_pricing(_local_model_cost_map):
"""
Test that openrouter/google/gemini-3.1-flash-lite (stable, no -preview suffix)
has a pricing entry.
@ -3505,8 +3419,6 @@ def test_openrouter_gemini_3_1_flash_lite_stable_pricing():
Pricing matches the existing -preview entry one-for-one (input $0.25/M, output
$1.50/M, cache-read $0.025/M) — Google did not change costs at the GA cutover.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model_name = "openrouter/google/gemini-3.1-flash-lite"
model_info = litellm.model_cost.get(model_name)
@ -3520,7 +3432,7 @@ def test_openrouter_gemini_3_1_flash_lite_stable_pricing():
assert model_info["max_output_tokens"] == 65536
def test_completion_cost_logs_reasoning_and_cache_breakdown():
def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_map):
"""
completion_cost must surface explicit reasoning and cache-read costs into the
cost_breakdown stored on the logging object, so they end up in the spend logs
@ -3531,8 +3443,6 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown():
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.utils import Choices, CompletionTokensDetailsWrapper, Message
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
logging_obj = Logging(
model="gemini-2.5-flash",
@ -3750,13 +3660,11 @@ def test_combine_usage_objects_sums_mirrored_cache_write_fields_once():
assert combined_pair.prompt_tokens_details.cache_creation_tokens == 100
def test_completion_cost_prices_anthropic_shaped_cache_read_tokens():
def test_completion_cost_prices_anthropic_shaped_cache_read_tokens(_local_model_cost_map):
"""Regression: an Anthropic /v1/messages response reports cache reads as top-level
cache_read_input_tokens with input_tokens excluding them. Reading that usage as
Responses API usage dropped the cache tokens and billed the whole prompt at the
uncached input rate, overstating spend on cache hits."""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
response = {
"id": "msg_1",

3
tests/test_litellm/test_main.py
View file

@ -2836,7 +2836,10 @@ def _priced_at(prompt_tokens, completion_tokens):
@pytest.fixture
def local_cost_map(monkeypatch):
"""The prices these tests assert are the checked-in ones. Setting the environment
variable alone does not reload the map, so pin the map itself."""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
def test_a_streamed_response_bills_the_usage_the_provider_reported(local_cost_map):

125
tests/test_litellm/test_router_model_cost_isolation.py
View file

@ -8,6 +8,7 @@ should still use the built-in pricing.
"""
import copy
import logging
import os
import re
import sys
@ -2007,3 +2008,127 @@ def test_a_falsy_id_is_still_scanned_for_collisions():
},
]
)
# --- a reservation declared while the feature is off says so ------------------------
def _ptu_warnings(caplog):
return tuple(
record.getMessage()
for record in caplog.records
if record.name == "LiteLLM Router" and record.levelno == logging.WARNING and "PTU" in record.getMessage()
)
def test_a_reservation_declared_while_the_feature_is_off_is_warned_about(caplog):
"""The deployment serves and bills per token, so without this the operator believes they
reserved capacity and sees no signal anywhere that nothing accrues."""
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
_ptu_router(ptu_enabled=False)
warnings = _ptu_warnings(caplog)
assert len(warnings) == 1
assert "gpt-4o-ptu" in warnings[0]
assert "LITELLM_ENABLE_PTU_COST_ATTRIBUTION" in warnings[0]
def test_a_reservation_is_not_warned_about_while_the_feature_is_on(caplog):
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
_ptu_router()
assert _ptu_warnings(caplog) == ()
def test_a_deployment_carrying_no_ptu_field_is_not_warned_about(caplog):
"""Most of every config.yaml, so warning here would fire on proxies that never asked."""
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
_ptu_router(model_info={"team_id": "team-alpha"}, ptu_enabled=False)
assert _ptu_warnings(caplog) == ()
def test_a_half_written_reservation_is_warned_about(caplog):
"""A count with no rate is not a chargeable reservation, but the operator still meant to
declare one, so what they wrote is what decides whether they hear about it."""
half_written = {k: v for k, v in _PTU_MODEL_INFO.items() if k != "cost_per_ptu_per_hour"}
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
_ptu_router(model_info=half_written, ptu_enabled=False)
assert len(_ptu_warnings(caplog)) == 1
@pytest.mark.parametrize(
"typo",
[
{"ptu_count": 0},
{"ptu_count": 0, "cost_per_ptu_per_hour": 0, "ptu_effective_from": None},
],
ids=["count out of range", "every value still a zero placeholder"],
)
def test_a_reservation_dropped_by_a_typo_is_warned_about(caplog, typo):
"""An out-of-range value fails ModelInfo before the flag is ever consulted, so the
deployment stops serving on a proxy that never enabled PTU. The warning is what tells the
operator which feature the entry that vanished belonged to.
Built the way proxy_server builds it, since dropping rather than raising is what
``ignore_invalid_deployments`` does and config.yaml is loaded with it on.
"""
with patch.dict(os.environ, {"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": ""}, clear=False):
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
router = Router(
ignore_invalid_deployments=True,
model_list=[
{
"model_name": "gpt-4o-ptu",
"litellm_params": {"model": "azure/gpt-4o", "api_key": "k", "api_base": "https://e.azure.com"},
"model_info": {**_PTU_MODEL_INFO, **typo},
}
],
)
assert router.model_list == []
assert len(_ptu_warnings(caplog)) == 1
def test_a_db_backed_reservation_is_not_warned_about(caplog):
"""/model/new already answered the caller with a 400, so repeating it on every reload
would report the operator's own rejected write back to them as a standing problem."""
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
_ptu_router(model_info={**_PTU_MODEL_INFO, "db_model": True}, ptu_enabled=False)
assert _ptu_warnings(caplog) == ()
def test_every_declaring_deployment_is_named(caplog):
"""One line naming all of them, so a reload does not bury the config in repeats."""
with patch.dict(os.environ, {"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": ""}, clear=False):
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
Router(
model_list=[
{
"model_name": "azure-ptu-east",
"litellm_params": {"model": "azure/gpt-4o", "api_key": "k", "api_base": "https://e.azure.com"},
"model_info": dict(_PTU_MODEL_INFO),
},
{
"model_name": "azure-ptu-west",
"litellm_params": {"model": "azure/gpt-4o", "api_key": "k", "api_base": "https://w.azure.com"},
"model_info": {**_PTU_MODEL_INFO, "id": "ptu-alpha-westus"},
},
{
"model_name": "plain-gpt-4o",
"litellm_params": {"model": "azure/gpt-4o", "api_key": "k", "api_base": "https://p.azure.com"},
"model_info": {"id": "plain"},
},
]
)
warnings = _ptu_warnings(caplog)
assert len(warnings) == 1
assert "azure-ptu-east" in warnings[0]
assert "azure-ptu-west" in warnings[0]
assert "plain-gpt-4o" not in warnings[0]