mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge pull request #33308 from BerriAI/litellm_internal_staging
chore(ci): promote internal staging to main
This commit is contained in:
commit
6a797f97b2
329 changed files with 19507 additions and 5643 deletions
2
.github/CODEOWNERS
vendored
Normal file
2
.github/CODEOWNERS
vendored
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
/ui/ @yuneng-jiang @ryan-crabbe-berri
|
||||
/litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri
|
||||
47
.github/actions/setup-uv-with-retries/action.yml
vendored
Normal file
47
.github/actions/setup-uv-with-retries/action.yml
vendored
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
name: "Set up uv with retries"
|
||||
description: >-
|
||||
Install uv via astral-sh/setup-uv, retrying on transient failures. Even with
|
||||
an exact pinned version, the action resolves the artifact URL by fetching
|
||||
https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a
|
||||
single request with no retry, timeout, or fallback, so one connection-level
|
||||
network error ("fetch failed") fails the whole job before any test runs.
|
||||
Retrying the full step covers the manifest fetch and the binary download.
|
||||
|
||||
inputs:
|
||||
version:
|
||||
description: "uv version to install"
|
||||
required: true
|
||||
|
||||
runs:
|
||||
using: composite
|
||||
steps:
|
||||
- name: Set up uv (attempt 1)
|
||||
id: attempt-1
|
||||
continue-on-error: true
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
|
||||
with:
|
||||
version: ${{ inputs.version }}
|
||||
|
||||
- name: Wait before attempt 2
|
||||
if: steps.attempt-1.outcome == 'failure'
|
||||
shell: bash
|
||||
run: sleep 15
|
||||
|
||||
- name: Set up uv (attempt 2)
|
||||
id: attempt-2
|
||||
if: steps.attempt-1.outcome == 'failure'
|
||||
continue-on-error: true
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
|
||||
with:
|
||||
version: ${{ inputs.version }}
|
||||
|
||||
- name: Wait before attempt 3
|
||||
if: steps.attempt-2.outcome == 'failure'
|
||||
shell: bash
|
||||
run: sleep 30
|
||||
|
||||
- name: Set up uv (attempt 3)
|
||||
if: steps.attempt-2.outcome == 'failure'
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
|
||||
with:
|
||||
version: ${{ inputs.version }}
|
||||
24
.github/pull_request_template.md
vendored
24
.github/pull_request_template.md
vendored
|
|
@ -41,3 +41,27 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
|
|||
✅ Test
|
||||
|
||||
## Changes
|
||||
|
||||
## QA runbook
|
||||
|
||||
<!-- Only needed when your PR edits tests/e2e; delete this section otherwise
|
||||
|
||||
For each e2e test you added or changed, list the manual steps a reviewer can follow to reproduce it by hand against a live proxy, mapping 1:1 to what the test asserts: one top-level bullet per test giving its pytest node id followed by what it proves in plain words, then a nested "- [ ]" checklist where each item is a concrete action (route, request body, expected response) and the final item is the sanity-check step shown in the examples. Note environment prerequisites (provider credentials, config flags) and any nuances a manual run will hit. See PRs #32914 and #32963 for full examples
|
||||
|
||||
Example checklists:
|
||||
|
||||
- tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py::TestKeyRateLimits::test_rpm_limit_blocks_over_limit - a key allowed 2 requests a minute serves exactly 2 and refuses the 3rd
|
||||
- [ ] Generate a limited key: curl -X POST http://localhost:4000/key/generate -H "Authorization: Bearer sk-1234" -d '{"rpm_limit": 2}'
|
||||
- [ ] Send three /v1/chat/completions requests with that key inside one minute
|
||||
- [ ] Expect the first two to return 200 and the third to return 429 naming the rpm limit
|
||||
- [ ] Sanity check: this test makes sense to add and is not hand-wavey (e.g., assert actual expected spend instead of just spend > 0) or potentially flaky
|
||||
|
||||
- tests/e2e/management/test_management_e2e.py::TestModelRoutes::test_model_create_appears_in_ui - a deployment created through the API shows up on the Admin UI models page
|
||||
- [ ] POST /model/new with the master key, a bedrock model, and aws_region_name (needs STORE_MODEL_IN_DB=True and AWS credentials)
|
||||
- [ ] Open http://localhost:4000/ui/?page=models and expect a deployment row showing the returned model id
|
||||
- [ ] Sanity check: this test makes sense to add and is not hand-wavey (e.g., assert actual expected spend instead of just spend > 0) or potentially flaky
|
||||
-->
|
||||
|
||||
### Final Attestation
|
||||
|
||||
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
|
||||
|
|
|
|||
2
.github/workflows/_test-unit-base.yml
vendored
2
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -63,7 +63,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ jobs:
|
|||
with:
|
||||
persist-credentials: false
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
- name: Update JSON Data
|
||||
|
|
|
|||
2
.github/workflows/check-ui-api-types.yml
vendored
2
.github/workflows/check-ui-api-types.yml
vendored
|
|
@ -31,7 +31,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
2
.github/workflows/codspeed.yml
vendored
2
.github/workflows/codspeed.yml
vendored
|
|
@ -37,7 +37,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
2
.github/workflows/mutation-test.yml
vendored
2
.github/workflows/mutation-test.yml
vendored
|
|
@ -39,7 +39,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
2
.github/workflows/oss_daily_guardrails.yml
vendored
2
.github/workflows/oss_daily_guardrails.yml
vendored
|
|
@ -35,7 +35,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
2
.github/workflows/test-code-quality.yml
vendored
2
.github/workflows/test-code-quality.yml
vendored
|
|
@ -38,7 +38,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
4
.github/workflows/test-linting.yml
vendored
4
.github/workflows/test-linting.yml
vendored
|
|
@ -33,7 +33,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
@ -172,7 +172,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
2
.github/workflows/test-litellm-ui-lint.yml
vendored
2
.github/workflows/test-litellm-ui-lint.yml
vendored
|
|
@ -89,4 +89,4 @@ jobs:
|
|||
|
||||
- name: Check for dead code (knip)
|
||||
if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }}
|
||||
run: npm run knip
|
||||
run: npm run knip:ci
|
||||
|
|
|
|||
2
.github/workflows/test-mcp.yml
vendored
2
.github/workflows/test-mcp.yml
vendored
|
|
@ -32,7 +32,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
2
.github/workflows/test-semgrep.yml
vendored
2
.github/workflows/test-semgrep.yml
vendored
|
|
@ -31,7 +31,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
|
|
@ -74,7 +74,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
2
.github/workflows/test-unit-proxy-legacy.yml
vendored
2
.github/workflows/test-unit-proxy-legacy.yml
vendored
|
|
@ -59,7 +59,7 @@ jobs:
|
|||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
|
|
|
|||
|
|
@ -39,6 +39,8 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a
|
|||
|
||||
Python max line length is 120, not 88
|
||||
|
||||
On a fresh worktree or clone, run `make bootstrap` before anything else. It provisions everything tests, `make pre-commit`, and a local proxy need
|
||||
|
||||
Run tests before you commit. Also, run `make pre-commit` right before each commit, which generates types (as needed) and formats/lints your code. Any errors found must be fixed. It only runs when there are staged frontend and/or backend changes and calculates violations, generates types, etc. based on the worktree, so stage what you need or stash/delete unwanted files in litellm/ or ui/ (where backend and frontend lint run, respectively) before running it. If it fails because dashboard api types are stale, it already regenerated them for you. You just need to stage the schema.d.ts, re-run `make pre-commit` to confirm it passes, and commit
|
||||
|
||||
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
|
||||
|
|
|
|||
15
Makefile
15
Makefile
|
|
@ -9,11 +9,12 @@
|
|||
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
|
||||
install-dev install-proxy-dev install-test-deps install-hooks \
|
||||
install-helm-unittest check-circular-imports check-import-safety pre-commit \
|
||||
lint-install lint-fetch-base
|
||||
lint-install lint-fetch-base bootstrap
|
||||
|
||||
# Default target
|
||||
help:
|
||||
@echo "Available commands:"
|
||||
@echo " make bootstrap - Provision a fresh clone/worktree"
|
||||
@echo " make install-dev - Install development dependencies"
|
||||
@echo " make install-proxy-dev - Install proxy development dependencies"
|
||||
@echo " make install-dev-ci - Install dev dependencies (CI-compatible, pins OpenAI)"
|
||||
|
|
@ -71,6 +72,18 @@ info:
|
|||
install-dev:
|
||||
$(UV) sync --inexact --frozen
|
||||
|
||||
bootstrap:
|
||||
$(UV) sync --inexact --frozen --extra proxy --group proxy-dev --group e2e-dev
|
||||
$(UV_RUN) python scripts/prisma_generate_if_needed.py
|
||||
cd ui/litellm-dashboard && npm ci --no-audit --no-fund
|
||||
@main_root=$$(git worktree list --porcelain | head -1 | sed 's/^worktree //'); \
|
||||
if [ "$$main_root" != "$$(git rev-parse --show-toplevel)" ] && [ -f "$$main_root/.env" ] && [ ! -f .env ]; then \
|
||||
cp "$$main_root/.env" .env && echo "bootstrap: copied .env from $$main_root"; \
|
||||
else \
|
||||
echo "bootstrap: .env left untouched"; \
|
||||
fi
|
||||
@echo "bootstrap: done"
|
||||
|
||||
install-proxy-dev:
|
||||
$(UV) sync --frozen --group proxy-dev --extra proxy
|
||||
|
||||
|
|
|
|||
13
README.md
13
README.md
|
|
@ -552,17 +552,12 @@ The Terraform modules live at [`terraform/litellm/aws/`](./terraform/litellm/aws
|
|||
2. Run dependent services `docker-compose up db prometheus`
|
||||
|
||||
#### Backend
|
||||
1. (In root) create virtual environment `python -m venv .venv`
|
||||
2. Activate virtual environment `source .venv/bin/activate`
|
||||
3. Install dependencies `uv sync --all-extras --group proxy-dev`
|
||||
4. `uv run prisma generate`
|
||||
5. `prisma generate`
|
||||
6. Start proxy backend `python litellm/proxy/proxy_cli.py`
|
||||
1. Run `make bootstrap`
|
||||
2. Start proxy backend: `uv run python litellm/proxy/proxy_cli.py`
|
||||
|
||||
#### Frontend
|
||||
1. Navigate to `ui/litellm-dashboard`
|
||||
2. Install dependencies `npm install`
|
||||
3. Run `npm run dev` to start the dashboard
|
||||
1. Navigate to `ui/litellm-dashboard` (dependencies were already installed w/ `make bootstrap`)
|
||||
2. Start dashboard: `npm run dev`
|
||||
|
||||
### Verify Docker Image Signatures
|
||||
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ class CheckBatchCost:
|
|||
proxy_logging_obj: "ProxyLogging",
|
||||
prisma_client: "PrismaClient",
|
||||
llm_router: "Router",
|
||||
track_unmanaged_vertex_batch_cost: bool = False,
|
||||
track_unmanaged_batch_cost: bool = False,
|
||||
):
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
|
@ -37,7 +37,7 @@ class CheckBatchCost:
|
|||
self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
|
||||
self.prisma_client: PrismaClient = prisma_client
|
||||
self.llm_router: Router = llm_router
|
||||
self._track_unmanaged_vertex_batch_cost = track_unmanaged_vertex_batch_cost
|
||||
self._track_unmanaged_batch_cost = track_unmanaged_batch_cost
|
||||
# Cached after the first poll cycle. Once we know the column is absent we skip
|
||||
# the guaranteed-failing primary query on every subsequent cycle.
|
||||
self._has_batch_processed_column: bool = True
|
||||
|
|
@ -118,11 +118,11 @@ class CheckBatchCost:
|
|||
Resolve (model_id, batch_id) for a managed-object row, where model_id is a router
|
||||
deployment id and batch_id is the raw provider batch id.
|
||||
|
||||
Managed batches encode both in a base64 unified id. Unmanaged Vertex batches, created with
|
||||
a raw gs:// input_file_id, store the raw provider job id as unified_object_id; when
|
||||
track_unmanaged_vertex_batch_cost is enabled the model is derived from the gs:// path and
|
||||
mapped to a configured vertex_ai deployment. Returns None (recording a metric) when the row
|
||||
can't be routed.
|
||||
Managed batches encode both in a base64 unified id. Unmanaged batches (created outside
|
||||
LiteLLM's own /v1/batches with a raw input_file_id) store the raw provider job id as
|
||||
unified_object_id instead; when track_unmanaged_batch_cost is enabled the model is derived
|
||||
from the provider-specific input_file_id layout (Vertex gs:// or Bedrock s3://) and mapped
|
||||
to a matching deployment. Returns None (recording a metric) when the row can't be routed.
|
||||
"""
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_is_base64_encoded_unified_file_id,
|
||||
|
|
@ -142,8 +142,43 @@ class CheckBatchCost:
|
|||
return None
|
||||
return model_id, get_batch_id_from_unified_batch_id(decoded)
|
||||
|
||||
if self._track_unmanaged_vertex_batch_cost:
|
||||
return self._resolve_unmanaged_vertex_routing(job, prom_logger)
|
||||
if self._track_unmanaged_batch_cost:
|
||||
from litellm.llms.bedrock.batches.transformation import (
|
||||
BedrockBatchesConfig,
|
||||
)
|
||||
from litellm.llms.vertex_ai.batches.transformation import (
|
||||
VertexAIBatchTransformation,
|
||||
)
|
||||
|
||||
input_file_id = self._get_input_file_id(job)
|
||||
if VertexAIBatchTransformation.is_unmanaged_gcs_batch_input_file_id(
|
||||
input_file_id
|
||||
):
|
||||
assert input_file_id is not None # narrowed by is_unmanaged_gcs_batch_input_file_id
|
||||
return self._resolve_unmanaged_provider_routing(
|
||||
job=job,
|
||||
prom_logger=prom_logger,
|
||||
llm_provider="vertex_ai",
|
||||
bare_model_name=VertexAIBatchTransformation.get_bare_model_name_from_gcs_file(
|
||||
input_file_id
|
||||
),
|
||||
)
|
||||
if BedrockBatchesConfig.is_unmanaged_s3_batch_input_file_id(input_file_id):
|
||||
assert input_file_id is not None # narrowed by is_unmanaged_s3_batch_input_file_id
|
||||
return self._resolve_unmanaged_provider_routing(
|
||||
job=job,
|
||||
prom_logger=prom_logger,
|
||||
llm_provider="bedrock",
|
||||
bare_model_name=BedrockBatchesConfig.get_bare_model_name_from_s3_file(
|
||||
input_file_id
|
||||
),
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
f"Skipping job {unified_object_id}: not a recognized unmanaged batch "
|
||||
"(no gs:// or s3:// input_file_id with an embedded model)"
|
||||
)
|
||||
self._record_error(prom_logger, "invalid_unified_id")
|
||||
return None
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"Skipping job {unified_object_id} because it is not a valid unified object id"
|
||||
|
|
@ -151,36 +186,17 @@ class CheckBatchCost:
|
|||
self._record_error(prom_logger, "invalid_unified_id")
|
||||
return None
|
||||
|
||||
def _resolve_unmanaged_vertex_routing(
|
||||
def _resolve_unmanaged_provider_routing(
|
||||
self,
|
||||
job: "LiteLLM_ManagedObjectTable",
|
||||
prom_logger: Optional["PrometheusLogger"],
|
||||
llm_provider: str,
|
||||
bare_model_name: str,
|
||||
) -> Optional[Tuple[str, str]]:
|
||||
from litellm.llms.vertex_ai.batches.transformation import (
|
||||
VertexAIBatchTransformation,
|
||||
)
|
||||
|
||||
input_file_id = self._get_input_file_id(job)
|
||||
if not VertexAIBatchTransformation.is_unmanaged_gcs_batch_input_file_id(
|
||||
input_file_id
|
||||
):
|
||||
verbose_proxy_logger.info(
|
||||
f"Skipping job {job.unified_object_id}: not an unmanaged vertex batch "
|
||||
"(no gs:// input_file_id with a publishers/ model path)"
|
||||
)
|
||||
self._record_error(prom_logger, "invalid_unified_id")
|
||||
return None
|
||||
assert input_file_id is not None # narrowed by is_unmanaged_gcs_batch_input_file_id
|
||||
|
||||
bare_model_name = VertexAIBatchTransformation.get_bare_model_name_from_gcs_file(
|
||||
input_file_id
|
||||
)
|
||||
deployment_id = self._get_vertex_ai_deployment_id_for_bare_model(
|
||||
bare_model_name
|
||||
)
|
||||
deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider)
|
||||
if deployment_id is None:
|
||||
verbose_proxy_logger.info(
|
||||
f"Skipping unmanaged vertex batch {job.unified_object_id}: no vertex_ai "
|
||||
f"Skipping unmanaged {llm_provider} batch {job.unified_object_id}: no {llm_provider} "
|
||||
f"deployment configured for model {bare_model_name}"
|
||||
)
|
||||
self._record_error(prom_logger, "unmanaged_no_matching_deployment")
|
||||
|
|
@ -188,22 +204,22 @@ class CheckBatchCost:
|
|||
|
||||
return deployment_id, job.unified_object_id
|
||||
|
||||
def _get_vertex_ai_deployment_id_for_bare_model(
|
||||
self, bare_model_name: str
|
||||
def _get_deployment_id_for_bare_model(
|
||||
self, bare_model_name: str, llm_provider: str
|
||||
) -> Optional[str]:
|
||||
model_group = self.llm_router.resolve_model_name_from_model_id(bare_model_name)
|
||||
deployment_id = (
|
||||
self._get_vertex_ai_deployment_id(model_group) if model_group else None
|
||||
self._get_deployment_id_for_provider(model_group, llm_provider) if model_group else None
|
||||
)
|
||||
if deployment_id is not None:
|
||||
return deployment_id
|
||||
|
||||
return self._get_vertex_ai_deployment_id_from_matching_deployments(
|
||||
bare_model_name
|
||||
return self._get_deployment_id_from_matching_deployments(
|
||||
bare_model_name, llm_provider
|
||||
)
|
||||
|
||||
def _get_vertex_ai_deployment_id_from_matching_deployments(
|
||||
self, bare_model_name: str
|
||||
def _get_deployment_id_from_matching_deployments(
|
||||
self, bare_model_name: str, llm_provider: str
|
||||
) -> Optional[str]:
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
|
|
@ -215,13 +231,13 @@ class CheckBatchCost:
|
|||
if not self._is_bare_model_match(actual_model, bare_model_name):
|
||||
continue
|
||||
try:
|
||||
_, llm_provider, _, _ = get_llm_provider(
|
||||
_, deployment_llm_provider, _, _ = get_llm_provider(
|
||||
model=actual_model,
|
||||
custom_llm_provider=litellm_params.get("custom_llm_provider"),
|
||||
)
|
||||
except Exception:
|
||||
continue
|
||||
if llm_provider != "vertex_ai":
|
||||
if deployment_llm_provider != llm_provider:
|
||||
continue
|
||||
model_info = deployment.get("model_info") or {}
|
||||
deployment_id = model_info.get("id")
|
||||
|
|
@ -231,15 +247,21 @@ class CheckBatchCost:
|
|||
|
||||
@staticmethod
|
||||
def _is_bare_model_match(actual_model: str, bare_model_name: str) -> bool:
|
||||
# Bedrock model ids may have ":" replaced with "-" in the S3 object key (see
|
||||
# BedrockBatchesConfig.get_bare_model_name_from_s3_file), so normalize both sides;
|
||||
# a no-op for providers like vertex_ai whose model ids never contain a colon.
|
||||
normalized_actual = actual_model.replace(":", "-")
|
||||
normalized_bare = bare_model_name.replace(":", "-")
|
||||
return (
|
||||
actual_model == bare_model_name
|
||||
or actual_model.endswith(f"/{bare_model_name}")
|
||||
or actual_model.endswith(f":{bare_model_name}")
|
||||
normalized_actual == normalized_bare
|
||||
or normalized_actual.endswith(f"/{normalized_bare}")
|
||||
)
|
||||
|
||||
def _get_vertex_ai_deployment_id(self, model_group: str) -> Optional[str]:
|
||||
def _get_deployment_id_for_provider(
|
||||
self, model_group: str, llm_provider: str
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Returns the first deployment id for `model_group` whose provider is vertex_ai,
|
||||
Returns the first deployment id for `model_group` whose provider is `llm_provider`,
|
||||
skipping deployments from other providers that happen to share the model group name.
|
||||
"""
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
|
@ -249,13 +271,13 @@ class CheckBatchCost:
|
|||
if deployment_info is None:
|
||||
continue
|
||||
try:
|
||||
_, llm_provider, _, _ = get_llm_provider(
|
||||
_, deployment_llm_provider, _, _ = get_llm_provider(
|
||||
model=deployment_info.litellm_params.model,
|
||||
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
|
||||
)
|
||||
except Exception:
|
||||
continue
|
||||
if llm_provider == "vertex_ai":
|
||||
if deployment_llm_provider == llm_provider:
|
||||
return deployment_id
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.49"
|
||||
version = "0.1.50"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.49"
|
||||
version = "0.1.50"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,6 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "key_type" TEXT;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "key_type" TEXT;
|
||||
|
||||
|
|
@ -422,6 +422,7 @@ model LiteLLM_VerificationToken {
|
|||
budget_reset_at DateTime?
|
||||
allowed_cache_controls String[] @default([])
|
||||
allowed_routes String[] @default([])
|
||||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
|
|
@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
budget_reset_at DateTime?
|
||||
allowed_cache_controls String[] @default([])
|
||||
allowed_routes String[] @default([])
|
||||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.76"
|
||||
version = "0.4.77"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.76"
|
||||
version = "0.4.77"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -118,6 +118,7 @@ def _batch_cost_calculator(
|
|||
total_cost = _get_batch_job_cost_from_file_content(
|
||||
file_content_dictionary=file_content_dictionary,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
model_info=model_info,
|
||||
)
|
||||
verbose_logger.debug("total_cost=%s", total_cost)
|
||||
|
|
@ -363,6 +364,7 @@ def _count_entry_tokens(
|
|||
def _get_batch_job_cost_from_file_content(
|
||||
file_content_dictionary: List[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
|
||||
model_name: Optional[str] = None,
|
||||
model_info: Optional[ModelInfo] = None,
|
||||
) -> float:
|
||||
"""
|
||||
|
|
@ -377,9 +379,15 @@ def _get_batch_job_cost_from_file_content(
|
|||
for _item in file_content_dictionary:
|
||||
if _batch_response_was_successful(_item, custom_llm_provider):
|
||||
_response_body = _get_response_from_batch_job_output_file(_item, custom_llm_provider)
|
||||
if model_info is not None or custom_llm_provider == "anthropic":
|
||||
if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"):
|
||||
usage = _get_batch_job_usage_from_response_body(_response_body, custom_llm_provider)
|
||||
model = _response_body.get("model", "")
|
||||
# Bedrock batch output lines report a short internal model id
|
||||
# (e.g. "claude-sonnet-4-6") that is not in the cost map; use the
|
||||
# deployment model name for pricing when available.
|
||||
if custom_llm_provider == "bedrock" and model_name:
|
||||
model = model_name
|
||||
else:
|
||||
model = _response_body.get("model") or model_name or ""
|
||||
prompt_cost, completion_cost = batch_cost_calculator(
|
||||
usage=usage,
|
||||
model=model,
|
||||
|
|
@ -485,7 +493,7 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov
|
|||
"""
|
||||
Get the tokens of a batch job from the response body
|
||||
"""
|
||||
if custom_llm_provider == "anthropic":
|
||||
if custom_llm_provider in ("anthropic", "bedrock"):
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
return AnthropicConfig().calculate_usage(
|
||||
|
|
@ -513,6 +521,8 @@ def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom
|
|||
"""
|
||||
if custom_llm_provider == "anthropic":
|
||||
return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("message", None) or {}
|
||||
if custom_llm_provider == "bedrock":
|
||||
return batch_job_output_file.get("modelOutput", None) or {}
|
||||
_response: dict = batch_job_output_file.get("response", None) or {}
|
||||
_response_body = _response.get("body", None) or {}
|
||||
return _response_body
|
||||
|
|
@ -523,9 +533,12 @@ def _batch_response_was_successful(batch_job_output_file: dict, custom_llm_provi
|
|||
Check if the batch job response was successful
|
||||
|
||||
OpenAI-shaped output rows report ``response.status_code == 200``; Anthropic
|
||||
message batch results lines report ``result.type == "succeeded"``.
|
||||
message batch results lines report ``result.type == "succeeded"``; Bedrock
|
||||
batch output lines report ``modelOutput`` (and no ``error``).
|
||||
"""
|
||||
if custom_llm_provider == "anthropic":
|
||||
return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("type") == "succeeded"
|
||||
if custom_llm_provider == "bedrock":
|
||||
return batch_job_output_file.get("modelOutput") is not None and batch_job_output_file.get("error") is None
|
||||
_response: dict = batch_job_output_file.get("response", None) or {}
|
||||
return _response.get("status_code", None) == 200
|
||||
|
|
|
|||
|
|
@ -515,6 +515,22 @@ class CustomGuardrail(CustomLogger):
|
|||
return True
|
||||
return False
|
||||
|
||||
def uses_apply_guardrail_interface(self) -> bool:
|
||||
return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
|
||||
|
||||
def _deployment_pre_call_target(self) -> "CustomLogger":
|
||||
if not self.uses_apply_guardrail_interface():
|
||||
return self
|
||||
try:
|
||||
from litellm.proxy.utils import unified_guardrail
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
f"Guardrail {self.guardrail_name or type(self).__name__} implements apply_guardrail, which needs "
|
||||
"the litellm proxy dependencies to run at the deployment level. "
|
||||
"Install them with: pip install 'litellm[proxy]'"
|
||||
) from e
|
||||
return unified_guardrail
|
||||
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
|
||||
) -> Optional[dict]:
|
||||
|
|
@ -533,7 +549,10 @@ class CustomGuardrail(CustomLogger):
|
|||
|
||||
# CHECK IF GUARDRAIL REJECTS THE REQUEST
|
||||
if call_type == CallTypes.completion or call_type == CallTypes.acompletion:
|
||||
result = await self.async_pre_call_hook(
|
||||
target = self._deployment_pre_call_target()
|
||||
if target is not self:
|
||||
kwargs["guardrail_to_apply"] = self
|
||||
result = await target.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_id=kwargs.get("user_api_key_user_id"),
|
||||
team_id=kwargs.get("user_api_key_team_id"),
|
||||
|
|
@ -543,7 +562,7 @@ class CustomGuardrail(CustomLogger):
|
|||
),
|
||||
cache=dc,
|
||||
data=kwargs,
|
||||
call_type=call_type.value or "acompletion", # type: ignore
|
||||
call_type="completion" if call_type == CallTypes.completion else "acompletion",
|
||||
)
|
||||
|
||||
if result is not None and isinstance(result, dict):
|
||||
|
|
|
|||
|
|
@ -239,6 +239,18 @@ class PrometheusLogger(CustomLogger):
|
|||
labelnames=self.get_labels_for_metric("litellm_output_audio_tokens_metric"),
|
||||
)
|
||||
|
||||
self.litellm_video_duration_seconds_metric = self._counter_factory(
|
||||
"litellm_video_duration_seconds_metric",
|
||||
"Seconds of video generated, from usage.duration_seconds on video generation calls",
|
||||
labelnames=self.get_labels_for_metric("litellm_video_duration_seconds_metric"),
|
||||
)
|
||||
|
||||
self.litellm_images_generated_metric = self._counter_factory(
|
||||
"litellm_images_generated_metric",
|
||||
"Number of images generated, from the image generation response",
|
||||
labelnames=self.get_labels_for_metric("litellm_images_generated_metric"),
|
||||
)
|
||||
|
||||
# Remaining Budget for Team
|
||||
self.litellm_remaining_team_budget_metric = self._gauge_factory(
|
||||
"litellm_remaining_team_budget_metric",
|
||||
|
|
@ -1336,6 +1348,12 @@ class PrometheusLogger(CustomLogger):
|
|||
label_context=label_context,
|
||||
)
|
||||
|
||||
self._increment_media_generation_metrics(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
enum_values=enum_values,
|
||||
label_context=label_context,
|
||||
)
|
||||
|
||||
# MCP tool call metrics
|
||||
self._increment_mcp_tool_call_metrics(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
|
|
@ -1459,8 +1477,65 @@ class PrometheusLogger(CustomLogger):
|
|||
),
|
||||
]
|
||||
|
||||
for counter, metric_name, value in detail_metrics:
|
||||
if not isinstance(value, (int, float)) or value <= 0:
|
||||
PrometheusLogger._inc_sparse_usage_counters(
|
||||
self,
|
||||
detail_metrics,
|
||||
enum_values=enum_values,
|
||||
label_context=label_context,
|
||||
)
|
||||
|
||||
def _increment_media_generation_metrics(
|
||||
self,
|
||||
standard_logging_payload: StandardLoggingPayload,
|
||||
enum_values: UserAPIKeyLabelValues,
|
||||
label_context: PrometheusLabelFactoryContext | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Increment video-seconds and images-generated counters from
|
||||
``standard_logging_payload["metadata"]["usage_object"]``. Video
|
||||
providers report ``duration_seconds`` there; image generation calls
|
||||
report ``output_image_count``. Both are sparse: only emitted when the
|
||||
value is present and > 0, so token-only call types are unaffected.
|
||||
"""
|
||||
metadata = standard_logging_payload.get("metadata") or {}
|
||||
usage_object = metadata.get("usage_object") if isinstance(metadata, dict) else None
|
||||
if not isinstance(usage_object, dict):
|
||||
return
|
||||
|
||||
media_metrics: list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]] = [
|
||||
(
|
||||
self.litellm_video_duration_seconds_metric,
|
||||
"litellm_video_duration_seconds_metric",
|
||||
usage_object.get("duration_seconds"),
|
||||
),
|
||||
(
|
||||
self.litellm_images_generated_metric,
|
||||
"litellm_images_generated_metric",
|
||||
usage_object.get("output_image_count"),
|
||||
),
|
||||
]
|
||||
|
||||
PrometheusLogger._inc_sparse_usage_counters(
|
||||
self,
|
||||
media_metrics,
|
||||
enum_values=enum_values,
|
||||
label_context=label_context,
|
||||
)
|
||||
|
||||
def _inc_sparse_usage_counters(
|
||||
self,
|
||||
counters_with_values: list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]],
|
||||
enum_values: UserAPIKeyLabelValues,
|
||||
label_context: PrometheusLabelFactoryContext | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Increment each ``(counter, metric_name, value)`` entry whose value is
|
||||
a positive number. Non-numeric values (including booleans from
|
||||
malformed provider usage dicts) and values <= 0 are skipped, keeping
|
||||
scrape output sparse.
|
||||
"""
|
||||
for counter, metric_name, value in counters_with_values:
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0:
|
||||
continue
|
||||
PrometheusLogger._inc_labeled_counter(
|
||||
self,
|
||||
|
|
@ -1716,6 +1791,35 @@ class PrometheusLogger(CustomLogger):
|
|||
amount=float(response_cost),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_remaining_from_v3_rate_limit_headers(
|
||||
standard_logging_payload: StandardLoggingPayload | None,
|
||||
rate_limit_type: Literal["requests", "tokens"],
|
||||
) -> int | None:
|
||||
"""
|
||||
Read the per-(key, model) remaining value emitted by the v3 rate
|
||||
limiter (``parallel_request_limiter_v3.py``), which writes
|
||||
``x-ratelimit-model_per_key-remaining-{requests,tokens}`` into
|
||||
``standard_logging_object.hidden_params.additional_headers`` instead
|
||||
of the ``litellm-key-remaining-*`` metadata keys the legacy limiter
|
||||
sets. The header carries no model group; it always refers to this
|
||||
request's model group, which is what the gauges are labeled with.
|
||||
Values are written in-process as plain ints (never HTTP-serialized
|
||||
strings), so anything else is rejected rather than coerced.
|
||||
"""
|
||||
if standard_logging_payload is None:
|
||||
return None
|
||||
hidden_params = standard_logging_payload.get("hidden_params")
|
||||
if hidden_params is None:
|
||||
return None
|
||||
additional_headers = hidden_params.get("additional_headers")
|
||||
if additional_headers is None:
|
||||
return None
|
||||
value = dict(additional_headers).get(f"x-ratelimit-model_per_key-remaining-{rate_limit_type}")
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
return None
|
||||
return value
|
||||
|
||||
def _set_virtual_key_rate_limit_metrics(
|
||||
self,
|
||||
user_api_key: Optional[str],
|
||||
|
|
@ -1733,11 +1837,20 @@ class PrometheusLogger(CustomLogger):
|
|||
model_group = get_model_group_from_litellm_kwargs(kwargs)
|
||||
remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}"
|
||||
remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}"
|
||||
standard_logging_payload: StandardLoggingPayload | None = kwargs.get("standard_logging_object")
|
||||
|
||||
remaining_requests = metadata.get(remaining_requests_variable_name)
|
||||
if remaining_requests is None:
|
||||
remaining_requests = self._get_remaining_from_v3_rate_limit_headers(
|
||||
standard_logging_payload=standard_logging_payload, rate_limit_type="requests"
|
||||
)
|
||||
if remaining_requests is None:
|
||||
remaining_requests = sys.maxsize
|
||||
remaining_tokens = metadata.get(remaining_tokens_variable_name)
|
||||
if remaining_tokens is None:
|
||||
remaining_tokens = self._get_remaining_from_v3_rate_limit_headers(
|
||||
standard_logging_payload=standard_logging_payload, rate_limit_type="tokens"
|
||||
)
|
||||
if remaining_tokens is None:
|
||||
remaining_tokens = sys.maxsize
|
||||
|
||||
|
|
|
|||
|
|
@ -161,8 +161,13 @@ def get_s3_object_key(
|
|||
start_time: datetime,
|
||||
s3_file_name: str,
|
||||
) -> str:
|
||||
sanitized_s3_file_name = s3_file_name.replace("/", "_")
|
||||
s3_object_key = (
|
||||
(s3_path.rstrip("/") + "/" if s3_path else "") + prefix + start_time.strftime("%Y-%m-%d") + "/" + s3_file_name
|
||||
(s3_path.rstrip("/") + "/" if s3_path else "")
|
||||
+ prefix
|
||||
+ start_time.strftime("%Y-%m-%d")
|
||||
+ "/"
|
||||
+ sanitized_s3_file_name
|
||||
) # we need the s3 key to include the time, so we log cache hits too
|
||||
s3_object_key += ".json"
|
||||
return s3_object_key
|
||||
|
|
|
|||
|
|
@ -2,26 +2,8 @@ from typing import Optional
|
|||
|
||||
from litellm.llms.openai.data_residency import infer_openai_data_residency
|
||||
|
||||
# Pre-define optional kwargs keys as frozenset for O(1) lookups
|
||||
# These are extracted from kwargs only if present, avoiding unnecessary .get() calls
|
||||
OPTIONAL_KWARGS_KEYS = frozenset(
|
||||
AWS_CREDENTIAL_KWARGS_KEYS = frozenset(
|
||||
{
|
||||
"azure_ad_token",
|
||||
"tenant_id",
|
||||
"client_id",
|
||||
"client_secret",
|
||||
"azure_username",
|
||||
"azure_password",
|
||||
"azure_scope",
|
||||
"timeout",
|
||||
"gcs_bucket_name",
|
||||
"bucket_name",
|
||||
"vertex_credentials",
|
||||
"vertex_project",
|
||||
"vertex_location",
|
||||
"vertex_ai_project",
|
||||
"vertex_ai_location",
|
||||
"vertex_ai_credentials",
|
||||
"aws_region_name",
|
||||
"aws_access_key_id",
|
||||
"aws_secret_access_key",
|
||||
|
|
@ -34,14 +16,40 @@ OPTIONAL_KWARGS_KEYS = frozenset(
|
|||
"aws_external_id",
|
||||
"aws_bedrock_runtime_endpoint",
|
||||
"aws_bedrock_project_id",
|
||||
"tpm",
|
||||
"rpm",
|
||||
"itpm",
|
||||
"otpm",
|
||||
"use_xai_oauth",
|
||||
}
|
||||
)
|
||||
|
||||
# Pre-define optional kwargs keys as frozenset for O(1) lookups
|
||||
# These are extracted from kwargs only if present, avoiding unnecessary .get() calls
|
||||
OPTIONAL_KWARGS_KEYS = (
|
||||
frozenset(
|
||||
{
|
||||
"azure_ad_token",
|
||||
"tenant_id",
|
||||
"client_id",
|
||||
"client_secret",
|
||||
"azure_username",
|
||||
"azure_password",
|
||||
"azure_scope",
|
||||
"timeout",
|
||||
"gcs_bucket_name",
|
||||
"bucket_name",
|
||||
"vertex_credentials",
|
||||
"vertex_project",
|
||||
"vertex_location",
|
||||
"vertex_ai_project",
|
||||
"vertex_ai_location",
|
||||
"vertex_ai_credentials",
|
||||
"tpm",
|
||||
"rpm",
|
||||
"itpm",
|
||||
"otpm",
|
||||
"use_xai_oauth",
|
||||
}
|
||||
)
|
||||
| AWS_CREDENTIAL_KWARGS_KEYS
|
||||
)
|
||||
|
||||
# Backward-compatible alias for existing imports/tests.
|
||||
_OPTIONAL_KWARGS_KEYS = OPTIONAL_KWARGS_KEYS
|
||||
|
||||
|
|
|
|||
|
|
@ -73,6 +73,7 @@ from litellm.litellm_core_utils.model_param_helper import ModelParamHelper
|
|||
from litellm.litellm_core_utils.redact_messages import (
|
||||
redact_message_input_output_from_custom_logger,
|
||||
redact_message_input_output_from_logging,
|
||||
redact_streaming_responses_for_custom_logger,
|
||||
)
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.llms.base_llm.search.transformation import SearchResponse
|
||||
|
|
@ -2576,6 +2577,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
model_call_details = callback.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details=model_call_details
|
||||
)
|
||||
model_call_details = redact_streaming_responses_for_custom_logger(
|
||||
model_call_details=model_call_details, custom_logger=callback
|
||||
)
|
||||
##################################
|
||||
if self.stream is True:
|
||||
if "async_complete_streaming_response" in model_call_details:
|
||||
|
|
@ -5208,10 +5212,15 @@ def get_standard_logging_object_payload(
|
|||
call_type = kwargs.get("call_type")
|
||||
cache_hit = kwargs.get("cache_hit", False)
|
||||
# Extract usage as a plain dict, avoiding Pydantic round-trip
|
||||
usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict(
|
||||
raw_usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict(
|
||||
response_obj=response_obj,
|
||||
combined_usage_object=cast(Optional[Usage], kwargs.get("combined_usage_object")),
|
||||
)
|
||||
usage_dict = (
|
||||
{**raw_usage_dict, "output_image_count": len(init_response_obj.data)}
|
||||
if isinstance(init_response_obj, ImageResponse) and init_response_obj.data
|
||||
else raw_usage_dict
|
||||
)
|
||||
|
||||
id = response_obj.get("id", kwargs.get("litellm_call_id"))
|
||||
|
||||
|
|
|
|||
|
|
@ -445,6 +445,7 @@ class PromptTokensDetailsResult(TypedDict):
|
|||
text_tokens: int
|
||||
audio_tokens: int
|
||||
image_tokens: int
|
||||
video_tokens: int
|
||||
character_count: int
|
||||
image_count: int
|
||||
video_length_seconds: float
|
||||
|
|
@ -473,6 +474,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
|||
)
|
||||
audio_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0
|
||||
image_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0
|
||||
video_tokens = _coerce_token_count(getattr(usage.prompt_tokens_details, "video_tokens", 0))
|
||||
character_count = (
|
||||
cast(
|
||||
Optional[int],
|
||||
|
|
@ -503,6 +505,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
|||
text_tokens=text_tokens,
|
||||
audio_tokens=audio_tokens,
|
||||
image_tokens=image_tokens,
|
||||
video_tokens=video_tokens,
|
||||
character_count=character_count,
|
||||
image_count=image_count,
|
||||
video_length_seconds=float(video_length_seconds),
|
||||
|
|
@ -515,6 +518,7 @@ class CompletionTokensDetailsResult(TypedDict):
|
|||
text_tokens: int
|
||||
reasoning_tokens: int
|
||||
image_tokens: int
|
||||
video_tokens: int
|
||||
|
||||
|
||||
def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult:
|
||||
|
|
@ -546,12 +550,14 @@ def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsRes
|
|||
)
|
||||
or 0
|
||||
)
|
||||
video_tokens = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0))
|
||||
|
||||
return CompletionTokensDetailsResult(
|
||||
audio_tokens=audio_tokens,
|
||||
text_tokens=text_tokens,
|
||||
reasoning_tokens=reasoning_tokens,
|
||||
image_tokens=image_tokens,
|
||||
video_tokens=video_tokens,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -586,6 +592,13 @@ def _calculate_input_cost(
|
|||
image_token_cost_key = "input_cost_per_token"
|
||||
prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"])
|
||||
|
||||
### VIDEO TOKEN COST
|
||||
if prompt_tokens_details["video_tokens"]:
|
||||
video_token_cost_key = "input_cost_per_video_token"
|
||||
if model_info.get(video_token_cost_key) is None:
|
||||
video_token_cost_key = "input_cost_per_token"
|
||||
prompt_cost += calculate_cost_component(model_info, video_token_cost_key, prompt_tokens_details["video_tokens"])
|
||||
|
||||
### CACHE WRITING COST - Now uses tiered pricing
|
||||
if (
|
||||
prompt_tokens_details["cache_creation_tokens"]
|
||||
|
|
@ -698,6 +711,7 @@ def generic_cost_per_token(
|
|||
text_tokens=usage.prompt_tokens,
|
||||
audio_tokens=0,
|
||||
image_tokens=0,
|
||||
video_tokens=0,
|
||||
character_count=0,
|
||||
image_count=0,
|
||||
video_length_seconds=0.0,
|
||||
|
|
@ -716,13 +730,14 @@ def generic_cost_per_token(
|
|||
audio_tokens = prompt_tokens_details["audio_tokens"]
|
||||
cache_creation = prompt_tokens_details["cache_creation_tokens"]
|
||||
image_tokens = prompt_tokens_details["image_tokens"]
|
||||
video_tokens = prompt_tokens_details["video_tokens"]
|
||||
|
||||
# Check for double-counting: sum of details > prompt_tokens means overlap
|
||||
total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens
|
||||
total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + video_tokens
|
||||
has_double_counting = cache_hit > 0 and total_details > usage.prompt_tokens
|
||||
|
||||
if (text_tokens == 0 and prompt_tokens_details["image_count"] == 0) or has_double_counting:
|
||||
text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens
|
||||
text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens
|
||||
# Clamp to zero: inconsistent streaming usage
|
||||
if text_tokens < 0:
|
||||
text_tokens = 0
|
||||
|
|
@ -751,6 +766,7 @@ def generic_cost_per_token(
|
|||
audio_tokens = 0
|
||||
reasoning_tokens = 0
|
||||
image_tokens = 0
|
||||
video_tokens = 0
|
||||
is_text_tokens_total = False
|
||||
if usage.completion_tokens_details is not None:
|
||||
completion_tokens_details = _parse_completion_tokens_details(usage)
|
||||
|
|
@ -758,19 +774,20 @@ def generic_cost_per_token(
|
|||
text_tokens = completion_tokens_details["text_tokens"]
|
||||
reasoning_tokens = completion_tokens_details["reasoning_tokens"]
|
||||
image_tokens = completion_tokens_details["image_tokens"]
|
||||
video_tokens = completion_tokens_details["video_tokens"]
|
||||
|
||||
# Handle text_tokens calculation:
|
||||
# 1. If text_tokens is explicitly provided and > 0, use it
|
||||
# 2. If there's a breakdown (reasoning/audio/image tokens), calculate text_tokens as the remainder
|
||||
# 2. If there's a breakdown (reasoning/audio/image/video tokens), calculate text_tokens as the remainder
|
||||
# 3. If no breakdown at all, assume all completion_tokens are text_tokens
|
||||
has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0
|
||||
has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0 or video_tokens > 0
|
||||
if text_tokens == 0:
|
||||
if has_token_breakdown:
|
||||
# Calculate text tokens as remainder when we have a breakdown
|
||||
# This handles cases like OpenAI's reasoning models where text_tokens isn't provided
|
||||
text_tokens = max(
|
||||
0,
|
||||
usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens,
|
||||
usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens - video_tokens,
|
||||
)
|
||||
else:
|
||||
# No breakdown at all, all tokens are text tokens
|
||||
|
|
@ -803,6 +820,14 @@ def generic_cost_per_token(
|
|||
)
|
||||
completion_cost += float(image_tokens) * _output_cost_per_image_token
|
||||
|
||||
## VIDEO COST
|
||||
if not is_text_tokens_total and video_tokens and video_tokens > 0:
|
||||
_output_cost_per_video_token = _get_cost_per_unit(model_info, "output_cost_per_video_token", None)
|
||||
_output_cost_per_video_token = (
|
||||
_output_cost_per_video_token if _output_cost_per_video_token is not None else completion_base_cost
|
||||
)
|
||||
completion_cost += float(video_tokens) * _output_cost_per_video_token
|
||||
|
||||
## REGIONAL DATA-RESIDENCY UPLIFT
|
||||
# Applied as a flat multiplier across all token costs for the request
|
||||
# when the upstream is a regionalized OpenAI host (eu./us.api.openai.com).
|
||||
|
|
|
|||
|
|
@ -5494,3 +5494,56 @@ def has_tool_with_name(tools: Any, tool_name: str) -> bool:
|
|||
elif tool.get("name") == tool_name:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def resolve_structured_messages(
|
||||
messages: list[dict[str, Any]] | None,
|
||||
request_kwargs: dict[str, Any],
|
||||
) -> list[dict[str, Any]] | None:
|
||||
"""
|
||||
Normalize a request's messages to OpenAI-spec chat-completions shape,
|
||||
regardless of which API surface produced them (chat completions,
|
||||
Anthropic /v1/messages, Responses API ``input``, etc).
|
||||
|
||||
Returns ``messages`` unchanged if already present. Otherwise dispatches
|
||||
through the guardrail translation handlers (the same per-surface
|
||||
conversion logic guardrails use) to convert e.g. Responses API ``input``
|
||||
into a message list. Returns ``None`` if no messages could be resolved.
|
||||
"""
|
||||
if messages:
|
||||
return messages
|
||||
|
||||
from litellm.litellm_core_utils.api_route_to_call_types import (
|
||||
get_call_types_for_route,
|
||||
)
|
||||
from litellm.llms import load_guardrail_translation_mappings
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
mappings = load_guardrail_translation_mappings()
|
||||
call_type: CallTypes | None = None
|
||||
|
||||
# 1. Try route-based inference from proxy metadata
|
||||
route = request_kwargs.get("litellm_metadata", {}).get("user_api_key_request_route")
|
||||
if route:
|
||||
call_types_list = get_call_types_for_route(route)
|
||||
if call_types_list:
|
||||
for ct in call_types_list:
|
||||
if ct in mappings:
|
||||
call_type = ct
|
||||
break
|
||||
|
||||
# 2. Fallback: try each mapped handler until one produces messages
|
||||
handlers_to_try: list[Any] = []
|
||||
if call_type is not None and call_type in mappings:
|
||||
handlers_to_try.append(mappings[call_type]())
|
||||
else:
|
||||
handlers_to_try.extend(handler_cls() for handler_cls in mappings.values())
|
||||
|
||||
for handler in handlers_to_try:
|
||||
structured = handler.get_structured_messages(request_kwargs)
|
||||
if structured:
|
||||
return [
|
||||
msg if isinstance(msg, dict) else msg.model_dump() # type: ignore
|
||||
for msg in structured
|
||||
]
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -38,10 +38,45 @@ def redact_message_input_output_from_custom_logger(
|
|||
litellm_logging_obj: LiteLLMLoggingObject, result, custom_logger: CustomLogger
|
||||
):
|
||||
if hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True:
|
||||
return perform_redaction(litellm_logging_obj.model_call_details, result)
|
||||
return perform_redaction(litellm_logging_obj.model_call_details, result, redact_streaming_responses=False)
|
||||
return result
|
||||
|
||||
|
||||
def redact_streaming_responses_for_custom_logger(model_call_details: dict, custom_logger: CustomLogger) -> dict:
|
||||
"""
|
||||
Returns a copy of model_call_details whose streaming response entries are redacted deepcopies
|
||||
when the custom logger has opted out of message logging. The shared model_call_details is left
|
||||
untouched so other callbacks still receive the unredacted response.
|
||||
"""
|
||||
if not (hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True):
|
||||
return model_call_details
|
||||
redacted_entries = {
|
||||
streaming_key: _redacted_streaming_response_copy(model_call_details[streaming_key])
|
||||
for streaming_key in ("complete_streaming_response", "async_complete_streaming_response")
|
||||
if model_call_details.get(streaming_key) is not None
|
||||
}
|
||||
if not redacted_entries:
|
||||
return model_call_details
|
||||
return {**model_call_details, **redacted_entries}
|
||||
|
||||
|
||||
def _redacted_streaming_response_copy(streaming_response):
|
||||
redacted_response = copy.deepcopy(streaming_response)
|
||||
_redact_streaming_response(redacted_response)
|
||||
return redacted_response
|
||||
|
||||
|
||||
def _redact_streaming_response(streaming_response):
|
||||
if hasattr(streaming_response, "choices"):
|
||||
for choice in streaming_response.choices:
|
||||
_redact_choice_content(choice)
|
||||
redact_vertex_ai_metadata_from_logged_object(streaming_response)
|
||||
elif hasattr(streaming_response, "output"):
|
||||
_redact_responses_api_output(streaming_response.output)
|
||||
if hasattr(streaming_response, "reasoning") and streaming_response.reasoning is not None:
|
||||
streaming_response.reasoning = None
|
||||
|
||||
|
||||
def _redact_choice_content(choice):
|
||||
"""Helper to redact content in a choice (message or delta)."""
|
||||
if isinstance(choice, litellm.Choices):
|
||||
|
|
@ -150,9 +185,13 @@ def _redact_model_response_dict_choices(choices, redacted_str: str):
|
|||
_redact_choice_content(choice)
|
||||
|
||||
|
||||
def perform_redaction(model_call_details: dict, result):
|
||||
def perform_redaction(model_call_details: dict, result, redact_streaming_responses: bool = True):
|
||||
"""
|
||||
Performs the actual redaction on the logging object and result.
|
||||
|
||||
redact_streaming_responses=False skips the in-place redaction of the shared streaming
|
||||
response entries; per-callback redaction hands each opted-out callback its own redacted
|
||||
copy via redact_streaming_responses_for_custom_logger instead.
|
||||
"""
|
||||
# Redact model_call_details
|
||||
model_call_details["messages"] = [{"role": "user", "content": "redacted-by-litellm"}]
|
||||
|
|
@ -162,17 +201,9 @@ def perform_redaction(model_call_details: dict, result):
|
|||
redact_vertex_ai_metadata_from_litellm_params(model_call_details)
|
||||
|
||||
# Redact streaming response
|
||||
if model_call_details.get("stream", False) is True and "complete_streaming_response" in model_call_details:
|
||||
_streaming_response = model_call_details["complete_streaming_response"]
|
||||
if hasattr(_streaming_response, "choices"):
|
||||
for choice in _streaming_response.choices:
|
||||
_redact_choice_content(choice)
|
||||
redact_vertex_ai_metadata_from_logged_object(_streaming_response)
|
||||
elif hasattr(_streaming_response, "output"):
|
||||
_redact_responses_api_output(_streaming_response.output)
|
||||
# Redact reasoning field in ResponsesAPIResponse
|
||||
if hasattr(_streaming_response, "reasoning") and _streaming_response.reasoning is not None:
|
||||
_streaming_response.reasoning = None
|
||||
if redact_streaming_responses and model_call_details.get("stream", False) is True:
|
||||
for _streaming_key in ("complete_streaming_response", "async_complete_streaming_response"):
|
||||
_redact_streaming_response(model_call_details.get(_streaming_key))
|
||||
|
||||
# Redact result
|
||||
if result is not None:
|
||||
|
|
|
|||
|
|
@ -227,6 +227,10 @@ DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING = (
|
|||
"Sonnet 4.6+, and Mythos Preview."
|
||||
)
|
||||
|
||||
DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING = (
|
||||
"Dropping adaptive `thinking` for model=%s: max_tokens is too small to fit the minimum thinking budget."
|
||||
)
|
||||
|
||||
DROP_UNSUPPORTED_SPEED_WARNING = (
|
||||
"Dropping unsupported `speed` for model=%s (drop_params=True). Fast mode is only supported on select Opus models."
|
||||
)
|
||||
|
|
@ -1220,6 +1224,23 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
llm_provider=llm_provider,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _cap_thinking_budget_to_max_tokens(
|
||||
thinking: AnthropicThinkingParam, max_tokens: Optional[int]
|
||||
) -> Optional[AnthropicThinkingParam]:
|
||||
"""Cap a legacy ``thinking.budget_tokens`` below ``max_tokens`` (Anthropic
|
||||
requires ``max_tokens > budget_tokens``). Returns the (possibly capped)
|
||||
thinking dict, or ``None`` when ``max_tokens`` is too small to fit even the
|
||||
minimum thinking budget and thinking should be dropped."""
|
||||
budget = thinking.get("budget_tokens")
|
||||
if max_tokens is None or not isinstance(budget, int):
|
||||
return thinking
|
||||
if max_tokens <= ANTHROPIC_MIN_THINKING_BUDGET_TOKENS:
|
||||
return None
|
||||
if budget < max_tokens:
|
||||
return thinking
|
||||
return AnthropicThinkingParam(type=thinking.get("type", "enabled"), budget_tokens=max_tokens - 1)
|
||||
|
||||
def _extract_json_schema_from_response_format(self, value: Optional[dict]) -> Optional[dict]:
|
||||
if value is None:
|
||||
return None
|
||||
|
|
@ -1420,24 +1441,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
output_key=param,
|
||||
)
|
||||
elif param == "response_format" and isinstance(value, dict):
|
||||
if any(
|
||||
substring in model
|
||||
for substring in {
|
||||
"sonnet-4.5",
|
||||
"sonnet-4-5",
|
||||
"opus-4.1",
|
||||
"opus-4-1",
|
||||
"opus-4.5",
|
||||
"opus-4-5",
|
||||
"opus-4.6",
|
||||
"opus-4-6",
|
||||
"opus-4.7",
|
||||
"opus-4-7",
|
||||
"sonnet-4.6",
|
||||
"sonnet-4-6",
|
||||
"sonnet_4.6",
|
||||
"sonnet_4_6",
|
||||
}
|
||||
if AnthropicConfig._supports_model_capability(
|
||||
model,
|
||||
"supports_native_structured_output",
|
||||
self._resolved_provider,
|
||||
):
|
||||
_output_format = self.map_response_format_to_anthropic_output_format(value)
|
||||
if _output_format is not None:
|
||||
|
|
@ -1463,7 +1470,38 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
):
|
||||
optional_params["metadata"] = {"user_id": value}
|
||||
elif param == "thinking":
|
||||
optional_params["thinking"] = value
|
||||
if (
|
||||
isinstance(value, dict)
|
||||
and value.get("type") == "adaptive"
|
||||
and not AnthropicConfig._is_adaptive_thinking_model(model, self._resolved_provider)
|
||||
):
|
||||
# Callers (e.g. Claude Code) send adaptive thinking
|
||||
# unconditionally; translate it down to the legacy
|
||||
# `thinking={type: enabled, budget_tokens}` interface a
|
||||
# pre-4.6 model actually supports instead of forwarding a
|
||||
# shape the model will reject.
|
||||
max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens")
|
||||
legacy_thinking = AnthropicConfig._map_reasoning_effort(
|
||||
reasoning_effort="medium",
|
||||
model=model,
|
||||
custom_llm_provider=self._resolved_provider,
|
||||
llm_provider=self._resolved_provider,
|
||||
)
|
||||
capped_thinking = (
|
||||
AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
|
||||
if legacy_thinking is not None
|
||||
else None
|
||||
)
|
||||
if capped_thinking is not None:
|
||||
optional_params["thinking"] = capped_thinking
|
||||
else:
|
||||
litellm.verbose_logger.warning(
|
||||
DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING,
|
||||
model,
|
||||
)
|
||||
optional_params.pop("thinking", None)
|
||||
else:
|
||||
optional_params["thinking"] = value
|
||||
elif param == "reasoning_effort":
|
||||
# Accept both string ("low") and dict ({"effort": "low",
|
||||
# "summary": "concise"}). The Responses->Chat parser keeps the
|
||||
|
|
|
|||
|
|
@ -340,11 +340,15 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
def _get_model_capability(model: str, key: str) -> Optional[bool]:
|
||||
"""Read boolean capability ``key`` from the model map, or None when
|
||||
no entry declares it."""
|
||||
from litellm.utils import _get_bundled_model_cost_map
|
||||
|
||||
try:
|
||||
for cand in AnthropicModelInfo._model_map_lookup_candidates(model):
|
||||
value = litellm.model_cost.get(cand, {}).get(key)
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
candidates = AnthropicModelInfo._model_map_lookup_candidates(model)
|
||||
for model_cost in (litellm.model_cost, _get_bundled_model_cost_map()):
|
||||
for cand in candidates:
|
||||
value = model_cost.get(cand, {}).get(key)
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ from typing import Any, AsyncIterator, Dict, List, Optional, Tuple
|
|||
import httpx
|
||||
|
||||
from litellm.constants import (
|
||||
ANTHROPIC_MIN_THINKING_BUDGET_TOKENS,
|
||||
DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
|
||||
DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
|
||||
|
|
@ -358,7 +357,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
except _BadRequestError as e:
|
||||
raise AnthropicError(message=str(e.message), status_code=400)
|
||||
capped_thinking = (
|
||||
AnthropicMessagesConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
|
||||
AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
|
||||
if legacy_thinking is not None
|
||||
else None
|
||||
)
|
||||
|
|
@ -377,19 +376,34 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
optional_params.pop("output_config", None)
|
||||
|
||||
@staticmethod
|
||||
def _cap_thinking_budget_to_max_tokens(thinking: Dict, max_tokens: Optional[int]) -> Optional[Dict]:
|
||||
"""Cap a legacy ``thinking.budget_tokens`` below ``max_tokens`` (Anthropic
|
||||
requires ``max_tokens > budget_tokens``). Returns the (possibly capped)
|
||||
thinking dict, or ``None`` when ``max_tokens`` is too small to fit even the
|
||||
minimum thinking budget and thinking should be dropped."""
|
||||
budget = thinking.get("budget_tokens")
|
||||
if max_tokens is None or not isinstance(budget, int):
|
||||
return thinking
|
||||
if max_tokens <= ANTHROPIC_MIN_THINKING_BUDGET_TOKENS:
|
||||
return None
|
||||
if budget < max_tokens:
|
||||
return thinking
|
||||
return {**thinking, "budget_tokens": max_tokens - 1}
|
||||
def _drop_incompatible_temperature_for_thinking(
|
||||
model: str, optional_params: dict, custom_llm_provider: str
|
||||
) -> None:
|
||||
"""Anthropic rejects any ``temperature`` other than 1 while extended thinking
|
||||
is enabled ("temperature may only be set to 1 when thinking is enabled").
|
||||
|
||||
Clients like Claude Code send ``thinking``/``output_config.effort`` together
|
||||
with a pinned ``temperature`` (e.g. the safety classifier uses ``temperature=0``
|
||||
for determinism). When the request lands on a non-adaptive model, the effort
|
||||
interface is reshaped above into legacy ``thinking={type: enabled}`` (or kept
|
||||
as ``output_config.effort`` on Opus 4.5), and the leftover ``temperature`` would
|
||||
400. Preserving the thinking the caller asked for wins over an unhonorable
|
||||
sampling value (Anthropic forces ``temperature=1`` under thinking regardless),
|
||||
so drop it and let the API default apply.
|
||||
|
||||
Adaptive models (4.6+) own this natively and are left untouched.
|
||||
"""
|
||||
if AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider):
|
||||
return
|
||||
temperature = optional_params.get("temperature")
|
||||
if temperature is None or temperature == 1:
|
||||
return
|
||||
thinking = optional_params.get("thinking")
|
||||
output_config = optional_params.get("output_config")
|
||||
thinking_enabled = isinstance(thinking, dict) and thinking.get("type") == "enabled"
|
||||
effort_enabled = isinstance(output_config, dict) and output_config.get("effort") is not None
|
||||
if thinking_enabled or effort_enabled:
|
||||
optional_params.pop("temperature", None)
|
||||
|
||||
def transform_anthropic_messages_request(
|
||||
self,
|
||||
|
|
@ -431,6 +445,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
custom_llm_provider=self._resolved_provider,
|
||||
)
|
||||
|
||||
self._drop_incompatible_temperature_for_thinking(
|
||||
model=model,
|
||||
optional_params=anthropic_messages_optional_request_params,
|
||||
custom_llm_provider=self._resolved_provider,
|
||||
)
|
||||
|
||||
system_param = anthropic_messages_optional_request_params.get("system")
|
||||
if self.should_strip_billing_metadata() and system_param is not None:
|
||||
filtered_system = self._filter_billing_headers_from_system(system_param)
|
||||
|
|
|
|||
|
|
@ -198,14 +198,16 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
@staticmethod
|
||||
def translate_tool_choice_to_responses_api(
|
||||
tool_choice: AnthropicMessagesToolChoice,
|
||||
) -> Dict[str, Any]:
|
||||
) -> Union[str, dict[str, Any]]:
|
||||
"""Convert Anthropic tool_choice to Responses API tool_choice."""
|
||||
tc_type = tool_choice.get("type")
|
||||
if tc_type == "any":
|
||||
return {"type": "required"}
|
||||
return "required"
|
||||
elif tc_type == "tool":
|
||||
return {"type": "function", "name": tool_choice.get("name", "")}
|
||||
return {"type": "auto"}
|
||||
elif tc_type == "none":
|
||||
return "none"
|
||||
return "auto"
|
||||
|
||||
@staticmethod
|
||||
def translate_context_management_to_responses_api(
|
||||
|
|
|
|||
|
|
@ -877,6 +877,15 @@ class BaseAWSLLM:
|
|||
"Resource": "*",
|
||||
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
|
||||
},
|
||||
{
|
||||
"Sid": "BedrockMantleLiteLLM",
|
||||
"Effect": "Allow",
|
||||
"Action": [
|
||||
"bedrock-mantle:CreateInference",
|
||||
],
|
||||
"Resource": "*",
|
||||
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
|
||||
},
|
||||
],
|
||||
}
|
||||
assume_role_params = {
|
||||
|
|
|
|||
|
|
@ -5,6 +5,9 @@ from typing import Any, Dict, List, Literal, Optional, Union, cast
|
|||
|
||||
from httpx import Headers, Response
|
||||
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
BEDROCK_MANAGED_S3_BATCH_PREFIX,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
|
@ -26,6 +29,15 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders
|
|||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..common_utils import CommonBatchFilesUtils
|
||||
|
||||
# Bedrock batch input files are uploaded as
|
||||
# s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see
|
||||
# BedrockFilesTransformation._get_s3_object_name). A uuid4 is always 36 hex/dash
|
||||
# characters, so it can be stripped off the end unambiguously even though the
|
||||
# model name itself may contain dashes.
|
||||
_S3_BATCH_FILE_UUID_SUFFIX_PATTERN = re.compile(
|
||||
r"-[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.jsonl$"
|
||||
)
|
||||
|
||||
|
||||
class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
||||
"""
|
||||
|
|
@ -40,6 +52,41 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
def custom_llm_provider(self) -> LlmProviders:
|
||||
return LlmProviders.BEDROCK
|
||||
|
||||
@classmethod
|
||||
def _get_bare_model_name_from_s3_key(cls, object_key: str) -> Optional[str]:
|
||||
if not object_key.startswith(BEDROCK_MANAGED_S3_BATCH_PREFIX):
|
||||
return None
|
||||
model_part = object_key[len(BEDROCK_MANAGED_S3_BATCH_PREFIX) :]
|
||||
match = _S3_BATCH_FILE_UUID_SUFFIX_PATTERN.search(model_part)
|
||||
if not match or match.start() == 0:
|
||||
return None
|
||||
return model_part[: match.start()]
|
||||
|
||||
@classmethod
|
||||
def is_unmanaged_s3_batch_input_file_id(cls, input_file_id: Optional[str]) -> bool:
|
||||
"""
|
||||
Returns True if `input_file_id` is a raw s3:// Bedrock batch input file (i.e. not a
|
||||
LiteLLM-managed unified file id) whose object key embeds the model name in the
|
||||
`litellm-bedrock-files-{model}-{uuid}.jsonl` layout.
|
||||
"""
|
||||
if input_file_id is None or not input_file_id.startswith("s3://"):
|
||||
return False
|
||||
object_key = input_file_id.rsplit("/", 1)[-1]
|
||||
return cls._get_bare_model_name_from_s3_key(object_key) is not None
|
||||
|
||||
@classmethod
|
||||
def get_bare_model_name_from_s3_file(cls, input_file_id: str) -> str:
|
||||
"""
|
||||
Extracts the bare model name (e.g. "us.anthropic.claude-sonnet-4-20250514-v1-0") from
|
||||
an unmanaged batch's s3:// input file id. Note any ":" in the original model id was
|
||||
replaced with "-" at upload time, so callers must fuzzy-match against configured
|
||||
deployments rather than expect an exact string match.
|
||||
"""
|
||||
object_key = input_file_id.rsplit("/", 1)[-1]
|
||||
bare_model_name = cls._get_bare_model_name_from_s3_key(object_key)
|
||||
assert bare_model_name is not None # narrowed by is_unmanaged_s3_batch_input_file_id
|
||||
return bare_model_name
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
make_valid_bedrock_tool_name,
|
||||
)
|
||||
from litellm.llms.anthropic.chat.transformation import (
|
||||
DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING,
|
||||
DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING,
|
||||
REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT,
|
||||
AnthropicConfig,
|
||||
|
|
@ -899,7 +900,28 @@ class AmazonConverseConfig(BaseConfig):
|
|||
"tool_choice": {"disable_parallel_tool_use": disable_parallel}
|
||||
}
|
||||
if param == "thinking":
|
||||
optional_params["thinking"] = value
|
||||
if (
|
||||
isinstance(value, dict)
|
||||
and value.get("type") == "adaptive"
|
||||
and not AnthropicConfig._is_adaptive_thinking_model(model, "bedrock")
|
||||
):
|
||||
max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens")
|
||||
legacy_thinking = AnthropicConfig._map_reasoning_effort(
|
||||
reasoning_effort="medium",
|
||||
model=model,
|
||||
custom_llm_provider="bedrock",
|
||||
)
|
||||
capped = (
|
||||
AnthropicConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens)
|
||||
if legacy_thinking is not None
|
||||
else None
|
||||
)
|
||||
if capped is not None:
|
||||
optional_params["thinking"] = capped
|
||||
else:
|
||||
litellm.verbose_logger.warning(DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, model)
|
||||
else:
|
||||
optional_params["thinking"] = value
|
||||
elif param == "reasoning_effort" and isinstance(value, str):
|
||||
self._handle_reasoning_effort_parameter(
|
||||
model=model, reasoning_effort=value, optional_params=optional_params
|
||||
|
|
|
|||
|
|
@ -20,6 +20,8 @@ from litellm.types.utils import LlmProviders
|
|||
|
||||
from ..common_utils import OpenAIError
|
||||
|
||||
OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS = 16
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
||||
|
|
@ -59,6 +61,19 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
key="supports_none_reasoning_effort",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _enforce_min_max_output_tokens(max_output_tokens: "int | None") -> "int | None":
|
||||
"""Raise sub-minimum max_output_tokens up to the OpenAI Responses API minimum.
|
||||
|
||||
OpenAI's Responses API rejects max_output_tokens below 16 for every model
|
||||
(not gpt-5 specific), so a client like Claude Code that sends a max_tokens=1
|
||||
warmup probe on model switch would otherwise 400. Values that are None or
|
||||
already at/above the minimum are returned unchanged.
|
||||
"""
|
||||
if isinstance(max_output_tokens, int) and max_output_tokens < OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS:
|
||||
return OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS
|
||||
return max_output_tokens
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
"""
|
||||
All OpenAI Responses API params are supported
|
||||
|
|
@ -92,6 +107,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
"""
|
||||
params = dict(response_api_optional_params)
|
||||
|
||||
if "max_output_tokens" in params:
|
||||
params["max_output_tokens"] = self._enforce_min_max_output_tokens(params.get("max_output_tokens"))
|
||||
|
||||
if self._is_gpt_5_model(model=model):
|
||||
temperature = params.get("temperature")
|
||||
if temperature is not None and temperature != 1:
|
||||
|
|
|
|||
|
|
@ -998,6 +998,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
response_modalities.append("IMAGE")
|
||||
elif modality == "audio":
|
||||
response_modalities.append("AUDIO")
|
||||
elif modality == "video":
|
||||
response_modalities.append("VIDEO")
|
||||
else:
|
||||
response_modalities.append("MODALITY_UNSPECIFIED")
|
||||
return response_modalities
|
||||
|
|
|
|||
|
|
@ -92,7 +92,10 @@ from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
|
|||
from litellm.litellm_core_utils.request_timeout_resolver import (
|
||||
get_configured_request_timeout,
|
||||
)
|
||||
from litellm.litellm_core_utils.get_litellm_params import OPTIONAL_KWARGS_KEYS
|
||||
from litellm.litellm_core_utils.get_litellm_params import (
|
||||
AWS_CREDENTIAL_KWARGS_KEYS,
|
||||
OPTIONAL_KWARGS_KEYS,
|
||||
)
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.get_provider_specific_headers import (
|
||||
ProviderSpecificHeaderUtils,
|
||||
|
|
@ -5322,7 +5325,7 @@ def completion( # type: ignore
|
|||
tpm=kwargs.get("tpm"),
|
||||
rpm=kwargs.get("rpm"),
|
||||
use_xai_oauth=kwargs.get("use_xai_oauth", False),
|
||||
aws_bedrock_project_id=kwargs.get("aws_bedrock_project_id"),
|
||||
**{key: kwargs[key] for key in AWS_CREDENTIAL_KWARGS_KEYS if key in kwargs},
|
||||
)
|
||||
cast(LiteLLMLoggingObj, logging).update_environment_variables(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -11331,6 +11331,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
|
|
@ -11362,6 +11363,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
|
|
@ -11424,6 +11426,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -11479,6 +11482,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
|
|
@ -11506,6 +11510,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
|
|
@ -11559,6 +11564,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_output_config": true
|
||||
|
|
@ -11586,6 +11592,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_output_config": true
|
||||
|
|
@ -11614,6 +11621,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"provider_specific_entry": {
|
||||
|
|
@ -11648,6 +11656,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"provider_specific_entry": {
|
||||
|
|
@ -11682,6 +11691,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -11718,6 +11728,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -11788,6 +11799,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -18747,6 +18759,49 @@
|
|||
},
|
||||
"supports_image_size": false
|
||||
},
|
||||
"gemini/gemini-3-pro-image": {
|
||||
"input_cost_per_image": 0.0011,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"input_cost_per_token_batches": 1e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 65536,
|
||||
"max_output_tokens": 32768,
|
||||
"max_tokens": 32768,
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.134,
|
||||
"output_cost_per_image_token": 0.00012,
|
||||
"output_cost_per_token": 1.2e-05,
|
||||
"rpm": 1000,
|
||||
"tpm": 4000000,
|
||||
"output_cost_per_token_batches": 6e-06,
|
||||
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/completions",
|
||||
"/v1/batch"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supports_function_calling": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_low": 0.014,
|
||||
"search_context_size_medium": 0.014,
|
||||
"search_context_size_high": 0.014
|
||||
},
|
||||
"web_search_billing_unit": "per_query"
|
||||
},
|
||||
"gemini/gemini-3-pro-image-preview": {
|
||||
"input_cost_per_image": 0.0011,
|
||||
"input_cost_per_token": 2e-06,
|
||||
|
|
@ -18778,6 +18833,49 @@
|
|||
],
|
||||
"supports_function_calling": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_low": 0.014,
|
||||
"search_context_size_medium": 0.014,
|
||||
"search_context_size_high": 0.014
|
||||
},
|
||||
"web_search_billing_unit": "per_query"
|
||||
},
|
||||
"gemini/gemini-3.1-flash-image": {
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
"input_cost_per_token_batches": 1.25e-07,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 65536,
|
||||
"max_output_tokens": 32768,
|
||||
"max_tokens": 32768,
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.045,
|
||||
"output_cost_per_image_token": 6e-05,
|
||||
"output_cost_per_token": 1.5e-06,
|
||||
"output_cost_per_token_batches": 7.5e-07,
|
||||
"rpm": 1000,
|
||||
"tpm": 4000000,
|
||||
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/completions",
|
||||
"/v1/batch"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supports_function_calling": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -18819,6 +18917,7 @@
|
|||
],
|
||||
"supports_function_calling": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -19513,6 +19612,39 @@
|
|||
},
|
||||
"web_search_billing_unit": "per_query"
|
||||
},
|
||||
"gemini/gemini-omni-flash-preview": {
|
||||
"input_cost_per_audio_token": 1.5e-06,
|
||||
"input_cost_per_token": 1.5e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 65535,
|
||||
"max_tokens": 65535,
|
||||
"mode": "chat",
|
||||
"output_cost_per_reasoning_token": 9e-06,
|
||||
"output_cost_per_token": 9e-06,
|
||||
"output_cost_per_video_token": 1.75e-05,
|
||||
"rpm": 2000,
|
||||
"source": "https://ai.google.dev/gemini-api/docs/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image",
|
||||
"audio",
|
||||
"video"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text",
|
||||
"video"
|
||||
],
|
||||
"supports_audio_input": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"tpm": 800000
|
||||
},
|
||||
"gemini/gemini-3.1-pro-preview": {
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
|
||||
|
|
@ -19677,6 +19809,37 @@
|
|||
},
|
||||
"web_search_billing_unit": "per_query"
|
||||
},
|
||||
"gemini-omni-flash-preview": {
|
||||
"input_cost_per_audio_token": 1.5e-06,
|
||||
"input_cost_per_token": 1.5e-06,
|
||||
"litellm_provider": "vertex_ai-language-models",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 65535,
|
||||
"max_tokens": 65535,
|
||||
"mode": "chat",
|
||||
"output_cost_per_reasoning_token": 9e-06,
|
||||
"output_cost_per_token": 9e-06,
|
||||
"output_cost_per_video_token": 1.75e-05,
|
||||
"source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/omni-flash-preview",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image",
|
||||
"audio",
|
||||
"video"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text",
|
||||
"video"
|
||||
],
|
||||
"supports_audio_input": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"gemini-3.5-flash": {
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"input_cost_per_audio_token": 1e-06,
|
||||
|
|
@ -36777,6 +36940,7 @@
|
|||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -36786,6 +36950,22 @@
|
|||
"tpm": 8000000,
|
||||
"supports_image_size": false
|
||||
},
|
||||
"vertex_ai/gemini-3-pro-image": {
|
||||
"input_cost_per_image": 0.0011,
|
||||
"input_cost_per_token": 2e-06,
|
||||
"input_cost_per_token_batches": 1e-06,
|
||||
"litellm_provider": "vertex_ai-language-models",
|
||||
"max_input_tokens": 65536,
|
||||
"max_output_tokens": 32768,
|
||||
"max_tokens": 32768,
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.134,
|
||||
"output_cost_per_image_token": 0.00012,
|
||||
"output_cost_per_token": 1.2e-05,
|
||||
"output_cost_per_token_batches": 6e-06,
|
||||
"supports_reasoning": false,
|
||||
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
|
||||
},
|
||||
"vertex_ai/gemini-3-pro-image-preview": {
|
||||
"input_cost_per_image": 0.0011,
|
||||
"input_cost_per_token": 2e-06,
|
||||
|
|
@ -36799,8 +36979,23 @@
|
|||
"output_cost_per_image_token": 0.00012,
|
||||
"output_cost_per_token": 1.2e-05,
|
||||
"output_cost_per_token_batches": 6e-06,
|
||||
"supports_reasoning": false,
|
||||
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
|
||||
},
|
||||
"vertex_ai/gemini-3.1-flash-image": {
|
||||
"input_cost_per_image": 0.00056,
|
||||
"input_cost_per_token": 5e-07,
|
||||
"litellm_provider": "vertex_ai-language-models",
|
||||
"max_input_tokens": 65536,
|
||||
"max_output_tokens": 32768,
|
||||
"max_tokens": 32768,
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.0672,
|
||||
"output_cost_per_image_token": 6e-05,
|
||||
"output_cost_per_token": 3e-06,
|
||||
"supports_reasoning": false,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models"
|
||||
},
|
||||
"vertex_ai/gemini-3.1-flash-image-preview": {
|
||||
"input_cost_per_image": 0.00056,
|
||||
"input_cost_per_token": 5e-07,
|
||||
|
|
@ -36812,6 +37007,7 @@
|
|||
"output_cost_per_image": 0.0672,
|
||||
"output_cost_per_image_token": 6e-05,
|
||||
"output_cost_per_token": 3e-06,
|
||||
"supports_reasoning": false,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models"
|
||||
},
|
||||
"vertex_ai/gemini-3.1-flash-lite-preview": {
|
||||
|
|
@ -45051,8 +45247,8 @@
|
|||
"rules": [
|
||||
{
|
||||
"name": "bedrock-claude-ids",
|
||||
"pattern": "anthropic\\.claude-",
|
||||
"description": "Any Bedrock-syntax Claude id: the dotted anthropic.claude- segment appears in bare (anthropic.claude-...), region-prefixed (us./eu./au./jp./apac.) and global.-prefixed ids, for every version. Routes these to bedrock before the bare-id Anthropic rule is consulted.",
|
||||
"pattern": "^(?:[a-z-]+\\.)?anthropic\\.claude-",
|
||||
"description": "A Bedrock-syntax Claude id, for every version: anthropic.claude- at the start of the name, optionally behind a single dotted geo segment (us./eu./au./jp./apac./global./us-gov.). Anchored to the start because routing rules see the raw request string and provider inference feeds the proxy's provider/* wildcard access checks: an id under an unrecognized namespace such as bedrockz/anthropic.claude-... must stay unroutable rather than resolve to bedrock and slip through a bedrock/* key. Routes to bedrock before the bare-id Anthropic rule is consulted.",
|
||||
"model_info": {
|
||||
"litellm_provider": "bedrock"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
|
|||
budget_reset_at: Optional[datetime] = None
|
||||
allowed_cache_controls: Optional[list] = []
|
||||
allowed_routes: Optional[list] = []
|
||||
key_type: str | None = None
|
||||
permissions: Dict = {}
|
||||
model_spend: Dict = {}
|
||||
model_max_budget: Dict = {}
|
||||
|
|
|
|||
|
|
@ -18,6 +18,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti
|
|||
is_bridge_envelope_shaped,
|
||||
resolve_bridge_envelope,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
|
||||
EnvelopeIdentity,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
UI_TEAM_ID,
|
||||
LiteLLM_TeamTable,
|
||||
|
|
@ -543,7 +546,7 @@ class MCPRequestHandler:
|
|||
header_key = server.alias or server.server_name
|
||||
if header_key is None:
|
||||
raise HTTPException(status_code=500, detail="Server misconfigured: MCP server has no routable name")
|
||||
admitted = await MCPRequestHandler._reload_admitted_key(result.identity.key_hash)
|
||||
admitted = await MCPRequestHandler._reload_admitted_principal(result.identity)
|
||||
await MCPRequestHandler._enforce_admitted_live_policy(admitted=admitted, request=request, route=route)
|
||||
injected = {header_key: {"Authorization": result.upstream_authorization.get_secret_value()}}
|
||||
new_headers = {**(mcp_server_auth_headers or {}), **injected}
|
||||
|
|
@ -572,6 +575,89 @@ class MCPRequestHandler:
|
|||
route=route,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _reload_admitted_principal(identity: EnvelopeIdentity) -> UserAPIKeyAuth:
|
||||
"""Reload the live litellm record the envelope's subject references.
|
||||
|
||||
Dispatches on the sealed subject type: a ``key_hash`` reloads the virtual key that
|
||||
minted the envelope (the scripted two-header client that presents a litellm key at the
|
||||
token endpoint), a ``user_id`` reloads the user that authenticated interactively (the
|
||||
DCR client, whose SSO login at the bridged authorize yields a user, not a key). Both
|
||||
return a ``UserAPIKeyAuth`` the caller runs through the centralized policy gate, so
|
||||
team/project/org/budget/SCIM enforcement is identical to the principal presenting
|
||||
itself directly."""
|
||||
match identity.subject_type:
|
||||
case "key_hash":
|
||||
return await MCPRequestHandler._reload_admitted_key(identity.subject)
|
||||
case "user_id":
|
||||
return await MCPRequestHandler._reload_admitted_user(identity.subject)
|
||||
case _:
|
||||
assert_never(identity.subject_type)
|
||||
|
||||
@staticmethod
|
||||
async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth:
|
||||
"""Reload the live user an interactively-minted envelope references and admit them as
|
||||
themselves.
|
||||
|
||||
The DCR client authenticates via SSO at the bridged authorize, which yields a user
|
||||
subject rather than a virtual key, so the envelope admits under the user's own
|
||||
identity: the reloaded ``user_id`` and the user's own MCP object permission ride on the
|
||||
returned ``UserAPIKeyAuth``, and the SAME ``get_allowed_mcp_servers`` the key path uses then
|
||||
computes which servers the user may reach, so the user's litellm MCP grants and access groups
|
||||
gate the request exactly as a key's do. Only the user's OWN object permission is bound: a
|
||||
``UserAPIKeyAuth`` carries a single ``team_id`` while a user may belong to many teams, so
|
||||
team-inherited MCP grants for a user are a follow-up (they need a many-teams union
|
||||
``get_allowed_mcp_servers`` does not do off one auth object). The caller's centralized policy
|
||||
gate enforces the user's live budget and org state, and a SCIM-deactivated owner fails closed.
|
||||
|
||||
Error handling mirrors the key path's retryable-503 contract, but ``get_user_object`` defeats a
|
||||
type-based check: where ``get_key_object`` raises a typed ``ProxyException`` for a missing key
|
||||
and lets a DB outage propagate raw, ``get_user_object`` catches every DB failure and re-raises a
|
||||
bare ``ValueError``, so a missing user and a real outage look identical and the original error
|
||||
survives only as ``__context__``. ``_raise_503_if_db_unavailable`` therefore walks the cause
|
||||
chain: a transient DB outage still surfaces as a retryable 503, while a missing user, or any
|
||||
other non-outage resolution failure, fails closed as a 401 rather than an opaque 500. The
|
||||
object-permission load shares this one boundary, so an outage there is classified the same
|
||||
way (``get_object_permission`` itself swallows a failed load to ``None``, matching how
|
||||
``get_key_object`` best-effort-loads a key's object permission)."""
|
||||
from litellm.proxy.auth.auth_checks import get_object_permission, get_user_object
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Server misconfigured: no database connection")
|
||||
try:
|
||||
user_object = await get_user_object(
|
||||
user_id=user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
)
|
||||
# Resolve the user's own MCP object permission (get_user_object does not load it) so the shared
|
||||
# get_allowed_mcp_servers can grant the user their litellm-granted servers. Reuses the same
|
||||
# get_object_permission resolver the key and team paths use; no permission logic is duplicated.
|
||||
object_permission = user_object.object_permission if user_object is not None else None
|
||||
if user_object is not None and object_permission is None and user_object.object_permission_id:
|
||||
object_permission = await get_object_permission(
|
||||
object_permission_id=user_object.object_permission_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
except (ProxyException, HTTPException):
|
||||
raise HTTPException(status_code=401, detail="Invalid or expired credential") from None
|
||||
except Exception as e: # noqa: BLE001 # a DB outage anywhere in the resolution is a retryable 503, not an opaque 500; anything else fails closed as 401
|
||||
MCPRequestHandler._raise_503_if_db_unavailable(e)
|
||||
raise HTTPException(status_code=401, detail="Invalid or expired credential") from None
|
||||
if user_object is None:
|
||||
raise HTTPException(status_code=401, detail="Invalid or expired credential")
|
||||
if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False:
|
||||
raise HTTPException(status_code=401, detail="Invalid or expired credential")
|
||||
return UserAPIKeyAuth(
|
||||
user_id=user_object.user_id,
|
||||
user_role=user_object.user_role,
|
||||
object_permission=object_permission,
|
||||
object_permission_id=user_object.object_permission_id,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth:
|
||||
"""Reload the live key record an admitted envelope references and re-check live policy.
|
||||
|
|
@ -615,10 +701,14 @@ class MCPRequestHandler:
|
|||
"""Raise a retryable 503 when ``e`` means the auth database is unreachable, else return so the
|
||||
caller applies its own fail-closed mapping. A DB outage must not masquerade as an auth failure
|
||||
(401) or surface as an opaque 500; the caller retries. Mirrors ``UserAPIKeyAuthExceptionHandler``,
|
||||
which renders a service-unavailable database error as 503 on the standard pipeline."""
|
||||
which renders a service-unavailable database error as 503 on the standard pipeline.
|
||||
|
||||
Classifies across the ``__cause__``/``__context__`` chain, not just ``e`` itself: ``get_user_object``
|
||||
re-raises every DB failure as a bare ``ValueError``, so a type-based check on the top exception
|
||||
would miss a real outage wrapped inside it."""
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e):
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.",
|
||||
|
|
|
|||
694
litellm/proxy/_experimental/mcp_server/bridge_token_flow.py
Normal file
694
litellm/proxy/_experimental/mcp_server/bridge_token_flow.py
Normal file
|
|
@ -0,0 +1,694 @@
|
|||
"""Bridge token flow: litellm identity resolution and the DCR-bridge oauth_delegate mint/refresh pipeline."""
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Literal, Optional
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import SecretStr
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _BridgeAuthorizationCode
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
|
||||
EnvelopeIdentity,
|
||||
EnvelopeKeys,
|
||||
RefreshCredential,
|
||||
UpstreamTokenGrant,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
||||
def _litellm_key_from_request(request: Request) -> Optional[str]:
|
||||
"""Return the LiteLLM API key presented on the request, or ``None``.
|
||||
|
||||
Accepts the key from ``x-litellm-api-key`` (what MCP clients such as Claude Desktop/Code
|
||||
send) as well as ``Authorization``; either may carry a bare token or ``Bearer <token>``.
|
||||
``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry
|
||||
an OAuth/upstream bearer.
|
||||
"""
|
||||
for header_value in (
|
||||
request.headers.get("x-litellm-api-key"),
|
||||
request.headers.get("Authorization") or request.headers.get("authorization"),
|
||||
):
|
||||
if not header_value:
|
||||
continue
|
||||
value = header_value.strip()
|
||||
if value.lower().startswith("bearer "):
|
||||
value = value[7:].strip()
|
||||
if value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _key_is_active(key_obj: "UserAPIKeyAuth") -> bool:
|
||||
"""``True`` when the presented key is neither blocked nor past its expiry.
|
||||
|
||||
The OAuth token endpoint is unauthenticated, so the presented key is validated here before it is
|
||||
trusted; a revoked or expired key must not mint a bridge envelope or write a stored credential.
|
||||
``get_key_object`` resolves a row without these checks (the main ``user_api_key_auth`` pipeline
|
||||
enforces them downstream, which this endpoint bypasses), so they are applied here. Deleted keys
|
||||
are already rejected upstream, where ``get_key_object`` raises on a row that no longer exists.
|
||||
|
||||
This is an active-state gate only; it deliberately does not require a ``user_id``. A valid
|
||||
team-scoped or service-account key has no ``user_id`` yet is a legitimate credential, so gating
|
||||
on ``user_id`` presence would wrongly reject it. Callers that need the user (the per-user token
|
||||
store) derive it separately via :func:`_active_key_user_id`.
|
||||
|
||||
Total by design: ``expires`` is typed ``str | datetime``, and an unparseable string would make
|
||||
``datetime.fromisoformat`` raise. Since the callers run this outside their key-resolution
|
||||
``try``, an uncaught parse error would surface as a 500 instead of the endpoint's fail-closed
|
||||
behavior, so a malformed expiry is treated as inactive (return ``False``) rather than raising.
|
||||
"""
|
||||
if key_obj.blocked is True:
|
||||
return False
|
||||
expires = key_obj.expires
|
||||
if expires is not None:
|
||||
if isinstance(expires, datetime):
|
||||
expiry = expires
|
||||
else:
|
||||
try:
|
||||
expiry = datetime.fromisoformat(expires)
|
||||
except (ValueError, TypeError):
|
||||
return False
|
||||
if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None:
|
||||
expiry = expiry.replace(tzinfo=timezone.utc)
|
||||
if expiry < datetime.now(timezone.utc):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> str | None:
|
||||
"""The active key's ``user_id``, or ``None`` when the key is blocked/expired or simply has no
|
||||
``user_id`` (a team-scoped or service-account key). Used only by the per-user token store, which
|
||||
needs a user to key the stored credential; the bridge mint uses the key hash and does not."""
|
||||
return key_obj.user_id if _key_is_active(key_obj) else None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ResolvedKey:
|
||||
"""An active litellm key resolved from the token request: its hash (the value ``get_key_object``
|
||||
and the cache/DB layer key the record by) and the live record."""
|
||||
|
||||
key_hash: str
|
||||
key: "UserAPIKeyAuth"
|
||||
|
||||
|
||||
_KeyResolutionFailure = Literal["no_active_key", "unavailable", "unresolvable"]
|
||||
"""Why a token request yielded no active litellm key, kept distinct so a caller statuses each truthfully
|
||||
instead of blaming the client for a gateway problem:
|
||||
- ``no_active_key``: none was presented, or the presented key is unknown / blocked / expired (the
|
||||
caller's request is at fault)
|
||||
- ``unavailable``: the auth database was transiently unreachable while resolving (retryable)
|
||||
- ``unresolvable``: the gateway cannot resolve identity right now (no DB connection, or an unexpected
|
||||
error) -- a gateway fault, not the caller's
|
||||
The classification mirrors admission's ``_reload_admitted_key`` so the mint (ingress) and admission
|
||||
(egress) never disagree on the status of the same outage."""
|
||||
|
||||
|
||||
async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyResolutionFailure":
|
||||
"""Resolve the presented litellm key to an active key record, or say precisely why not.
|
||||
|
||||
Single resolution path the OAuth token endpoint reuses, resolving authoritatively via
|
||||
``get_key_object`` (cache first, then DB). The failure is a value, not a bare ``None``, so a caller
|
||||
can tell "the client sent no usable credential" (a request error) apart from "the gateway could not
|
||||
check" (an infrastructure error) and status each truthfully; collapsing both to ``None`` is what let
|
||||
a DB outage read as a 400. A resolved key is still gated by ``_key_is_active``, so a blocked or
|
||||
expired key is ``no_active_key`` while a valid team-scoped or service-account key (no ``user_id``)
|
||||
resolves. Classification mirrors admission's ``_reload_admitted_key``: no DB connection is a gateway
|
||||
fault, a ``ProxyException`` / ``HTTPException`` from ``get_key_object`` is an unknown or invalid key,
|
||||
a database-service-unavailable error is a retryable outage, and anything else is an unexpected
|
||||
gateway fault."""
|
||||
token = _litellm_key_from_request(request)
|
||||
if not token:
|
||||
return "no_active_key"
|
||||
from litellm.proxy._types import hash_token # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
|
||||
return await _reload_active_key_by_hash(hash_token(token))
|
||||
|
||||
|
||||
async def _reload_active_key_by_hash(key_hash: str) -> "_ResolvedKey | _KeyResolutionFailure":
|
||||
"""Reload the live key record for ``key_hash`` (cache first, then DB) and gate it on active state,
|
||||
returning the resolved key or a precise failure. Shared by the token request's presented-key
|
||||
resolution (:func:`_resolve_active_litellm_key`, which hashes the presented key) and the refresh
|
||||
path (which already holds the hash sealed in the refresh envelope), so both re-validate identity
|
||||
through one active-key gate and one failure classification. Classification mirrors admission's
|
||||
``_reload_admitted_key``: no DB connection is a gateway fault, a ``ProxyException`` / ``HTTPException``
|
||||
from ``get_key_object`` is an unknown or invalid key, a database-service-unavailable error is a
|
||||
retryable outage, and anything else is an unexpected gateway fault. A blocked or expired key is
|
||||
``no_active_key``, so a revoked key can neither mint nor refresh a bridge envelope."""
|
||||
from litellm.proxy._types import (
|
||||
ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
get_key_object,
|
||||
)
|
||||
from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
PrismaDBExceptionHandler,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
return "unresolvable"
|
||||
try:
|
||||
key_obj = await get_key_object(
|
||||
hashed_token=key_hash,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
except (ProxyException, HTTPException):
|
||||
return "no_active_key"
|
||||
except Exception as exc: # noqa: BLE001 # classify: a DB outage is retryable, anything else is an opaque gateway fault
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error(exc):
|
||||
return "unavailable"
|
||||
verbose_logger.debug(
|
||||
"_reload_active_key_by_hash: unexpected key-resolution error (%s)",
|
||||
type(exc).__name__,
|
||||
)
|
||||
return "unresolvable"
|
||||
if not _key_is_active(key_obj):
|
||||
return "no_active_key"
|
||||
return _ResolvedKey(key_hash=key_hash, key=key_obj)
|
||||
|
||||
|
||||
async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | None":
|
||||
"""Re-validate a live litellm user by id, returning ``None`` when the user is active or a precise
|
||||
failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a
|
||||
user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a
|
||||
deactivated user cannot keep refreshing, mirroring how admission re-validates the same user subject on
|
||||
the egress side. No DB connection is a gateway fault (``unresolvable``) and a
|
||||
database-service-unavailable error is a retryable outage (``unavailable``). Everything else fails
|
||||
closed as ``no_active_key`` (the caller maps it to invalid_grant): a ``ProxyException`` /
|
||||
``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object``
|
||||
catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look
|
||||
identical, the original error surviving only as ``__context__``), so the outage check walks the cause
|
||||
chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault."""
|
||||
from litellm.proxy._types import (
|
||||
ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
get_user_object,
|
||||
)
|
||||
from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
PrismaDBExceptionHandler,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
return "unresolvable"
|
||||
try:
|
||||
user_object = await get_user_object(
|
||||
user_id=user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
)
|
||||
except (ProxyException, HTTPException):
|
||||
return "no_active_key"
|
||||
except Exception as exc: # noqa: BLE001 # a DB outage is retryable; a missing user (get_user_object's wrapped ValueError) or any other resolution failure fails closed as no_active_key, never a 500
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(exc):
|
||||
return "unavailable"
|
||||
verbose_logger.debug("_reload_active_user_by_id: user-resolution error (%s)", type(exc).__name__)
|
||||
return "no_active_key"
|
||||
if user_object is None:
|
||||
return "no_active_key"
|
||||
if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False:
|
||||
return "no_active_key"
|
||||
return None
|
||||
|
||||
|
||||
async def _key_owner_scim_deactivated(key: "UserAPIKeyAuth") -> bool:
|
||||
"""True only when the key's owning user was explicitly SCIM-deactivated, so a refresh revokes an
|
||||
offboarded owner's key exactly as admission does via ``_reject_if_admitted_owner_scim_deactivated``.
|
||||
A key with no owner, a missing owner record, or a failed lookup fails OPEN (returns ``False``),
|
||||
matching admission and the standard builder: a key may outlive its owner record, and a transient DB
|
||||
blip must not revoke a live key. Only an explicit ``scim_active`` of ``False`` gates renewal."""
|
||||
if key.user_id is None:
|
||||
return False
|
||||
from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
get_user_object,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
return False
|
||||
try:
|
||||
owner = await get_user_object(
|
||||
user_id=key.user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 # fail open: a missing owner (get_user_object's wrapped ValueError) or a DB blip must not revoke a live key
|
||||
verbose_logger.debug("refresh: key-owner SCIM lookup failed, not revoking (%s)", type(exc).__name__)
|
||||
return False
|
||||
return owner is not None and isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False
|
||||
|
||||
|
||||
async def _revalidate_active_subject(identity: "EnvelopeIdentity") -> "_KeyResolutionFailure | None":
|
||||
"""Re-validate that the subject sealed in a refresh envelope is still live, dispatching on its type:
|
||||
a key_hash reloads the virtual key, a user_id reloads the user. Returns ``None`` when the subject is
|
||||
active or a precise failure otherwise, so revocation gates renewal for either identity source the same
|
||||
way admission gates the egress: a blocked or expired key, a SCIM-deactivated key owner (mirroring
|
||||
admission's owner check, so an offboarded user cannot keep renewing a still-active key), and a
|
||||
deactivated or deleted user all fail closed to ``no_active_key``."""
|
||||
match identity.subject_type:
|
||||
case "key_hash":
|
||||
reloaded = await _reload_active_key_by_hash(identity.subject)
|
||||
if not isinstance(reloaded, _ResolvedKey):
|
||||
return reloaded
|
||||
if await _key_owner_scim_deactivated(reloaded.key):
|
||||
return "no_active_key"
|
||||
return None
|
||||
case "user_id":
|
||||
return await _reload_active_user_by_id(identity.subject)
|
||||
case _:
|
||||
assert_never(identity.subject_type)
|
||||
|
||||
|
||||
async def _extract_user_id_from_request(request: Request) -> str | None:
|
||||
"""The litellm ``user_id`` for the token request, so a per-user token is stored under the same
|
||||
identity the egress later reads it by. Storage is best-effort, so every non-resolved outcome
|
||||
(including a transient DB outage) collapses to ``None`` here and the caller simply skips the store;
|
||||
the bridge mint, which must status those outcomes differently, consumes
|
||||
:func:`_resolve_active_litellm_key` directly."""
|
||||
resolved = await _resolve_active_litellm_key(request)
|
||||
if not isinstance(resolved, _ResolvedKey):
|
||||
return None
|
||||
return _active_key_user_id(resolved.key)
|
||||
|
||||
|
||||
_UpstreamGrantRejection = Literal["no_access_token", "expired_lifetime"]
|
||||
"""Why an upstream token response cannot back a bridge envelope:
|
||||
- ``no_access_token``: the response carries no usable ``access_token``
|
||||
- ``expired_lifetime``: the response reports a parseable, non-positive ``expires_in``, i.e. an upstream
|
||||
token that is already dead, so sealing it would forward a bearer the edge cannot use
|
||||
An absent or unparseable ``expires_in`` is NOT a rejection; the lifetime is merely unknown and the
|
||||
envelope caps it, the by-design behaviour for an upstream that omits the field."""
|
||||
|
||||
|
||||
def _classify_upstream_lifetime(raw_expires_in: object) -> "int | Literal['unspecified', 'expired']":
|
||||
"""Classify an upstream ``expires_in`` into a positive number of seconds, ``"unspecified"`` (absent
|
||||
or unparseable, so the envelope caps it), or ``"expired"`` (a non-positive value the upstream reports
|
||||
as already elapsed). Telling "we do not know the lifetime" apart from "the upstream says it is
|
||||
already dead" is what stops an explicitly-expired token from silently receiving the envelope's 1h
|
||||
cap. The expired decision is made on the parsed numeric value, not on ``int(...)`` of it, so a
|
||||
positive sub-second lifetime in ``(0, 1)`` is not truncated to ``0`` and misread as elapsed; the
|
||||
envelope works in whole seconds, so such a lifetime clamps up to its 1s floor. ``bool`` is excluded
|
||||
(an ``int`` subclass but never a real lifetime), and the conversions can raise on ``NaN`` /
|
||||
``Infinity`` / oversized input, which reads as unparseable rather than surfacing as a 500."""
|
||||
if raw_expires_in is None or isinstance(raw_expires_in, bool) or not isinstance(raw_expires_in, (int, float, str)):
|
||||
return "unspecified"
|
||||
try:
|
||||
numeric = float(raw_expires_in)
|
||||
seconds = int(numeric)
|
||||
except (ValueError, TypeError, OverflowError):
|
||||
return "unspecified"
|
||||
if numeric <= 0:
|
||||
return "expired"
|
||||
return max(1, seconds)
|
||||
|
||||
|
||||
def _bridge_grant_from_token_response(token_response: object) -> "UpstreamTokenGrant | _UpstreamGrantRejection":
|
||||
"""Validate an upstream OAuth token response into a typed grant, or say why it cannot back an
|
||||
envelope. Each field is isinstance-checked so nothing untyped from ``response.json()`` reaches the
|
||||
grant. ``expires_in`` is read three ways (see :func:`_classify_upstream_lifetime`): an unknown
|
||||
lifetime leaves the grant ``expires_in`` ``None`` for the envelope to cap, a positive value is
|
||||
honoured, and an explicit already-elapsed value is a rejection rather than a silent fall-through to
|
||||
the cap."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
UpstreamTokenGrant,
|
||||
)
|
||||
|
||||
if not isinstance(token_response, dict):
|
||||
return "no_access_token"
|
||||
access = token_response.get("access_token")
|
||||
if not isinstance(access, str) or not access:
|
||||
return "no_access_token"
|
||||
lifetime = _classify_upstream_lifetime(token_response.get("expires_in"))
|
||||
if lifetime == "expired":
|
||||
return "expired_lifetime"
|
||||
token_type = token_response.get("token_type")
|
||||
scope = token_response.get("scope")
|
||||
return UpstreamTokenGrant(
|
||||
access_token=SecretStr(access),
|
||||
token_type=token_type if isinstance(token_type, str) and token_type else "Bearer",
|
||||
# The upstream refresh_token is deliberately NOT sealed: the edge never consumes it (it forwards
|
||||
# only token_type + access_token), so it would be dead weight embedding a long-lived upstream
|
||||
# credential in the client-held bearer, and it enlarges the envelope. Refresh support is a
|
||||
# follow-up (a dedicated refresh-envelope); the client re-runs authorization_code at the cap.
|
||||
refresh_token=None,
|
||||
scope=scope if isinstance(scope, str) and scope else None,
|
||||
expires_in=lifetime if isinstance(lifetime, int) else None,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DCR-bridge oauth_delegate mint: a three-phase pipeline whose failures are values.
|
||||
#
|
||||
# prepare (before the upstream exchange) -> validate every precondition and resolve identity+keys
|
||||
# exchange (the single-use upstream code is consumed here, in exchange_token_with_server)
|
||||
# finish (after the exchange) -> seal the upstream grant into the client-held envelope
|
||||
#
|
||||
# Every precondition lives in ``prepare``, which runs BEFORE the exchange, so no failure can burn the
|
||||
# single-use code or rotate a refresh token, for either grant type -- that whole class of bug is gone
|
||||
# by construction rather than guarded case by case. Failures are values mapped to an OAuth-shaped
|
||||
# response in one place (``_bridge_mint_error_response``), so status codes and the RFC 6749 §5.2 body
|
||||
# shape are uniform. Adding a failure mode is a new literal plus a match arm the type checker forces.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_BridgeMintError = Literal[
|
||||
"no_identity",
|
||||
"invalid_refresh",
|
||||
"identity_unavailable",
|
||||
"identity_unresolvable",
|
||||
"not_configured",
|
||||
"no_upstream_token",
|
||||
"upstream_token_expired",
|
||||
"too_large",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BridgeMintReady:
|
||||
"""Everything the seal needs, resolved once before the exchange: the identity to bind the envelope
|
||||
to and the master-key-derived envelope keys. The identity is a key_hash subject for the scripted
|
||||
two-header client (resolved from the litellm key it presents) or a user_id subject for the
|
||||
interactive SSO client (the user recovered from the gateway authorization code), so one phase-3 seal
|
||||
serves both. Resolving identity here means ``_finish_bridge_mint`` has no preconditions left to
|
||||
fail."""
|
||||
|
||||
identity: "EnvelopeIdentity"
|
||||
keys: "EnvelopeKeys"
|
||||
|
||||
|
||||
def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse:
|
||||
"""Map a bridge-mint failure value to its token-endpoint response: one place, RFC 6749 §5.2 shape
|
||||
(top-level ``error``, no-store headers) for every case, with a status truthful about where the
|
||||
failure is. The caller's request is 400, a transient gateway outage is 503, a gateway
|
||||
misconfiguration is 500, and an upstream problem is 502. The identity-resolution statuses match how
|
||||
admission statuses the same conditions on the egress side, so mint and admit never disagree under
|
||||
one outage."""
|
||||
match error:
|
||||
case "no_identity":
|
||||
status, code, desc = (
|
||||
400,
|
||||
"invalid_request",
|
||||
"this server issues a gateway-bound credential; complete the interactive sign-in, or "
|
||||
"send a litellm credential (x-litellm-api-key or Authorization) on the token request",
|
||||
)
|
||||
case "invalid_refresh":
|
||||
status, code, desc = (
|
||||
400,
|
||||
"invalid_grant",
|
||||
"the refresh credential is not a valid, live refresh envelope for this server; "
|
||||
"re-run authorization_code to obtain a new one",
|
||||
)
|
||||
case "identity_unavailable":
|
||||
status, code, desc = (
|
||||
503,
|
||||
"temporarily_unavailable",
|
||||
"the authentication database is temporarily unreachable; retry shortly",
|
||||
)
|
||||
case "identity_unresolvable":
|
||||
status, code, desc = (
|
||||
500,
|
||||
"server_error",
|
||||
"the gateway could not resolve the litellm identity for this request",
|
||||
)
|
||||
case "not_configured":
|
||||
status, code, desc = (
|
||||
500,
|
||||
"server_error",
|
||||
"the gateway is not configured to mint a gateway-bound credential (master_key is not set)",
|
||||
)
|
||||
case "no_upstream_token":
|
||||
status, code, desc = (
|
||||
502,
|
||||
"server_error",
|
||||
"the upstream token response has no usable access_token",
|
||||
)
|
||||
case "upstream_token_expired":
|
||||
status, code, desc = (
|
||||
502,
|
||||
"server_error",
|
||||
"the upstream token response reports an already-expired lifetime",
|
||||
)
|
||||
case "too_large":
|
||||
status, code, desc = (
|
||||
502,
|
||||
"server_error",
|
||||
"the upstream token is too large to seal into a gateway-bound credential",
|
||||
)
|
||||
case _:
|
||||
assert_never(error)
|
||||
return JSONResponse(
|
||||
status_code=status, content={"error": code, "error_description": desc}, headers=TOKEN_NO_CACHE_HEADERS
|
||||
)
|
||||
|
||||
|
||||
def _key_resolution_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError:
|
||||
"""Lift an identity-resolution failure into the mint taxonomy, preserving origin so the status stays
|
||||
truthful: the caller's missing credential is 400, a transient DB outage is 503, and a gateway that
|
||||
cannot resolve identity is 500."""
|
||||
match failure:
|
||||
case "no_active_key":
|
||||
return "no_identity"
|
||||
case "unavailable":
|
||||
return "identity_unavailable"
|
||||
case "unresolvable":
|
||||
return "identity_unresolvable"
|
||||
case _:
|
||||
assert_never(failure)
|
||||
|
||||
|
||||
def _upstream_rejection_to_mint_error(rejection: _UpstreamGrantRejection) -> _BridgeMintError:
|
||||
"""Lift an upstream-response rejection into the mint taxonomy; both are upstream faults (502)."""
|
||||
match rejection:
|
||||
case "no_access_token":
|
||||
return "no_upstream_token"
|
||||
case "expired_lifetime":
|
||||
return "upstream_token_expired"
|
||||
case _:
|
||||
assert_never(rejection)
|
||||
|
||||
|
||||
async def _prepare_bridge_mint(
|
||||
request: Request,
|
||||
mcp_server: MCPServer,
|
||||
bridge_identity: "_BridgeAuthorizationCode | None" = None,
|
||||
) -> "_BridgeMintReady | _BridgeMintError":
|
||||
"""Phase 1 for the authorization_code grant, BEFORE the upstream exchange: confirm the gateway can
|
||||
mint (master_key set), resolve the litellm identity, and derive the envelope keys. Returns a ready
|
||||
context or a precise failure value. Running before the exchange is what makes every failure here fail
|
||||
closed without consuming the single-use code.
|
||||
|
||||
Two identity sources, one envelope. The interactive DCR client authenticates via SSO at the bridged
|
||||
authorize, so its identity arrives as ``bridge_identity`` (the user recovered from the gateway
|
||||
authorization code) and mints a user subject. The scripted two-header client presents a litellm key
|
||||
on the token request instead, so its identity is the active key's hash and mints a key_hash subject.
|
||||
A missing or invalid presented key keeps its resolution origin so the mapper statuses it truthfully;
|
||||
neither source present is ``no_identity``. The refresh_token grant has its own phase-1
|
||||
(:func:`_prepare_bridge_refresh`), which recovers identity from the presented refresh envelope."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
envelope_keys_from_master_key,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
key_hash_identity,
|
||||
user_identity,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
master_key,
|
||||
)
|
||||
|
||||
if not master_key:
|
||||
return "not_configured"
|
||||
keys = envelope_keys_from_master_key(master_key)
|
||||
if bridge_identity is not None:
|
||||
identity = user_identity(server_id=mcp_server.server_id, user_id=bridge_identity.litellm_user_id)
|
||||
return _BridgeMintReady(identity=identity, keys=keys)
|
||||
resolved = await _resolve_active_litellm_key(request)
|
||||
if not isinstance(resolved, _ResolvedKey):
|
||||
return _key_resolution_failure_to_mint_error(resolved)
|
||||
identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=resolved.key_hash)
|
||||
return _BridgeMintReady(identity=identity, keys=keys)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BridgeRefreshReady:
|
||||
"""A validated refresh request: the identity+keys to mint the renewed pair under, the upstream refresh
|
||||
token (unwrapped from the client's refresh envelope) to exchange with the upstream IdP, and the scope
|
||||
sealed alongside it at mint. The upstream refresh token is a ``SecretStr`` like every other credential
|
||||
in this layer, so a repr or a traceback that captures this value never exposes the raw upstream refresh
|
||||
token in plaintext. ``upstream_scope`` carries the originally-granted scope so the renewal re-requests
|
||||
it when the client (a DCR/MCP client that typically omits scope on refresh) sends none, keeping the
|
||||
renewed token's scope stable against an upstream that would otherwise narrow or drop it."""
|
||||
|
||||
ready: "_BridgeMintReady"
|
||||
upstream_refresh_token: SecretStr
|
||||
upstream_scope: str | None = None
|
||||
|
||||
|
||||
def _refresh_key_failure_to_mint_error(failure: _KeyResolutionFailure) -> _BridgeMintError:
|
||||
"""Lift an identity-resolution failure on the refresh path into the mint taxonomy. Unlike the mint
|
||||
path, a resolved-but-inactive (or unknown) key is ``invalid_grant`` rather than ``invalid_request``:
|
||||
the client did present an identity (sealed in the refresh envelope), but it is no longer live, so the
|
||||
refresh is invalid and the client must re-authenticate. A transient outage is still 503 and a gateway
|
||||
fault still 500, matching the mint path and admission."""
|
||||
match failure:
|
||||
case "no_active_key":
|
||||
return "invalid_refresh"
|
||||
case "unavailable":
|
||||
return "identity_unavailable"
|
||||
case "unresolvable":
|
||||
return "identity_unresolvable"
|
||||
case _:
|
||||
assert_never(failure)
|
||||
|
||||
|
||||
async def _prepare_bridge_refresh(
|
||||
mcp_server: MCPServer, refresh_value: str | None
|
||||
) -> "_BridgeRefreshReady | _BridgeMintError":
|
||||
"""Phase 1 for the refresh_token grant, BEFORE the upstream exchange: open the client's refresh
|
||||
envelope, re-validate the sealed litellm identity so a revoked key cannot keep refreshing, and
|
||||
recover the upstream refresh token to exchange. Identity comes entirely from the sealed envelope, not
|
||||
the HTTP request, so the request object is not needed here. The client presents a refresh envelope,
|
||||
never a raw upstream refresh token, so a missing value, a non-envelope, an unopenable envelope, or one
|
||||
minted for another server is ``invalid_grant``. Running before the exchange means a rejected refresh
|
||||
never consumes or rotates the upstream refresh token."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
BridgeRefreshOpened,
|
||||
envelope_keys_from_master_key,
|
||||
open_bridge_refresh_envelope,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
master_key,
|
||||
)
|
||||
|
||||
if not master_key:
|
||||
return "not_configured"
|
||||
if not refresh_value:
|
||||
return "invalid_refresh"
|
||||
keys = envelope_keys_from_master_key(master_key)
|
||||
opened = open_bridge_refresh_envelope(refresh_value, keys, datetime.now(timezone.utc), mcp_server.server_id)
|
||||
if not isinstance(opened, BridgeRefreshOpened):
|
||||
return "invalid_refresh"
|
||||
failure = await _revalidate_active_subject(opened.identity)
|
||||
if failure is not None:
|
||||
return _refresh_key_failure_to_mint_error(failure)
|
||||
return _BridgeRefreshReady(
|
||||
ready=_BridgeMintReady(identity=opened.identity, keys=keys),
|
||||
upstream_refresh_token=opened.refresh.refresh_token,
|
||||
upstream_scope=opened.refresh.scope,
|
||||
)
|
||||
|
||||
|
||||
def _finish_bridge_mint(
|
||||
ready: "_BridgeMintReady", mcp_server: MCPServer, token_response: object, now: datetime
|
||||
) -> "JSONResponse | _BridgeMintError":
|
||||
"""Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held access envelope
|
||||
using the pre-resolved identity and keys, and, when the upstream returned a refresh token, seal a
|
||||
long-lived refresh envelope alongside it so the client can renew without re-authenticating. Shared by
|
||||
the authorization_code and refresh_token paths, so a renewal that the upstream rotates re-issues a
|
||||
fresh refresh envelope. The only hard failures here are properties of the upstream access token (no
|
||||
usable token, an already-expired lifetime, or a token too large to seal); a refresh token that cannot
|
||||
be sealed degrades to an access-only response rather than failing the whole exchange."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
build_bridge_token_response,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
SealedEnvelope,
|
||||
UpstreamTokenGrant,
|
||||
)
|
||||
|
||||
grant = _bridge_grant_from_token_response(token_response)
|
||||
if not isinstance(grant, UpstreamTokenGrant):
|
||||
return _upstream_rejection_to_mint_error(grant)
|
||||
sealed = build_bridge_token_response(ready.identity, grant, ready.keys, now)
|
||||
if not isinstance(sealed, SealedEnvelope):
|
||||
return "too_large"
|
||||
# Report expires_in from the JWT's own second-truncated exp, rounding the elapsed portion up, so the
|
||||
# client is never told the bearer lives past the point admission (which uses that exp) rejects it.
|
||||
expires_in = max(0, int(sealed.expires_at.timestamp()) - math.ceil(now.timestamp()))
|
||||
refresh_envelope = _mint_refresh_envelope_value(ready.identity, token_response, ready.keys, now, mcp_server)
|
||||
body = {
|
||||
"access_token": sealed.token.get_secret_value(),
|
||||
"token_type": "Bearer",
|
||||
"expires_in": expires_in,
|
||||
# A refresh envelope rides along only when the upstream returned a refresh token to seal; when it
|
||||
# rotates on renewal, the client receives the new one and the old envelope's upstream token dies.
|
||||
**({"refresh_token": refresh_envelope} if refresh_envelope is not None else {}),
|
||||
}
|
||||
return JSONResponse(body, headers=TOKEN_NO_CACHE_HEADERS)
|
||||
|
||||
|
||||
def _upstream_refresh_credential(token_response: object) -> "RefreshCredential | None":
|
||||
"""Extract the upstream refresh grant from a token response, or ``None`` when there is none to seal.
|
||||
Each field is isinstance-checked so nothing untyped reaches the refresh envelope; ``refresh_expires_in``
|
||||
(the refresh token's own lifetime, when the upstream reports it) is classified like ``expires_in`` and
|
||||
bounds the refresh envelope's TTL. An upstream that reports the refresh token itself as already elapsed
|
||||
(``refresh_expires_in`` non-positive) yields ``None`` rather than a refresh envelope: sealing a dead
|
||||
token would hand the client a full-TTL-capped envelope the IdP will reject, so the exchange degrades to
|
||||
an access-only response (the client re-authenticates at access expiry), mirroring how
|
||||
:func:`_bridge_grant_from_token_response` refuses an already-elapsed access token instead of capping it."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
RefreshCredential,
|
||||
)
|
||||
|
||||
if not isinstance(token_response, dict):
|
||||
return None
|
||||
refresh = token_response.get("refresh_token")
|
||||
if not isinstance(refresh, str) or not refresh:
|
||||
return None
|
||||
lifetime = _classify_upstream_lifetime(token_response.get("refresh_expires_in"))
|
||||
if lifetime == "expired":
|
||||
return None
|
||||
scope = token_response.get("scope")
|
||||
return RefreshCredential(
|
||||
refresh_token=SecretStr(refresh),
|
||||
scope=scope if isinstance(scope, str) and scope else None,
|
||||
expires_in=lifetime if isinstance(lifetime, int) else None,
|
||||
)
|
||||
|
||||
|
||||
def _mint_refresh_envelope_value(
|
||||
identity: "EnvelopeIdentity", token_response: object, keys: "EnvelopeKeys", now: datetime, mcp_server: MCPServer
|
||||
) -> str | None:
|
||||
"""Seal the upstream refresh grant (if any) into a refresh envelope and return its bearer string, or
|
||||
``None`` when the upstream returned no refresh token or the refresh token is too large to seal. A
|
||||
too-large refresh token degrades to an access-only response (logged) rather than failing an exchange
|
||||
that already succeeded upstream: the client simply re-authenticates when the access envelope expires."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
build_bridge_refresh_token_response,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
SealedEnvelope,
|
||||
)
|
||||
|
||||
refresh_credential = _upstream_refresh_credential(token_response)
|
||||
if refresh_credential is None:
|
||||
return None
|
||||
sealed = build_bridge_refresh_token_response(identity, refresh_credential, keys, now)
|
||||
if isinstance(sealed, SealedEnvelope):
|
||||
return sealed.token.get_secret_value()
|
||||
verbose_logger.warning(
|
||||
"bridge mint: the upstream refresh token is too large to seal into a refresh envelope for "
|
||||
"server=%s; issuing an access-only response, so the client re-authenticates at access expiry",
|
||||
mcp_server.server_id,
|
||||
)
|
||||
return None
|
||||
|
|
@ -10,7 +10,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
|
|||
import httpx
|
||||
from fastapi import APIRouter, Form, HTTPException, Request
|
||||
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
|
|
@ -21,6 +21,24 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import (
|
|||
TokenEndpointAuthConfigError,
|
||||
build_token_endpoint_client_auth,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.bridge_token_flow import (
|
||||
_bridge_mint_error_response,
|
||||
_BridgeMintReady,
|
||||
_BridgeRefreshReady,
|
||||
_extract_user_id_from_request,
|
||||
_finish_bridge_mint,
|
||||
_prepare_bridge_mint,
|
||||
_prepare_bridge_refresh,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.faults import (
|
||||
CallerRejected,
|
||||
CredentialSource,
|
||||
UpstreamProtocolFault,
|
||||
classify_upstream_dcr_rejection,
|
||||
classify_upstream_token_rejection,
|
||||
dcr_fault_detail,
|
||||
render_token_fault,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.oauth_utils import (
|
||||
TOKEN_NO_CACHE_HEADERS,
|
||||
get_request_base_url,
|
||||
|
|
@ -37,7 +55,7 @@ from litellm.types.mcp import MCPAuth, MCPCredentials
|
|||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import LiteLLM_MCPServerTable, UserAPIKeyAuth
|
||||
from litellm.proxy._types import LiteLLM_MCPServerTable
|
||||
|
||||
# TTL cache for upstream OAuth metadata fetched from pass-through MCP servers.
|
||||
# Keeps us from hammering the upstream IdP on each discovery request.
|
||||
|
|
@ -91,6 +109,8 @@ def encode_state_with_base_url(
|
|||
code_challenge: Optional[str] = None,
|
||||
code_challenge_method: Optional[str] = None,
|
||||
client_redirect_uri: Optional[str] = None,
|
||||
litellm_user_id: str | None = None,
|
||||
mcp_server_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Encode the base_url, original state, and PKCE parameters using encryption.
|
||||
|
|
@ -101,6 +121,11 @@ def encode_state_with_base_url(
|
|||
code_challenge: PKCE code challenge from client
|
||||
code_challenge_method: PKCE code challenge method from client
|
||||
client_redirect_uri: Original redirect_uri from client
|
||||
litellm_user_id: The SSO-authenticated litellm user captured at the bridge authorize
|
||||
(interactive dcr_bridge oauth_delegate only); the callback seals it into the gateway
|
||||
authorization code so the token mint can bind the envelope to this user
|
||||
mcp_server_id: The bridge server the interactive flow targets, sealed alongside
|
||||
litellm_user_id so the gateway code cannot be replayed against another server
|
||||
|
||||
Returns:
|
||||
An encrypted string that encodes all values
|
||||
|
|
@ -111,6 +136,8 @@ def encode_state_with_base_url(
|
|||
"code_challenge": code_challenge,
|
||||
"code_challenge_method": code_challenge_method,
|
||||
"client_redirect_uri": client_redirect_uri,
|
||||
"litellm_user_id": litellm_user_id,
|
||||
"mcp_server_id": mcp_server_id,
|
||||
}
|
||||
state_json = json.dumps(state_data, sort_keys=True)
|
||||
encrypted_state = encrypt_value_helper(state_json)
|
||||
|
|
@ -138,6 +165,68 @@ def decode_state_hash(encrypted_state: str) -> dict:
|
|||
return state_data
|
||||
|
||||
|
||||
_BRIDGE_AUTH_CODE_PREFIX = "llm_bcode_"
|
||||
|
||||
|
||||
class _BridgeAuthorizationCode(BaseModel):
|
||||
"""The identity and upstream code the gateway seals into the authorization code it hands a DCR
|
||||
client for an interactive dcr_bridge oauth_delegate sign-in, recovered at the token endpoint."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
upstream_code: str = Field(min_length=1)
|
||||
litellm_user_id: str = Field(min_length=1)
|
||||
mcp_server_id: str = Field(min_length=1)
|
||||
|
||||
|
||||
def is_bridge_authorization_code(code: str) -> bool:
|
||||
"""Cheap prefix check that ``code`` is a gateway-sealed bridge authorization code rather than a
|
||||
raw upstream code, so the token endpoint can route without decrypting."""
|
||||
return code.startswith(_BRIDGE_AUTH_CODE_PREFIX)
|
||||
|
||||
|
||||
def seal_bridge_authorization_code(upstream_code: str, litellm_user_id: str, mcp_server_id: str) -> str:
|
||||
"""Seal the upstream authorization code and the SSO-captured litellm user into a gateway
|
||||
authorization code. The DCR client only echoes this opaque value back at the token endpoint; the
|
||||
gateway decrypts it there to recover the user (to bind the envelope) and the upstream code (to
|
||||
exchange with the upstream), so a litellm identity captured in the browser at authorize survives
|
||||
to the back-channel token call with nothing stored server-side. Encrypted with the repo's
|
||||
authenticated symmetric helper (the same family the OAuth state uses), so the client can neither
|
||||
read nor forge it."""
|
||||
payload = json.dumps(
|
||||
{"upstream_code": upstream_code, "litellm_user_id": litellm_user_id, "mcp_server_id": mcp_server_id},
|
||||
sort_keys=True,
|
||||
)
|
||||
return _BRIDGE_AUTH_CODE_PREFIX + encrypt_value_helper(payload)
|
||||
|
||||
|
||||
def open_bridge_authorization_code(code: str) -> _BridgeAuthorizationCode | None:
|
||||
"""Recover the sealed identity and upstream code, or ``None`` when ``code`` is not a gateway
|
||||
bridge code or does not decrypt / validate. Total over hostile input: a raw upstream code (the
|
||||
scripted two-header path) returns ``None`` and the caller falls through to the existing
|
||||
behavior."""
|
||||
if not is_bridge_authorization_code(code):
|
||||
return None
|
||||
decrypted = decrypt_value_helper(
|
||||
code[len(_BRIDGE_AUTH_CODE_PREFIX) :], "bridge_authorization_code", return_original_value=False
|
||||
)
|
||||
if not isinstance(decrypted, str):
|
||||
return None
|
||||
try:
|
||||
return _BridgeAuthorizationCode.model_validate_json(decrypted)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _redirect_to_litellm_login(request: Request) -> RedirectResponse:
|
||||
"""Send an unauthenticated browser through litellm login before the interactive bridge authorize
|
||||
can capture its identity. The bridge oauth_delegate flow seals the SSO user into the gateway code,
|
||||
so a session is required; without one there is nothing to bind. After login the user re-initiates
|
||||
the connection, which then finds the session cookie (the seamless return-to round-trip, which is
|
||||
origin-validated against the control-plane URL, is a follow-up)."""
|
||||
base_url = get_request_base_url(request)
|
||||
return RedirectResponse(f"{base_url}/sso/key/generate")
|
||||
|
||||
|
||||
# LIT-4197: some upstream authorization servers reject an over-long ``state``
|
||||
# (the encrypted OAuth session blob routinely exceeds their limit). The upstream
|
||||
# only needs an opaque value it echoes back on ``/callback``, so we forward a
|
||||
|
|
@ -304,90 +393,6 @@ def _validate_token_response(
|
|||
)
|
||||
|
||||
|
||||
def _litellm_key_from_request(request: Request) -> Optional[str]:
|
||||
"""Return the LiteLLM API key presented on the request, or ``None``.
|
||||
|
||||
Accepts the key from ``x-litellm-api-key`` (what MCP clients such as Claude Desktop/Code
|
||||
send) as well as ``Authorization``; either may carry a bare token or ``Bearer <token>``.
|
||||
``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry
|
||||
an OAuth/upstream bearer.
|
||||
"""
|
||||
for header_value in (
|
||||
request.headers.get("x-litellm-api-key"),
|
||||
request.headers.get("Authorization") or request.headers.get("authorization"),
|
||||
):
|
||||
if not header_value:
|
||||
continue
|
||||
value = header_value.strip()
|
||||
if value.lower().startswith("bearer "):
|
||||
value = value[7:].strip()
|
||||
if value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> Optional[str]:
|
||||
"""The key's ``user_id``, or ``None`` if the key is blocked or expired.
|
||||
|
||||
The OAuth token endpoint is unauthenticated, so the presented key is validated here before its
|
||||
identity is trusted to key a stored credential; a revoked or expired key must not be able to
|
||||
write or overwrite the per-user OAuth token. ``get_key_object`` resolves a row without these
|
||||
checks (the main ``user_api_key_auth`` pipeline enforces them downstream, which this endpoint
|
||||
bypasses), so they are applied here. Deleted keys are already rejected upstream, where
|
||||
``get_key_object`` raises on a row that no longer exists.
|
||||
"""
|
||||
if key_obj.blocked is True:
|
||||
return None
|
||||
expires = key_obj.expires
|
||||
if expires is not None:
|
||||
expiry = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires)
|
||||
if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None:
|
||||
expiry = expiry.replace(tzinfo=timezone.utc)
|
||||
if expiry < datetime.now(timezone.utc):
|
||||
return None
|
||||
return key_obj.user_id
|
||||
|
||||
|
||||
async def _extract_user_id_from_request(request: Request) -> Optional[str]:
|
||||
"""Resolve the LiteLLM ``user_id`` at the OAuth token endpoint so a per-user token is stored
|
||||
under the same identity the egress later reads it by (``user_api_key_auth.user_id``).
|
||||
|
||||
Resolves authoritatively via ``get_key_object`` (cache first, then DB) instead of a raw cache
|
||||
peek. On a multi-replica gateway the token-exchange request can land on a worker whose in-memory
|
||||
cache never saw the key, and a cross-replica Redis hit deserializes to a plain ``dict`` rather
|
||||
than a ``UserAPIKeyAuth``; the previous code read only ``Authorization`` and did
|
||||
``getattr(cached, "user_id")`` with no ``model_type`` rehydration and no DB fallback, so it
|
||||
silently returned ``None`` and the token was never persisted, which makes the egress 401 on every
|
||||
reconnect. The resolved key is validated (``_active_key_user_id``) before its identity is trusted,
|
||||
so a blocked or expired key cannot write. Returns ``None`` when no key is present, the key cannot
|
||||
be resolved, or it is blocked/expired.
|
||||
"""
|
||||
token = _litellm_key_from_request(request)
|
||||
if not token:
|
||||
return None
|
||||
try:
|
||||
from litellm.proxy._types import hash_token # noqa: PLC0415
|
||||
from litellm.proxy.auth.auth_checks import get_key_object # noqa: PLC0415
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
key_obj = await get_key_object(
|
||||
hashed_token=hash_token(token),
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
return _active_key_user_id(key_obj)
|
||||
except Exception as exc:
|
||||
verbose_logger.debug(
|
||||
"_extract_user_id_from_request: could not resolve a LiteLLM user_id for the presented "
|
||||
"key (%s); per-user token will not be stored server-side.",
|
||||
type(exc).__name__,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def _store_per_user_token_server_side(
|
||||
server: MCPServer,
|
||||
user_id: str,
|
||||
|
|
@ -620,12 +625,31 @@ async def authorize_with_server(
|
|||
parsed = urlparse(redirect_uri)
|
||||
base_url = urlunparse(parsed._replace(query=""))
|
||||
request_base_url = get_request_base_url(request)
|
||||
|
||||
# Interactive dcr_bridge oauth_delegate sign-in: this arm runs the gateway /callback and /token in
|
||||
# the loop, so the gateway can capture the litellm user here (from the browser's UI session) and
|
||||
# carry it to the back-channel token mint. Seal the SSO user and the target server into the state;
|
||||
# the callback reads them back to mint the gateway authorization code. A DCR client cannot present a
|
||||
# litellm key, so the browser session is the only identity source; without one there is nothing to
|
||||
# bind, so send the user through login first. Every other oauth2 server keeps the identity-less state.
|
||||
litellm_user_id: str | None = None
|
||||
if mcp_server.is_dcr_bridge and mcp_server.is_oauth_delegate:
|
||||
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
_user_id_from_session_cookie,
|
||||
)
|
||||
|
||||
litellm_user_id = _user_id_from_session_cookie(request)
|
||||
if litellm_user_id is None:
|
||||
return _redirect_to_litellm_login(request)
|
||||
|
||||
encoded_state = encode_state_with_base_url(
|
||||
base_url=base_url,
|
||||
original_state=state,
|
||||
code_challenge=code_challenge,
|
||||
code_challenge_method=code_challenge_method,
|
||||
client_redirect_uri=redirect_uri,
|
||||
litellm_user_id=litellm_user_id,
|
||||
mcp_server_id=mcp_server.server_id if litellm_user_id else None,
|
||||
)
|
||||
relay_state = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES)
|
||||
|
||||
|
|
@ -654,6 +678,13 @@ async def authorize_with_server(
|
|||
return response
|
||||
|
||||
|
||||
def _token_credential_source(mcp_server: MCPServer) -> CredentialSource:
|
||||
"""Mirrors the resolved-client rule in :func:`exchange_token_with_server`: when the server has a
|
||||
stored client_id the gateway presents its own credentials upstream, so a credential rejection is
|
||||
the operator's fault, not the caller's."""
|
||||
return "gateway_stored" if mcp_server.client_id else "caller_supplied"
|
||||
|
||||
|
||||
async def exchange_token_with_server(
|
||||
request: Request,
|
||||
mcp_server: MCPServer,
|
||||
|
|
@ -688,25 +719,61 @@ async def exchange_token_with_server(
|
|||
except TokenEndpointAuthConfigError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
|
||||
bridge_identity: _BridgeAuthorizationCode | None = None
|
||||
bridge_mint_ready: _BridgeMintReady | None = None
|
||||
bridge_upstream_refresh: SecretStr | None = None
|
||||
bridge_upstream_scope: str | None = None
|
||||
refresh_request_scope: str | None = None
|
||||
is_bridge = mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge
|
||||
|
||||
if grant_type == "refresh_token":
|
||||
if not refresh_token:
|
||||
# Phase 1 for a bridge refresh: open the client's refresh envelope, re-validate the sealed
|
||||
# identity, and unwrap the real upstream refresh token BEFORE building token_data, so the exchange
|
||||
# sends the upstream token and never the envelope. A failure returns without touching the upstream.
|
||||
if is_bridge:
|
||||
prepared_refresh = await _prepare_bridge_refresh(mcp_server, refresh_token)
|
||||
if not isinstance(prepared_refresh, _BridgeRefreshReady):
|
||||
return _bridge_mint_error_response(prepared_refresh)
|
||||
bridge_mint_ready = prepared_refresh.ready
|
||||
bridge_upstream_refresh = prepared_refresh.upstream_refresh_token
|
||||
bridge_upstream_scope = prepared_refresh.upstream_scope
|
||||
# A bridge server sends the unwrapped upstream refresh token recovered from the client's refresh
|
||||
# envelope above; every other server sends the client's own refresh token verbatim.
|
||||
upstream_refresh_token = (
|
||||
bridge_upstream_refresh.get_secret_value() if bridge_upstream_refresh is not None else refresh_token
|
||||
)
|
||||
if not upstream_refresh_token:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="refresh_token is required for refresh_token grant",
|
||||
)
|
||||
token_data: dict = {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refresh_token,
|
||||
"refresh_token": upstream_refresh_token,
|
||||
**client_auth.body,
|
||||
}
|
||||
if scope:
|
||||
token_data["scope"] = scope
|
||||
refresh_request_scope = scope or bridge_upstream_scope
|
||||
if refresh_request_scope:
|
||||
token_data["scope"] = refresh_request_scope
|
||||
else:
|
||||
if not code:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="code is required for authorization_code grant",
|
||||
)
|
||||
# Interactive dcr_bridge oauth_delegate: the client presents the gateway authorization code the
|
||||
# callback sealed. Recover the SSO user and the real upstream code from it; the upstream exchange
|
||||
# below uses the upstream code, and the mint binds the envelope to the recovered user. Bind the
|
||||
# sealed server to this request so a code minted for one bridge server cannot be spent at another.
|
||||
# A raw upstream code (scripted path) opens to None and the code is used as-is.
|
||||
bridge_identity = open_bridge_authorization_code(code)
|
||||
if bridge_identity is not None:
|
||||
if bridge_identity.mcp_server_id != mcp_server.server_id:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Authorization code was issued for a different MCP server",
|
||||
)
|
||||
code = bridge_identity.upstream_code
|
||||
bridge_token_relay = _dcr_bridge_relays_client_registration(mcp_server)
|
||||
if bridge_token_relay and not redirect_uri:
|
||||
raise HTTPException(
|
||||
|
|
@ -726,32 +793,49 @@ async def exchange_token_with_server(
|
|||
}
|
||||
if code_verifier:
|
||||
token_data["code_verifier"] = code_verifier
|
||||
|
||||
# Phase 1 for a bridge authorization_code mint: resolve identity (the SSO user recovered above, or
|
||||
# the presented litellm key) and the envelope keys BEFORE the exchange consumes the single-use code.
|
||||
if is_bridge:
|
||||
prepared = await _prepare_bridge_mint(request, mcp_server, bridge_identity)
|
||||
if not isinstance(prepared, _BridgeMintReady):
|
||||
return _bridge_mint_error_response(prepared)
|
||||
bridge_mint_ready = prepared
|
||||
async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
|
||||
response = await async_client.post(
|
||||
mcp_server.token_url,
|
||||
headers={"Accept": "application/json", **client_auth.headers},
|
||||
data=token_data,
|
||||
)
|
||||
try:
|
||||
response = await async_client.post(
|
||||
mcp_server.token_url,
|
||||
headers={"Accept": "application/json", **client_auth.headers},
|
||||
data=token_data,
|
||||
)
|
||||
if response is not None:
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as exc:
|
||||
fault = classify_upstream_token_rejection(
|
||||
exc.response,
|
||||
credential_source=_token_credential_source(mcp_server),
|
||||
log_context=mcp_server.server_id,
|
||||
)
|
||||
upstream_rejected_bridge_refresh = (
|
||||
is_bridge
|
||||
and grant_type == "refresh_token"
|
||||
and isinstance(fault, CallerRejected)
|
||||
and fault.code == "invalid_grant"
|
||||
)
|
||||
if upstream_rejected_bridge_refresh:
|
||||
verbose_logger.info(
|
||||
"bridge refresh: the upstream rejected the sealed refresh token for server=%s with "
|
||||
"invalid_grant (revoked or expired at the IdP); returning invalid_grant so the client "
|
||||
"re-runs authorization_code rather than an opaque upstream error",
|
||||
mcp_server.server_id,
|
||||
)
|
||||
return _bridge_mint_error_response("invalid_refresh")
|
||||
return render_token_fault(fault)
|
||||
if response is None:
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail="MCP upstream token endpoint returned no response",
|
||||
)
|
||||
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as exc:
|
||||
if "invalid_target" in exc.response.text:
|
||||
verbose_logger.warning(
|
||||
"MCP server %s: the upstream authorization server rejected the token request with "
|
||||
"invalid_target; it may require RFC 8707 resource indicators, which the gateway "
|
||||
"does not send yet (tracked as LIT-4339)",
|
||||
mcp_server.server_id,
|
||||
)
|
||||
raise
|
||||
token_response = response.json()
|
||||
access_token = token_response["access_token"]
|
||||
|
||||
# Validate token response against server-configured rules before any storage.
|
||||
# This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc.
|
||||
|
|
@ -791,8 +875,23 @@ async def exchange_token_with_server(
|
|||
mcp_server.server_id,
|
||||
)
|
||||
|
||||
# A DCR-bridge oauth_delegate server hands the client a gateway-bound envelope (identity plus the
|
||||
# upstream token) instead of the raw upstream token, so the one bearer both admits the caller and
|
||||
# forwards the upstream credential. Only this mode mints; every other server returns the raw token.
|
||||
if bridge_mint_ready is not None:
|
||||
if refresh_request_scope and isinstance(token_response, dict) and not token_response.get("scope"):
|
||||
token_response = {**token_response, "scope": refresh_request_scope}
|
||||
# Phase 3: seal the upstream grant into the client-held envelope; failures map through the same
|
||||
# OAuth-shaped response as the phase-1 preconditions.
|
||||
minted = _finish_bridge_mint(bridge_mint_ready, mcp_server, token_response, datetime.now(timezone.utc))
|
||||
return minted if isinstance(minted, JSONResponse) else _bridge_mint_error_response(minted)
|
||||
|
||||
raw_access_token = token_response.get("access_token") if isinstance(token_response, dict) else None
|
||||
if not isinstance(raw_access_token, str) or not raw_access_token:
|
||||
return render_token_fault(UpstreamProtocolFault(note="the upstream token response has no usable access_token"))
|
||||
|
||||
result = {
|
||||
"access_token": access_token,
|
||||
"access_token": raw_access_token,
|
||||
"token_type": token_response.get("token_type", "Bearer"),
|
||||
}
|
||||
|
||||
|
|
@ -1048,21 +1147,6 @@ async def _persist_dcr_client_registration(
|
|||
return "failed"
|
||||
|
||||
|
||||
_MAX_UPSTREAM_ERROR_CHARS = 500
|
||||
|
||||
|
||||
def _safe_upstream_error_detail(response: httpx.Response) -> str:
|
||||
"""Bounded plaintext summary of an upstream registration failure for the client.
|
||||
|
||||
RFC 7591 error bodies are small JSON objects (``error`` / ``error_description``); relaying the
|
||||
text lets the client read the real reason instead of a bare 500, and the length bound keeps a
|
||||
hostile or oversized upstream body from bloating the gateway response."""
|
||||
body = response.text
|
||||
if not body:
|
||||
return response.reason_phrase or "upstream registration failed"
|
||||
return body[:_MAX_UPSTREAM_ERROR_CHARS]
|
||||
|
||||
|
||||
async def register_client_with_server(
|
||||
request: Request,
|
||||
mcp_server: MCPServer,
|
||||
|
|
@ -1122,19 +1206,24 @@ async def register_client_with_server(
|
|||
}
|
||||
|
||||
async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Register)
|
||||
response = await async_client.post(
|
||||
mcp_server.registration_url,
|
||||
headers=headers,
|
||||
json=register_data,
|
||||
)
|
||||
try:
|
||||
response = await async_client.post(
|
||||
mcp_server.registration_url,
|
||||
headers=headers,
|
||||
json=register_data,
|
||||
)
|
||||
if response is not None:
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as exc:
|
||||
status_code, detail = dcr_fault_detail(
|
||||
classify_upstream_dcr_rejection(exc.response, log_context=mcp_server.server_id)
|
||||
)
|
||||
raise HTTPException(status_code=status_code, detail=detail) from exc
|
||||
if response is None:
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail="MCP upstream registration endpoint returned no response",
|
||||
)
|
||||
if bridge_relay and response.status_code >= 400:
|
||||
raise HTTPException(status_code=response.status_code, detail=_safe_upstream_error_detail(response))
|
||||
response.raise_for_status()
|
||||
|
||||
token_response = response.json()
|
||||
|
||||
|
|
@ -1362,7 +1451,20 @@ async def callback(
|
|||
# states while permitting same-origin / allowlisted clients.
|
||||
redirect_uri = _get_validated_client_redirect_uri(request, state_data)
|
||||
|
||||
params = {"code": code, "state": original_state}
|
||||
# Interactive dcr_bridge oauth_delegate: the state carries the litellm user the authorize step
|
||||
# captured. Instead of forwarding the raw upstream code (which the client would present at the
|
||||
# token endpoint with no way to prove who signed in), seal the user and the upstream code into a
|
||||
# gateway authorization code and forward THAT. The token endpoint decrypts it to bind the
|
||||
# envelope to this user. Every other flow forwards the raw code unchanged.
|
||||
litellm_user_id = state_data.get("litellm_user_id")
|
||||
mcp_server_id = state_data.get("mcp_server_id")
|
||||
forwarded_code = code
|
||||
if isinstance(litellm_user_id, str) and litellm_user_id and isinstance(mcp_server_id, str) and mcp_server_id:
|
||||
forwarded_code = seal_bridge_authorization_code(
|
||||
upstream_code=code, litellm_user_id=litellm_user_id, mcp_server_id=mcp_server_id
|
||||
)
|
||||
|
||||
params = {"code": forwarded_code, "state": original_state}
|
||||
complete_returned_url = _append_query_params(redirect_uri, params)
|
||||
response = RedirectResponse(url=complete_returned_url, status_code=302)
|
||||
_clear_oauth_state_cookie(response, request, state)
|
||||
|
|
|
|||
38
litellm/proxy/_experimental/mcp_server/faults/__init__.py
Normal file
38
litellm/proxy/_experimental/mcp_server/faults/__init__.py
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
"""Typed fault values for upstream OAuth/DCR failures (phase 1 of the MCP error-handling framework).
|
||||
|
||||
The invariant this package exists to enforce: an upstream failure is classified ONCE into a single
|
||||
fault value, and the response status, wire error code, and prose are all derived from that value.
|
||||
Deriving all three from one classification makes contradictory pairings (a caller-fault error code on
|
||||
a server-fault status) unrepresentable, and gives the trust-boundary rule one enforcement point:
|
||||
spec-defined machine fields may cross to callers, upstream prose and raw bodies go to server logs.
|
||||
"""
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.faults.classify import (
|
||||
classify_upstream_dcr_rejection,
|
||||
classify_upstream_token_rejection,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.faults.render_oauth import (
|
||||
dcr_fault_detail,
|
||||
render_token_fault,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.faults.types import (
|
||||
CallerRejected,
|
||||
CredentialSource,
|
||||
GatewayRejected,
|
||||
UpstreamOAuthFault,
|
||||
UpstreamProtocolFault,
|
||||
UpstreamReportedFault,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CallerRejected",
|
||||
"CredentialSource",
|
||||
"GatewayRejected",
|
||||
"UpstreamOAuthFault",
|
||||
"UpstreamProtocolFault",
|
||||
"UpstreamReportedFault",
|
||||
"classify_upstream_dcr_rejection",
|
||||
"classify_upstream_token_rejection",
|
||||
"dcr_fault_detail",
|
||||
"render_token_fault",
|
||||
]
|
||||
133
litellm/proxy/_experimental/mcp_server/faults/classify.py
Normal file
133
litellm/proxy/_experimental/mcp_server/faults/classify.py
Normal file
|
|
@ -0,0 +1,133 @@
|
|||
"""The single place that reads upstream OAuth/DCR failure responses.
|
||||
|
||||
Every accessor here is total: an upstream that lies about its content encoding, sends an undecodable
|
||||
body, or omits the spec fields yields a classified fault, never an exception. Nothing outside this
|
||||
module should touch a failed upstream response's body.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._experimental.mcp_server.faults.types import (
|
||||
GATEWAY_CAPABILITY_CODES,
|
||||
GATEWAY_CREDENTIAL_CODES,
|
||||
MAX_WIRE_FIELD_CHARS,
|
||||
CallerRejected,
|
||||
CredentialSource,
|
||||
GatewayRejected,
|
||||
UpstreamOAuthFault,
|
||||
UpstreamProtocolFault,
|
||||
UpstreamReportedFault,
|
||||
)
|
||||
|
||||
|
||||
def _safe_text(response: httpx.Response) -> str:
|
||||
try:
|
||||
return response.text
|
||||
except Exception:
|
||||
return ""
|
||||
|
||||
|
||||
def _safe_json(response: httpx.Response) -> object:
|
||||
try:
|
||||
return response.json()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _bounded_field(value: object) -> str | None:
|
||||
if not isinstance(value, str) or not value:
|
||||
return None
|
||||
return value[:MAX_WIRE_FIELD_CHARS]
|
||||
|
||||
|
||||
def _log_out_of_contract(endpoint_kind: str, response: httpx.Response, log_context: str) -> None:
|
||||
verbose_logger.warning(
|
||||
"MCP upstream %s endpoint (%s) returned HTTP %s outside the OAuth error contract (first %s chars): %s",
|
||||
endpoint_kind,
|
||||
log_context,
|
||||
response.status_code,
|
||||
MAX_WIRE_FIELD_CHARS,
|
||||
_safe_text(response)[:MAX_WIRE_FIELD_CHARS],
|
||||
)
|
||||
|
||||
|
||||
def _classify_oauth_error_code(
|
||||
code: str,
|
||||
description: str | None,
|
||||
error_uri: str | None,
|
||||
credential_source: CredentialSource,
|
||||
log_context: str,
|
||||
) -> UpstreamOAuthFault:
|
||||
"""Blame assignment for a contract-conformant OAuth error code, shared by the token and DCR
|
||||
classifiers. Codes by which the upstream blames itself keep that blame; ``invalid_target`` is a
|
||||
gateway capability gap (RFC 8707 resource indicators, LIT-4339) no matter whose credentials were
|
||||
presented; credential-indicting codes follow the credential source; everything else, including
|
||||
codes we do not recognize, is the caller's to act on. The upstream's HTTP status is deliberately
|
||||
never consulted: status derives from this classification at render time, which is what keeps
|
||||
status and code from contradicting each other."""
|
||||
if code == "server_error" or code == "temporarily_unavailable":
|
||||
return UpstreamReportedFault(code=code)
|
||||
if code in GATEWAY_CAPABILITY_CODES:
|
||||
verbose_logger.warning(
|
||||
"MCP server %s: the upstream authorization server rejected the request with "
|
||||
"invalid_target; it may require RFC 8707 resource indicators, which the gateway "
|
||||
"does not send yet (tracked as LIT-4339)",
|
||||
log_context,
|
||||
)
|
||||
return GatewayRejected(code=code)
|
||||
if credential_source == "gateway_stored" and code in GATEWAY_CREDENTIAL_CODES:
|
||||
verbose_logger.warning(
|
||||
"MCP server %s: upstream authorization server rejected the gateway's configured client "
|
||||
"credentials (%s): %s",
|
||||
log_context,
|
||||
code,
|
||||
description or "<no description>",
|
||||
)
|
||||
return GatewayRejected(code=code)
|
||||
return CallerRejected(code=code, description=description, error_uri=error_uri)
|
||||
|
||||
|
||||
def classify_upstream_token_rejection(
|
||||
response: httpx.Response,
|
||||
credential_source: CredentialSource,
|
||||
log_context: str,
|
||||
) -> UpstreamOAuthFault:
|
||||
"""Classify a token-endpoint rejection into exactly one fault: a body with an RFC 6749 §5.2
|
||||
``error`` field goes through blame assignment (:func:`_classify_oauth_error_code`); anything
|
||||
without a usable ``error`` field is an upstream protocol fault."""
|
||||
parsed = _safe_json(response)
|
||||
fields = parsed if isinstance(parsed, dict) else {}
|
||||
code = _bounded_field(fields.get("error"))
|
||||
if code is None:
|
||||
_log_out_of_contract("token", response, log_context)
|
||||
return UpstreamProtocolFault(note=f"upstream token endpoint returned HTTP {response.status_code}")
|
||||
return _classify_oauth_error_code(
|
||||
code,
|
||||
description=_bounded_field(fields.get("error_description")),
|
||||
error_uri=_bounded_field(fields.get("error_uri")),
|
||||
credential_source=credential_source,
|
||||
log_context=log_context,
|
||||
)
|
||||
|
||||
|
||||
def classify_upstream_dcr_rejection(response: httpx.Response, log_context: str) -> UpstreamOAuthFault:
|
||||
"""Classify a dynamic-client-registration rejection. RFC 7591 §3.2.2 errors carry
|
||||
``error`` / ``error_description`` and go through the same blame assignment as token errors
|
||||
(registration sends no client credentials, so credential codes stay caller-actionable); anything
|
||||
without a usable ``error`` field is an upstream protocol fault."""
|
||||
parsed = _safe_json(response)
|
||||
fields = parsed if isinstance(parsed, dict) else {}
|
||||
code = _bounded_field(fields.get("error"))
|
||||
if code is None:
|
||||
_log_out_of_contract("registration", response, log_context)
|
||||
return UpstreamProtocolFault(note=f"upstream registration failed with HTTP {response.status_code}")
|
||||
return _classify_oauth_error_code(
|
||||
code,
|
||||
description=_bounded_field(fields.get("error_description")),
|
||||
error_uri=None,
|
||||
credential_source="caller_supplied",
|
||||
log_context=log_context,
|
||||
)
|
||||
|
|
@ -0,0 +1,89 @@
|
|||
"""Render upstream OAuth/DCR faults onto the wire. The only place that chooses statuses and bodies
|
||||
for these faults, so every consumer emits the same contract: RFC 6749 §5.2-shaped JSON with the §5.1
|
||||
no-store headers on token endpoints, HTTPException details on registration. Status, code, and prose
|
||||
all derive from the fault tag; exhaustive matches keep a new fault arm from shipping unrendered.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi.responses import JSONResponse
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.faults.types import UpstreamOAuthFault
|
||||
from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS
|
||||
|
||||
|
||||
def _gateway_rejected_description(code: str) -> str:
|
||||
if code == "invalid_target":
|
||||
return (
|
||||
"the upstream authorization server rejected the request (invalid_target); "
|
||||
"it may require RFC 8707 resource indicators, which the gateway does not send yet"
|
||||
)
|
||||
return (
|
||||
f"the upstream authorization server rejected the gateway's configured client credentials "
|
||||
f"({code}); verify the MCP server's client_id and client_secret"
|
||||
)
|
||||
|
||||
|
||||
def _upstream_reported_status_and_description(code: str) -> tuple[int, str]:
|
||||
if code == "temporarily_unavailable":
|
||||
return 503, "the upstream authorization server is temporarily unavailable; retry shortly"
|
||||
return 502, "the upstream authorization server reported an internal error"
|
||||
|
||||
|
||||
def render_token_fault(fault: UpstreamOAuthFault) -> JSONResponse:
|
||||
"""RFC 6749 §5.2 response for a token-endpoint fault. Caller-actionable rejections relay the
|
||||
upstream's code on the status that code implies (401 for invalid_client per §5.2, else 400);
|
||||
gateway-side faults are 502 ``server_error`` with gateway-authored prose so a caller is never
|
||||
blamed for, or shown the internals of, a failure only the operator can fix."""
|
||||
match fault.tag:
|
||||
case "caller_rejected":
|
||||
content = {
|
||||
"error": fault.code,
|
||||
**({"error_description": fault.description} if fault.description else {}),
|
||||
**({"error_uri": fault.error_uri} if fault.error_uri else {}),
|
||||
}
|
||||
status_code = 401 if fault.code == "invalid_client" else 400
|
||||
return JSONResponse(status_code=status_code, content=content, headers=TOKEN_NO_CACHE_HEADERS)
|
||||
case "gateway_rejected":
|
||||
return JSONResponse(
|
||||
status_code=502,
|
||||
content={
|
||||
"error": "server_error",
|
||||
"error_description": _gateway_rejected_description(fault.code),
|
||||
},
|
||||
headers=TOKEN_NO_CACHE_HEADERS,
|
||||
)
|
||||
case "upstream_reported_fault":
|
||||
status_code, description = _upstream_reported_status_and_description(fault.code)
|
||||
return JSONResponse(
|
||||
status_code=status_code,
|
||||
content={"error": fault.code, "error_description": description},
|
||||
headers=TOKEN_NO_CACHE_HEADERS,
|
||||
)
|
||||
case "upstream_protocol_fault":
|
||||
return JSONResponse(
|
||||
status_code=502,
|
||||
content={"error": "server_error", "error_description": fault.note},
|
||||
headers=TOKEN_NO_CACHE_HEADERS,
|
||||
)
|
||||
case _:
|
||||
assert_never(fault.tag)
|
||||
|
||||
|
||||
def dcr_fault_detail(fault: UpstreamOAuthFault) -> tuple[int, str]:
|
||||
"""Status and detail string for a registration fault, raised as HTTPException by the caller.
|
||||
RFC 7591 §3.2.2 defines registration errors as 400, so a contract-conformant rejection is 400
|
||||
regardless of the status the upstream chose; everything else is a 502 upstream fault."""
|
||||
match fault.tag:
|
||||
case "caller_rejected":
|
||||
detail = f"{fault.code}: {fault.description}" if fault.description else fault.code
|
||||
return 400, detail
|
||||
case "gateway_rejected":
|
||||
return 502, _gateway_rejected_description(fault.code)
|
||||
case "upstream_reported_fault":
|
||||
return _upstream_reported_status_and_description(fault.code)
|
||||
case "upstream_protocol_fault":
|
||||
return 502, fault.note
|
||||
case _:
|
||||
assert_never(fault.tag)
|
||||
79
litellm/proxy/_experimental/mcp_server/faults/types.py
Normal file
79
litellm/proxy/_experimental/mcp_server/faults/types.py
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
"""Fault taxonomy for upstream OAuth token and DCR registration failures.
|
||||
|
||||
Each fault is a frozen model on a ``tag`` literal. The tag alone decides the HTTP status, the wire
|
||||
error code, and whose prose the caller sees, so those three facts can never disagree the way they can
|
||||
when an upstream's status and error code are relayed independently.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
MAX_WIRE_FIELD_CHARS = 500
|
||||
"""Bound on every upstream-derived string that crosses to a caller or into a log line."""
|
||||
|
||||
CredentialSource: TypeAlias = Literal["gateway_stored", "caller_supplied"]
|
||||
"""Whose client credentials the gateway presented upstream: the MCP server's stored configuration or
|
||||
credentials the caller supplied on the request. Decides whether a credential rejection is the
|
||||
caller's problem to fix or the gateway operator's."""
|
||||
|
||||
GATEWAY_CREDENTIAL_CODES: frozenset[str] = frozenset({"invalid_client", "unauthorized_client"})
|
||||
"""RFC 6749 error codes that indict the OAuth client's credentials or grant authorization. When the
|
||||
gateway presented its own stored credentials, these are gateway-side faults the caller cannot act on;
|
||||
when the caller supplied the credentials, they are the caller's to fix."""
|
||||
|
||||
GATEWAY_CAPABILITY_CODES: frozenset[str] = frozenset({"invalid_target"})
|
||||
"""Codes that indict a gateway capability regardless of whose credentials were presented:
|
||||
``invalid_target`` means the upstream wants RFC 8707 resource indicators, which the gateway does not
|
||||
send yet (LIT-4339). Never the caller's fault."""
|
||||
|
||||
UPSTREAM_FAULT_CODES: frozenset[str] = frozenset({"server_error", "temporarily_unavailable"})
|
||||
"""Codes by which the upstream blames itself. Relaying them as caller faults would invert blame, so
|
||||
they classify as upstream-reported faults and render on the 5xx their meaning implies."""
|
||||
|
||||
|
||||
class CallerRejected(BaseModel):
|
||||
"""The upstream spoke the OAuth error contract and the failure is actionable by our caller
|
||||
(e.g. ``invalid_grant``: re-run authorization). The code and its bounded prose relay on the
|
||||
4xx status the code itself implies."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["caller_rejected"] = "caller_rejected"
|
||||
code: str
|
||||
description: str | None = None
|
||||
error_uri: str | None = None
|
||||
|
||||
|
||||
class GatewayRejected(BaseModel):
|
||||
"""The upstream rejected the request for a cause only the gateway operator can address: the
|
||||
server's stored client credentials or a gateway capability gap. Not actionable by the caller:
|
||||
rendered as 502 with gateway-authored prose naming the code; the upstream's prose goes to
|
||||
server logs only."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["gateway_rejected"] = "gateway_rejected"
|
||||
code: str
|
||||
|
||||
|
||||
class UpstreamReportedFault(BaseModel):
|
||||
"""The upstream blamed itself in the OAuth vocabulary. Rendered on the 5xx the code implies
|
||||
(``server_error`` 502, ``temporarily_unavailable`` 503) so blame and status agree."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["upstream_reported_fault"] = "upstream_reported_fault"
|
||||
code: Literal["server_error", "temporarily_unavailable"]
|
||||
|
||||
|
||||
class UpstreamProtocolFault(BaseModel):
|
||||
"""The upstream broke the error contract: no JSON ``error`` field, an undecodable body, or a
|
||||
success response without a usable token. Rendered as 502 with a gateway-authored note; the
|
||||
upstream body never crosses to the caller."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["upstream_protocol_fault"] = "upstream_protocol_fault"
|
||||
note: str
|
||||
|
||||
|
||||
UpstreamOAuthFault: TypeAlias = CallerRejected | GatewayRejected | UpstreamReportedFault | UpstreamProtocolFault
|
||||
|
|
@ -21,11 +21,16 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import
|
|||
EnvelopeKeys,
|
||||
EnvelopeMintError,
|
||||
OpenedEnvelope,
|
||||
OpenedRefreshEnvelope,
|
||||
RefreshCredential,
|
||||
SealedEnvelope,
|
||||
UpstreamTokenGrant,
|
||||
is_envelope,
|
||||
is_refresh_envelope,
|
||||
mint_envelope,
|
||||
mint_refresh_envelope,
|
||||
open_envelope,
|
||||
open_refresh_envelope,
|
||||
)
|
||||
|
||||
_SIGNING_KEY_DOMAIN = b"litellm-mcp-bridge:envelope-signing:"
|
||||
|
|
@ -92,6 +97,67 @@ def build_bridge_token_response(
|
|||
return mint_envelope(identity, grant, keys, now)
|
||||
|
||||
|
||||
def build_bridge_refresh_token_response(
|
||||
identity: EnvelopeIdentity,
|
||||
refresh: RefreshCredential,
|
||||
keys: EnvelopeKeys,
|
||||
now: datetime,
|
||||
) -> SealedEnvelope | EnvelopeMintError:
|
||||
"""Seal ``refresh`` for ``identity`` into the long-lived refresh envelope the token endpoint returns
|
||||
alongside the access envelope, so the client can renew without re-authenticating. A thin, pure
|
||||
wrapper over :func:`mint_refresh_envelope`; returns the mint error as a value for the caller to map.
|
||||
"""
|
||||
return mint_refresh_envelope(identity, refresh, keys, now)
|
||||
|
||||
|
||||
class BridgeRefreshOpened(BaseModel):
|
||||
"""A valid refresh envelope presented to the token endpoint: the identity to re-validate and renew
|
||||
under, and the upstream refresh grant to exchange."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["opened"] = "opened"
|
||||
identity: EnvelopeIdentity
|
||||
refresh: RefreshCredential
|
||||
|
||||
|
||||
class BridgeRefreshInvalid(BaseModel):
|
||||
"""The presented refresh grant is not a valid refresh envelope for this server (not refresh-shaped,
|
||||
will not open, or minted for a different server); the token endpoint fails the refresh closed."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["invalid"] = "invalid"
|
||||
|
||||
|
||||
BridgeRefreshResult: TypeAlias = BridgeRefreshOpened | BridgeRefreshInvalid
|
||||
|
||||
|
||||
def open_bridge_refresh_envelope(
|
||||
refresh_value: str,
|
||||
keys: EnvelopeKeys,
|
||||
now: datetime,
|
||||
expected_server_id: str,
|
||||
) -> BridgeRefreshResult:
|
||||
"""Open a refresh envelope a bridge ``oauth_delegate`` client presented on a refresh_token grant.
|
||||
|
||||
The token-endpoint mirror of :func:`resolve_bridge_envelope`: strips an optional ``Bearer`` scheme,
|
||||
then returns ``BridgeRefreshOpened`` with the recovered identity and upstream refresh grant, or
|
||||
``BridgeRefreshInvalid`` for anything that is not a valid refresh envelope for this server. Never
|
||||
raises; total over hostile input via :func:`open_refresh_envelope`. ``expected_server_id`` binds the
|
||||
envelope to the server the request targets, so a refresh envelope minted for one server cannot renew
|
||||
against another. A raw upstream refresh token (not envelope-shaped) is ``BridgeRefreshInvalid``: this
|
||||
mode never hands the client a bare upstream refresh token, so it must never accept one.
|
||||
"""
|
||||
candidate = _strip_bearer(refresh_value)
|
||||
if not is_refresh_envelope(candidate):
|
||||
return BridgeRefreshInvalid()
|
||||
opened = open_refresh_envelope(candidate, keys, now)
|
||||
if not isinstance(opened, OpenedRefreshEnvelope):
|
||||
return BridgeRefreshInvalid()
|
||||
if opened.identity.server_id != expected_server_id:
|
||||
return BridgeRefreshInvalid()
|
||||
return BridgeRefreshOpened(identity=opened.identity, refresh=opened.refresh)
|
||||
|
||||
|
||||
class NotBridgeEnvelope(BaseModel):
|
||||
"""The bearer is not an envelope; admission continues on its normal path."""
|
||||
|
||||
|
|
@ -128,10 +194,12 @@ def _strip_bearer(value: str) -> str:
|
|||
|
||||
|
||||
def is_bridge_envelope_shaped(authorization_value: str) -> bool:
|
||||
"""Cheap, keyless test that an ``Authorization`` value carries an envelope (optional
|
||||
``Bearer`` scheme stripped). The admission edge engages the bridge arm only for an
|
||||
envelope, so a plain upstream bearer falls through to normal oauth2 admission."""
|
||||
return is_envelope(_strip_bearer(authorization_value))
|
||||
"""Cheap, keyless test that an ``Authorization`` value carries an envelope of either kind (optional
|
||||
``Bearer`` scheme stripped). The admission edge engages the bridge arm for an access envelope (to
|
||||
admit) and for a refresh envelope (to reject it explicitly, since a refresh credential is never
|
||||
usable at the tool-call edge); a plain upstream bearer falls through to normal oauth2 admission."""
|
||||
candidate = _strip_bearer(authorization_value)
|
||||
return is_envelope(candidate) or is_refresh_envelope(candidate)
|
||||
|
||||
|
||||
def resolve_bridge_envelope(
|
||||
|
|
@ -148,6 +216,10 @@ def resolve_bridge_envelope(
|
|||
envelope, and ``BridgeEnvelopeInvalid`` for an envelope-shaped bearer that will not
|
||||
open. Never raises: it is total over hostile input via :func:`open_envelope`.
|
||||
|
||||
A refresh envelope is ``BridgeEnvelopeInvalid`` here: it is a valid gateway credential but only ever
|
||||
presented back to the token endpoint, never usable to authenticate a tool call, so admission must
|
||||
fail it closed rather than let it fall through to another arm.
|
||||
|
||||
``expected_server_id`` is the ``server_id`` of the MCP server the request targets; an
|
||||
opened envelope whose sealed ``server_id`` does not match is rejected as
|
||||
``BridgeEnvelopeInvalid``. Binding here (rather than leaving it to the caller) prevents
|
||||
|
|
@ -157,6 +229,8 @@ def resolve_bridge_envelope(
|
|||
unlike ``hmac.compare_digest`` on ``str``, does not raise on a non-ASCII server_id.
|
||||
"""
|
||||
candidate = _strip_bearer(authorization_value)
|
||||
if is_refresh_envelope(candidate):
|
||||
return BridgeEnvelopeInvalid()
|
||||
if not is_envelope(candidate):
|
||||
return NotBridgeEnvelope()
|
||||
opened = open_envelope(candidate, keys, now)
|
||||
|
|
|
|||
|
|
@ -44,18 +44,33 @@ from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
|
|||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value
|
||||
|
||||
ENVELOPE_PREFIX = "llm_env_"
|
||||
"""Marker prefix on every serialized envelope so the edge can cheaply tell an envelope
|
||||
"""Marker prefix on every serialized ACCESS envelope so the edge can cheaply tell an envelope
|
||||
from a raw upstream token before doing any cryptography."""
|
||||
|
||||
REFRESH_ENVELOPE_PREFIX = "llm_refresh_"
|
||||
"""Marker prefix on every serialized REFRESH envelope. A distinct prefix keeps the two credentials
|
||||
routable without crypto and, together with the signed ``kind`` claim, stops one from being presented
|
||||
where the other is expected: a refresh envelope carries a long-lived upstream refresh token and is only
|
||||
ever presented back to the token endpoint, never forwarded upstream on a tool call."""
|
||||
|
||||
ENVELOPE_ISSUER = "litellm-mcp-bridge"
|
||||
"""``iss`` claim stamped into every envelope and required back on open."""
|
||||
|
||||
MAX_ENVELOPE_TTL_SECONDS = 3600
|
||||
"""Hard ceiling on envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)``
|
||||
"""Hard ceiling on ACCESS envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)``
|
||||
(the cap alone when the upstream omits ``expires_in``), matching the 1h lifetime of the
|
||||
BYOK session bearer this module's signing approach is borrowed from: a client-held
|
||||
credential should never outlive a bounded window even when the upstream token does."""
|
||||
|
||||
MAX_REFRESH_ENVELOPE_TTL_SECONDS = 1209600
|
||||
"""Hard ceiling on REFRESH envelope lifetime (14 days). A refresh envelope only renews the short-lived
|
||||
access envelope, and each renewal re-validates the sealed litellm key (revocation gates it) and is
|
||||
re-minted with a fresh window, so the practical bound is idle time, not a fixed session. ``exp`` is
|
||||
``min(upstream refresh_expires_in, this cap)`` (the cap alone when the upstream omits it); if the
|
||||
upstream refresh token dies first, the next renewal simply fails at the upstream and the client
|
||||
re-authenticates. The value is deliberately far shorter than a typical upstream refresh-token lifetime
|
||||
so a leaked refresh envelope is bounded even if the upstream would have honoured it for longer."""
|
||||
|
||||
MAX_ENVELOPE_BYTES = 12288
|
||||
"""Size cap on the final serialized envelope (prefix + JWT, in bytes). Upstream JWTs
|
||||
commonly run 2-4KB; base64 plus encryption overhead roughly doubles that inside the
|
||||
|
|
@ -66,21 +81,48 @@ typed error, never truncated."""
|
|||
|
||||
_ENVELOPE_JWT_ALGORITHM = "HS256"
|
||||
|
||||
EnvelopeKind = Literal["access", "refresh"]
|
||||
"""Which credential an envelope is. Stamped into the signed claims and required to match on open, so a
|
||||
signature-valid envelope of one kind cannot be replayed as the other even if its wire prefix is swapped
|
||||
(the prefix is not part of the signed payload; this claim is)."""
|
||||
|
||||
|
||||
EnvelopeSubjectType: TypeAlias = Literal["key_hash", "user_id"]
|
||||
"""Discriminator for what litellm principal the envelope binds the grant to.
|
||||
|
||||
``key_hash`` is a hashed virtual key (the scripted two-header client mints under the key it
|
||||
presents at the token endpoint); ``user_id`` is a litellm user subject (the interactive DCR
|
||||
client mints under the SSO-authenticated user, which is the only identity that browser login
|
||||
yields). Admission reloads a key record for the first and a user record for the second, then
|
||||
runs both through the same live-policy gate, so team/org/budget/revocation enforcement is
|
||||
identical either way."""
|
||||
|
||||
|
||||
class EnvelopeIdentity(BaseModel):
|
||||
"""The litellm identity the envelope binds the inner grant to.
|
||||
"""The litellm principal the envelope binds the inner grant to.
|
||||
|
||||
``key_hash`` is the hashed litellm key that authorized the mint, never a raw
|
||||
credential (and the edge rejects a bare hash presented as a bearer). Admission
|
||||
reloads the live key record by it, so the key's current team/org/object-permission
|
||||
restrictions and its revocation state are enforced at use time rather than frozen at
|
||||
mint time. ``server_id`` binds the envelope to one MCP server so it cannot be replayed
|
||||
across a server boundary.
|
||||
``subject`` is the principal identifier and ``subject_type`` says how to resolve it: a
|
||||
hashed litellm key (``key_hash``) or a litellm user id (``user_id``), never a raw
|
||||
credential (and the edge rejects a bare hash or id presented as a bearer). Admission
|
||||
reloads the live record by it, so the principal's current team/org restrictions and its
|
||||
revocation state are enforced at use time rather than frozen at mint time. ``server_id``
|
||||
binds the envelope to one MCP server so it cannot be replayed across a server boundary.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
server_id: str = Field(min_length=1)
|
||||
key_hash: str = Field(min_length=1)
|
||||
subject_type: EnvelopeSubjectType
|
||||
subject: str = Field(min_length=1)
|
||||
|
||||
|
||||
def key_hash_identity(server_id: str, key_hash: str) -> EnvelopeIdentity:
|
||||
"""The identity for the scripted client that mints under a presented virtual key."""
|
||||
return EnvelopeIdentity(server_id=server_id, subject_type="key_hash", subject=key_hash)
|
||||
|
||||
|
||||
def user_identity(server_id: str, user_id: str) -> EnvelopeIdentity:
|
||||
"""The identity for the interactive DCR client that mints under its SSO user subject."""
|
||||
return EnvelopeIdentity(server_id=server_id, subject_type="user_id", subject=user_id)
|
||||
|
||||
|
||||
class UpstreamTokenGrant(BaseModel):
|
||||
|
|
@ -99,6 +141,21 @@ class UpstreamTokenGrant(BaseModel):
|
|||
expires_in: int | None = Field(default=None, gt=0)
|
||||
|
||||
|
||||
class RefreshCredential(BaseModel):
|
||||
"""The upstream refresh grant sealed inside a refresh envelope.
|
||||
|
||||
Only the refresh token (plus the scope to re-request and the refresh token's own lifetime, when the
|
||||
upstream reports it) is sealed; the access token is never in a refresh envelope. ``refresh_token`` is
|
||||
a ``SecretStr`` so reprs never leak it, and ``expires_in`` (the refresh token's lifetime, not the
|
||||
access token's) must be positive when present.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
refresh_token: SecretStr = Field(min_length=1)
|
||||
scope: str | None = None
|
||||
expires_in: int | None = Field(default=None, gt=0)
|
||||
|
||||
|
||||
class EnvelopeKeys(BaseModel):
|
||||
"""Injected key material: the HS256 signing key and the symmetric encryption key.
|
||||
|
||||
|
|
@ -121,13 +178,21 @@ class SealedEnvelope(BaseModel):
|
|||
|
||||
|
||||
class OpenedEnvelope(BaseModel):
|
||||
"""A validated envelope: the identity it was minted for and the recovered grant."""
|
||||
"""A validated access envelope: the identity it was minted for and the recovered grant."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
identity: EnvelopeIdentity
|
||||
grant: UpstreamTokenGrant
|
||||
|
||||
|
||||
class OpenedRefreshEnvelope(BaseModel):
|
||||
"""A validated refresh envelope: the identity it was minted for and the recovered refresh grant."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
identity: EnvelopeIdentity
|
||||
refresh: RefreshCredential
|
||||
|
||||
|
||||
class EnvelopeTooLarge(BaseModel):
|
||||
"""The serialized envelope exceeded ``MAX_ENVELOPE_BYTES``; carries sizes only."""
|
||||
|
||||
|
|
@ -199,8 +264,10 @@ class _EnvelopeClaims(BaseModel):
|
|||
iss: str
|
||||
iat: int
|
||||
exp: int
|
||||
kind: EnvelopeKind
|
||||
server_id: str = Field(min_length=1)
|
||||
key_hash: str = Field(min_length=1)
|
||||
subject_type: EnvelopeSubjectType
|
||||
subject: str = Field(min_length=1)
|
||||
grant: str = Field(min_length=1)
|
||||
|
||||
|
||||
|
|
@ -213,11 +280,25 @@ class _GrantWire(BaseModel):
|
|||
expires_in: int | None = None
|
||||
|
||||
|
||||
class _RefreshWire(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
refresh_token: str
|
||||
scope: str | None = None
|
||||
expires_in: int | None = None
|
||||
|
||||
|
||||
def is_envelope(candidate: str) -> bool:
|
||||
"""Cheap prefix check so the edge can route envelopes vs raw tokens without crypto."""
|
||||
"""Cheap prefix check for an ACCESS envelope so the edge can route envelopes vs raw tokens without
|
||||
crypto. A refresh envelope has a different prefix and is not an access envelope."""
|
||||
return candidate.startswith(ENVELOPE_PREFIX)
|
||||
|
||||
|
||||
def is_refresh_envelope(candidate: str) -> bool:
|
||||
"""Cheap prefix check for a REFRESH envelope so the token endpoint can route a refresh grant that
|
||||
carries an envelope vs a raw upstream refresh token without crypto."""
|
||||
return candidate.startswith(REFRESH_ENVELOPE_PREFIX)
|
||||
|
||||
|
||||
def mint_envelope(
|
||||
identity: EnvelopeIdentity,
|
||||
grant: UpstreamTokenGrant,
|
||||
|
|
@ -231,23 +312,15 @@ def mint_envelope(
|
|||
serialized envelope exceeds ``MAX_ENVELOPE_BYTES``.
|
||||
"""
|
||||
expires_at = now + timedelta(seconds=_envelope_ttl_seconds(grant.expires_in))
|
||||
claims = _EnvelopeClaims(
|
||||
iss=ENVELOPE_ISSUER,
|
||||
iat=int(now.timestamp()),
|
||||
exp=int(expires_at.timestamp()),
|
||||
server_id=identity.server_id,
|
||||
key_hash=identity.key_hash,
|
||||
grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key),
|
||||
return _seal(
|
||||
kind="access",
|
||||
prefix=ENVELOPE_PREFIX,
|
||||
identity=identity,
|
||||
grant_blob=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key),
|
||||
expires_at=expires_at,
|
||||
signing_key=keys.signing_key,
|
||||
now=now,
|
||||
)
|
||||
token = ENVELOPE_PREFIX + jwt.encode(
|
||||
claims.model_dump(),
|
||||
keys.signing_key.get_secret_value(),
|
||||
algorithm=_ENVELOPE_JWT_ALGORITHM,
|
||||
)
|
||||
size_bytes = len(token.encode("utf-8"))
|
||||
if size_bytes > MAX_ENVELOPE_BYTES:
|
||||
return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES)
|
||||
return SealedEnvelope(token=SecretStr(token), expires_at=expires_at)
|
||||
|
||||
|
||||
def open_envelope(
|
||||
|
|
@ -263,35 +336,136 @@ def open_envelope(
|
|||
re-derived, so it is stale by up to the envelope's lifetime; callers that need a
|
||||
live remaining lifetime should use ``now`` against the upstream, not this field.
|
||||
"""
|
||||
if not is_envelope(candidate):
|
||||
return NotAnEnvelope()
|
||||
# UTF-8 byte length is never below character length, so a character count already over the
|
||||
# cap rejects an oversize candidate in O(1) without encoding it; the exact byte check then
|
||||
# runs only on candidates already bounded to <= MAX_ENVELOPE_BYTES characters.
|
||||
if len(candidate) > MAX_ENVELOPE_BYTES:
|
||||
return MalformedPayload()
|
||||
if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES:
|
||||
return MalformedPayload()
|
||||
claims = _decode_claims(candidate.removeprefix(ENVELOPE_PREFIX), keys.signing_key)
|
||||
claims = _open_claims(candidate, prefix=ENVELOPE_PREFIX, expected_kind="access", keys=keys, now=now)
|
||||
if not isinstance(claims, _EnvelopeClaims):
|
||||
return claims
|
||||
if now.timestamp() >= claims.exp:
|
||||
return Expired()
|
||||
grant = _decrypt_grant(claims.grant, keys.encryption_key)
|
||||
if not isinstance(grant, UpstreamTokenGrant):
|
||||
return grant
|
||||
return OpenedEnvelope(
|
||||
identity=EnvelopeIdentity(server_id=claims.server_id, key_hash=claims.key_hash),
|
||||
identity=EnvelopeIdentity(server_id=claims.server_id, subject_type=claims.subject_type, subject=claims.subject),
|
||||
grant=grant,
|
||||
)
|
||||
|
||||
|
||||
def mint_refresh_envelope(
|
||||
identity: EnvelopeIdentity,
|
||||
refresh: RefreshCredential,
|
||||
keys: EnvelopeKeys,
|
||||
now: datetime,
|
||||
) -> SealedEnvelope | EnvelopeMintError:
|
||||
"""Seal ``refresh`` for ``identity`` into a long-lived, client-held refresh envelope.
|
||||
|
||||
``exp`` is ``min(refresh.expires_in, MAX_REFRESH_ENVELOPE_TTL_SECONDS)`` seconds from ``now`` (the
|
||||
cap alone when the upstream omits the refresh lifetime). Sealing a distinct ``kind="refresh"`` claim
|
||||
is what keeps a refresh envelope from ever opening as an access credential at the MCP edge. Returns
|
||||
``EnvelopeTooLarge`` when the serialized envelope exceeds ``MAX_ENVELOPE_BYTES``.
|
||||
"""
|
||||
expires_at = now + timedelta(seconds=_refresh_ttl_seconds(refresh.expires_in))
|
||||
return _seal(
|
||||
kind="refresh",
|
||||
prefix=REFRESH_ENVELOPE_PREFIX,
|
||||
identity=identity,
|
||||
grant_blob=_encrypt_grant_blob(_refresh_plaintext(refresh), keys.encryption_key),
|
||||
expires_at=expires_at,
|
||||
signing_key=keys.signing_key,
|
||||
now=now,
|
||||
)
|
||||
|
||||
|
||||
def open_refresh_envelope(
|
||||
candidate: str,
|
||||
keys: EnvelopeKeys,
|
||||
now: datetime,
|
||||
) -> OpenedRefreshEnvelope | EnvelopeOpenError:
|
||||
"""Validate a refresh ``candidate`` and recover the identity and inner refresh grant.
|
||||
|
||||
Total over hostile input exactly like :func:`open_envelope`: every invalid, expired, tampered,
|
||||
wrong-kind, or undecryptable candidate maps to a distinct ``EnvelopeOpenError`` variant, never a
|
||||
raise. The ``kind="refresh"`` claim is required, so an access envelope re-prefixed as a refresh one
|
||||
is rejected as ``MalformedPayload``.
|
||||
"""
|
||||
claims = _open_claims(candidate, prefix=REFRESH_ENVELOPE_PREFIX, expected_kind="refresh", keys=keys, now=now)
|
||||
if not isinstance(claims, _EnvelopeClaims):
|
||||
return claims
|
||||
refresh = _decrypt_refresh(claims.grant, keys.encryption_key)
|
||||
if not isinstance(refresh, RefreshCredential):
|
||||
return refresh
|
||||
return OpenedRefreshEnvelope(
|
||||
identity=EnvelopeIdentity(server_id=claims.server_id, subject_type=claims.subject_type, subject=claims.subject),
|
||||
refresh=refresh,
|
||||
)
|
||||
|
||||
|
||||
def _seal(
|
||||
kind: EnvelopeKind,
|
||||
prefix: str,
|
||||
identity: EnvelopeIdentity,
|
||||
grant_blob: str,
|
||||
expires_at: datetime,
|
||||
signing_key: SecretStr,
|
||||
now: datetime,
|
||||
) -> SealedEnvelope | EnvelopeTooLarge:
|
||||
"""Sign the claims for either envelope kind and enforce the size cap. Shared by both mints so the
|
||||
JWT shape, issuer, and size guard cannot drift between access and refresh envelopes."""
|
||||
claims = _EnvelopeClaims(
|
||||
iss=ENVELOPE_ISSUER,
|
||||
iat=int(now.timestamp()),
|
||||
exp=int(expires_at.timestamp()),
|
||||
kind=kind,
|
||||
server_id=identity.server_id,
|
||||
subject_type=identity.subject_type,
|
||||
subject=identity.subject,
|
||||
grant=grant_blob,
|
||||
)
|
||||
token = prefix + jwt.encode(claims.model_dump(), signing_key.get_secret_value(), algorithm=_ENVELOPE_JWT_ALGORITHM)
|
||||
size_bytes = len(token.encode("utf-8"))
|
||||
if size_bytes > MAX_ENVELOPE_BYTES:
|
||||
return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES)
|
||||
return SealedEnvelope(token=SecretStr(token), expires_at=expires_at)
|
||||
|
||||
|
||||
def _open_claims(
|
||||
candidate: str,
|
||||
prefix: str,
|
||||
expected_kind: EnvelopeKind,
|
||||
keys: EnvelopeKeys,
|
||||
now: datetime,
|
||||
) -> _EnvelopeClaims | EnvelopeOpenError:
|
||||
"""Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an attacker-controlled
|
||||
candidate, shared by both openers so the security gate is identical for access and refresh. Returns
|
||||
the validated claims or a distinct ``EnvelopeOpenError``; never raises."""
|
||||
if not candidate.startswith(prefix):
|
||||
return NotAnEnvelope()
|
||||
# UTF-8 byte length is never below character length, so a character count already over the cap
|
||||
# rejects an oversize candidate in O(1) without encoding it; the exact byte check then runs only on
|
||||
# candidates already bounded to <= MAX_ENVELOPE_BYTES characters.
|
||||
if len(candidate) > MAX_ENVELOPE_BYTES:
|
||||
return MalformedPayload()
|
||||
if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES:
|
||||
return MalformedPayload()
|
||||
claims = _decode_claims(candidate.removeprefix(prefix), keys.signing_key)
|
||||
if not isinstance(claims, _EnvelopeClaims):
|
||||
return claims
|
||||
if claims.kind != expected_kind:
|
||||
return MalformedPayload()
|
||||
if now.timestamp() >= claims.exp:
|
||||
return Expired()
|
||||
return claims
|
||||
|
||||
|
||||
def _envelope_ttl_seconds(upstream_expires_in: int | None) -> int:
|
||||
if upstream_expires_in is None:
|
||||
return MAX_ENVELOPE_TTL_SECONDS
|
||||
return min(upstream_expires_in, MAX_ENVELOPE_TTL_SECONDS)
|
||||
|
||||
|
||||
def _refresh_ttl_seconds(upstream_refresh_expires_in: int | None) -> int:
|
||||
if upstream_refresh_expires_in is None:
|
||||
return MAX_REFRESH_ENVELOPE_TTL_SECONDS
|
||||
return min(upstream_refresh_expires_in, MAX_REFRESH_ENVELOPE_TTL_SECONDS)
|
||||
|
||||
|
||||
def _grant_plaintext(grant: UpstreamTokenGrant) -> str:
|
||||
wire = _GrantWire(
|
||||
access_token=grant.access_token.get_secret_value(),
|
||||
|
|
@ -303,6 +477,15 @@ def _grant_plaintext(grant: UpstreamTokenGrant) -> str:
|
|||
return wire.model_dump_json(exclude_none=True)
|
||||
|
||||
|
||||
def _refresh_plaintext(refresh: RefreshCredential) -> str:
|
||||
wire = _RefreshWire(
|
||||
refresh_token=refresh.refresh_token.get_secret_value(),
|
||||
scope=refresh.scope,
|
||||
expires_in=refresh.expires_in,
|
||||
)
|
||||
return wire.model_dump_json(exclude_none=True)
|
||||
|
||||
|
||||
def _decode_claims(
|
||||
compact: str,
|
||||
signing_key: SecretStr,
|
||||
|
|
@ -364,3 +547,22 @@ def _decrypt_grant(
|
|||
return UpstreamTokenGrant.model_validate_json(plaintext)
|
||||
except ValidationError:
|
||||
return MalformedPayload()
|
||||
|
||||
|
||||
def _decrypt_refresh(
|
||||
blob: str,
|
||||
encryption_key: SecretStr,
|
||||
) -> RefreshCredential | DecryptFailed | MalformedPayload:
|
||||
from nacl.exceptions import CryptoError
|
||||
|
||||
try:
|
||||
plaintext = decrypt_value(
|
||||
value=base64.urlsafe_b64decode(blob),
|
||||
signing_key=encryption_key.get_secret_value(),
|
||||
)
|
||||
except (CryptoError, ValueError):
|
||||
return DecryptFailed()
|
||||
try:
|
||||
return RefreshCredential.model_validate_json(plaintext)
|
||||
except ValidationError:
|
||||
return MalformedPayload()
|
||||
|
|
|
|||
|
|
@ -3719,8 +3719,15 @@ if MCP_AVAILABLE:
|
|||
headers={"www-authenticate": upstream_www_authenticate},
|
||||
)
|
||||
|
||||
def _get_authorization_header_from_scope(scope: Scope) -> Optional[str]:
|
||||
"""First ``Authorization`` header value in the ASGI scope, or None."""
|
||||
for key, value in scope.get("headers", []):
|
||||
if key.lower() == b"authorization":
|
||||
return value.decode("latin-1")
|
||||
return None
|
||||
|
||||
def _scope_has_authorization_header(scope: Scope) -> bool:
|
||||
return any(key.lower() == b"authorization" for key, _ in scope.get("headers", []))
|
||||
return _get_authorization_header_from_scope(scope) is not None
|
||||
|
||||
def _get_forwarded_auth_from_scope(scope: Scope) -> Optional[str]:
|
||||
"""Return the upstream-bound ``Authorization`` header value, or None.
|
||||
|
|
@ -3733,17 +3740,24 @@ if MCP_AVAILABLE:
|
|||
``MCPRequestHandler.process_mcp_request``), and forwarding it upstream
|
||||
would leak the proxy key to a third-party MCP server.
|
||||
"""
|
||||
authorization = None
|
||||
has_litellm_key_header = False
|
||||
for key, value in scope.get("headers", []):
|
||||
key_lower = key.lower()
|
||||
if key_lower == b"authorization":
|
||||
authorization = value.decode("latin-1")
|
||||
elif key_lower == b"x-litellm-api-key":
|
||||
has_litellm_key_header = True
|
||||
has_litellm_key_header = any(key.lower() == b"x-litellm-api-key" for key, _ in scope.get("headers", []))
|
||||
if not has_litellm_key_header:
|
||||
return None
|
||||
return authorization
|
||||
return _get_authorization_header_from_scope(scope)
|
||||
|
||||
def _is_delegate_upstream_probe_target(server: MCPServer) -> bool:
|
||||
"""Whether ``server`` is an interactive delegate-auth server whose client-supplied
|
||||
token should be preflighted upstream.
|
||||
|
||||
Mirrors the anonymous-delegate gate in ``get_allowed_mcp_servers``: the flow is
|
||||
resolved via ``effective_oauth2_flow`` so an unstamped M2M-shape row fails closed
|
||||
(its stored client credentials drive egress; the caller's bearer is irrelevant).
|
||||
"""
|
||||
return (
|
||||
server.auth_type == MCPAuth.oauth2
|
||||
and server.delegate_auth_to_upstream is True
|
||||
and MCPServerManager.effective_oauth2_flow(server) != "client_credentials"
|
||||
)
|
||||
|
||||
async def _probe_upstream_auth(
|
||||
url: str,
|
||||
|
|
@ -3805,7 +3819,7 @@ if MCP_AVAILABLE:
|
|||
mcp_servers: Optional[List[str]],
|
||||
client_ip: Optional[str],
|
||||
) -> None:
|
||||
"""Probe pass-through upstream servers in parallel before the MCP session starts.
|
||||
"""Probe pass-through and delegate-auth upstream servers in parallel before the MCP session starts.
|
||||
|
||||
Only servers the caller's key is already authorized to reach are probed —
|
||||
the list is derived from _get_allowed_mcp_servers so that a user cannot
|
||||
|
|
@ -3813,11 +3827,42 @@ if MCP_AVAILABLE:
|
|||
|
||||
The MCP SDK commits HTTP 200 headers before invoking handlers, so a 401
|
||||
can only be returned before that point. This function raises HTTPException(401)
|
||||
with a WWW-Authenticate header if any upstream rejects the client token.
|
||||
with a WWW-Authenticate header if any upstream rejects the client token, or 403
|
||||
if the upstream accepts it but forbids the caller.
|
||||
Fails-open: network errors are logged and the request is allowed through.
|
||||
|
||||
Delegate-auth servers (``auth_type=oauth2`` + ``delegate_auth_to_upstream``)
|
||||
are probed with the caller's bare ``Authorization`` bearer. That bearer is only
|
||||
an upstream token (never a LiteLLM key) when admission took the delegate bypass,
|
||||
so the delegate target is resolved through ``get_mcp_server_by_name`` -- the same
|
||||
resolver admission used -- rather than the wider allowed-server prefix/access-group
|
||||
matching. A name that only reaches a delegate server via server_id or an access
|
||||
group would have been admitted as a real LiteLLM key, so probing it would leak that
|
||||
key upstream; requiring the admission-resolver match closes that gap. Without the
|
||||
probe a rejected token is absorbed by the tools/list handler and masked as an empty
|
||||
tool list. Gated to single-server routes so one rejected token cannot 401 a
|
||||
multi-server aggregate connect, matching the OBO preflight gating; the challenge
|
||||
echoes the requested name so aliased routes get the same resource_metadata URL as
|
||||
the tokenless preemptive challenge.
|
||||
"""
|
||||
forwarded_auth = _get_forwarded_auth_from_scope(scope)
|
||||
if not forwarded_auth:
|
||||
requested_single_target = mcp_servers[0] if mcp_servers is not None and len(mcp_servers) == 1 else None
|
||||
# The bare Authorization header (no x-litellm-api-key) is a valid upstream token
|
||||
# only when admission classified it as one, i.e. the single requested name resolves
|
||||
# to a delegate server under admission's own resolver. Resolve it the same way here
|
||||
# so a server_id- or access-group-named delegate (which admission would have treated
|
||||
# as a LiteLLM key) is never probed with that key.
|
||||
delegate_server = (
|
||||
global_mcp_server_manager.get_mcp_server_by_name(requested_single_target, client_ip=client_ip)
|
||||
if requested_single_target
|
||||
else None
|
||||
)
|
||||
delegate_auth = (
|
||||
_get_authorization_header_from_scope(scope)
|
||||
if delegate_server is not None and _is_delegate_upstream_probe_target(delegate_server)
|
||||
else None
|
||||
)
|
||||
if not forwarded_auth and not delegate_auth:
|
||||
return
|
||||
|
||||
# Use the authorized server set, not the raw user-supplied names, so that
|
||||
|
|
@ -3827,33 +3872,49 @@ if MCP_AVAILABLE:
|
|||
mcp_servers=mcp_servers,
|
||||
client_ip=client_ip,
|
||||
)
|
||||
passthrough_servers = [
|
||||
srv
|
||||
for srv in allowed_servers
|
||||
# Restrict to genuine OAuth pass-through servers (auth_type none +
|
||||
# Authorization in extra_headers). Gateway-managed OAuth2 servers
|
||||
# must not receive the ``resource_metadata=`` challenge emitted
|
||||
# below — they require ``authorization_uri=`` pointing at the
|
||||
# gateway AS metadata. ``is_oauth_passthrough`` already requires
|
||||
# ``auth_type in (None, MCPAuth.none)``, which is mutually
|
||||
# exclusive with ``has_client_credentials`` (oauth2 + M2M flow),
|
||||
# so M2M servers are implicitly excluded here.
|
||||
if srv.is_oauth_passthrough
|
||||
]
|
||||
if not passthrough_servers:
|
||||
passthrough_targets: Tuple[Tuple[MCPServer, str, str], ...] = (
|
||||
tuple(
|
||||
(srv, forwarded_auth, srv.name)
|
||||
for srv in allowed_servers
|
||||
# Restrict to genuine OAuth pass-through servers (auth_type none +
|
||||
# Authorization in extra_headers). Gateway-managed OAuth2 servers
|
||||
# must not receive the ``resource_metadata=`` challenge emitted
|
||||
# below — they require ``authorization_uri=`` pointing at the
|
||||
# gateway AS metadata. ``is_oauth_passthrough`` already requires
|
||||
# ``auth_type in (None, MCPAuth.none)``, which is mutually
|
||||
# exclusive with ``has_client_credentials`` (oauth2 + M2M flow),
|
||||
# so M2M servers are implicitly excluded here.
|
||||
if srv.is_oauth_passthrough
|
||||
)
|
||||
if forwarded_auth
|
||||
else ()
|
||||
)
|
||||
# Probe the admission-resolved delegate server only when the caller is actually
|
||||
# authorized for it (present in the IP-filtered allowed set), keyed by server_id.
|
||||
delegate_targets: Tuple[Tuple[MCPServer, str, str], ...] = (
|
||||
tuple(
|
||||
(srv, delegate_auth, requested_single_target)
|
||||
for srv in allowed_servers
|
||||
if delegate_server is not None and srv.server_id == delegate_server.server_id
|
||||
)
|
||||
if delegate_auth and requested_single_target
|
||||
else ()
|
||||
)
|
||||
probe_targets = passthrough_targets + delegate_targets
|
||||
if not probe_targets:
|
||||
return
|
||||
|
||||
probe_results = await asyncio.gather(
|
||||
*[_probe_upstream_auth(srv.url or "", forwarded_auth) for srv in passthrough_servers]
|
||||
*[_probe_upstream_auth(srv.url or "", auth_header) for srv, auth_header, _ in probe_targets]
|
||||
)
|
||||
for srv, (probe_status, _) in zip(passthrough_servers, probe_results):
|
||||
for (srv, _, challenge_server_name), (probe_status, _) in zip(probe_targets, probe_results):
|
||||
if probe_status == 401:
|
||||
# Token is missing or expired: keep pass-through clients on the
|
||||
# protected-resource discovery flow so they re-authorize against
|
||||
# the upstream IdP metadata proxied by LiteLLM.
|
||||
www_authenticate = _get_passthrough_www_authenticate(
|
||||
scope=scope,
|
||||
server_name=srv.name,
|
||||
server_name=challenge_server_name,
|
||||
invalid_token=True,
|
||||
)
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -1118,6 +1118,7 @@ class GenerateKeyRequest(KeyRequestBase):
|
|||
class GenerateKeyResponse(KeyRequestBase):
|
||||
key: str # type: ignore
|
||||
key_name: Optional[str] = None
|
||||
key_type: str | None = None
|
||||
expires: Optional[datetime] = None
|
||||
user_id: Optional[str] = None
|
||||
token_id: Optional[str] = None
|
||||
|
|
@ -2421,6 +2422,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
"is active as a reminder that hard enforcement is relaxed."
|
||||
),
|
||||
)
|
||||
skip_user_budget_on_team_key: bool | None = Field(
|
||||
None,
|
||||
description=(
|
||||
"If True, restores the legacy behavior where a user's personal "
|
||||
"max_budget is NOT enforced when their key belongs to a team; only "
|
||||
"the team (and team-member) budgets apply. Defaults to False, meaning "
|
||||
"the user's personal max_budget is always enforced regardless of "
|
||||
"whether the key belongs to a team (see GitHub issue #12905)."
|
||||
),
|
||||
)
|
||||
user_url_validation: Optional[bool] = Field(
|
||||
None,
|
||||
description=(
|
||||
|
|
|
|||
|
|
@ -626,26 +626,29 @@ async def common_checks(
|
|||
)
|
||||
|
||||
async def _user_max_budget_check() -> None:
|
||||
# 4.1 personal budget, if personal key
|
||||
if (
|
||||
(team_object is None or team_object.team_id is None)
|
||||
and user_object is not None
|
||||
and user_object.max_budget is not None
|
||||
):
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
if user_object is None or user_object.max_budget is None:
|
||||
return
|
||||
skip_for_team = (
|
||||
general_settings.get("skip_user_budget_on_team_key") is True
|
||||
and team_object is not None
|
||||
and team_object.team_id is not None
|
||||
)
|
||||
if skip_for_team:
|
||||
return
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
user_budget = user_object.max_budget
|
||||
user_spend = await get_current_spend(
|
||||
counter_key=f"spend:user:{user_object.user_id}",
|
||||
fallback_spend=user_object.spend or 0.0,
|
||||
user_budget = user_object.max_budget
|
||||
user_spend = await get_current_spend(
|
||||
counter_key=f"spend:user:{user_object.user_id}",
|
||||
fallback_spend=user_object.spend or 0.0,
|
||||
max_budget=user_budget,
|
||||
)
|
||||
if math.isfinite(user_budget) and user_spend >= user_budget:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=user_spend,
|
||||
max_budget=user_budget,
|
||||
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
|
||||
)
|
||||
if math.isfinite(user_budget) and user_spend >= user_budget:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=user_spend,
|
||||
max_budget=user_budget,
|
||||
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
|
||||
)
|
||||
|
||||
# Each scope reads a distinct counter key with no cross-scope ordering
|
||||
# dependency, so the per-scope Redis-first reads run concurrently instead
|
||||
|
|
@ -4383,14 +4386,23 @@ def _model_custom_llm_provider_matches_wildcard_pattern(model: str, allowed_mode
|
|||
or
|
||||
- `model=claude-3-5-sonnet-20240620`
|
||||
- `allowed_model_pattern=anthropic/*`
|
||||
|
||||
A model that already carries a namespace get_llm_provider did not consume
|
||||
(e.g. `bedrockz/anthropic.claude-...`) is never granted here: its provider was
|
||||
inferred from a fragment of the full string, so rebuilding
|
||||
`{provider}/{model}` would produce `bedrock/bedrockz/...` and slip an
|
||||
unrecognized namespace through a `bedrock/*` key.
|
||||
"""
|
||||
try:
|
||||
model, custom_llm_provider, _, _ = get_llm_provider(model=model)
|
||||
stripped_model, custom_llm_provider, _, _ = get_llm_provider(model=model)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
if stripped_model == model and "/" in model:
|
||||
return False
|
||||
|
||||
return is_model_allowed_by_pattern(
|
||||
model=f"{custom_llm_provider}/{model}",
|
||||
model=f"{custom_llm_provider}/{stripped_model}",
|
||||
allowed_model_pattern=allowed_model_pattern,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1191,13 +1191,15 @@ async def _user_api_key_auth_builder(
|
|||
return await handle_oauth2_proxy_request(request=request)
|
||||
|
||||
if general_settings.get("enable_jwt_auth", False) is True:
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
||||
if premium_user is not True:
|
||||
raise ValueError(f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}")
|
||||
is_jwt = jwt_handler.is_jwt(token=api_key)
|
||||
verbose_proxy_logger.debug("is_jwt: %s", is_jwt)
|
||||
if is_jwt:
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
||||
if premium_user is not True:
|
||||
raise ValueError(
|
||||
f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}"
|
||||
)
|
||||
# Try JWT-to-Virtual-Key mapping first to avoid
|
||||
# unnecessary DB queries in auth_builder
|
||||
do_standard_jwt_auth = True
|
||||
|
|
@ -2442,6 +2444,7 @@ async def _reserve_budget_after_common_checks(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
end_user_id=end_user_id,
|
||||
end_user_object=end_user_object,
|
||||
skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,10 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
# Bounds the __cause__/__context__ walk in is_database_service_unavailable_error_in_chain.
|
||||
# Real exception chains are a few links deep; the cap also makes the walk cycle-safe.
|
||||
_MAX_EXCEPTION_CHAIN_DEPTH = 20
|
||||
|
||||
|
||||
class PrismaDBExceptionHandler:
|
||||
"""
|
||||
|
|
@ -218,6 +222,32 @@ class PrismaDBExceptionHandler:
|
|||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def is_database_service_unavailable_error_in_chain(e: BaseException) -> bool:
|
||||
"""Like ``is_database_service_unavailable_error`` but also walks the
|
||||
``__cause__`` / ``__context__`` chain.
|
||||
|
||||
``is_database_service_unavailable_error`` classifies a single exception
|
||||
by type, which a caller that catches a raw DB failure and re-raises a
|
||||
domain exception of a different type defeats. ``get_user_object`` in
|
||||
``litellm/proxy/auth/auth_checks.py`` is the concrete case: it wraps
|
||||
every DB error, a genuine outage included, in a bare ``ValueError``
|
||||
whose original error survives only as ``__context__``. A type check on
|
||||
the ``ValueError`` misses the outage, so the caller would mistake an
|
||||
infrastructure fault for an auth failure. Walking the chain recovers the
|
||||
real signal, which is the PEP 3134 way to inspect a wrapped cause.
|
||||
|
||||
The walk is depth-bounded, which also makes it cycle-safe.
|
||||
"""
|
||||
current: BaseException | None = e
|
||||
for _ in range(_MAX_EXCEPTION_CHAIN_DEPTH):
|
||||
if not isinstance(current, Exception):
|
||||
return False
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error(current):
|
||||
return True
|
||||
current = current.__cause__ or current.__context__
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def handle_db_exception(e: Exception):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -199,9 +199,10 @@ model_list:
|
|||
|
||||
general_settings:
|
||||
master_key: sk-1234
|
||||
# Opt-in: let CheckBatchCost track cost for unmanaged Vertex batches created with a raw gs:// input_file_id.
|
||||
# Requires a vertex_ai deployment configured for the batched model. Defaults to false.
|
||||
# track_unmanaged_vertex_batch_cost: true
|
||||
# Opt-in: let CheckBatchCost track cost for unmanaged batches created with a raw
|
||||
# gs:// (Vertex) or s3:// (Bedrock) input_file_id. Requires a matching deployment
|
||||
# configured for the batched model. Defaults to false.
|
||||
# track_unmanaged_batch_cost: true
|
||||
|
||||
sandbox_tools:
|
||||
- sandbox_tool_name: e2b_sandbox
|
||||
|
|
|
|||
|
|
@ -34,6 +34,12 @@ def is_text_content_call_type(call_type: str) -> bool:
|
|||
|
||||
TEXT_PART_TYPES: FrozenSet[str] = frozenset({"text", "input_text", "output_text"})
|
||||
|
||||
# Responses-API item types whose ``output`` field carries user/tool text
|
||||
# that guardrails should inspect. ``function_call_output`` is the
|
||||
# built-in shape; ``custom_tool_call_output`` is the custom-tool
|
||||
# counterpart (see ``ChatCompletionCustomToolCallOutput``).
|
||||
_OUTPUT_ITEM_TYPES: frozenset[str] = frozenset({"function_call_output", "custom_tool_call_output"})
|
||||
|
||||
|
||||
def _iter_text_parts_in_content(content: Any) -> Iterator[str]:
|
||||
"""Yield text fragments from a ``message.content`` value (string or
|
||||
|
|
@ -72,7 +78,7 @@ def _coerce_input_to_messages(input_value: Any) -> List[Dict[str, Any]]:
|
|||
messages.append({"role": item.get("role") or "user", "content": [item]})
|
||||
elif "content" in item:
|
||||
messages.append({"role": item.get("role") or "user", "content": item["content"]})
|
||||
elif item.get("type") == "function_call_output" and "output" in item:
|
||||
elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item:
|
||||
messages.append({"role": item.get("role") or "tool", "content": item["output"]})
|
||||
return messages
|
||||
|
||||
|
|
@ -157,7 +163,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int:
|
|||
input_value[idx] = {**item, "text": visit(item["text"])}
|
||||
elif "content" in item:
|
||||
item["content"] = _rewrite_content(item["content"])
|
||||
elif item.get("type") == "function_call_output" and "output" in item:
|
||||
elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item:
|
||||
item["output"] = _rewrite_content(item["output"])
|
||||
return visited
|
||||
|
||||
|
|
|
|||
|
|
@ -772,7 +772,13 @@ class LassoGuardrail(CustomGuardrail):
|
|||
data: Request data (used for conversation_id generation and tools extraction)
|
||||
cache: Cache instance for storing conversation_id (optional for post-call)
|
||||
"""
|
||||
payload: Dict[str, Any] = {"messages": messages, "messageType": message_type}
|
||||
payload: Dict[str, Any] = {
|
||||
"messages": messages,
|
||||
"messageType": message_type,
|
||||
# Drives the "Used By" badge on Lasso Application API Keys: every call from this
|
||||
# integration is attributed as "litellm" on the keys list.
|
||||
"source": {"type": "litellm"},
|
||||
}
|
||||
|
||||
# Add optional parameters if available
|
||||
if self.user_id:
|
||||
|
|
|
|||
|
|
@ -44,12 +44,18 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
|
||||
from litellm.types.utils import (
|
||||
GenericGuardrailAPIInputs,
|
||||
GuardrailStatus,
|
||||
StandardLoggingGuardrailInformation,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
|
|
@ -60,6 +66,13 @@ if TYPE_CHECKING:
|
|||
)
|
||||
|
||||
|
||||
def _sanitize_scan_result_for_logging(scan_result: dict) -> dict:
|
||||
without_secrets = {key: value for key, value in scan_result.items() if key != "secret_fields"}
|
||||
redacted = redact_nested_match_and_regex_keys(without_secrets)
|
||||
masked = mask_credentials_in_payload(redacted if isinstance(redacted, dict) else without_secrets)
|
||||
return masked if isinstance(masked, dict) else without_secrets
|
||||
|
||||
|
||||
_DEFAULT_API_BASE = "https://api-xecguard.cycraft.ai"
|
||||
_SCAN_ENDPOINT = "/xecguard/v1/scan"
|
||||
_GROUNDING_ENDPOINT = "/xecguard/v1/grounding"
|
||||
|
|
@ -246,16 +259,21 @@ class XecGuardGuardrail(CustomGuardrail):
|
|||
"guardrail_intervened" if scan_result.get("decision") == "UNSAFE" else "success"
|
||||
)
|
||||
end_time = datetime.now()
|
||||
kwargs["standard_logging_object"]["guardrail_information"] = {
|
||||
"duration": (end_time - start_time).total_seconds(),
|
||||
"end_time": end_time.timestamp(),
|
||||
"guardrail_mode": "logging_only",
|
||||
"guardrail_name": "xecguard",
|
||||
"guardrail_response": scan_result,
|
||||
"guardrail_status": guardrail_status,
|
||||
"masked_entity_count": None,
|
||||
"start_time": start_time.timestamp(),
|
||||
}
|
||||
slg = StandardLoggingGuardrailInformation(
|
||||
guardrail_name=self.guardrail_name or "xecguard",
|
||||
guardrail_mode=GuardrailEventHooks.logging_only,
|
||||
guardrail_response=_sanitize_scan_result_for_logging(scan_result),
|
||||
guardrail_status=guardrail_status,
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=end_time.timestamp(),
|
||||
duration=(end_time - start_time).total_seconds(),
|
||||
masked_entity_count=None,
|
||||
)
|
||||
existing = kwargs["standard_logging_object"].get("guardrail_information")
|
||||
if isinstance(existing, list):
|
||||
existing.append(slg)
|
||||
else:
|
||||
kwargs["standard_logging_object"]["guardrail_information"] = [slg]
|
||||
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
|
|||
|
|
@ -28,11 +28,20 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
|||
)
|
||||
from litellm.proxy.utils import ProxyUpdateSpend
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
StandardLoggingPayload,
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
)
|
||||
from litellm.utils import get_end_user_id_for_cost_tracking
|
||||
|
||||
_PASS_THROUGH_CALL_TYPES: frozenset[str] = frozenset(
|
||||
{
|
||||
CallTypes.pass_through.value,
|
||||
CallTypes.llm_passthrough_route.value,
|
||||
CallTypes.allm_passthrough_route.value,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _ProxyDBLogger(CustomLogger):
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
|
|
@ -219,11 +228,13 @@ class _ProxyDBLogger(CustomLogger):
|
|||
verbose_proxy_logger.debug(
|
||||
f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}"
|
||||
)
|
||||
call_type: Optional[str] = kwargs.get("call_type")
|
||||
if _should_track_cost_callback(
|
||||
user_api_key=user_api_key,
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
end_user_id=end_user_id,
|
||||
call_type=call_type,
|
||||
):
|
||||
## UPDATE DATABASE
|
||||
await _update_database_and_spend_counters(
|
||||
|
|
@ -412,9 +423,15 @@ def _should_track_cost_callback(
|
|||
user_id: Optional[str],
|
||||
team_id: Optional[str],
|
||||
end_user_id: Optional[str],
|
||||
call_type: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Determine if the cost callback should be tracked based on the kwargs
|
||||
|
||||
Pass-through endpoints can be configured with ``auth=false``, which leaves
|
||||
the request with no key/user/team/end-user to attribute spend to. Those
|
||||
requests still forward real provider traffic that operators expect to see
|
||||
in request/usage logs, so they are tracked even when unauthenticated.
|
||||
"""
|
||||
|
||||
# don't run track cost callback if user opted into disabling spend
|
||||
|
|
@ -423,7 +440,7 @@ def _should_track_cost_callback(
|
|||
|
||||
if user_api_key is not None or user_id is not None or team_id is not None or end_user_id is not None:
|
||||
return True
|
||||
return False
|
||||
return call_type in _PASS_THROUGH_CALL_TYPES
|
||||
|
||||
|
||||
def _get_budget_reservation_from_metadata(metadata: dict) -> Optional[dict]:
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ from litellm.proxy._experimental.mcp_server.db import (
|
|||
rotate_mcp_user_env_vars_master_key,
|
||||
)
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken, hash_token
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
_delete_cache_key_object,
|
||||
can_team_access_model,
|
||||
|
|
@ -468,7 +468,10 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict:
|
|||
Handle the key type.
|
||||
"""
|
||||
key_type = data.key_type
|
||||
data_json.pop("key_type", None)
|
||||
if key_type is None:
|
||||
data_json.pop("key_type", None)
|
||||
return data_json
|
||||
data_json["key_type"] = key_type.value
|
||||
if key_type == LiteLLMKeyType.LLM_API:
|
||||
data_json["allowed_routes"] = ["llm_api_routes"]
|
||||
elif key_type == LiteLLMKeyType.MANAGEMENT:
|
||||
|
|
@ -3566,6 +3569,7 @@ async def generate_key_helper_fn(
|
|||
created_by: Optional[str] = None,
|
||||
updated_by: Optional[str] = None,
|
||||
allowed_routes: Optional[list] = None,
|
||||
key_type: str | None = None,
|
||||
sso_user_id: Optional[str] = None,
|
||||
object_permission_id: Optional[str] = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable
|
||||
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None,
|
||||
|
|
@ -3706,6 +3710,7 @@ async def generate_key_helper_fn(
|
|||
"created_by": created_by,
|
||||
"updated_by": updated_by,
|
||||
"allowed_routes": allowed_routes or [],
|
||||
"key_type": key_type,
|
||||
"object_permission_id": object_permission_id,
|
||||
"router_settings": router_settings_json,
|
||||
"access_group_ids": access_group_ids or [],
|
||||
|
|
@ -3772,7 +3777,10 @@ async def generate_key_helper_fn(
|
|||
return user_data
|
||||
|
||||
## CREATE KEY
|
||||
verbose_proxy_logger.debug("prisma_client: Creating Key= %s", key_data)
|
||||
verbose_proxy_logger.debug(
|
||||
"prisma_client: Creating Key= %s",
|
||||
{**key_data, "token": hash_token(token=token)},
|
||||
)
|
||||
create_key_response = await prisma_client.insert_data(data=key_data, table_name="key")
|
||||
|
||||
key_data["token_id"] = getattr(create_key_response, "token", None)
|
||||
|
|
|
|||
|
|
@ -3550,7 +3550,7 @@ async def team_info(
|
|||
try:
|
||||
team_info: Optional[BaseModel] = await TeamRepository(prisma_client).table.find_unique(
|
||||
where={"team_id": team_id},
|
||||
include={"object_permission": True},
|
||||
include={"litellm_model_table": True, "object_permission": True},
|
||||
)
|
||||
if team_info is None:
|
||||
raise Exception
|
||||
|
|
|
|||
|
|
@ -4034,30 +4034,41 @@ class MicrosoftSSOHandler:
|
|||
base_url = MicrosoftSSOHandler.get_graph_api_base_url()
|
||||
# Endpoint to get app role assignments for the given service principal
|
||||
endpoint = f"/servicePrincipals/{service_principal_id}/appRoleAssignedTo"
|
||||
url = base_url + endpoint
|
||||
next_link: str | None = base_url + endpoint
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
response = await async_client.get(url, headers=headers)
|
||||
response_json = response.json()
|
||||
verbose_proxy_logger.debug(f"Response from service principal app role assigned to: {response_json}")
|
||||
group_ids: List[str] = []
|
||||
service_principal_teams: List[MicrosoftServicePrincipalTeam] = []
|
||||
page_count = 0
|
||||
|
||||
for _object in response_json.get("value", []):
|
||||
if _object.get("principalType") == "Group":
|
||||
# Append the group ID to the list
|
||||
group_ids.append(_object.get("principalId"))
|
||||
# Append the service principal team to the list
|
||||
service_principal_teams.append(
|
||||
MicrosoftServicePrincipalTeam(
|
||||
principalDisplayName=_object.get("principalDisplayName"),
|
||||
principalId=_object.get("principalId"),
|
||||
while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES:
|
||||
response = await async_client.get(next_link, headers=headers)
|
||||
response_json = response.json()
|
||||
verbose_proxy_logger.debug(f"Response from service principal app role assigned to: {response_json}")
|
||||
|
||||
for _object in response_json.get("value", []):
|
||||
if _object.get("principalType") == "Group":
|
||||
# Append the group ID to the list
|
||||
group_ids.append(_object.get("principalId"))
|
||||
# Append the service principal team to the list
|
||||
service_principal_teams.append(
|
||||
MicrosoftServicePrincipalTeam(
|
||||
principalDisplayName=_object.get("principalDisplayName"),
|
||||
principalId=_object.get("principalId"),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
next_link = response_json.get("@odata.nextLink")
|
||||
page_count += 1
|
||||
|
||||
if next_link is not None and page_count >= MicrosoftSSOHandler.MAX_GRAPH_API_PAGES:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Reached maximum page limit of {MicrosoftSSOHandler.MAX_GRAPH_API_PAGES}. Some service principal group assignments may not be included."
|
||||
)
|
||||
|
||||
return group_ids, service_principal_teams
|
||||
|
||||
|
|
|
|||
|
|
@ -7882,7 +7882,7 @@ class ProxyStartupEvent:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
prisma_client=prisma_client,
|
||||
llm_router=llm_router,
|
||||
track_unmanaged_vertex_batch_cost=general_settings.get("track_unmanaged_vertex_batch_cost", False),
|
||||
track_unmanaged_batch_cost=general_settings.get("track_unmanaged_batch_cost", False),
|
||||
)
|
||||
scheduler.add_job(
|
||||
check_batch_cost_job.check_batch_cost,
|
||||
|
|
@ -14805,6 +14805,7 @@ async def get_config_list(
|
|||
"forward_client_headers_to_llm_api": {"type": "Boolean"},
|
||||
"mcp_required_fields": {"type": "List"},
|
||||
"cancel_on_disconnect": {"type": "Boolean"},
|
||||
"skip_user_budget_on_team_key": {"type": "Boolean"},
|
||||
}
|
||||
|
||||
return_val = []
|
||||
|
|
|
|||
|
|
@ -422,6 +422,7 @@ model LiteLLM_VerificationToken {
|
|||
budget_reset_at DateTime?
|
||||
allowed_cache_controls String[] @default([])
|
||||
allowed_routes String[] @default([])
|
||||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
|
|
@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
budget_reset_at DateTime?
|
||||
allowed_cache_controls String[] @default([])
|
||||
allowed_routes String[] @default([])
|
||||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
|
|
|
|||
|
|
@ -123,6 +123,7 @@ async def reserve_budget_for_request(
|
|||
proxy_logging_obj: ProxyLogging,
|
||||
end_user_id: Optional[str] = None,
|
||||
end_user_object: Optional[Any] = None,
|
||||
skip_user_budget_on_team_key: bool = False,
|
||||
) -> Optional[dict]:
|
||||
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
|
||||
return None
|
||||
|
|
@ -141,6 +142,7 @@ async def reserve_budget_for_request(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
end_user_id=end_user_id,
|
||||
end_user_object=end_user_object,
|
||||
skip_user_budget_on_team_key=skip_user_budget_on_team_key,
|
||||
)
|
||||
if not counters:
|
||||
return None
|
||||
|
|
@ -296,6 +298,7 @@ async def _get_budget_counters(
|
|||
proxy_logging_obj: ProxyLogging,
|
||||
end_user_id: Optional[str] = None,
|
||||
end_user_object: Optional[Any] = None,
|
||||
skip_user_budget_on_team_key: bool = False,
|
||||
) -> List[_BudgetCounter]:
|
||||
counters: List[_BudgetCounter] = []
|
||||
|
||||
|
|
@ -344,8 +347,9 @@ async def _get_budget_counters(
|
|||
)
|
||||
)
|
||||
|
||||
is_team_key = team_object is not None and team_object.team_id is not None
|
||||
if (
|
||||
(team_object is None or team_object.team_id is None)
|
||||
not (is_team_key and skip_user_budget_on_team_key)
|
||||
and user_object is not None
|
||||
and user_object.user_id is not None
|
||||
and user_object.max_budget is not None
|
||||
|
|
|
|||
|
|
@ -3592,7 +3592,10 @@ class PrismaClient:
|
|||
"""
|
||||
start_time = time.time()
|
||||
try:
|
||||
verbose_proxy_logger.debug("PrismaClient: insert_data: %s", data)
|
||||
verbose_proxy_logger.debug(
|
||||
"PrismaClient: insert_data: %s",
|
||||
{**data, "token": self.hash_token(token=data["token"])} if data.get("token") is not None else data,
|
||||
)
|
||||
if table_name == "key":
|
||||
token = data["token"]
|
||||
hashed_token = self.hash_token(token=token)
|
||||
|
|
|
|||
|
|
@ -261,7 +261,7 @@ async def aresponses_api_with_mcp(
|
|||
pre_processed_mcp_tools=original_mcp_tools,
|
||||
)
|
||||
|
||||
return LiteLLM_Proxy_MCP_Handler._create_mcp_streaming_response(
|
||||
mcp_streaming_response = LiteLLM_Proxy_MCP_Handler._create_mcp_streaming_response(
|
||||
input=input,
|
||||
model=model,
|
||||
all_tools=all_tools,
|
||||
|
|
@ -272,6 +272,10 @@ async def aresponses_api_with_mcp(
|
|||
tool_server_map=tool_server_map,
|
||||
**kwargs,
|
||||
)
|
||||
await mcp_streaming_response._create_initial_response_iterator()
|
||||
if mcp_streaming_response._initial_creation_error is not None:
|
||||
raise mcp_streaming_response._initial_creation_error
|
||||
return mcp_streaming_response
|
||||
|
||||
# Determine if we should auto-execute tools
|
||||
should_auto_execute = bool(mcp_tools_with_litellm_proxy) and LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ from litellm._uuid import uuid
|
|||
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
|
||||
from litellm.types.llms.openai import (
|
||||
BaseLiteLLMOpenAIResponseObject,
|
||||
ErrorEvent,
|
||||
ErrorEventError,
|
||||
MCPCallArgumentsDeltaEvent,
|
||||
MCPCallArgumentsDoneEvent,
|
||||
MCPCallCompletedEvent,
|
||||
|
|
@ -316,6 +318,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
# Cache the response ID to ensure consistency across all events
|
||||
self._cached_response_id: Optional[str] = None
|
||||
|
||||
self._initial_creation_error: Exception | None = None
|
||||
self._stream_error: Exception | None = None
|
||||
self._error_event_emitted = False
|
||||
self._last_sequence_number = 0
|
||||
|
||||
def _extract_mcp_headers_from_params(self) -> None:
|
||||
"""Extract MCP headers from original request params to pass to tool calls"""
|
||||
from typing import Dict, Optional
|
||||
|
|
@ -380,10 +387,31 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
|
||||
return LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(self.mcp_tools_with_litellm_proxy)
|
||||
|
||||
def _make_stream_error_event(self) -> ResponsesAPIStreamingResponse:
|
||||
err = self._stream_error
|
||||
status_code = getattr(err, "status_code", None)
|
||||
return ErrorEvent(
|
||||
type=ResponsesAPIStreamEvents.ERROR,
|
||||
sequence_number=self._last_sequence_number + 1,
|
||||
error=ErrorEventError(
|
||||
type="mcp_gateway_error",
|
||||
code=str(status_code) if status_code is not None else "internal_error",
|
||||
message=str(err) if err is not None else "MCP gateway stream failed",
|
||||
param=None,
|
||||
),
|
||||
)
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> ResponsesAPIStreamingResponse:
|
||||
chunk = await self._anext_impl()
|
||||
sequence_number = getattr(chunk, "sequence_number", None)
|
||||
if isinstance(sequence_number, int) and sequence_number > self._last_sequence_number:
|
||||
self._last_sequence_number = sequence_number
|
||||
return chunk
|
||||
|
||||
async def _anext_impl(self) -> ResponsesAPIStreamingResponse:
|
||||
"""
|
||||
Phase-based streaming:
|
||||
1. initial_response - Stream the first LLM response (includes response.created, response.in_progress, response.output_item.added)
|
||||
|
|
@ -438,10 +466,16 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
self.phase = "continue_initial_response"
|
||||
return await self.__anext__()
|
||||
self.phase = "finished"
|
||||
if self._stream_error is not None and not self._error_event_emitted:
|
||||
self._error_event_emitted = True
|
||||
return self._make_stream_error_event()
|
||||
raise StopAsyncIteration
|
||||
|
||||
# Phase 6: Finished
|
||||
if self.phase == "finished":
|
||||
if self._stream_error is not None and not self._error_event_emitted:
|
||||
self._error_event_emitted = True
|
||||
return self._make_stream_error_event()
|
||||
raise StopAsyncIteration
|
||||
|
||||
# Should not reach here
|
||||
|
|
@ -530,6 +564,12 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
|
||||
chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined]
|
||||
|
||||
if self._cached_response_id is None and hasattr(chunk, "response"):
|
||||
new_response = getattr(chunk, "response", None)
|
||||
new_response_id = getattr(new_response, "id", None) if new_response is not None else None
|
||||
if new_response_id:
|
||||
self._cached_response_id = new_response_id
|
||||
|
||||
# Ensure response ID consistency - update chunk if needed
|
||||
if self._cached_response_id and hasattr(chunk, "response"):
|
||||
response_obj = getattr(chunk, "response", None)
|
||||
|
|
@ -589,6 +629,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
|
||||
traceback.print_exc()
|
||||
self.base_iterator = None
|
||||
self._initial_creation_error = e
|
||||
self._stream_error = e
|
||||
# Don't set phase to "finished" here — let __anext__ emit any
|
||||
# pre-generated MCP discovery events before ending the iteration.
|
||||
|
||||
|
|
@ -761,6 +803,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
if hasattr(follow_up_response, "__aiter__"):
|
||||
self.base_iterator = follow_up_response
|
||||
self.collected_response = None
|
||||
self._cached_response_id = None
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.error(f"Error creating follow-up iterator: {e}")
|
||||
|
|
@ -768,6 +811,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
|
||||
traceback.print_exc()
|
||||
self.base_iterator = None
|
||||
self._stream_error = e
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
|
|
|||
|
|
@ -179,7 +179,9 @@ from litellm.types.router import (
|
|||
RouterModelGroupAliasItem,
|
||||
RouterRateLimitError,
|
||||
RouterRateLimitErrorBasic,
|
||||
RoutingContext,
|
||||
RoutingGroup,
|
||||
RoutingPlugin,
|
||||
RoutingStrategy,
|
||||
SearchToolTypedDict,
|
||||
)
|
||||
|
|
@ -299,6 +301,7 @@ class Router:
|
|||
enable_pre_call_checks: bool = False,
|
||||
enable_tag_filtering: bool = False,
|
||||
tag_filtering_match_any: bool = True,
|
||||
plugins: list[RoutingPlugin] | None = None,
|
||||
retry_after: int = 0, # min time to wait before retrying a failed request
|
||||
retry_policy: Optional[Union[RetryPolicy, dict]] = None, # set custom retries for different exceptions
|
||||
model_group_retry_policy: Dict[str, RetryPolicy] = {}, # set custom retry policies based on model group
|
||||
|
|
@ -477,6 +480,7 @@ class Router:
|
|||
self.complexity_routers: Dict[str, "ComplexityRouter"] = {}
|
||||
self.adaptive_routers: Dict[str, "AdaptiveRouter"] = {}
|
||||
self.quality_routers: Dict[str, "QualityRouter"] = {}
|
||||
self.routing_plugins: list[RoutingPlugin] = list(plugins) if plugins else []
|
||||
|
||||
# Initialize model_group_alias early since it's used in set_model_list
|
||||
self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = (
|
||||
|
|
@ -7552,7 +7556,11 @@ class Router:
|
|||
if default_model is None and complexity_router_config:
|
||||
tiers = complexity_router_config.get("tiers", {})
|
||||
# Use MEDIUM tier as fallback default
|
||||
default_model = tiers.get("MEDIUM") or tiers.get("SIMPLE")
|
||||
medium = tiers.get("MEDIUM") or tiers.get("SIMPLE")
|
||||
if isinstance(medium, list):
|
||||
default_model = medium[0] if medium else None
|
||||
else:
|
||||
default_model = medium
|
||||
|
||||
if default_model is None:
|
||||
raise ValueError(
|
||||
|
|
@ -7589,15 +7597,6 @@ class Router:
|
|||
AdaptiveRouterPostCallHook,
|
||||
)
|
||||
|
||||
for _cb_list in (
|
||||
litellm.callbacks,
|
||||
litellm.success_callback,
|
||||
litellm.failure_callback,
|
||||
litellm._async_success_callback,
|
||||
litellm._async_failure_callback,
|
||||
):
|
||||
litellm.logging_callback_manager.remove_callbacks_by_type(_cb_list, AdaptiveRouterPostCallHook)
|
||||
|
||||
for entry in self.model_list or []:
|
||||
lp = entry.get("litellm_params") if isinstance(entry, dict) else entry.litellm_params
|
||||
lp_model = (lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None
|
||||
|
|
@ -7606,15 +7605,29 @@ class Router:
|
|||
model_name = entry.get("model_name") if isinstance(entry, dict) else entry.model_name
|
||||
if not model_name or not lp:
|
||||
continue
|
||||
if model_name in self.adaptive_routers:
|
||||
continue
|
||||
deployment = Deployment(
|
||||
model_name=model_name,
|
||||
litellm_params=(lp if not isinstance(lp, dict) else LiteLLM_Params(**lp)),
|
||||
model_info=(entry.get("model_info") if isinstance(entry, dict) else entry.model_info),
|
||||
)
|
||||
if model_name in self.adaptive_routers:
|
||||
continue
|
||||
self.init_adaptive_router_deployment(deployment=deployment)
|
||||
|
||||
for model_name, complexity_router in self.complexity_routers.items():
|
||||
if not complexity_router.config.adaptive or model_name in self.adaptive_routers:
|
||||
continue
|
||||
adaptive_router = complexity_router._ensure_adaptive_router()
|
||||
if adaptive_router is not None:
|
||||
self.adaptive_routers[model_name] = adaptive_router
|
||||
|
||||
for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook):
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(callback)
|
||||
for adaptive_router in self.adaptive_routers.values():
|
||||
litellm.logging_callback_manager.add_litellm_callback(
|
||||
AdaptiveRouterPostCallHook(adaptive_router=adaptive_router)
|
||||
)
|
||||
|
||||
def init_adaptive_router_deployment(self, deployment: Deployment) -> None:
|
||||
"""
|
||||
Build an AdaptiveRouter instance for this deployment and register its
|
||||
|
|
@ -10321,6 +10334,12 @@ class Router:
|
|||
metadata_variable_name=self._get_metadata_variable_name_from_kwargs(request_kwargs),
|
||||
)
|
||||
|
||||
# narrow to whatever `self.routing_plugins` left in candidate_models
|
||||
healthy_deployments = self._filter_by_routing_plugin_candidates(
|
||||
healthy_deployments=healthy_deployments,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2)
|
||||
_target_order = (request_kwargs or {}).pop("_target_order", None)
|
||||
healthy_deployments = litellm.utils._get_order_filtered_deployments(
|
||||
|
|
@ -10596,6 +10615,76 @@ class Router:
|
|||
)
|
||||
raise e
|
||||
|
||||
async def _run_routing_plugins(
|
||||
self,
|
||||
model: str,
|
||||
request_kwargs: dict,
|
||||
messages: list[dict[str, Any]] | None,
|
||||
) -> RoutingContext:
|
||||
"""
|
||||
Build a RoutingContext for `model`, run it through `self.routing_plugins`
|
||||
in order, then stash the narrowed candidate list and accumulated signals
|
||||
onto `request_kwargs["metadata"]` so `_filter_by_routing_plugin_candidates`
|
||||
(called later, during healthy-deployment filtering) can consume them.
|
||||
"""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
resolve_structured_messages,
|
||||
)
|
||||
|
||||
deployments = self.get_model_list(model_name=model) or []
|
||||
candidate_models = [
|
||||
d["litellm_params"]["model"] for d in deployments if d.get("litellm_params", {}).get("model")
|
||||
]
|
||||
|
||||
metadata_key = self._get_metadata_variable_name_from_kwargs(request_kwargs)
|
||||
metadata = request_kwargs.setdefault(metadata_key, {})
|
||||
|
||||
context = RoutingContext(
|
||||
raw_messages=messages or [],
|
||||
structured_messages=resolve_structured_messages(messages=messages, request_kwargs=request_kwargs) or [],
|
||||
candidate_models=candidate_models,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
for plugin in self.routing_plugins:
|
||||
context = await plugin.run(context)
|
||||
|
||||
metadata["routing_plugin_signals"] = context.signals
|
||||
if len(context.candidate_models) < len(candidate_models):
|
||||
metadata["_routing_plugin_candidate_models"] = context.candidate_models
|
||||
|
||||
return context
|
||||
|
||||
def _filter_by_routing_plugin_candidates(
|
||||
self,
|
||||
healthy_deployments: Union[list[dict], dict],
|
||||
request_kwargs: dict,
|
||||
) -> Union[list[dict], dict]:
|
||||
"""
|
||||
Narrow `healthy_deployments` to whatever `self.routing_plugins` left in
|
||||
`context.candidate_models`. Raises rather than silently falling back to
|
||||
the unfiltered pool -- a plugin narrowing to nothing is a policy decision
|
||||
(e.g. no model this tenant's budget allows), not something to bypass.
|
||||
"""
|
||||
if not self.routing_plugins or not isinstance(healthy_deployments, list):
|
||||
return healthy_deployments
|
||||
|
||||
metadata_key = self._get_metadata_variable_name_from_kwargs(request_kwargs)
|
||||
candidate_models = (request_kwargs.get(metadata_key) or {}).get("_routing_plugin_candidate_models")
|
||||
# `is None` (not falsy-check): a plugin narrowing to an empty list must
|
||||
# still hit the "no deployments left" raise below, not be treated the
|
||||
# same as "no plugin ever set this key".
|
||||
if candidate_models is None:
|
||||
return healthy_deployments
|
||||
|
||||
candidate_set = set(candidate_models)
|
||||
filtered = [d for d in healthy_deployments if d.get("litellm_params", {}).get("model") in candidate_set]
|
||||
|
||||
if not filtered:
|
||||
raise ValueError(f"No deployments left after routing-plugin filtering. candidate_models={candidate_models}")
|
||||
|
||||
return filtered
|
||||
|
||||
async def async_pre_routing_hook(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -10610,55 +10699,47 @@ class Router:
|
|||
Used for the litellm auto-router to modify the request before the routing decision is made.
|
||||
"""
|
||||
#########################################################
|
||||
# Check if any auto-router should be used
|
||||
# Run the routing-plugin pipeline, if any plugins are configured.
|
||||
# Plugins narrow the candidate deployment pool (consumed later by
|
||||
# `_filter_by_routing_plugin_candidates`) and may attach signals for
|
||||
# downstream strategies (auto-router, complexity-router, ...) to read.
|
||||
#########################################################
|
||||
if model in self.auto_routers:
|
||||
return await self.auto_routers[model].async_pre_routing_hook(
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
)
|
||||
if self.routing_plugins:
|
||||
await self._run_routing_plugins(model=model, request_kwargs=request_kwargs, messages=messages)
|
||||
|
||||
#########################################################
|
||||
# Check if any complexity-router should be used
|
||||
#########################################################
|
||||
if model in self.complexity_routers:
|
||||
return await self.complexity_routers[model].async_pre_routing_hook(
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
)
|
||||
router_strategy = (
|
||||
self.auto_routers.get(model)
|
||||
or self.complexity_routers.get(model)
|
||||
or self.adaptive_routers.get(model)
|
||||
or self.quality_routers.get(model)
|
||||
)
|
||||
if router_strategy is None:
|
||||
return None
|
||||
|
||||
#########################################################
|
||||
# Check if an adaptive-router should be used
|
||||
#########################################################
|
||||
adaptive_router = self.adaptive_routers.get(model)
|
||||
if adaptive_router is not None:
|
||||
return await adaptive_router.async_pre_routing_hook(
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
)
|
||||
pre_routing_hook_response = await router_strategy.async_pre_routing_hook(
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
)
|
||||
|
||||
#########################################################
|
||||
# Check if any quality-router should be used
|
||||
#########################################################
|
||||
if model in self.quality_routers:
|
||||
return await self.quality_routers[model].async_pre_routing_hook(
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
)
|
||||
# `model` (the alias, e.g. "smart-router") is never the deployment actually
|
||||
# called - apply the alias's own litellm_params (besides `model` itself,
|
||||
# which is just the alias marker) to the request, since the tier/route
|
||||
# deployment the hook selected won't have them. Router-only fields
|
||||
# (tpm, rpm, weight, complexity_router_config, ...) are excluded from the
|
||||
# actual outbound LLM call downstream by litellm.types.utils.all_litellm_params,
|
||||
# not here.
|
||||
if pre_routing_hook_response is not None:
|
||||
alias_index = self.model_name_to_deployment_indices.get(model, [])
|
||||
if alias_index:
|
||||
alias_litellm_params = self.model_list[alias_index[0]].get("litellm_params", {})
|
||||
for key, value in alias_litellm_params.items():
|
||||
if key != "model" and value is not None:
|
||||
request_kwargs.setdefault(key, value)
|
||||
|
||||
return None
|
||||
return pre_routing_hook_response
|
||||
|
||||
def get_available_deployment(
|
||||
self,
|
||||
|
|
@ -10671,6 +10752,18 @@ class Router:
|
|||
"""
|
||||
Returns the deployment based on routing strategy
|
||||
"""
|
||||
if self.routing_plugins:
|
||||
raise ValueError(
|
||||
"Router(plugins=[...]) is configured, but this call resolved to the synchronous "
|
||||
"deployment-selection path, which never runs the routing-plugin pipeline. This "
|
||||
"happens for sync Router methods (e.g. Router.completion()) and for async calls "
|
||||
"with a routing_strategy that has no async-native selector (e.g. legacy "
|
||||
"'usage-based-routing', v1). Silently skipping "
|
||||
"configured plugins would let a policy plugin (e.g. a deny-all rule) be bypassed. "
|
||||
"Use an async Router method with a supported routing_strategy (simple-shuffle, "
|
||||
"usage-based-routing-v2, cost-based-routing, latency-based-routing, least-busy), "
|
||||
"or remove `plugins` from the Router config."
|
||||
)
|
||||
# users need to explicitly call a specific deployment, by setting `specific_deployment = True` as completion()/embedding() kwarg
|
||||
# When this was no explicit we had several issues with fallbacks timing out
|
||||
|
||||
|
|
|
|||
|
|
@ -56,11 +56,10 @@ Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key
|
|||
- **Per-request decision.** Sample once per eligible model, score with
|
||||
`quality_weight·sample + cost_weight·normalized_cost`, pick the argmax.
|
||||
Routing is stateless per-turn — no sticky lookup. Each call resamples.
|
||||
- **Owner-cache attribution.** Post-call, the conversation's first picked
|
||||
model claims an "owner slot" for `OWNER_CACHE_TTL_SECONDS` (24h). Later
|
||||
turns of the same conversation only fire bandit/state updates if the
|
||||
same model handled them — mismatches are dropped (no attribution) and
|
||||
counted in `skipped_updates_total`. Conversation identity is the
|
||||
- **Previous-response attribution.** Post-call, feedback from the current user
|
||||
message is attributed to the model that produced the previous response, while
|
||||
response signals are attributed to the current model. Contexts expire after
|
||||
24 hours and the in-memory cache is capped at 1,024 sessions. Conversation identity is the
|
||||
client-supplied `litellm_session_id` if present, otherwise a sha256 over
|
||||
caller identity (api key hash, team, user, end-user) + the first message.
|
||||
- **Per-turn updates.** `satisfaction → +α`. `misalignment, stagnation,
|
||||
|
|
@ -76,12 +75,6 @@ Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key
|
|||
model can still be picked.
|
||||
- **Hard sample cap at 200.** Once `α + β > 200`, deltas are silently dropped.
|
||||
No rescaling — drift is a v1 concern.
|
||||
- **24h owner-cache TTL.** No explicit eviction below TTL. The in-memory map
|
||||
can grow if traffic patterns produce many one-shot sessions.
|
||||
- **Owner-recovery skew.** If model A "owns" a conversation but is then
|
||||
dethroned in the bandit, later turns served by model B are dropped — so
|
||||
bandit updates for that conversation flatline until A's TTL expires.
|
||||
Tracked via `skipped_updates_total`.
|
||||
- **Signals are regex + tool-call only.** No LLM-judge, no embedding similarity,
|
||||
no exemplar storage. Signals are best-effort and biased toward English.
|
||||
- **One AdaptiveRouter per `Router`.** Multiple `adaptive_router/*` deployments
|
||||
|
|
|
|||
|
|
@ -3,25 +3,21 @@ Main adaptive router strategy. See README.md for design overview.
|
|||
|
||||
One AdaptiveRouter instance per router_name. Holds in-memory caches:
|
||||
- _cells: Beta(alpha, beta) bandit posteriors per (request_type, model)
|
||||
- _owner_cache: session_key -> (owner_model, expires_at) — the first model
|
||||
picked for a conversation owns its bandit-update slot
|
||||
- _session_states: (session_key, model) -> SessionState for incremental signal updates
|
||||
|
||||
Owns the AdaptiveRouterUpdateQueue used by the proxy's flusher to persist
|
||||
state and session snapshots back to Postgres.
|
||||
|
||||
Routing is stateless per-turn (Thompson sample fresh on every call). The
|
||||
owner cache is consulted only at post-call time to decide whether a turn's
|
||||
signals should fire a bandit update — turns served by a different model than
|
||||
the conversation's owner are skipped to avoid cross-model misattribution.
|
||||
Routing is stateless per-turn (Thompson sample fresh on every call).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from dataclasses import asdict
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union, cast
|
||||
from collections import OrderedDict
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any, Union, cast
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
|
|
@ -38,13 +34,18 @@ from litellm.router_strategy.adaptive_router.config import (
|
|||
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
|
||||
MIN_QUALITY_TIER_HEADER,
|
||||
MIN_QUALITY_TIER_METADATA_KEY,
|
||||
MIN_TURNS_FOR_CLEAN_CREDIT,
|
||||
OWNER_CACHE_TTL_SECONDS,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.signals import (
|
||||
SessionState,
|
||||
SignalDelta,
|
||||
Turn,
|
||||
apply_turn,
|
||||
advance_session_state,
|
||||
apply_signal_delta,
|
||||
detect_response_signals,
|
||||
detect_user_feedback,
|
||||
merge_signal_deltas,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.update_queue import (
|
||||
AdaptiveRouterUpdateQueue,
|
||||
|
|
@ -53,8 +54,7 @@ from litellm.router_strategy.adaptive_router.update_queue import (
|
|||
# Sweep session-state cache when it exceeds this many live entries. Expired
|
||||
# entries are dropped in bulk; amortizes to O(1) per insert.
|
||||
_SESSION_STATE_SWEEP_THRESHOLD: int = 1024
|
||||
# Same pattern for the owner cache.
|
||||
_OWNER_CACHE_SWEEP_THRESHOLD: int = 1024
|
||||
_FEEDBACK_CONTEXT_MAX_ENTRIES: int = 1024
|
||||
from litellm.repositories.table_repositories import AdaptiveRouterStateRepository
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.router import (
|
||||
|
|
@ -70,6 +70,17 @@ def _default_prefs() -> AdaptiveRouterPreferences:
|
|||
return AdaptiveRouterPreferences(quality_tier=2, strengths=[])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _FeedbackContext:
|
||||
model_name: str
|
||||
request_type: RequestType
|
||||
user_content: str | None
|
||||
assistant_content: str | None
|
||||
turn_count: int
|
||||
clean_credit_awarded: bool
|
||||
expires_at: float
|
||||
|
||||
|
||||
class AdaptiveRouter:
|
||||
"""One instance per router_name. Holds in-memory caches + the update queue."""
|
||||
|
||||
|
|
@ -77,8 +88,8 @@ class AdaptiveRouter:
|
|||
self,
|
||||
router_name: str,
|
||||
config: AdaptiveRouterConfig,
|
||||
model_to_prefs: Dict[str, AdaptiveRouterPreferences],
|
||||
model_to_cost: Dict[str, float],
|
||||
model_to_prefs: dict[str, AdaptiveRouterPreferences],
|
||||
model_to_cost: dict[str, float],
|
||||
) -> None:
|
||||
self.router_name = router_name
|
||||
self.config = config
|
||||
|
|
@ -86,13 +97,14 @@ class AdaptiveRouter:
|
|||
self.model_to_cost = model_to_cost
|
||||
self.queue = AdaptiveRouterUpdateQueue()
|
||||
|
||||
self._cells: Dict[Tuple[RequestType, str], BanditCell] = {}
|
||||
self._owner_cache: Dict[str, Tuple[str, float]] = {}
|
||||
self._session_states: Dict[Tuple[str, str], SessionState] = {}
|
||||
# Parallel expiry map for _session_states, same TTL as _owner_cache.
|
||||
# Evicted opportunistically in `get_or_create_session_state`.
|
||||
self._session_states_expiry: Dict[Tuple[str, str], float] = {}
|
||||
self._skipped_updates_total: int = 0
|
||||
self._cells: dict[tuple[RequestType, str], BanditCell] = {}
|
||||
self._session_states: dict[tuple[str, str], SessionState] = {}
|
||||
self._feedback_contexts: OrderedDict[str, _FeedbackContext] = OrderedDict()
|
||||
self._session_states_expiry: dict[tuple[str, str], float] = {}
|
||||
self._feedback_attributed_total: int = 0
|
||||
self._feedback_without_context_total: int = 0
|
||||
self._cross_model_feedback_total: int = 0
|
||||
self._response_signal_updates_total: int = 0
|
||||
# Set to True once the proxy flusher has loaded persisted priors from
|
||||
# Postgres. Checked to support lazy-load on hot-reloaded routers.
|
||||
self._state_loaded: bool = False
|
||||
|
|
@ -145,11 +157,11 @@ class AdaptiveRouter:
|
|||
async def async_pre_routing_hook(
|
||||
self,
|
||||
model: str,
|
||||
request_kwargs: Dict[str, Any],
|
||||
messages: Optional[List[Dict[str, Any]]] = None,
|
||||
input: Optional[Union[str, List]] = None,
|
||||
specific_deployment: Optional[bool] = False,
|
||||
) -> Optional[PreRoutingHookResponse]:
|
||||
request_kwargs: dict[str, Any],
|
||||
messages: list[dict[str, Any]] | None = None,
|
||||
input: Union[str, list] | None = None,
|
||||
specific_deployment: bool | None = False,
|
||||
) -> PreRoutingHookResponse | None:
|
||||
"""
|
||||
Plugin entry point invoked by `Router.async_pre_routing_hook` when the
|
||||
inbound `model` matches this adaptive router's `router_name`.
|
||||
|
|
@ -159,11 +171,9 @@ class AdaptiveRouter:
|
|||
post-call hook can surface it as a response header.
|
||||
|
||||
Routing is stateless per-turn: every call Thompson-samples fresh,
|
||||
regardless of any prior pick for the same session. Cross-turn
|
||||
attribution is enforced post-call via the owner cache (see
|
||||
`claim_or_check_owner`).
|
||||
regardless of any prior pick for the same session.
|
||||
"""
|
||||
user_text = get_last_user_message(cast(List[AllMessageValues], messages or [])) or ""
|
||||
user_text = get_last_user_message(cast(list[AllMessageValues], messages or [])) or ""
|
||||
|
||||
request_type = classify_prompt(user_text)
|
||||
min_quality_tier = self._extract_min_quality_tier(request_kwargs)
|
||||
|
|
@ -190,7 +200,7 @@ class AdaptiveRouter:
|
|||
async def pick_model(
|
||||
self,
|
||||
request_type: RequestType,
|
||||
min_quality_tier: Optional[int] = None,
|
||||
min_quality_tier: int | None = None,
|
||||
) -> str:
|
||||
"""Thompson-sample across eligible models. Stateless per-turn."""
|
||||
eligible = self._eligible_models(min_quality_tier)
|
||||
|
|
@ -206,44 +216,7 @@ class AdaptiveRouter:
|
|||
cost_weight=self.config.weights.cost,
|
||||
)
|
||||
|
||||
def claim_or_check_owner(self, session_key: str, current_model: str) -> bool:
|
||||
"""Resolve attribution for a turn under stateless routing.
|
||||
|
||||
Returns True iff this turn should fire a bandit/state update. The
|
||||
first call for a `session_key` claims ownership for `current_model`
|
||||
and returns True. Subsequent calls return True only if the owner is
|
||||
still live AND matches `current_model`. Mismatches (a different
|
||||
model handled this turn) and expired owners both increment
|
||||
`_skipped_updates_total` and return False — no attribution.
|
||||
"""
|
||||
now = time.time()
|
||||
existing = self._owner_cache.get(session_key)
|
||||
if existing is not None and existing[1] > now:
|
||||
owner_model, _ = existing
|
||||
if owner_model == current_model:
|
||||
return True
|
||||
self._skipped_updates_total += 1
|
||||
return False
|
||||
|
||||
# Opportunistic bulk sweep — sessions that never come back would
|
||||
# otherwise pile up here forever. Same threshold pattern as the
|
||||
# session-state cache.
|
||||
if len(self._owner_cache) >= _OWNER_CACHE_SWEEP_THRESHOLD:
|
||||
self._evict_expired_owner_cache(now)
|
||||
|
||||
# No live owner -> claim for current_model.
|
||||
self._owner_cache[session_key] = (
|
||||
current_model,
|
||||
now + OWNER_CACHE_TTL_SECONDS,
|
||||
)
|
||||
return True
|
||||
|
||||
def _evict_expired_owner_cache(self, now: float) -> None:
|
||||
expired = [k for k, (_, exp) in self._owner_cache.items() if exp <= now]
|
||||
for k in expired:
|
||||
self._owner_cache.pop(k, None)
|
||||
|
||||
async def get_state_snapshot(self) -> Dict[str, Any]:
|
||||
async def get_state_snapshot(self) -> dict[str, Any]:
|
||||
"""In-memory snapshot for the introspection endpoint. Cheap; no DB hit."""
|
||||
cells = []
|
||||
for (rt, model), cell in sorted(self._cells.items(), key=lambda kv: (kv[0][0].value, kv[0][1])):
|
||||
|
|
@ -264,7 +237,7 @@ class AdaptiveRouter:
|
|||
)
|
||||
queue = await self.queue.queue_size()
|
||||
now = time.time()
|
||||
owner_cache_live = sum(1 for _, exp in self._owner_cache.values() if exp > now)
|
||||
feedback_contexts_live = sum(1 for context in self._feedback_contexts.values() if context.expires_at > now)
|
||||
return {
|
||||
"router_name": self.router_name,
|
||||
"available_models": list(self.config.available_models),
|
||||
|
|
@ -274,15 +247,18 @@ class AdaptiveRouter:
|
|||
},
|
||||
"model_costs": dict(self.model_to_cost),
|
||||
"cells": cells,
|
||||
"owner_cache_live": owner_cache_live,
|
||||
"skipped_updates_total": self._skipped_updates_total,
|
||||
"feedback_contexts_live": feedback_contexts_live,
|
||||
"feedback_attributed_total": self._feedback_attributed_total,
|
||||
"feedback_without_context_total": self._feedback_without_context_total,
|
||||
"cross_model_feedback_total": self._cross_model_feedback_total,
|
||||
"response_signal_updates_total": self._response_signal_updates_total,
|
||||
"queue": queue,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _extract_min_quality_tier(
|
||||
request_kwargs: Dict[str, Any],
|
||||
) -> Optional[int]:
|
||||
request_kwargs: dict[str, Any],
|
||||
) -> int | None:
|
||||
"""Pull `min_quality_tier` from request headers or metadata.
|
||||
|
||||
Precedence: headers (`x-litellm-min-quality-tier`) over metadata
|
||||
|
|
@ -310,7 +286,7 @@ class AdaptiveRouter:
|
|||
return None
|
||||
return None
|
||||
|
||||
def _eligible_models(self, min_quality_tier: Optional[int]) -> List[str]:
|
||||
def _eligible_models(self, min_quality_tier: int | None) -> list[str]:
|
||||
if min_quality_tier is None:
|
||||
return list(self.config.available_models)
|
||||
return [
|
||||
|
|
@ -363,17 +339,131 @@ class AdaptiveRouter:
|
|||
request_type: RequestType,
|
||||
turn: Turn,
|
||||
) -> SignalDelta:
|
||||
"""Apply one turn, push session snapshot + bandit deltas to the queue."""
|
||||
state = self.get_or_create_session_state(session_id, model_name, request_type)
|
||||
delta = apply_turn(state, turn)
|
||||
verbose_router_logger.debug("AdaptiveRouter[%s]: record_turn delta=%s", self.router_name, delta)
|
||||
"""Attribute feedback to the previous response and response signals to the current model."""
|
||||
async with self._lock:
|
||||
now = time.time()
|
||||
while self._feedback_contexts:
|
||||
oldest_context = next(iter(self._feedback_contexts.values()))
|
||||
if oldest_context.expires_at > now:
|
||||
break
|
||||
self._feedback_contexts.popitem(last=False)
|
||||
previous = self._feedback_contexts.pop(session_id, None)
|
||||
|
||||
# Strip the raw conversation content before persisting. The
|
||||
# last_user/assistant_content and tool_call_history fields are only
|
||||
# needed in-memory for the next turn's incremental signal detection;
|
||||
# writing user prompts and tool payloads to the DB would store PII
|
||||
# for every adaptive-router conversation. Counts + bookkeeping is
|
||||
# all the persisted row needs.
|
||||
effective_request_type = (
|
||||
previous.request_type if previous is not None and request_type == RequestType.GENERAL else request_type
|
||||
)
|
||||
current_state = self.get_or_create_session_state(
|
||||
session_id,
|
||||
model_name,
|
||||
effective_request_type,
|
||||
)
|
||||
feedback_delta = detect_user_feedback(
|
||||
previous.user_content if previous else None,
|
||||
turn.user_content,
|
||||
turn.tool_results,
|
||||
allow_satisfaction=(
|
||||
previous is not None
|
||||
and not previous.clean_credit_awarded
|
||||
and previous.turn_count + 1 >= MIN_TURNS_FOR_CLEAN_CREDIT
|
||||
),
|
||||
)
|
||||
previous_assistant = previous.assistant_content if previous else None
|
||||
response_delta = detect_response_signals(
|
||||
previous_assistant,
|
||||
turn.assistant_content,
|
||||
current_state.tool_call_history,
|
||||
turn.tool_calls,
|
||||
turn.tool_results,
|
||||
turn.response_status,
|
||||
)
|
||||
states_to_persist: dict[str, SessionState] = {model_name: current_state}
|
||||
bandit_deltas: dict[tuple[RequestType, str], SignalDelta] = {}
|
||||
|
||||
if previous is not None:
|
||||
feedback_state = self.get_or_create_session_state(
|
||||
session_id,
|
||||
previous.model_name,
|
||||
previous.request_type,
|
||||
)
|
||||
apply_signal_delta(feedback_state, feedback_delta)
|
||||
if feedback_delta.satisfaction:
|
||||
feedback_state.clean_credit_awarded = True
|
||||
states_to_persist[previous.model_name] = feedback_state
|
||||
if feedback_delta.any_fired():
|
||||
self._feedback_attributed_total += 1
|
||||
if previous.model_name != model_name:
|
||||
self._cross_model_feedback_total += 1
|
||||
bandit_deltas[(previous.request_type, previous.model_name)] = feedback_delta
|
||||
else:
|
||||
if feedback_delta.any_fired():
|
||||
self._feedback_without_context_total += 1
|
||||
initial_failure = SignalDelta(failure=feedback_delta.failure)
|
||||
apply_signal_delta(current_state, initial_failure)
|
||||
bandit_deltas[(effective_request_type, model_name)] = initial_failure
|
||||
|
||||
apply_signal_delta(current_state, response_delta)
|
||||
if self._compute_bandit_delta(response_delta) != (0.0, 0.0):
|
||||
self._response_signal_updates_total += 1
|
||||
current_key = (effective_request_type, model_name)
|
||||
bandit_deltas[current_key] = merge_signal_deltas(
|
||||
bandit_deltas.get(current_key, SignalDelta()),
|
||||
response_delta,
|
||||
)
|
||||
advance_session_state(current_state, turn)
|
||||
|
||||
next_turn_count = (previous.turn_count if previous else 0) + 1
|
||||
clean_credit_awarded = bool((previous and previous.clean_credit_awarded) or feedback_delta.satisfaction)
|
||||
if len(self._feedback_contexts) >= _FEEDBACK_CONTEXT_MAX_ENTRIES:
|
||||
self._feedback_contexts.popitem(last=False)
|
||||
self._feedback_contexts[session_id] = _FeedbackContext(
|
||||
model_name=model_name,
|
||||
request_type=effective_request_type,
|
||||
user_content=turn.user_content,
|
||||
assistant_content=turn.assistant_content,
|
||||
turn_count=next_turn_count,
|
||||
clean_credit_awarded=clean_credit_awarded,
|
||||
expires_at=now + OWNER_CACHE_TTL_SECONDS,
|
||||
)
|
||||
|
||||
for state_model, state in states_to_persist.items():
|
||||
await self.queue.add_session_state(
|
||||
session_id,
|
||||
self.router_name,
|
||||
state_model,
|
||||
self._persistable_session_snapshot(state),
|
||||
)
|
||||
|
||||
combined_delta = SignalDelta()
|
||||
for (attribution_type, target_model), delta in bandit_deltas.items():
|
||||
combined_delta = merge_signal_deltas(combined_delta, delta)
|
||||
d_alpha, d_beta = self._compute_bandit_delta(delta)
|
||||
if d_alpha == 0 and d_beta == 0:
|
||||
continue
|
||||
cell_key = (attribution_type, target_model)
|
||||
self._cells[cell_key] = apply_delta(
|
||||
self._cells[cell_key],
|
||||
d_alpha,
|
||||
d_beta,
|
||||
)
|
||||
await self.queue.add_state_delta(
|
||||
self.router_name,
|
||||
attribution_type.value,
|
||||
target_model,
|
||||
d_alpha,
|
||||
d_beta,
|
||||
)
|
||||
|
||||
verbose_router_logger.debug(
|
||||
"AdaptiveRouter[%s]: feedback_target=%s current_model=%s delta=%s",
|
||||
self.router_name,
|
||||
previous.model_name if previous else None,
|
||||
model_name,
|
||||
combined_delta,
|
||||
)
|
||||
return combined_delta
|
||||
|
||||
@staticmethod
|
||||
def _persistable_session_snapshot(state: SessionState) -> dict[str, Any]:
|
||||
snapshot = asdict(state)
|
||||
for sensitive in (
|
||||
"last_user_content",
|
||||
|
|
@ -382,38 +472,10 @@ class AdaptiveRouter:
|
|||
"pending_tool_calls",
|
||||
):
|
||||
snapshot.pop(sensitive, None)
|
||||
await self.queue.add_session_state(session_id, self.router_name, model_name, snapshot)
|
||||
|
||||
d_alpha, d_beta = self._compute_bandit_delta(delta)
|
||||
verbose_router_logger.debug(
|
||||
"AdaptiveRouter[%s]: bandit delta alpha=%.2f beta=%.2f",
|
||||
self.router_name,
|
||||
d_alpha,
|
||||
d_beta,
|
||||
)
|
||||
if d_alpha != 0 or d_beta != 0:
|
||||
# For non-GENERAL turns, attribute to the current-turn classification
|
||||
# so genuine mid-session topic shifts (e.g. code → math) update the
|
||||
# correct cell. For GENERAL turns ("thanks!", "ok", "sounds good"), fall
|
||||
# back to the session's original type so closing pleasantries don't
|
||||
# misattribute the reward.
|
||||
attribution_type = (
|
||||
request_type if request_type != RequestType.GENERAL else RequestType(state.classified_type)
|
||||
)
|
||||
cell_key = (attribution_type, model_name)
|
||||
self._cells[cell_key] = apply_delta(self._cells[cell_key], d_alpha, d_beta)
|
||||
await self.queue.add_state_delta(
|
||||
self.router_name,
|
||||
attribution_type.value,
|
||||
model_name,
|
||||
d_alpha,
|
||||
d_beta,
|
||||
)
|
||||
|
||||
return delta
|
||||
return snapshot
|
||||
|
||||
@staticmethod
|
||||
def _compute_bandit_delta(delta: SignalDelta) -> Tuple[float, float]:
|
||||
def _compute_bandit_delta(delta: SignalDelta) -> tuple[float, float]:
|
||||
"""
|
||||
Translate per-turn signal deltas into bandit-cell deltas.
|
||||
|
||||
|
|
|
|||
|
|
@ -214,10 +214,6 @@ class AdaptiveRouterPostCallHook(CustomLogger):
|
|||
) -> None:
|
||||
try:
|
||||
messages = kwargs.get("messages") or []
|
||||
if len(messages) < SIGNAL_GATE_MIN_MESSAGES:
|
||||
# Too few turns for any signal to be meaningful — skip.
|
||||
return
|
||||
|
||||
session_key = _resolve_session_key(kwargs)
|
||||
if not session_key:
|
||||
return
|
||||
|
|
@ -233,10 +229,6 @@ class AdaptiveRouterPostCallHook(CustomLogger):
|
|||
if not current_model:
|
||||
return
|
||||
|
||||
if not self.adaptive_router.claim_or_check_owner(session_key, current_model):
|
||||
# A different model owns this conversation — skip attribution.
|
||||
return
|
||||
|
||||
user_text = _last_user_content(messages)
|
||||
assistant_text, tool_calls = _assistant_content_and_tool_calls(response_obj)
|
||||
tool_results = _recent_tool_results(messages)
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from __future__ import annotations
|
|||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Set
|
||||
from typing import Any
|
||||
|
||||
from litellm.router_strategy.adaptive_router.config import (
|
||||
LOOP_REPEAT_THRESHOLD,
|
||||
|
|
@ -74,26 +74,26 @@ class SessionState:
|
|||
loop_count: int = 0
|
||||
exhaustion_count: int = 0
|
||||
|
||||
last_user_content: Optional[str] = None
|
||||
last_assistant_content: Optional[str] = None
|
||||
tool_call_history: List[str] = field(default_factory=list)
|
||||
pending_tool_calls: Dict[str, str] = field(default_factory=dict)
|
||||
last_user_content: str | None = None
|
||||
last_assistant_content: str | None = None
|
||||
tool_call_history: list[str] = field(default_factory=list)
|
||||
pending_tool_calls: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
turn_count: int = 0
|
||||
last_processed_turn: int = -1
|
||||
clean_credit_awarded: bool = False
|
||||
terminal_status: Optional[int] = None
|
||||
terminal_status: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Turn:
|
||||
"""One turn of input. Caller assembles this from the request/response."""
|
||||
|
||||
user_content: Optional[str] = None
|
||||
assistant_content: Optional[str] = None
|
||||
tool_calls: List[Dict[str, Any]] = field(default_factory=list)
|
||||
tool_results: List[Dict[str, Any]] = field(default_factory=list)
|
||||
response_status: Optional[int] = None
|
||||
user_content: str | None = None
|
||||
assistant_content: str | None = None
|
||||
tool_calls: list[dict[str, Any]] = field(default_factory=list)
|
||||
tool_results: list[dict[str, Any]] = field(default_factory=list)
|
||||
response_status: int | None = None
|
||||
|
||||
|
||||
# ---- Detection helpers ----------------------------------------------------
|
||||
|
|
@ -101,13 +101,13 @@ class Turn:
|
|||
_TOKEN_RE = re.compile(r"[A-Za-z0-9]+")
|
||||
|
||||
|
||||
def _tokens(text: Optional[str]) -> Set[str]:
|
||||
def _tokens(text: str | None) -> set[str]:
|
||||
if not text:
|
||||
return set()
|
||||
return {t.lower() for t in _TOKEN_RE.findall(text)}
|
||||
|
||||
|
||||
def _jaccard(a: Set[str], b: Set[str]) -> float:
|
||||
def _jaccard(a: set[str], b: set[str]) -> float:
|
||||
union = a | b
|
||||
if not union:
|
||||
return 0.0
|
||||
|
|
@ -130,7 +130,7 @@ _SATISFACTION_PATTERNS = [
|
|||
]
|
||||
|
||||
|
||||
def _detect_misalignment(prev_user: Optional[str], curr_user: Optional[str]) -> bool:
|
||||
def _detect_misalignment(prev_user: str | None, curr_user: str | None) -> bool:
|
||||
"""Fires when consecutive user messages share *some* topic (jaccard > 0)
|
||||
but are sufficiently different (jaccard < threshold) — i.e. user is
|
||||
rephrasing, not changing topic, not repeating."""
|
||||
|
|
@ -140,7 +140,7 @@ def _detect_misalignment(prev_user: Optional[str], curr_user: Optional[str]) ->
|
|||
return 0.0 < j < MISALIGNMENT_JACCARD_THRESHOLD
|
||||
|
||||
|
||||
def _detect_stagnation(prev_asst: Optional[str], curr_asst: Optional[str]) -> bool:
|
||||
def _detect_stagnation(prev_asst: str | None, curr_asst: str | None) -> bool:
|
||||
"""Fires when consecutive assistant messages are near-duplicates."""
|
||||
if not prev_asst or not curr_asst:
|
||||
return False
|
||||
|
|
@ -148,19 +148,19 @@ def _detect_stagnation(prev_asst: Optional[str], curr_asst: Optional[str]) -> bo
|
|||
return j >= STAGNATION_JACCARD_NEAR_DUP
|
||||
|
||||
|
||||
def _detect_disengagement(curr_user: Optional[str]) -> bool:
|
||||
def _detect_disengagement(curr_user: str | None) -> bool:
|
||||
if not curr_user:
|
||||
return False
|
||||
return any(p.search(curr_user) for p in _DISENGAGEMENT_PATTERNS)
|
||||
|
||||
|
||||
def _detect_satisfaction(curr_user: Optional[str]) -> bool:
|
||||
def _detect_satisfaction(curr_user: str | None) -> bool:
|
||||
if not curr_user:
|
||||
return False
|
||||
return any(p.search(curr_user) for p in _SATISFACTION_PATTERNS)
|
||||
|
||||
|
||||
def _detect_failure(tool_results: List[Dict[str, Any]]) -> bool:
|
||||
def _detect_failure(tool_results: list[dict[str, Any]]) -> bool:
|
||||
"""Any tool result explicitly flagged as an error.
|
||||
|
||||
We do NOT treat empty content as failure — many tools legitimately return
|
||||
|
|
@ -173,7 +173,7 @@ def _detect_failure(tool_results: List[Dict[str, Any]]) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _signature(call: Dict[str, Any]) -> str:
|
||||
def _signature(call: dict[str, Any]) -> str:
|
||||
"""Stable signature for loop detection: name + sorted JSON-ish args."""
|
||||
name = call.get("name") or call.get("function", {}).get("name", "")
|
||||
call_args = call.get("arguments")
|
||||
|
|
@ -184,7 +184,7 @@ def _signature(call: Dict[str, Any]) -> str:
|
|||
return f"{name}({call_args})"
|
||||
|
||||
|
||||
def _detect_loop(history: List[str], new_calls: List[Dict[str, Any]]) -> bool:
|
||||
def _detect_loop(history: list[str], new_calls: list[dict[str, Any]]) -> bool:
|
||||
"""Fires if any new call's signature appears >= LOOP_REPEAT_THRESHOLD-1 times
|
||||
in recent history (so this call would be the Nth)."""
|
||||
if not new_calls:
|
||||
|
|
@ -209,7 +209,7 @@ _EXHAUSTION_KEYWORDS = (
|
|||
)
|
||||
|
||||
|
||||
def _detect_exhaustion(status: Optional[int], tool_results: List[Dict[str, Any]]) -> bool:
|
||||
def _detect_exhaustion(status: int | None, tool_results: list[dict[str, Any]]) -> bool:
|
||||
if status is not None and status in _EXHAUSTION_STATUSES:
|
||||
return True
|
||||
for r in tool_results:
|
||||
|
|
@ -219,39 +219,53 @@ def _detect_exhaustion(status: Optional[int], tool_results: List[Dict[str, Any]]
|
|||
return False
|
||||
|
||||
|
||||
# ---- Public entrypoint ----------------------------------------------------
|
||||
def detect_user_feedback(
|
||||
previous_user_content: str | None,
|
||||
current_user_content: str | None,
|
||||
tool_results: list[dict[str, Any]],
|
||||
allow_satisfaction: bool,
|
||||
) -> SignalDelta:
|
||||
return SignalDelta(
|
||||
misalignment=int(_detect_misalignment(previous_user_content, current_user_content)),
|
||||
disengagement=int(_detect_disengagement(current_user_content)),
|
||||
satisfaction=int(allow_satisfaction and _detect_satisfaction(current_user_content)),
|
||||
failure=int(_detect_failure(tool_results)),
|
||||
)
|
||||
|
||||
|
||||
def apply_turn(state: SessionState, turn: Turn) -> SignalDelta:
|
||||
"""
|
||||
Detect signals on this turn, mutate state, return the delta.
|
||||
def detect_response_signals(
|
||||
previous_assistant_content: str | None,
|
||||
current_assistant_content: str | None,
|
||||
tool_call_history: list[str],
|
||||
tool_calls: list[dict[str, Any]],
|
||||
tool_results: list[dict[str, Any]],
|
||||
response_status: int | None,
|
||||
) -> SignalDelta:
|
||||
return SignalDelta(
|
||||
stagnation=int(
|
||||
_detect_stagnation(
|
||||
previous_assistant_content,
|
||||
current_assistant_content,
|
||||
)
|
||||
),
|
||||
loop=int(_detect_loop(tool_call_history, tool_calls)),
|
||||
exhaustion=int(_detect_exhaustion(response_status, tool_results)),
|
||||
)
|
||||
|
||||
O(1) per turn (no full-history rescan). Only inspects last_*, recent tool history
|
||||
(which is bounded at TOOL_CALL_HISTORY_MAX), and the new turn payload.
|
||||
"""
|
||||
delta = SignalDelta()
|
||||
|
||||
if _detect_misalignment(state.last_user_content, turn.user_content):
|
||||
delta.misalignment = 1
|
||||
if _detect_stagnation(state.last_assistant_content, turn.assistant_content):
|
||||
delta.stagnation = 1
|
||||
if _detect_disengagement(turn.user_content):
|
||||
delta.disengagement = 1
|
||||
if _detect_satisfaction(turn.user_content):
|
||||
# Gate: only award satisfaction credit once per session, and only
|
||||
# after MIN_TURNS_FOR_CLEAN_CREDIT turns of context. Early "thanks"
|
||||
# on turn 1-2 is noise, not a validated quality signal.
|
||||
current_turn_index = state.turn_count + 1
|
||||
if not state.clean_credit_awarded and current_turn_index >= MIN_TURNS_FOR_CLEAN_CREDIT:
|
||||
delta.satisfaction = 1
|
||||
state.clean_credit_awarded = True
|
||||
if _detect_failure(turn.tool_results):
|
||||
delta.failure = 1
|
||||
if _detect_loop(state.tool_call_history, turn.tool_calls):
|
||||
delta.loop = 1
|
||||
if _detect_exhaustion(turn.response_status, turn.tool_results):
|
||||
delta.exhaustion = 1
|
||||
def merge_signal_deltas(*deltas: SignalDelta) -> SignalDelta:
|
||||
return SignalDelta(
|
||||
misalignment=sum(delta.misalignment for delta in deltas),
|
||||
stagnation=sum(delta.stagnation for delta in deltas),
|
||||
disengagement=sum(delta.disengagement for delta in deltas),
|
||||
satisfaction=sum(delta.satisfaction for delta in deltas),
|
||||
failure=sum(delta.failure for delta in deltas),
|
||||
loop=sum(delta.loop for delta in deltas),
|
||||
exhaustion=sum(delta.exhaustion for delta in deltas),
|
||||
)
|
||||
|
||||
|
||||
def apply_signal_delta(state: SessionState, delta: SignalDelta) -> None:
|
||||
state.misalignment_count += delta.misalignment
|
||||
state.stagnation_count += delta.stagnation
|
||||
state.disengagement_count += delta.disengagement
|
||||
|
|
@ -260,6 +274,8 @@ def apply_turn(state: SessionState, turn: Turn) -> SignalDelta:
|
|||
state.loop_count += delta.loop
|
||||
state.exhaustion_count += delta.exhaustion
|
||||
|
||||
|
||||
def advance_session_state(state: SessionState, turn: Turn) -> None:
|
||||
if turn.user_content:
|
||||
state.last_user_content = turn.user_content
|
||||
if turn.assistant_content:
|
||||
|
|
@ -276,4 +292,38 @@ def apply_turn(state: SessionState, turn: Turn) -> SignalDelta:
|
|||
state.turn_count += 1
|
||||
state.last_processed_turn = state.turn_count
|
||||
|
||||
|
||||
# ---- Public entrypoint ----------------------------------------------------
|
||||
|
||||
|
||||
def apply_turn(state: SessionState, turn: Turn) -> SignalDelta:
|
||||
"""
|
||||
Detect signals on this turn, mutate state, return the delta.
|
||||
|
||||
O(1) per turn (no full-history rescan). Only inspects last_*, recent tool history
|
||||
(which is bounded at TOOL_CALL_HISTORY_MAX), and the new turn payload.
|
||||
"""
|
||||
feedback_delta = detect_user_feedback(
|
||||
state.last_user_content,
|
||||
turn.user_content,
|
||||
turn.tool_results,
|
||||
allow_satisfaction=(not state.clean_credit_awarded and state.turn_count + 1 >= MIN_TURNS_FOR_CLEAN_CREDIT),
|
||||
)
|
||||
response_delta = detect_response_signals(
|
||||
state.last_assistant_content,
|
||||
turn.assistant_content,
|
||||
state.tool_call_history,
|
||||
turn.tool_calls,
|
||||
turn.tool_results,
|
||||
turn.response_status,
|
||||
)
|
||||
delta = merge_signal_deltas(
|
||||
feedback_delta,
|
||||
response_delta,
|
||||
)
|
||||
apply_signal_delta(state, delta)
|
||||
if delta.satisfaction:
|
||||
state.clean_credit_awarded = True
|
||||
advance_session_state(state, turn)
|
||||
|
||||
return delta
|
||||
|
|
|
|||
|
|
@ -13,9 +13,12 @@ evaluated before either classification strategy and force a tier outright when m
|
|||
Inspired by ClawRouter: https://github.com/BlockRunAI/ClawRouter
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import random
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Literal, Union, cast
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -37,6 +40,7 @@ if TYPE_CHECKING:
|
|||
from semantic_router.routers import SemanticRouter
|
||||
|
||||
from litellm.router import Router
|
||||
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
else:
|
||||
Router = Any
|
||||
|
|
@ -62,7 +66,7 @@ Tiers:
|
|||
{prompt}"""
|
||||
|
||||
|
||||
def _append_custom_keywords(base_keywords: list[str], custom_keywords: Optional[list[str]]) -> list[str]:
|
||||
def _append_custom_keywords(base_keywords: list[str], custom_keywords: list[str] | None) -> list[str]:
|
||||
if not custom_keywords:
|
||||
return base_keywords
|
||||
base_lowered = frozenset(keyword.lower() for keyword in base_keywords)
|
||||
|
|
@ -94,7 +98,7 @@ def _sanitize_user_api_key_auth(auth: Any) -> Any:
|
|||
return auth
|
||||
|
||||
|
||||
def _classifier_call_metadata(metadata: Optional[dict[str, Any]]) -> Optional[dict[str, Any]]:
|
||||
def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any] | None:
|
||||
if not metadata:
|
||||
return metadata
|
||||
return {
|
||||
|
|
@ -109,7 +113,7 @@ class DimensionScore:
|
|||
|
||||
__slots__ = ("name", "score", "signal")
|
||||
|
||||
def __init__(self, name: str, score: float, signal: Optional[str] = None):
|
||||
def __init__(self, name: str, score: float, signal: str | None = None):
|
||||
self.name = name
|
||||
self.score = score
|
||||
self.signal = signal
|
||||
|
|
@ -133,9 +137,9 @@ class ComplexityRouter(CustomLogger):
|
|||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
litellm_router_instance: "Router",
|
||||
complexity_router_config: Optional[Dict[str, Any]] = None,
|
||||
default_model: Optional[str] = None,
|
||||
litellm_router_instance: Router,
|
||||
complexity_router_config: dict[str, Any] | None = None,
|
||||
default_model: str | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize ComplexityRouter.
|
||||
|
|
@ -172,7 +176,7 @@ class ComplexityRouter(CustomLogger):
|
|||
# embeddings are static, only the prompt is embedded per request). The lock
|
||||
# serializes the one-time build so concurrent cold-start requests don't each
|
||||
# construct the index and fire duplicate embedding calls.
|
||||
self._semantic_routelayer: Optional[SemanticRouter] = None
|
||||
self._semantic_routelayer: SemanticRouter | None = None
|
||||
self._semantic_routelayer_lock = asyncio.Lock()
|
||||
|
||||
# Pre-compile regex patterns for efficiency
|
||||
|
|
@ -184,6 +188,10 @@ class ComplexityRouter(CustomLogger):
|
|||
re.compile(r"[a-z]\)\s", re.IGNORECASE),
|
||||
]
|
||||
|
||||
self.adaptive_router: AdaptiveRouter | None = None
|
||||
self._model_tiers: dict[str, tuple[ComplexityTier, ...]] = {}
|
||||
self._adaptive_init_attempted = False
|
||||
|
||||
verbose_router_logger.debug(f"ComplexityRouter initialized for {model_name} with tiers: {self.config.tiers}")
|
||||
|
||||
def _estimate_tokens(self, text: str) -> int:
|
||||
|
|
@ -227,12 +235,12 @@ class ComplexityRouter(CustomLogger):
|
|||
def _score_keyword_match(
|
||||
self,
|
||||
text: str,
|
||||
keywords: List[str],
|
||||
keywords: list[str],
|
||||
name: str,
|
||||
signal_label: str,
|
||||
thresholds: Tuple[int, int], # (low, high)
|
||||
scores: Tuple[float, float, float], # (none, low, high)
|
||||
) -> Tuple[DimensionScore, int]:
|
||||
thresholds: tuple[int, int], # (low, high)
|
||||
scores: tuple[float, float, float], # (none, low, high)
|
||||
) -> tuple[DimensionScore, int]:
|
||||
"""Score based on keyword matches using word boundary matching.
|
||||
|
||||
Returns:
|
||||
|
|
@ -270,7 +278,7 @@ class ComplexityRouter(CustomLogger):
|
|||
return DimensionScore("questionComplexity", 0.5, f"{count} questions")
|
||||
return DimensionScore("questionComplexity", 0, None)
|
||||
|
||||
def classify(self, prompt: str, system_prompt: Optional[str] = None) -> Tuple[ComplexityTier, float, List[str]]:
|
||||
def classify(self, prompt: str, system_prompt: str | None = None) -> tuple[ComplexityTier, float, list[str]]:
|
||||
"""
|
||||
Classify a prompt by complexity.
|
||||
|
||||
|
|
@ -329,7 +337,7 @@ class ComplexityRouter(CustomLogger):
|
|||
(0, -1.0, -1.0),
|
||||
)
|
||||
|
||||
dimensions: List[DimensionScore] = [
|
||||
dimensions: list[DimensionScore] = [
|
||||
self._score_token_count(estimated_tokens),
|
||||
code_score,
|
||||
reasoning_score,
|
||||
|
|
@ -371,8 +379,8 @@ class ComplexityRouter(CustomLogger):
|
|||
async def aclassify(
|
||||
self,
|
||||
prompt: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
request_kwargs: Optional[dict[str, Any]] = None,
|
||||
system_prompt: str | None = None,
|
||||
request_kwargs: dict[str, Any] | None = None,
|
||||
) -> tuple[ComplexityTier, float, list[str]]:
|
||||
"""
|
||||
Classify a prompt by complexity, using the LLM classifier when configured.
|
||||
|
|
@ -395,8 +403,8 @@ class ComplexityRouter(CustomLogger):
|
|||
async def _classify_with_llm(
|
||||
self,
|
||||
prompt: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
request_kwargs: Optional[dict[str, Any]] = None,
|
||||
system_prompt: str | None = None,
|
||||
request_kwargs: dict[str, Any] | None = None,
|
||||
) -> ComplexityTier:
|
||||
"""Call the configured classifier model and parse its structured tier response."""
|
||||
llm_config = self.config.classifier_llm_config
|
||||
|
|
@ -437,23 +445,196 @@ class ComplexityRouter(CustomLogger):
|
|||
"""
|
||||
tier_key = tier.value if isinstance(tier, ComplexityTier) else tier
|
||||
|
||||
# Check config tiers mapping
|
||||
model = self.config.tiers.get(tier_key)
|
||||
if model:
|
||||
return model
|
||||
if tier_key in self.config.tiers:
|
||||
return self._pick_from_tier_value(self.config.tiers[tier_key], tier_key)
|
||||
|
||||
# Fallback to default model if configured
|
||||
if self.config.default_model:
|
||||
return self.config.default_model
|
||||
|
||||
# Last resort: return MEDIUM tier model or error
|
||||
medium_model = self.config.tiers.get(ComplexityTier.MEDIUM.value)
|
||||
if medium_model:
|
||||
return medium_model
|
||||
medium_key = ComplexityTier.MEDIUM.value
|
||||
if medium_key in self.config.tiers:
|
||||
return self._pick_from_tier_value(self.config.tiers[medium_key], medium_key)
|
||||
|
||||
raise ValueError(f"No model configured for tier {tier_key} and no default_model set")
|
||||
|
||||
def _lexical_tier_override(self, user_message: str) -> Optional[ComplexityTier]:
|
||||
@staticmethod
|
||||
def _pick_from_tier_value(model: str | list[str], tier_key: str) -> str:
|
||||
if isinstance(model, str):
|
||||
return model
|
||||
if not model:
|
||||
raise ValueError(f"Empty model pool for tier {tier_key}")
|
||||
return random.choice(model)
|
||||
|
||||
def _tier_pools(self) -> dict[str, list[str]]:
|
||||
return {tier: (models if isinstance(models, list) else [models]) for tier, models in self.config.tiers.items()}
|
||||
|
||||
def _ensure_adaptive_router(self) -> Any | None:
|
||||
if not self.config.adaptive:
|
||||
return None
|
||||
if self.adaptive_router is not None:
|
||||
return self.adaptive_router
|
||||
if self._adaptive_init_attempted:
|
||||
return self.adaptive_router
|
||||
self._adaptive_init_attempted = True
|
||||
|
||||
from litellm.router_strategy.adaptive_router.adaptive_router import (
|
||||
AdaptiveRouter,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.config import (
|
||||
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
|
||||
)
|
||||
from litellm.types.router import (
|
||||
AdaptiveRouterConfig,
|
||||
AdaptiveRouterPreferences,
|
||||
)
|
||||
|
||||
pools = self._tier_pools()
|
||||
available_models = list(dict.fromkeys(model for models in pools.values() for model in models))
|
||||
self._model_tiers = {
|
||||
model: tuple(ComplexityTier(tier_name) for tier_name, models in pools.items() if model in models)
|
||||
for model in available_models
|
||||
}
|
||||
|
||||
model_to_prefs: dict[str, AdaptiveRouterPreferences] = {}
|
||||
model_to_cost: dict[str, float] = {}
|
||||
model_list = getattr(self.litellm_router_instance, "model_list", None) or []
|
||||
name_to_indices = getattr(self.litellm_router_instance, "model_name_to_deployment_indices", {}) or {}
|
||||
for name in available_models:
|
||||
indices = name_to_indices.get(name, [])
|
||||
if not indices:
|
||||
model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[])
|
||||
model_to_cost[name] = 0.0
|
||||
continue
|
||||
deployment = model_list[indices[0]]
|
||||
mi = deployment.get("model_info") if isinstance(deployment, dict) else deployment.model_info
|
||||
mi_dict: dict[str, Any] = mi if isinstance(mi, dict) else (mi.model_dump() if mi else {})
|
||||
prefs_raw = mi_dict.get("adaptive_router_preferences")
|
||||
if prefs_raw is not None:
|
||||
model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw)
|
||||
else:
|
||||
model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[])
|
||||
|
||||
lp = deployment.get("litellm_params") if isinstance(deployment, dict) else deployment.litellm_params
|
||||
lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {})
|
||||
cost = lp_dict.get("input_cost_per_token")
|
||||
model_to_cost[name] = float(cost) if cost is not None else 0.0
|
||||
|
||||
self.adaptive_router = AdaptiveRouter(
|
||||
router_name=self.model_name,
|
||||
config=AdaptiveRouterConfig(
|
||||
available_models=available_models,
|
||||
weights=self.config.adaptive_weights,
|
||||
),
|
||||
model_to_prefs=model_to_prefs,
|
||||
model_to_cost=model_to_cost,
|
||||
)
|
||||
self._adaptive_chosen_model_key = ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY
|
||||
return self.adaptive_router
|
||||
|
||||
def _soft_floor_pick(
|
||||
self,
|
||||
classified_tier: ComplexityTier,
|
||||
user_message: str,
|
||||
request_kwargs: dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
from litellm.router_strategy.adaptive_router.bandit import (
|
||||
normalized_cost,
|
||||
thompson_sample,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
|
||||
|
||||
adaptive = self._ensure_adaptive_router()
|
||||
if adaptive is None:
|
||||
return self.get_model_for_tier(classified_tier)
|
||||
|
||||
request_type = classify_prompt(user_message)
|
||||
classified_idx = TIER_SEVERITY_ORDER.index(classified_tier)
|
||||
pools = self._tier_pools()
|
||||
classified_candidates = tuple(pools.get(classified_tier.value, ()))
|
||||
cold_start_candidates = tuple(
|
||||
model for model in classified_candidates if adaptive._cells[(request_type, model)].total_samples == 0
|
||||
)
|
||||
if cold_start_candidates:
|
||||
chosen_model = random.choice(cold_start_candidates)
|
||||
if request_kwargs is not None:
|
||||
metadata = request_kwargs.setdefault("metadata", {})
|
||||
if isinstance(metadata, dict):
|
||||
metadata["adaptive_router_decision"] = {
|
||||
"phase": "cold_start",
|
||||
"classified_tier": classified_tier.value,
|
||||
"request_type": request_type.value,
|
||||
"eligible_mode": "classified_tier",
|
||||
"quality_weight": self.config.adaptive_weights.quality,
|
||||
"cost_weight": self.config.adaptive_weights.cost,
|
||||
"tier_distance_penalty": self.config.tier_distance_penalty,
|
||||
"chosen_model": chosen_model,
|
||||
"candidates": [
|
||||
{
|
||||
"model": model,
|
||||
"total_samples": adaptive._cells[(request_type, model)].total_samples,
|
||||
}
|
||||
for model in cold_start_candidates
|
||||
],
|
||||
}
|
||||
return chosen_model
|
||||
if self.config.adaptive_eligible == "classified_tier":
|
||||
candidates = list(classified_candidates)
|
||||
if not candidates:
|
||||
return self.get_model_for_tier(classified_tier)
|
||||
else:
|
||||
candidates = list(adaptive.config.available_models)
|
||||
|
||||
all_costs = [adaptive.model_to_cost.get(m, 0.0) for m in candidates]
|
||||
quality_weight = self.config.adaptive_weights.quality
|
||||
cost_weight = self.config.adaptive_weights.cost
|
||||
penalty_weight = self.config.tier_distance_penalty
|
||||
|
||||
best_model: str | None = None
|
||||
best_score = float("-inf")
|
||||
candidate_scores: list[dict[str, Any]] = []
|
||||
for model in candidates:
|
||||
cell = adaptive._cells[(request_type, model)]
|
||||
quality_sample = thompson_sample(cell)
|
||||
cost_score = normalized_cost(adaptive.model_to_cost.get(model, 0.0), all_costs)
|
||||
if self.config.adaptive_eligible == "classified_tier":
|
||||
distance = 0
|
||||
else:
|
||||
model_tiers = self._model_tiers.get(model, (classified_tier,))
|
||||
distance = min(
|
||||
abs(TIER_SEVERITY_ORDER.index(model_tier) - classified_idx) for model_tier in model_tiers
|
||||
)
|
||||
score = quality_weight * quality_sample + cost_weight * cost_score - penalty_weight * distance
|
||||
candidate_scores.append(
|
||||
{
|
||||
"model": model,
|
||||
"quality_sample": quality_sample,
|
||||
"cost_score": cost_score,
|
||||
"tier_distance": distance,
|
||||
"score": score,
|
||||
}
|
||||
)
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_model = model
|
||||
if best_model is None:
|
||||
return self.get_model_for_tier(classified_tier)
|
||||
if request_kwargs is not None:
|
||||
metadata = request_kwargs.setdefault("metadata", {})
|
||||
if isinstance(metadata, dict):
|
||||
metadata["adaptive_router_decision"] = {
|
||||
"phase": "adaptive",
|
||||
"classified_tier": classified_tier.value,
|
||||
"request_type": request_type.value,
|
||||
"eligible_mode": self.config.adaptive_eligible,
|
||||
"quality_weight": quality_weight,
|
||||
"cost_weight": cost_weight,
|
||||
"tier_distance_penalty": penalty_weight,
|
||||
"chosen_model": best_model,
|
||||
"candidates": candidate_scores,
|
||||
}
|
||||
return best_model
|
||||
|
||||
def _lexical_tier_override(self, user_message: str) -> ComplexityTier | None:
|
||||
"""When keyword_tier_rules match literally, the most-severe matched tier wins.
|
||||
|
||||
Escalating to the highest tier (rather than the first rule in the list) keeps
|
||||
|
|
@ -471,7 +652,7 @@ class ComplexityRouter(CustomLogger):
|
|||
return None
|
||||
return max(matched_tiers, key=TIER_SEVERITY_ORDER.index)
|
||||
|
||||
def _get_or_create_semantic_routelayer(self) -> "SemanticRouter":
|
||||
def _get_or_create_semantic_routelayer(self) -> SemanticRouter:
|
||||
"""Build (once) a SemanticRouter with one route per tier, utterances = that tier's keywords."""
|
||||
if self._semantic_routelayer is not None:
|
||||
return self._semantic_routelayer
|
||||
|
|
@ -510,7 +691,7 @@ class ComplexityRouter(CustomLogger):
|
|||
self._semantic_routelayer = routelayer
|
||||
return routelayer
|
||||
|
||||
async def _ensure_semantic_routelayer(self) -> "SemanticRouter":
|
||||
async def _ensure_semantic_routelayer(self) -> SemanticRouter:
|
||||
"""Return the cached route layer, building it once under a lock if needed.
|
||||
|
||||
The build embeds the static route utterances via the encoder's synchronous path,
|
||||
|
|
@ -526,7 +707,7 @@ class ComplexityRouter(CustomLogger):
|
|||
routelayer = await asyncio.to_thread(self._get_or_create_semantic_routelayer)
|
||||
return routelayer
|
||||
|
||||
async def _semantic_tier_override(self, user_message: str, request_kwargs: Dict) -> Optional[ComplexityTier]:
|
||||
async def _semantic_tier_override(self, user_message: str, request_kwargs: dict) -> ComplexityTier | None:
|
||||
"""Match the prompt against keyword_tier_rules by embedding similarity.
|
||||
|
||||
Embeds the query ourselves (instead of letting SemanticRouter.acall embed it
|
||||
|
|
@ -566,7 +747,7 @@ class ComplexityRouter(CustomLogger):
|
|||
except ValueError:
|
||||
return None
|
||||
|
||||
async def _resolve_keyword_tier_override(self, user_message: str, request_kwargs: Dict) -> Optional[ComplexityTier]:
|
||||
async def _resolve_keyword_tier_override(self, user_message: str, request_kwargs: dict) -> ComplexityTier | None:
|
||||
"""Resolve a keyword_tier_rule override, semantically or lexically per config.
|
||||
|
||||
Returns None (no override -> fall through to the scorer) not only when no rule
|
||||
|
|
@ -587,60 +768,28 @@ class ComplexityRouter(CustomLogger):
|
|||
|
||||
def _resolve_messages(
|
||||
self,
|
||||
messages: Optional[List[Dict[str, Any]]],
|
||||
request_kwargs: Dict,
|
||||
) -> Optional[List[Dict[str, Any]]]:
|
||||
messages: list[dict[str, Any]] | None,
|
||||
request_kwargs: dict,
|
||||
) -> list[dict[str, Any]] | None:
|
||||
"""
|
||||
Resolve messages from the request, converting from other formats if needed.
|
||||
|
||||
Uses the guardrail translation handler dispatch to convert Responses API
|
||||
``input`` (or other non-chat-completions formats) into OpenAI-spec messages.
|
||||
"""
|
||||
if messages:
|
||||
return messages
|
||||
|
||||
from litellm.litellm_core_utils.api_route_to_call_types import (
|
||||
get_call_types_for_route,
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
resolve_structured_messages,
|
||||
)
|
||||
from litellm.llms import load_guardrail_translation_mappings
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
mappings = load_guardrail_translation_mappings()
|
||||
call_type: Optional[CallTypes] = None
|
||||
|
||||
# 1. Try route-based inference from proxy metadata
|
||||
route = request_kwargs.get("litellm_metadata", {}).get("user_api_key_request_route")
|
||||
if route:
|
||||
call_types_list = get_call_types_for_route(route)
|
||||
if call_types_list:
|
||||
for ct in call_types_list:
|
||||
if ct in mappings:
|
||||
call_type = ct
|
||||
break
|
||||
|
||||
# 2. Fallback: try each mapped handler until one produces messages
|
||||
handlers_to_try: List[Any] = []
|
||||
if call_type is not None and call_type in mappings:
|
||||
handlers_to_try.append(mappings[call_type]())
|
||||
else:
|
||||
handlers_to_try.extend(handler_cls() for handler_cls in mappings.values())
|
||||
|
||||
for handler in handlers_to_try:
|
||||
structured = handler.get_structured_messages(request_kwargs)
|
||||
if structured:
|
||||
return [
|
||||
msg if isinstance(msg, dict) else msg.model_dump() # type: ignore
|
||||
for msg in structured
|
||||
]
|
||||
return None
|
||||
return resolve_structured_messages(messages=messages, request_kwargs=request_kwargs)
|
||||
|
||||
@staticmethod
|
||||
def _extract_user_message_and_system_prompt(
|
||||
messages: List[Dict[str, Any]],
|
||||
) -> Tuple[Optional[str], Optional[str]]:
|
||||
messages: list[dict[str, Any]],
|
||||
) -> tuple[str | None, str | None]:
|
||||
"""Extract the last user message text and last system prompt from messages."""
|
||||
user_message: Optional[str] = None
|
||||
system_prompt: Optional[str] = None
|
||||
user_message: str | None = None
|
||||
system_prompt: str | None = None
|
||||
|
||||
for msg in reversed(messages):
|
||||
role = msg.get("role", "")
|
||||
|
|
@ -660,17 +809,115 @@ class ComplexityRouter(CustomLogger):
|
|||
|
||||
return user_message, system_prompt
|
||||
|
||||
@staticmethod
|
||||
def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]:
|
||||
"""Metadata may land on `metadata` or `litellm_metadata` depending on the
|
||||
endpoint, mirroring DeploymentAffinityCheck's precedence."""
|
||||
return [
|
||||
metadata
|
||||
for metadata_key in ("litellm_metadata", "metadata")
|
||||
if isinstance(metadata := request_kwargs.get(metadata_key), dict)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None:
|
||||
"""Resolve a client-supplied session_id."""
|
||||
for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs):
|
||||
session_id = metadata.get("session_id")
|
||||
if session_id is not None:
|
||||
return str(session_id)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_user_api_key_hash_from_request_kwargs(request_kwargs: dict) -> str | None:
|
||||
"""Resolve the proxy-derived API key hash, the same trust boundary
|
||||
DeploymentAffinityCheck uses for its own key-based affinity (not the
|
||||
client-supplied OpenAI `user` param, which isn't authenticated)."""
|
||||
for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs):
|
||||
user_key = metadata.get("user_api_key_hash")
|
||||
if user_key is not None:
|
||||
return str(user_key)
|
||||
return None
|
||||
|
||||
def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: dict) -> str:
|
||||
# Namespace by the caller's API key hash so two different callers reusing the
|
||||
# same client-supplied session_id can't poison each other's routing pin. Falls
|
||||
# back to "unscoped" only when there's no authenticated caller to scope by
|
||||
# (e.g. direct Router usage without the proxy layer).
|
||||
caller_scope = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped"
|
||||
return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}"
|
||||
|
||||
async def async_pre_routing_hook(
|
||||
self,
|
||||
model: str,
|
||||
request_kwargs: Dict,
|
||||
messages: Optional[List[Dict[str, Any]]] = None,
|
||||
input: Optional[Union[str, List]] = None,
|
||||
specific_deployment: Optional[bool] = False,
|
||||
) -> Optional["PreRoutingHookResponse"]:
|
||||
request_kwargs: dict,
|
||||
messages: list[dict[str, Any]] | None = None,
|
||||
input: Union[str, list] | None = None,
|
||||
specific_deployment: bool | None = False,
|
||||
) -> PreRoutingHookResponse | None:
|
||||
"""
|
||||
Pre-routing hook called before the routing decision.
|
||||
|
||||
When `session_affinity` is enabled and a session_id is resolvable on the request,
|
||||
pins the model chosen on the session's first turn and reuses it for every later
|
||||
turn, skipping classification entirely. Otherwise delegates to `_classify_and_route`.
|
||||
"""
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
|
||||
session_id = self._get_session_id_from_request_kwargs(request_kwargs) if self.config.session_affinity else None
|
||||
cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None
|
||||
|
||||
if cache_key is not None:
|
||||
pinned_model = await self.litellm_router_instance.cache.async_get_cache(key=cache_key)
|
||||
if isinstance(pinned_model, str):
|
||||
# Refresh the TTL on every hit so an active session doesn't lose its
|
||||
# pin mid-conversation just because it outlives the original write.
|
||||
await self.litellm_router_instance.cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=pinned_model,
|
||||
ttl=self.config.session_affinity_ttl_seconds,
|
||||
)
|
||||
if self.config.adaptive:
|
||||
from litellm.router_strategy.adaptive_router.config import (
|
||||
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
|
||||
)
|
||||
|
||||
kwargs_metadata = request_kwargs.setdefault("metadata", {})
|
||||
if isinstance(kwargs_metadata, dict):
|
||||
kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = pinned_model
|
||||
verbose_router_logger.info(
|
||||
f"ComplexityRouter: routing decision cause=session_affinity_pin, routed_model={pinned_model}"
|
||||
)
|
||||
has_original_messages = messages is not None and len(messages) > 0
|
||||
return PreRoutingHookResponse(
|
||||
model=pinned_model,
|
||||
messages=messages if has_original_messages else None,
|
||||
)
|
||||
|
||||
response = await self._classify_and_route(
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
)
|
||||
if cache_key is not None and response is not None:
|
||||
await self.litellm_router_instance.cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=response.model,
|
||||
ttl=self.config.session_affinity_ttl_seconds,
|
||||
)
|
||||
return response
|
||||
|
||||
async def _classify_and_route(
|
||||
self,
|
||||
model: str,
|
||||
request_kwargs: dict,
|
||||
messages: list[dict[str, Any]] | None = None,
|
||||
input: Union[str, list] | None = None,
|
||||
specific_deployment: bool | None = False,
|
||||
) -> PreRoutingHookResponse | None:
|
||||
"""
|
||||
Classifies the request by complexity and returns the appropriate model.
|
||||
Supports chat completions (messages), Responses API (input), and other
|
||||
formats via the guardrail translation handler dispatch.
|
||||
|
|
@ -719,12 +966,25 @@ class ComplexityRouter(CustomLogger):
|
|||
)
|
||||
|
||||
tier, score, signals = await self.aclassify(user_message, system_prompt, request_kwargs)
|
||||
routed_model = self.get_model_for_tier(tier)
|
||||
|
||||
verbose_router_logger.info(
|
||||
f"ComplexityRouter: routing decision cause=complexity_scorer, tier={tier.value}, "
|
||||
f"score={score:.3f}, signals={signals}, routed_model={routed_model}"
|
||||
)
|
||||
if self.config.adaptive:
|
||||
routed_model = self._soft_floor_pick(tier, user_message, request_kwargs)
|
||||
adaptive = self._ensure_adaptive_router()
|
||||
if adaptive is not None:
|
||||
kwargs_metadata = request_kwargs.setdefault("metadata", {})
|
||||
if isinstance(kwargs_metadata, dict):
|
||||
chosen_key = getattr(self, "_adaptive_chosen_model_key", "adaptive_router_chosen_model")
|
||||
kwargs_metadata[chosen_key] = routed_model
|
||||
verbose_router_logger.info(
|
||||
f"ComplexityRouter[adaptive]: routing decision cause=complexity_scorer, "
|
||||
f"tier={tier.value}, score={score:.3f}, "
|
||||
f"signals={signals}, routed_model={routed_model}"
|
||||
)
|
||||
else:
|
||||
routed_model = self.get_model_for_tier(tier)
|
||||
verbose_router_logger.info(
|
||||
f"ComplexityRouter: routing decision cause=complexity_scorer, tier={tier.value}, "
|
||||
f"score={score:.3f}, signals={signals}, routed_model={routed_model}"
|
||||
)
|
||||
|
||||
return PreRoutingHookResponse(
|
||||
model=routed_model,
|
||||
|
|
|
|||
|
|
@ -6,9 +6,11 @@ All values are configurable via proxy config.yaml.
|
|||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Dict, List, Literal, Optional
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
|
||||
from litellm.types.router import AdaptiveRouterWeights
|
||||
|
||||
|
||||
class ComplexityTier(str, Enum):
|
||||
|
|
@ -27,11 +29,13 @@ TIER_SEVERITY_ORDER: tuple[ComplexityTier, ...] = (
|
|||
ComplexityTier.REASONING,
|
||||
)
|
||||
|
||||
DEFAULT_TIER_DISTANCE_PENALTY: float = 0.5
|
||||
|
||||
|
||||
class KeywordTierRule(BaseModel):
|
||||
"""A deterministic override: if any keyword matches, route to this tier."""
|
||||
|
||||
keywords: List[str] = Field(
|
||||
keywords: list[str] = Field(
|
||||
min_length=1,
|
||||
description="Keywords/phrases that trigger this rule (lexical or semantic match)",
|
||||
)
|
||||
|
|
@ -56,7 +60,7 @@ class KeywordTierRule(BaseModel):
|
|||
# Note: Keywords should be full words/phrases to avoid substring false positives.
|
||||
# The matching logic uses word boundary detection for single-word keywords.
|
||||
|
||||
DEFAULT_CODE_KEYWORDS: List[str] = [
|
||||
DEFAULT_CODE_KEYWORDS: list[str] = [
|
||||
"function",
|
||||
"class",
|
||||
"def",
|
||||
|
|
@ -104,7 +108,7 @@ DEFAULT_CODE_KEYWORDS: List[str] = [
|
|||
"pull request",
|
||||
]
|
||||
|
||||
DEFAULT_REASONING_KEYWORDS: List[str] = [
|
||||
DEFAULT_REASONING_KEYWORDS: list[str] = [
|
||||
"step by step",
|
||||
"think through",
|
||||
"let's think",
|
||||
|
|
@ -126,7 +130,7 @@ DEFAULT_REASONING_KEYWORDS: List[str] = [
|
|||
"conclude",
|
||||
]
|
||||
|
||||
DEFAULT_TECHNICAL_KEYWORDS: List[str] = [
|
||||
DEFAULT_TECHNICAL_KEYWORDS: list[str] = [
|
||||
"architecture",
|
||||
"distributed",
|
||||
"scalable",
|
||||
|
|
@ -158,7 +162,7 @@ DEFAULT_TECHNICAL_KEYWORDS: List[str] = [
|
|||
# Note: "async", "kubernetes", "docker" are in DEFAULT_CODE_KEYWORDS
|
||||
]
|
||||
|
||||
DEFAULT_SIMPLE_KEYWORDS: List[str] = [
|
||||
DEFAULT_SIMPLE_KEYWORDS: list[str] = [
|
||||
"what is",
|
||||
"what's",
|
||||
"define",
|
||||
|
|
@ -191,7 +195,7 @@ DEFAULT_SIMPLE_KEYWORDS: List[str] = [
|
|||
|
||||
# ─── Default Dimension Weights ───
|
||||
|
||||
DEFAULT_DIMENSION_WEIGHTS: Dict[str, float] = {
|
||||
DEFAULT_DIMENSION_WEIGHTS: dict[str, float] = {
|
||||
"tokenCount": 0.10, # Reduced - length is less important than content
|
||||
"codePresence": 0.30, # High - code requests need capable models
|
||||
"reasoningMarkers": 0.25, # High - explicit reasoning requests
|
||||
|
|
@ -204,7 +208,7 @@ DEFAULT_DIMENSION_WEIGHTS: Dict[str, float] = {
|
|||
|
||||
# ─── Default Tier Boundaries ───
|
||||
|
||||
DEFAULT_TIER_BOUNDARIES: Dict[str, float] = {
|
||||
DEFAULT_TIER_BOUNDARIES: dict[str, float] = {
|
||||
"simple_medium": 0.15, # Lower threshold to catch more MEDIUM cases
|
||||
"medium_complex": 0.35, # Lower threshold to catch technical COMPLEX cases
|
||||
"complex_reasoning": 0.60, # Reasoning tier reserved for explicit reasoning markers
|
||||
|
|
@ -213,7 +217,7 @@ DEFAULT_TIER_BOUNDARIES: Dict[str, float] = {
|
|||
|
||||
# ─── Default Token Thresholds ───
|
||||
|
||||
DEFAULT_TOKEN_THRESHOLDS: Dict[str, int] = {
|
||||
DEFAULT_TOKEN_THRESHOLDS: dict[str, int] = {
|
||||
"simple": 15, # Only very short prompts (<15 tokens) are penalized
|
||||
"complex": 400, # Long prompts (>400 tokens) get complexity boost
|
||||
}
|
||||
|
|
@ -221,7 +225,7 @@ DEFAULT_TOKEN_THRESHOLDS: Dict[str, int] = {
|
|||
|
||||
# ─── Default Tier to Model Mapping ───
|
||||
|
||||
DEFAULT_TIER_MODELS: Dict[str, str] = {
|
||||
DEFAULT_TIER_MODELS: dict[str, str] = {
|
||||
"SIMPLE": "gpt-4o-mini",
|
||||
"MEDIUM": "gpt-4o",
|
||||
"COMPLEX": "claude-sonnet-4-20250514",
|
||||
|
|
@ -244,44 +248,47 @@ class ClassifierLLMConfig(BaseModel):
|
|||
class ComplexityRouterConfig(BaseModel):
|
||||
"""Configuration for the ComplexityRouter."""
|
||||
|
||||
# Tier to model mapping
|
||||
tiers: Dict[str, str] = Field(
|
||||
# string = pin; list = random pick when adaptive=False, soft-floor home pool when adaptive=True
|
||||
tiers: dict[str, str | list[str]] = Field(
|
||||
default_factory=lambda: DEFAULT_TIER_MODELS.copy(),
|
||||
description="Mapping of complexity tiers to model names",
|
||||
description=(
|
||||
"Mapping of complexity tiers to a model or model pool. "
|
||||
"A list is randomly picked from when adaptive=False, and used as a soft-floor home pool when adaptive=True"
|
||||
),
|
||||
)
|
||||
|
||||
# Tier boundaries (normalized scores)
|
||||
tier_boundaries: Dict[str, float] = Field(
|
||||
tier_boundaries: dict[str, float] = Field(
|
||||
default_factory=lambda: DEFAULT_TIER_BOUNDARIES.copy(),
|
||||
description="Score boundaries between tiers",
|
||||
)
|
||||
|
||||
# Token count thresholds
|
||||
token_thresholds: Dict[str, int] = Field(
|
||||
token_thresholds: dict[str, int] = Field(
|
||||
default_factory=lambda: DEFAULT_TOKEN_THRESHOLDS.copy(),
|
||||
description="Token count thresholds for simple/complex classification",
|
||||
)
|
||||
|
||||
# Dimension weights
|
||||
dimension_weights: Dict[str, float] = Field(
|
||||
dimension_weights: dict[str, float] = Field(
|
||||
default_factory=lambda: DEFAULT_DIMENSION_WEIGHTS.copy(),
|
||||
description="Weights for each scoring dimension",
|
||||
)
|
||||
|
||||
# Keyword lists (overridable)
|
||||
code_keywords: Optional[List[str]] = Field(
|
||||
code_keywords: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Keywords indicating code-related content",
|
||||
)
|
||||
reasoning_keywords: Optional[List[str]] = Field(
|
||||
reasoning_keywords: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Keywords indicating reasoning-required content",
|
||||
)
|
||||
technical_keywords: Optional[List[str]] = Field(
|
||||
technical_keywords: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Keywords indicating technical content",
|
||||
)
|
||||
custom_technical_keywords: Optional[list[str]] = Field(
|
||||
custom_technical_keywords: list[str] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Domain-specific technical keywords appended to the effective base list "
|
||||
|
|
@ -290,13 +297,13 @@ class ComplexityRouterConfig(BaseModel):
|
|||
"the base list and within this list."
|
||||
),
|
||||
)
|
||||
simple_keywords: Optional[List[str]] = Field(
|
||||
simple_keywords: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Keywords indicating simple/basic queries",
|
||||
)
|
||||
|
||||
# Default model if scoring fails
|
||||
default_model: Optional[str] = Field(
|
||||
default_model: str | None = Field(
|
||||
default=None,
|
||||
description="Default model to use if tier cannot be determined",
|
||||
)
|
||||
|
|
@ -306,13 +313,34 @@ class ComplexityRouterConfig(BaseModel):
|
|||
default="heuristic",
|
||||
description="Classification strategy: local regex/keyword scoring, or an LLM call",
|
||||
)
|
||||
classifier_llm_config: Optional[ClassifierLLMConfig] = Field(
|
||||
classifier_llm_config: ClassifierLLMConfig | None = Field(
|
||||
default=None,
|
||||
description="Configuration for the LLM classifier; required when classifier_type is 'llm'",
|
||||
)
|
||||
|
||||
adaptive: bool = Field(
|
||||
default=False,
|
||||
description="Enable adaptive bandit selection with soft complexity floors",
|
||||
)
|
||||
adaptive_weights: AdaptiveRouterWeights = Field(
|
||||
default_factory=lambda: AdaptiveRouterWeights(quality=0.3, cost=0.7),
|
||||
description="Quality vs cost weights for adaptive selection (used when adaptive=True)",
|
||||
)
|
||||
tier_distance_penalty: float = Field(
|
||||
default=DEFAULT_TIER_DISTANCE_PENALTY,
|
||||
ge=0.0,
|
||||
description="Score penalty per tier-step away from the classified tier when adaptive=True",
|
||||
)
|
||||
adaptive_eligible: Literal["all", "classified_tier"] = Field(
|
||||
default="all",
|
||||
description=(
|
||||
"When adaptive=True: 'all' scores every pool model with a tier-distance penalty (soft floors); "
|
||||
"'classified_tier' Thompson-samples only inside the classified tier's pool"
|
||||
),
|
||||
)
|
||||
|
||||
# Deterministic keyword -> tier overrides, evaluated before weighted scoring
|
||||
keyword_tier_rules: Optional[List[KeywordTierRule]] = Field(
|
||||
keyword_tier_rules: list[KeywordTierRule] | None = Field(
|
||||
default=None,
|
||||
description="Rules that force a specific tier when their keywords match the prompt",
|
||||
)
|
||||
|
|
@ -322,7 +350,7 @@ class ComplexityRouterConfig(BaseModel):
|
|||
default=False,
|
||||
description="Match keyword_tier_rules by embedding similarity instead of literal text",
|
||||
)
|
||||
embedding_model: Optional[str] = Field(
|
||||
embedding_model: str | None = Field(
|
||||
default=None,
|
||||
description="Embedding model (LiteLLM model name) used when semantic_keyword_matching is enabled",
|
||||
)
|
||||
|
|
@ -333,14 +361,56 @@ class ComplexityRouterConfig(BaseModel):
|
|||
description="Minimum cosine similarity for a semantic keyword match",
|
||||
)
|
||||
|
||||
# Session affinity: pin the first turn's routed model for the rest of the session
|
||||
session_affinity: bool = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"When True and a session_id is resolvable on the request, pin the model chosen on the "
|
||||
"session's first turn and reuse it for every later turn, skipping re-classification."
|
||||
),
|
||||
)
|
||||
session_affinity_ttl_seconds: int = Field(
|
||||
default=3600,
|
||||
gt=0,
|
||||
description="TTL for the session affinity pin; refreshed on every cache hit",
|
||||
)
|
||||
|
||||
model_config = ConfigDict(extra="allow") # Allow additional fields
|
||||
|
||||
@field_validator("tiers", mode="before")
|
||||
@classmethod
|
||||
def _coerce_tier_values(cls, value: object) -> object:
|
||||
if not isinstance(value, dict):
|
||||
return value
|
||||
coerced: dict[str, object] = {}
|
||||
for key, item in value.items():
|
||||
if isinstance(item, str):
|
||||
coerced[key] = item
|
||||
elif isinstance(item, (list, tuple)):
|
||||
coerced[key] = list(item)
|
||||
else:
|
||||
coerced[key] = item
|
||||
return coerced
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_llm_classifier_config(self) -> "ComplexityRouterConfig":
|
||||
if self.classifier_type == "llm" and self.classifier_llm_config is None:
|
||||
raise ValueError("classifier_llm_config is required when classifier_type is 'llm'")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_adaptive_pools(self) -> "ComplexityRouterConfig":
|
||||
if not self.adaptive:
|
||||
return self
|
||||
normalized = {tier: (models if isinstance(models, list) else [models]) for tier, models in self.tiers.items()}
|
||||
if not any(normalized.values()):
|
||||
raise ValueError("adaptive=True requires at least one non-empty tier pool")
|
||||
empty = [tier for tier, models in normalized.items() if not models]
|
||||
if empty:
|
||||
raise ValueError(f"adaptive=True tier pools must be non-empty; empty tiers: {empty}")
|
||||
self.tiers = normalized
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_semantic_matching(self) -> "ComplexityRouterConfig":
|
||||
if not self.semantic_keyword_matching:
|
||||
|
|
|
|||
|
|
@ -213,6 +213,8 @@ DEFINED_PROMETHEUS_METRICS = Literal[
|
|||
"litellm_input_audio_tokens_metric",
|
||||
"litellm_output_reasoning_tokens_metric",
|
||||
"litellm_output_audio_tokens_metric",
|
||||
"litellm_video_duration_seconds_metric",
|
||||
"litellm_images_generated_metric",
|
||||
"litellm_deployment_successful_fallbacks",
|
||||
"litellm_deployment_failed_fallbacks",
|
||||
"litellm_remaining_team_budget_metric",
|
||||
|
|
@ -506,6 +508,9 @@ class PrometheusMetricLabels:
|
|||
litellm_output_reasoning_tokens_metric = litellm_output_tokens_metric
|
||||
litellm_output_audio_tokens_metric = litellm_output_tokens_metric
|
||||
|
||||
litellm_video_duration_seconds_metric = litellm_output_tokens_metric
|
||||
litellm_images_generated_metric = litellm_output_tokens_metric
|
||||
|
||||
litellm_deployment_state = [
|
||||
UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value,
|
||||
UserAPIKeyLabelNames.MODEL_ID.value,
|
||||
|
|
@ -717,6 +722,8 @@ class PrometheusMetricLabels:
|
|||
"litellm_input_tokens_metric",
|
||||
"litellm_total_tokens_metric",
|
||||
"litellm_output_tokens_metric",
|
||||
"litellm_video_duration_seconds_metric",
|
||||
"litellm_images_generated_metric",
|
||||
}
|
||||
)
|
||||
# Managed batch metrics
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Union, get_type_hi
|
|||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
from typing_extensions import Required, TypedDict
|
||||
from typing_extensions import Protocol, Required, TypedDict
|
||||
|
||||
from litellm._uuid import uuid
|
||||
|
||||
|
|
@ -829,6 +829,35 @@ class PreRoutingHookResponse(BaseModel):
|
|||
messages: Optional[List[Dict[str, Any]]]
|
||||
|
||||
|
||||
class RoutingContext(BaseModel):
|
||||
"""
|
||||
Passed through a Router's `plugins` pipeline before the routing decision is made.
|
||||
|
||||
Each plugin reads and mutates this object; the next plugin sees the previous
|
||||
plugin's changes. `candidate_models` narrows as the pipeline runs -- Router
|
||||
only selects a deployment whose `litellm_params.model` survives the pipeline.
|
||||
|
||||
`raw_messages` and `structured_messages` mirror the pattern
|
||||
`CustomGuardrail.apply_guardrail` uses: the message shape differs by API
|
||||
surface (chat completions, Anthropic /v1/messages, Responses API `input`,
|
||||
...), so plugins that need a stable, provider-agnostic shape should read
|
||||
`structured_messages` (normalized to OpenAI chat-completions format);
|
||||
plugins that need the exact original payload can read `raw_messages`.
|
||||
"""
|
||||
|
||||
raw_messages: list[dict[str, Any]]
|
||||
structured_messages: list[dict[str, Any]]
|
||||
candidate_models: list[str]
|
||||
metadata: dict[str, Any] = Field(default_factory=dict)
|
||||
signals: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class RoutingPlugin(Protocol):
|
||||
"""Interface a custom routing plugin must implement to run in `Router(plugins=[...])`."""
|
||||
|
||||
async def run(self, context: RoutingContext) -> RoutingContext: ...
|
||||
|
||||
|
||||
class RequestType(str, enum.Enum):
|
||||
"""Fixed v0 taxonomy. User-extensible types come in v1."""
|
||||
|
||||
|
|
|
|||
|
|
@ -209,6 +209,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
input_cost_per_query: Optional[float] # only for rerank models
|
||||
input_cost_per_image: Optional[float] # only for vertex ai models
|
||||
input_cost_per_image_token: Optional[float] # for gpt-image-1 and similar models
|
||||
input_cost_per_video_token: Optional[float] # for gemini omni models with video input
|
||||
input_cost_per_audio_per_second: Optional[float] # only for vertex ai models
|
||||
input_cost_per_video_per_second: Optional[float] # only for vertex ai models
|
||||
input_cost_per_second: Optional[float] # for OpenAI Speech models
|
||||
|
|
@ -234,6 +235,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
output_cost_per_character_above_128k_tokens: Optional[float] # only for vertex ai models
|
||||
output_cost_per_image: Optional[float]
|
||||
output_cost_per_image_token: Optional[float]
|
||||
output_cost_per_video_token: Optional[float] # for gemini omni models with video output
|
||||
output_vector_size: Optional[int]
|
||||
output_cost_per_reasoning_token: Optional[float]
|
||||
output_cost_per_video_per_second: Optional[float] # only for vertex ai models
|
||||
|
|
@ -3046,6 +3048,7 @@ class CustomPricingLiteLLMParams(BaseModel):
|
|||
output_cost_per_character_above_128k_tokens: Optional[float] = None
|
||||
output_cost_per_image: Optional[float] = None
|
||||
output_cost_per_image_token: Optional[float] = None
|
||||
output_cost_per_video_token: Optional[float] = None
|
||||
output_cost_per_reasoning_token: Optional[float] = None
|
||||
output_cost_per_video_per_second: Optional[float] = None
|
||||
output_cost_per_audio_per_second: Optional[float] = None
|
||||
|
|
@ -3055,6 +3058,7 @@ class CustomPricingLiteLLMParams(BaseModel):
|
|||
cache_read_input_token_cost_above_272k_tokens: Optional[float] = None
|
||||
cache_read_input_token_cost_above_512k_tokens: Optional[float] = None
|
||||
input_cost_per_image_token: Optional[float] = None
|
||||
input_cost_per_video_token: Optional[float] = None
|
||||
input_cost_per_token_above_272k_tokens: Optional[float] = None
|
||||
input_cost_per_token_above_512k_tokens: Optional[float] = None
|
||||
output_cost_per_token_above_272k_tokens: Optional[float] = None
|
||||
|
|
@ -3210,6 +3214,16 @@ all_litellm_params = (
|
|||
"_litellm_tpm_reserved_model",
|
||||
"_litellm_tpm_reserved_scopes",
|
||||
"_litellm_tpm_reservation_released",
|
||||
"auto_router_config_path",
|
||||
"auto_router_config",
|
||||
"auto_router_default_model",
|
||||
"auto_router_embedding_model",
|
||||
"complexity_router_config",
|
||||
"complexity_router_default_model",
|
||||
"adaptive_router_config",
|
||||
"adaptive_router_default_model",
|
||||
"quality_router_config",
|
||||
"quality_router_default_model",
|
||||
]
|
||||
+ list(StandardCallbackDynamicParams.__annotations__.keys())
|
||||
+ list(CustomPricingLiteLLMParams.model_fields.keys())
|
||||
|
|
|
|||
|
|
@ -5437,6 +5437,7 @@ def _get_model_info_helper(
|
|||
input_cost_per_second=_model_info.get("input_cost_per_second", None),
|
||||
input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None),
|
||||
input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None),
|
||||
input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None),
|
||||
input_cost_per_image=_model_info.get("input_cost_per_image", None),
|
||||
input_cost_per_audio_per_second=_model_info.get("input_cost_per_audio_per_second", None),
|
||||
input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None),
|
||||
|
|
@ -5480,6 +5481,7 @@ def _get_model_info_helper(
|
|||
output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None),
|
||||
output_cost_per_image=_model_info.get("output_cost_per_image", None),
|
||||
output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None),
|
||||
output_cost_per_video_token=_model_info.get("output_cost_per_video_token", None),
|
||||
output_vector_size=_model_info.get("output_vector_size", None),
|
||||
citation_cost_per_token=_model_info.get("citation_cost_per_token", None),
|
||||
tiered_pricing=_model_info.get("tiered_pricing", None),
|
||||
|
|
|
|||
|
|
@ -11331,6 +11331,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
|
|
@ -11362,6 +11363,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
|
|
@ -11424,6 +11426,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -11479,6 +11482,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
|
|
@ -11506,6 +11510,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
|
|
@ -11559,6 +11564,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_output_config": true
|
||||
|
|
@ -11586,6 +11592,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_output_config": true
|
||||
|
|
@ -11614,6 +11621,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"provider_specific_entry": {
|
||||
|
|
@ -11648,6 +11656,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"provider_specific_entry": {
|
||||
|
|
@ -11682,6 +11691,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -11718,6 +11728,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -11788,6 +11799,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_sampling_params": false,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -18865,7 +18877,8 @@
|
|||
"search_context_size_medium": 0.014,
|
||||
"search_context_size_high": 0.014
|
||||
},
|
||||
"web_search_billing_unit": "per_query"
|
||||
"web_search_billing_unit": "per_query",
|
||||
"supports_reasoning": false
|
||||
},
|
||||
"gemini/gemini-3-pro-image-preview": {
|
||||
"input_cost_per_image": 0.0011,
|
||||
|
|
@ -18898,6 +18911,7 @@
|
|||
],
|
||||
"supports_function_calling": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -18939,6 +18953,7 @@
|
|||
],
|
||||
"supports_function_calling": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -18980,6 +18995,7 @@
|
|||
],
|
||||
"supports_function_calling": false,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
|
|
@ -19674,6 +19690,39 @@
|
|||
},
|
||||
"web_search_billing_unit": "per_query"
|
||||
},
|
||||
"gemini/gemini-omni-flash-preview": {
|
||||
"input_cost_per_audio_token": 1.5e-06,
|
||||
"input_cost_per_token": 1.5e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 65535,
|
||||
"max_tokens": 65535,
|
||||
"mode": "chat",
|
||||
"output_cost_per_reasoning_token": 9e-06,
|
||||
"output_cost_per_token": 9e-06,
|
||||
"output_cost_per_video_token": 1.75e-05,
|
||||
"rpm": 2000,
|
||||
"source": "https://ai.google.dev/gemini-api/docs/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image",
|
||||
"audio",
|
||||
"video"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text",
|
||||
"video"
|
||||
],
|
||||
"supports_audio_input": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"tpm": 800000
|
||||
},
|
||||
"gemini/gemini-3.1-pro-preview": {
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
|
||||
|
|
@ -19838,6 +19887,37 @@
|
|||
},
|
||||
"web_search_billing_unit": "per_query"
|
||||
},
|
||||
"gemini-omni-flash-preview": {
|
||||
"input_cost_per_audio_token": 1.5e-06,
|
||||
"input_cost_per_token": 1.5e-06,
|
||||
"litellm_provider": "vertex_ai-language-models",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 65535,
|
||||
"max_tokens": 65535,
|
||||
"mode": "chat",
|
||||
"output_cost_per_reasoning_token": 9e-06,
|
||||
"output_cost_per_token": 9e-06,
|
||||
"output_cost_per_video_token": 1.75e-05,
|
||||
"source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/omni-flash-preview",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image",
|
||||
"audio",
|
||||
"video"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text",
|
||||
"video"
|
||||
],
|
||||
"supports_audio_input": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"gemini-3.5-flash": {
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"input_cost_per_audio_token": 1e-06,
|
||||
|
|
@ -36951,6 +37031,7 @@
|
|||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -36973,6 +37054,7 @@
|
|||
"output_cost_per_image_token": 0.00012,
|
||||
"output_cost_per_token": 1.2e-05,
|
||||
"output_cost_per_token_batches": 6e-06,
|
||||
"supports_reasoning": false,
|
||||
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
|
||||
},
|
||||
"vertex_ai/gemini-3-pro-image-preview": {
|
||||
|
|
@ -36988,6 +37070,7 @@
|
|||
"output_cost_per_image_token": 0.00012,
|
||||
"output_cost_per_token": 1.2e-05,
|
||||
"output_cost_per_token_batches": 6e-06,
|
||||
"supports_reasoning": false,
|
||||
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
|
||||
},
|
||||
"vertex_ai/gemini-3.1-flash-image": {
|
||||
|
|
@ -37001,6 +37084,7 @@
|
|||
"output_cost_per_image": 0.0672,
|
||||
"output_cost_per_image_token": 6e-05,
|
||||
"output_cost_per_token": 3e-06,
|
||||
"supports_reasoning": false,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models"
|
||||
},
|
||||
"vertex_ai/gemini-3.1-flash-image-preview": {
|
||||
|
|
@ -37014,6 +37098,7 @@
|
|||
"output_cost_per_image": 0.0672,
|
||||
"output_cost_per_image_token": 6e-05,
|
||||
"output_cost_per_token": 3e-06,
|
||||
"supports_reasoning": false,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models"
|
||||
},
|
||||
"vertex_ai/gemini-3.1-flash-lite-preview": {
|
||||
|
|
@ -45284,8 +45369,8 @@
|
|||
"rules": [
|
||||
{
|
||||
"name": "bedrock-claude-ids",
|
||||
"pattern": "anthropic\\.claude-",
|
||||
"description": "Any Bedrock-syntax Claude id: the dotted anthropic.claude- segment appears in bare (anthropic.claude-...), region-prefixed (us./eu./au./jp./apac.) and global.-prefixed ids, for every version. Routes these to bedrock before the bare-id Anthropic rule is consulted.",
|
||||
"pattern": "^(?:[a-z-]+\\.)?anthropic\\.claude-",
|
||||
"description": "A Bedrock-syntax Claude id, for every version: anthropic.claude- at the start of the name, optionally behind a single dotted geo segment (us./eu./au./jp./apac./global./us-gov.). Anchored to the start because routing rules see the raw request string and provider inference feeds the proxy's provider/* wildcard access checks: an id under an unrecognized namespace such as bedrockz/anthropic.claude-... must stay unroutable rather than resolve to bedrock and slip through a bedrock/* key. Routes to bedrock before the bare-id Anthropic rule is consulted.",
|
||||
"model_info": {
|
||||
"litellm_provider": "bedrock"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm"
|
||||
version = "1.93.0"
|
||||
version = "1.94.0"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10, <3.14"
|
||||
|
|
@ -62,8 +62,8 @@ proxy = [
|
|||
"azure-identity>=1.25.2,<2.0",
|
||||
"azure-storage-blob>=12.28.0,<13.0",
|
||||
"mcp>=1.26.0,<2.0",
|
||||
"litellm-proxy-extras==0.4.76",
|
||||
"litellm-enterprise==0.1.49",
|
||||
"litellm-proxy-extras==0.4.77",
|
||||
"litellm-enterprise==0.1.50",
|
||||
"RestrictedPython>=8.1,<9.0",
|
||||
"rich>=13.9.4,<14.0",
|
||||
"polars>=1.38.1,<2.0",
|
||||
|
|
@ -205,7 +205,7 @@ ci = [
|
|||
# protobuf, Pillow is a compiled C extension).
|
||||
"tenacity==8.5.0",
|
||||
"google-generativeai==0.8.6",
|
||||
"Pillow==12.2.0",
|
||||
"Pillow==12.3.0",
|
||||
# Azure batch E2E tests still import psycopg2 directly.
|
||||
"psycopg2-binary==2.9.11",
|
||||
"pytest-codspeed==4.3.0",
|
||||
|
|
@ -264,6 +264,8 @@ constraint-dependencies = [
|
|||
"aiohttp>=3.14.1,<4.0",
|
||||
"packaging>=24.0",
|
||||
"soupsieve>=2.8.4",
|
||||
"httplib2>=0.32.0",
|
||||
"setuptools>=83.0.0",
|
||||
]
|
||||
override-dependencies = [
|
||||
# a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0.
|
||||
|
|
@ -284,7 +286,7 @@ members = ["enterprise", "litellm-proxy-extras"]
|
|||
profile = "black"
|
||||
|
||||
[tool.commitizen]
|
||||
version = "1.93.0"
|
||||
version = "1.94.0"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -306,7 +306,7 @@
|
|||
"limit": 9
|
||||
},
|
||||
"TID251": {
|
||||
"limit": 2710
|
||||
"limit": 2701
|
||||
},
|
||||
"TRY002": {
|
||||
"limit": 548
|
||||
|
|
@ -324,7 +324,7 @@
|
|||
"limit": 883
|
||||
},
|
||||
"UP006": {
|
||||
"limit": 12869
|
||||
"limit": 12792
|
||||
},
|
||||
"UP007": {
|
||||
"limit": 2570
|
||||
|
|
@ -354,7 +354,7 @@
|
|||
"limit": 4
|
||||
},
|
||||
"UP035": {
|
||||
"limit": 2295
|
||||
"limit": 2284
|
||||
},
|
||||
"UP036": {
|
||||
"limit": 4
|
||||
|
|
@ -363,6 +363,6 @@
|
|||
"limit": 105
|
||||
},
|
||||
"UP045": {
|
||||
"limit": 18517
|
||||
"limit": 18462
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -422,6 +422,7 @@ model LiteLLM_VerificationToken {
|
|||
budget_reset_at DateTime?
|
||||
allowed_cache_controls String[] @default([])
|
||||
allowed_routes String[] @default([])
|
||||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
|
|
@ -516,6 +517,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
budget_reset_at DateTime?
|
||||
allowed_cache_controls String[] @default([])
|
||||
allowed_routes String[] @default([])
|
||||
key_type String?
|
||||
policies String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
|
|
|
|||
|
|
@ -89,6 +89,11 @@ EOF
|
|||
|
||||
status=0
|
||||
|
||||
bootstrap_hint() {
|
||||
echo " This checkout looks unprovisioned (fresh worktree or clone)." >&2
|
||||
echo " Fix: make bootstrap" >&2
|
||||
}
|
||||
|
||||
if [ -n "$litellm_py_files" ]; then
|
||||
echo "pre-commit: linting Python (make lint)"
|
||||
make lint || { echo "✗ Python lint failed. Fix the reds above, then re-run make pre-commit." >&2; status=1; }
|
||||
|
|
@ -109,7 +114,13 @@ fi
|
|||
|
||||
if [ -n "$ui_prettier_files" ] || [ -n "$ui_eslint_files" ]; then
|
||||
echo "pre-commit: linting dashboard (prettier + eslint + lint budgets)"
|
||||
lint_dashboard || { echo "✗ Dashboard lint failed. See above; format with: (cd ui/litellm-dashboard && npm run format)." >&2; status=1; }
|
||||
if [ ! -d ui/litellm-dashboard/node_modules ]; then
|
||||
echo "✗ ui/litellm-dashboard/node_modules is missing; dashboard lint cannot run." >&2
|
||||
bootstrap_hint
|
||||
status=1
|
||||
else
|
||||
lint_dashboard || { echo "✗ Dashboard lint failed. See above; format with: (cd ui/litellm-dashboard && npm run format)." >&2; status=1; }
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -n "$spec_files" ]; then
|
||||
|
|
@ -118,7 +129,15 @@ if [ -n "$spec_files" ]; then
|
|||
# and an up-to-date Prisma client; check-ui-api-types.yml installs those and runs
|
||||
# prisma generate before gen:api, so mirror that here or a stale client can mask
|
||||
# drift that CI will still flag.
|
||||
if ! uv run --no-sync python scripts/prisma_generate_if_needed.py; then
|
||||
if [ ! -d ui/litellm-dashboard/node_modules ]; then
|
||||
echo "✗ ui/litellm-dashboard/node_modules is missing; the gen:api sync check cannot run." >&2
|
||||
bootstrap_hint
|
||||
status=1
|
||||
elif ! uv run --no-sync python -c "import orjson, prisma" 2>/dev/null; then
|
||||
echo "✗ The Python env lacks the proxy deps (orjson/prisma) that gen:api needs." >&2
|
||||
bootstrap_hint
|
||||
status=1
|
||||
elif ! uv run --no-sync python scripts/prisma_generate_if_needed.py; then
|
||||
echo "✗ Could not regenerate Prisma client (prisma generate failed)." >&2
|
||||
status=1
|
||||
elif ( cd ui/litellm-dashboard && LITELLM_PYTHON="uv run --no-sync python" npm run gen:api ); then
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
|
|||
- `embeddings/` - the `/embeddings` endpoint across providers
|
||||
- `batches/` - the `/batches` endpoint (placeholder until the first test lands)
|
||||
- `realtime/` - realtime websocket sessions, including the pipecat audio path
|
||||
- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window) and `spend_tracking/` (spend logging and cost attribution on `/spend/*`)
|
||||
- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`)
|
||||
- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials; also the dashboard UI behavior on top of them, driven through the proxy-served UI at /ui with playwright (optional dep behind importorskip)
|
||||
- `logging/` - logging-integration delivery (datadog and friends)
|
||||
- `security/` - secret handling and log-leak protection
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from __future__ import annotations
|
|||
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Callable
|
||||
|
||||
import pytest
|
||||
|
|
@ -49,7 +50,7 @@ from e2e_http import (
|
|||
unwrap,
|
||||
)
|
||||
from lifecycle import ResourceManager
|
||||
from models import KeyGenerateBody, SpendLogRow, SpendLogsParams
|
||||
from models import KeyGenerateBody, SpendLogRow
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
|
@ -349,6 +350,10 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
|
|||
the file-read path fires while the batch itself is not blocked.
|
||||
``resources.key()`` cannot set limits, so the key is minted on the gateway
|
||||
directly and its delete deferred.
|
||||
|
||||
Snapshots read /spend/logs/v2 over a bounded window around the test instead
|
||||
of the unpaginated /spend/logs whole-table read, which grows with the
|
||||
environment and OOMed the e2e runner on stage.
|
||||
"""
|
||||
user_id = f"e2e-batch-rl-{unique_marker()}"
|
||||
key = client.gateway.generate_key(
|
||||
|
|
@ -356,8 +361,13 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
|
|||
)
|
||||
resources.defer(lambda: client.gateway.delete_key(key))
|
||||
|
||||
window_start = datetime.now(timezone.utc) - timedelta(hours=1)
|
||||
window_end = window_start + timedelta(hours=2)
|
||||
before = frozenset(
|
||||
row.request_id for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams()))
|
||||
row.request_id
|
||||
for row in unattributed_rows(
|
||||
client.gateway.spend_logs_window(start=window_start, end=window_end)
|
||||
)
|
||||
)
|
||||
|
||||
file = unwrap(
|
||||
|
|
@ -379,7 +389,9 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
|
|||
|
||||
new_orphans = [
|
||||
row
|
||||
for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams()))
|
||||
for row in unattributed_rows(
|
||||
client.gateway.spend_logs_window(start=window_start, end=window_end)
|
||||
)
|
||||
if row.request_id not in before
|
||||
]
|
||||
assert not new_orphans, (
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@
|
|||
- {id: logging.datadog.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/datadog/datadog.py", rationale: "Powers dashboards/alerts; cardinality regressions common"}
|
||||
- {id: logging.datadog.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions], source: "integrations/datadog/datadog.py", rationale: "Failure metrics for alerting/SLO"}
|
||||
- {id: logging.prometheus.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/prometheus.py", rationale: "Standard OSS metrics; per-key cardinality (existing e2e)"}
|
||||
- {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"}
|
||||
- {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"}
|
||||
- {id: logging.otel.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions, messages], source: "integrations/otel/logger.py", rationale: "Error spans for observability continuity"}
|
||||
- {id: logging.braintrust.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/braintrust_logging.py", rationale: "Evals platform spend"}
|
||||
- {id: logging.langsmith.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/langsmith.py", rationale: "LangChain ecosystem"}
|
||||
|
|
|
|||
|
|
@ -18,6 +18,12 @@ configs:
|
|||
type: redis
|
||||
host: redis
|
||||
port: 6379
|
||||
# OTEL v2 trace destination for the logging suite's trace-completeness
|
||||
# tests: the arize_phoenix preset is OTLP with a configurable endpoint
|
||||
# (PHOENIX_COLLECTOR_HTTP_ENDPOINT below points it at the jaeger service),
|
||||
# so gen-AI spans export through a preset-owned provider - the code path
|
||||
# where trace splits actually happen - with no cloud credentials needed.
|
||||
callbacks: ["arize_phoenix"]
|
||||
|
||||
router_settings:
|
||||
routing_strategy: simple-shuffle
|
||||
|
|
@ -68,9 +74,14 @@ services:
|
|||
condition: service_healthy
|
||||
redis:
|
||||
condition: service_healthy
|
||||
jaeger:
|
||||
condition: service_healthy
|
||||
env_file: .env
|
||||
environment:
|
||||
LITELLM_MASTER_KEY: sk-1234
|
||||
LITELLM_OTEL_V2: "true"
|
||||
PHOENIX_COLLECTOR_HTTP_ENDPOINT: http://jaeger:4318/v1/traces
|
||||
PHOENIX_API_KEY: local-jaeger-noauth
|
||||
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
|
||||
UI_USERNAME: admin
|
||||
UI_PASSWORD: sk-1234
|
||||
|
|
@ -114,3 +125,15 @@ services:
|
|||
interval: 3s
|
||||
timeout: 3s
|
||||
retries: 20
|
||||
|
||||
# throwaway OTEL trace destination (OTLP ingest on 4318 inside the network,
|
||||
# query API on host 16686 for test read-back; see E2E_OTEL_QUERY_URL)
|
||||
jaeger:
|
||||
image: jaegertracing/all-in-one:1.62.0
|
||||
ports:
|
||||
- "16686:16686"
|
||||
healthcheck:
|
||||
test: ["CMD", "wget", "-qO-", "http://localhost:14269/"]
|
||||
interval: 3s
|
||||
timeout: 3s
|
||||
retries: 20
|
||||
|
|
|
|||
|
|
@ -24,6 +24,14 @@ CONTROL_PLANE_BASE_URL = os.environ.get(
|
|||
UI_USERNAME = os.environ.get("E2E_UI_USERNAME", "admin")
|
||||
UI_PASSWORD = os.environ.get("E2E_UI_PASSWORD", MASTER_KEY)
|
||||
|
||||
CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haiku-4-5")
|
||||
CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5")
|
||||
|
||||
# Jaeger query API of the compose stack's OTEL trace destination (the `jaeger`
|
||||
# service in docker-compose.yml maps it to host 16686). Trace-completeness tests
|
||||
# read exported spans back through it.
|
||||
OTEL_QUERY_URL = os.environ.get("E2E_OTEL_QUERY_URL", "http://localhost:16686").rstrip("/")
|
||||
|
||||
# Writes on the proxy are eventually consistent (e.g. spend rows flush on
|
||||
# proxy_batch_write_at, ~60s). Read-backs poll to this deadline, never sleep-once.
|
||||
POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120"))
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import time
|
|||
import warnings
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
|
||||
from e2e_http import (
|
||||
NoBody,
|
||||
|
|
@ -50,6 +51,8 @@ from models import (
|
|||
OcrResponse,
|
||||
SpendLogRow,
|
||||
SpendLogs,
|
||||
SpendLogsPage,
|
||||
SpendLogsPageParams,
|
||||
SpendLogsParams,
|
||||
)
|
||||
from e2e_config import (
|
||||
|
|
@ -255,6 +258,28 @@ class Gateway:
|
|||
case _:
|
||||
return []
|
||||
|
||||
def spend_logs_window(self, *, start: datetime, end: datetime) -> list[SpendLogRow]:
|
||||
def fetch(page: int) -> SpendLogsPage:
|
||||
return unwrap(
|
||||
self.transport.get(
|
||||
"/spend/logs/v2",
|
||||
headers=self.transport.master,
|
||||
params=SpendLogsPageParams(
|
||||
start_date=start.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
end_date=end.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
page=page,
|
||||
page_size=100,
|
||||
),
|
||||
response_type=SpendLogsPage,
|
||||
)
|
||||
)
|
||||
|
||||
first = fetch(1)
|
||||
return [
|
||||
*first.data,
|
||||
*(row for page in range(2, first.total_pages + 1) for row in fetch(page).data),
|
||||
]
|
||||
|
||||
def poll_logs_for_key(
|
||||
self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None
|
||||
) -> list[SpendLogRow]:
|
||||
|
|
|
|||
|
|
@ -109,14 +109,16 @@ class StreamingResponse(BaseModel):
|
|||
"""Raw outcome for calls whose body is provider-native or streamed: status, the
|
||||
x-litellm-call-id header, the x-litellm-response-cost header (StandardLogging
|
||||
response_cost), the content-type (which tells streaming `text/event-stream` from
|
||||
non-streaming `application/json`), and the body. SpendLogs.request_id is the
|
||||
completion body id, not call_id. Used by passthrough and streaming, where one
|
||||
validated JSON model does not fit."""
|
||||
non-streaming `application/json`), the response headers (lowercased names, e.g.
|
||||
the x-ratelimit-* pacing headers and retry-after on a 429), and the body.
|
||||
SpendLogs.request_id is the completion body id, not call_id. Used by passthrough
|
||||
and streaming, where one validated JSON model does not fit."""
|
||||
|
||||
status_code: int
|
||||
call_id: str | None = None # x-litellm-call-id header
|
||||
response_cost: float | None = None # x-litellm-response-cost header
|
||||
content_type: str | None = None
|
||||
headers: dict[str, str] = {}
|
||||
body: str
|
||||
chunks: int = 0 # streamed events (0 for non-streaming)
|
||||
|
||||
|
|
@ -276,12 +278,14 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon
|
|||
call_id = _hdr(resp, "x-litellm-call-id")
|
||||
response_cost = _parse_response_cost(resp)
|
||||
content_type = _hdr(resp, "content-type")
|
||||
headers = {name.lower(): value for name, value in resp.headers.items()}
|
||||
if not stream or not (200 <= resp.status_code < 300):
|
||||
return StreamingResponse(
|
||||
status_code=resp.status_code,
|
||||
call_id=call_id,
|
||||
response_cost=response_cost,
|
||||
content_type=content_type,
|
||||
headers=headers,
|
||||
body=resp.text,
|
||||
)
|
||||
lines = cast("Iterator[bytes]", resp.iter_lines())
|
||||
|
|
@ -291,6 +295,7 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon
|
|||
call_id=call_id,
|
||||
response_cost=response_cost,
|
||||
content_type=content_type,
|
||||
headers=headers,
|
||||
body="<streamed>",
|
||||
chunks=chunks,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import os
|
|||
import pytest
|
||||
|
||||
from logging_client import LangfuseCreds, LoggingClient, build_logging_client, load_langfuse_creds
|
||||
from otel_client import OtelReader, build_otel_reader
|
||||
|
||||
|
||||
def pytest_configure(config: pytest.Config) -> None:
|
||||
|
|
@ -28,6 +29,12 @@ def client() -> LoggingClient:
|
|||
return build_logging_client()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def otel_reader() -> OtelReader:
|
||||
"""Read-back client for the compose stack's Jaeger trace destination."""
|
||||
return build_otel_reader()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def datadog_creds() -> None:
|
||||
"""Require Datadog shipping credentials. Hard-fail when absent; never skip."""
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ from e2e_http import (
|
|||
unwrap,
|
||||
)
|
||||
from models import (
|
||||
AnthropicMessagesBody,
|
||||
ChatBody,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
|
|
@ -75,6 +76,14 @@ WEATHER_TOOL = ChatTool(
|
|||
)
|
||||
|
||||
|
||||
class ResponsesRequestBody(BaseModel):
|
||||
"""OpenAI Responses API /v1/responses request (non-streaming)."""
|
||||
|
||||
model: str
|
||||
input: str
|
||||
max_output_tokens: int
|
||||
|
||||
|
||||
class TeamCallbackBody(BaseModel):
|
||||
callback_name: Literal["langfuse_otel", "langfuse", "langsmith", "gcs"]
|
||||
callback_type: Literal["success", "failure", "success_and_failure"]
|
||||
|
|
@ -455,6 +464,32 @@ class LoggingClient:
|
|||
json=body,
|
||||
)
|
||||
|
||||
def messages_raw(self, key: str, model: str, text: str, *, max_tokens: int = 16) -> StreamingResponse:
|
||||
"""Non-streaming POST /v1/messages (Anthropic-native body): raw outcome
|
||||
judged by status/body/headers, for tests that need x-litellm-call-id."""
|
||||
return self.gateway.transport.send(
|
||||
"/v1/messages",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
json=AnthropicMessagesBody(
|
||||
model=model,
|
||||
max_tokens=max_tokens,
|
||||
messages=[ChatMessage(role="user", content=text)],
|
||||
),
|
||||
)
|
||||
|
||||
def responses_raw(
|
||||
self, key: str, model: str, text: str, *, max_output_tokens: int = 64
|
||||
) -> StreamingResponse:
|
||||
"""Non-streaming POST /v1/responses (OpenAI Responses API): raw outcome
|
||||
judged by status/body/headers, for tests that need x-litellm-call-id.
|
||||
max_output_tokens caps reasoning-model output cost; a capped response is
|
||||
still a 200 and still exports the trace."""
|
||||
return self.gateway.transport.send(
|
||||
"/v1/responses",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
json=ResponsesRequestBody(model=model, input=text, max_output_tokens=max_output_tokens),
|
||||
)
|
||||
|
||||
def scrape_metrics(self) -> str:
|
||||
return self.gateway.probe("/metrics", params=NoBody()).body
|
||||
|
||||
|
|
|
|||
138
tests/e2e/logging/otel_client.py
Normal file
138
tests/e2e/logging/otel_client.py
Normal file
|
|
@ -0,0 +1,138 @@
|
|||
"""Jaeger read-back for the OTEL trace-completeness tests: typed models over the
|
||||
Jaeger query API (the destination's own API - completeness is judged on what the
|
||||
backend actually holds, never on "export succeeded" proxy-side).
|
||||
|
||||
Traces are fetched server-side by the ``litellm.call_id`` tag the gen-AI span
|
||||
carries (the request's x-litellm-call-id response header), so read-back is
|
||||
immune to the query page filling up with unrelated traffic (background jobs,
|
||||
other suites sharing the stack). Jaeger returns every span of a matching trace,
|
||||
so the completeness assertions see the whole tree. A failed query is a hard
|
||||
failure, never an empty result - an unreachable destination must not read as
|
||||
"the trace never arrived".
|
||||
|
||||
External reads go through ``e2e_http`` (the only module allowed to call
|
||||
``requests.*``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from e2e_config import OTEL_QUERY_URL, POLL_INTERVAL, POLL_TIMEOUT
|
||||
from e2e_http import URL, NoBody, Success, get
|
||||
|
||||
#: OTEL resource service.name the proxy exports under (OTEL_SERVICE_NAME default).
|
||||
JAEGER_SERVICE = "litellm"
|
||||
#: Span tag carrying the request's x-litellm-call-id (stamped on the gen-AI span).
|
||||
CALL_ID_TAG = "litellm.call_id"
|
||||
|
||||
|
||||
class JaegerTag(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
key: str
|
||||
value: str | int | float | bool | None = None
|
||||
|
||||
|
||||
class JaegerReference(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", populate_by_name=True)
|
||||
|
||||
ref_type: str = Field(alias="refType")
|
||||
trace_id: str = Field(alias="traceID")
|
||||
span_id: str = Field(alias="spanID")
|
||||
|
||||
|
||||
class JaegerSpan(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", populate_by_name=True)
|
||||
|
||||
span_id: str = Field(alias="spanID")
|
||||
operation_name: str = Field(alias="operationName")
|
||||
start_time: int = Field(default=0, alias="startTime")
|
||||
references: list[JaegerReference] = []
|
||||
tags: list[JaegerTag] = []
|
||||
|
||||
@property
|
||||
def kind(self) -> str:
|
||||
for tag in self.tags:
|
||||
if tag.key == "span.kind":
|
||||
return str(tag.value)
|
||||
return ""
|
||||
|
||||
|
||||
class JaegerTrace(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", populate_by_name=True)
|
||||
|
||||
trace_id: str = Field(alias="traceID")
|
||||
spans: list[JaegerSpan] = []
|
||||
|
||||
def span_names(self) -> list[str]:
|
||||
return sorted(span.operation_name for span in self.spans)
|
||||
|
||||
|
||||
class JaegerTracesPage(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
data: list[JaegerTrace] = []
|
||||
|
||||
|
||||
class _TracesQuery(BaseModel):
|
||||
service: str
|
||||
tags: str
|
||||
limit: int = 20
|
||||
lookback: str = "1h"
|
||||
|
||||
|
||||
def _settled(trace: JaegerTrace, names: set[str], prefixes: set[str]) -> bool:
|
||||
present = set(trace.span_names())
|
||||
return names.issubset(present) and all(
|
||||
any(name.startswith(prefix) for name in present) for prefix in prefixes
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OtelReader:
|
||||
query_url: str
|
||||
|
||||
def traces_for_call(self, call_id: str) -> list[JaegerTrace]:
|
||||
"""Every trace holding a span tagged with this call id. Jaeger matches
|
||||
spans server-side and returns their full traces; more than one hit for
|
||||
one call IS the split-trace bug, so this never collapses to one."""
|
||||
result = get(
|
||||
URL(f"{self.query_url}/api/traces"),
|
||||
headers=NoBody(),
|
||||
params=_TracesQuery(service=JAEGER_SERVICE, tags=json.dumps({CALL_ID_TAG: call_id})),
|
||||
response_type=JaegerTracesPage,
|
||||
timeout=30.0,
|
||||
)
|
||||
match result:
|
||||
case Success(data=page):
|
||||
return page.data
|
||||
case failure:
|
||||
pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}")
|
||||
|
||||
def poll_traces_for_call(
|
||||
self, *, call_id: str, settled_names: set[str], settled_prefixes: set[str]
|
||||
) -> list[JaegerTrace]:
|
||||
"""Poll until exactly one trace holds the call and it carries every span
|
||||
name in ``settled_names`` plus at least one name per prefix in
|
||||
``settled_prefixes`` (spans flush in batches, the cost write lands after
|
||||
the response), then return the hits. At the deadline the last hits are
|
||||
returned as-is so the caller's assertions report the real final state -
|
||||
on a split trace this never settles and the orphan comes back."""
|
||||
deadline = time.monotonic() + POLL_TIMEOUT
|
||||
hits: list[JaegerTrace] = []
|
||||
while time.monotonic() < deadline:
|
||||
hits = self.traces_for_call(call_id)
|
||||
if len(hits) == 1 and _settled(hits[0], settled_names, settled_prefixes):
|
||||
return hits
|
||||
time.sleep(POLL_INTERVAL)
|
||||
return hits
|
||||
|
||||
|
||||
def build_otel_reader() -> OtelReader:
|
||||
return OtelReader(query_url=OTEL_QUERY_URL)
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue