Merge branch 'BerriAI:litellm_internal_staging' into fix/realtime-usage-detail-keys

This commit is contained in:
lmcdonald-godaddy 2026-05-12 13:04:04 -07:00 • committed by GitHub
commit e347ddeb75
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
195 changed files with 13811 additions and 4809 deletions

131
.github/workflows/mutation-test.yml vendored Normal file
View file

@ -0,0 +1,131 @@
name: "Mutation Test (manual)"
# Manually-triggered mutation testing. Runs mutmut against the scope
# configured in [tool.mutmut] in pyproject.toml (currently the
# litellm/proxy/management_endpoints/ folder). Intended cadence is roughly
# weekly — clicked from the Actions tab when someone wants a fresh report.
#
# Uploads a structured `mutation-report.md` (Meta ACH-style: original +
# mutated function with `# MUTANT START`/`# MUTANT END` delimiters + the
# existing tests + a task instruction) as a workflow artifact. Failures
# do not block anything because nothing depends on this workflow.
on:
workflow_dispatch:
permissions:
contents: read
concurrency:
group: mutation-test-${{ github.ref }}
cancel-in-progress: true
jobs:
mutation:
name: Run mutmut
runs-on: ubuntu-latest
# Whole-folder mutation against ~15 files / ~7.5k LOC can take hours.
# 350 minutes is just under the GitHub-hosted job cap of 360 minutes.
timeout-minutes: 350
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
with:
version: "0.10.9"
- name: Cache uv dependencies
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
- name: Install dependencies
run: |
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Generate Prisma client
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
# mutmut 3.x runs tests inside a `mutants/` sandbox where it injects
# mutation trampolines. uv installs the project as editable by default,
# which puts the original source dir on sys.path via a .pth file and
# shadows the sandbox copy — so tests would never exercise the mutated
# code. Reinstalling non-editable removes the .pth shadow.
- name: Reinstall litellm non-editable (so mutants/ is not shadowed)
run: |
uv pip uninstall litellm
uv pip install . --no-deps
# pytest-retry's pytest_configure hook crashes with
# `INTERNALERROR: no option named 'filtered_exceptions'` when invoked
# via mutmut's in-process pytest.main() call. The entry-point name
# doesn't normalize cleanly with `-p no:<name>`, so just remove the
# package outright. Reruns are wrong for mutation testing anyway —
# rerunning a "failed" mutant test would mask which mutants are killed.
- name: Remove pytest plugins that conflict with mutmut
run: |
uv pip uninstall pytest-retry || true
- name: Run mutmut
env:
# Make the mutants/ sandbox win over site-packages on sys.path so the
# trampolined files are imported instead of the installed copy.
PYTHONPATH: ${{ github.workspace }}/mutants
run: |
set -o pipefail
mkdir -p mutants
uv run --no-sync --with mutmut==3.5.0 mutmut run 2>&1 | tee mutmut-run.log
# Generate the structured report. The script embeds the enclosing
# function source for each survivor (via Python AST) and includes the
# existing test files, so an LLM agent has enough context to write
# killing tests without further file lookups. Modeled on Meta's ACH
# prompt template (arXiv 2501.12862).
- name: Generate detailed mutation report
if: always()
run: |
set +e
uv run --no-sync --with mutmut==3.5.0 mutmut export-cicd-stats > /dev/null 2>&1
uv run --no-sync --with mutmut==3.5.0 mutmut results > mutmut-results.txt 2>&1
uv run --no-sync python scripts/mutation_report.py
# The full report can be very long for big test files; the run-page
# summary cuts off at 1 MB. Append the head of the report (summary
# + survivor list) and link out to the artifact for the full body.
{
head -c 900000 mutation-report.md
echo ""
echo ""
echo "_Full report (with embedded function bodies and test files) is in the workflow artifact._"
} >> "$GITHUB_STEP_SUMMARY"
- name: Upload mutmut artifacts
if: always()
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: mutmut-${{ github.run_id }}-${{ github.run_attempt }}
path: |
mutation-report.md
mutmut-results.txt
mutmut-run.log
mutants/mutmut-stats.json
mutants/mutmut-cicd-stats.json
mutants/litellm/proxy/management_endpoints/**/*.py
if-no-files-found: warn
retention-days: 14

View file

@ -241,10 +241,27 @@ When opening issues or pull requests, follow these templates:
### Running the proxy server
Start the proxy with a config file:
Create a minimal config file and start the proxy:
```yaml
# config.yaml
model_list:
- model_name: fake-openai-endpoint
litellm_params:
model: openai/fake-model
api_key: fake-key
api_base: https://fake-api.example.com
general_settings:
master_key: sk-1234
litellm_settings:
drop_params: True
telemetry: False
```
```bash
uv run litellm --config dev_config.yaml --port 4000
uv run litellm --config config.yaml --port 4000
```
The proxy takes ~15-20 seconds to fully start (it runs Prisma migrations on boot). Wait for `/health` to return before sending requests. Without a PostgreSQL `DATABASE_URL`, the proxy connects to a default Neon dev database embedded in the `litellm-proxy-extras` package.

View file

@ -146,7 +146,7 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components:
- **Bound large result sets.** Prisma materializes full results in memory. For results over ~10 MB, paginate with `take`/`skip` or `cursor`/`take`, always with an explicit `order`. Prefer cursor-based pagination (`skip` is O(n)). Don't paginate naturally small result sets.
- **Limit fetched columns on wide tables.** Use `select` to fetch only needed fields — returns a partial object, so downstream code must not access unselected fields.
- **Check index coverage.** For new or modified queries, check `schema.prisma` for a supporting index. Prefer extending an existing index (e.g. `@@index([a])` → `@@index([a, b])`) over adding a new one, unless it's a `@@unique`. Only add indexes for large/frequent queries.
- **Keep schema files in sync.** Apply schema changes to all `schema.prisma` copies (`schema.prisma`, `litellm/proxy/`, `litellm-proxy-extras/`, `litellm-js/spend-logs/` for SpendLogs) with a migration under `litellm-proxy-extras/litellm_proxy_extras/migrations/`.
- **Keep schema files in sync.** Apply schema changes to all `schema.prisma` copies (`schema.prisma`, `litellm/proxy/`, `litellm-proxy-extras/`) with a migration under `litellm-proxy-extras/litellm_proxy_extras/migrations/`.
### Setup Wizard (`litellm/setup_wizard.py`)
- The wizard is implemented as a single `SetupWizard` class with `@staticmethod` methods — keep it that way. No module-level functions except `run_setup_wizard()` (the public entrypoint) and pure helpers (color, ANSI).

View file

@ -1,18 +0,0 @@
# Use the provided base image
FROM ghcr.io/berriai/litellm:main-latest@sha256:7c311546c25e7bb6e8cafede9fcd3d0d622ac636b5c9418befaa32e85dfb0186
# Set the working directory to /app
WORKDIR /app
# Copy the configuration file into the container at /app
COPY config.yaml .
# Make sure your docker/entrypoint.sh is executable
# Convert Windows line endings to Unix
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh
# Expose the necessary port
EXPOSE 4000/tcp
# Override the CMD instruction with your desired command and arguments
CMD ["--port", "4000", "--config", "config.yaml", "--detailed_debug", "--run_gunicorn"]

View file

@ -100,6 +100,16 @@ spec:
- name: DATABASE_URL
value: {{ .Values.db.url | quote }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }}
- name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaUrlKey }}
{{- else if .Values.db.readReplicaUrl }}
- name: DATABASE_URL_READ_REPLICA
value: {{ .Values.db.readReplicaUrl | quote }}
{{- end }}
- name: PROXY_MASTER_KEY
valueFrom:
secretKeyRef:

View file

@ -252,6 +252,26 @@ db:
passwordKey: password
# Optional: when set, DATABASE_HOST will be sourced from this secret key instead of db.endpoint
endpointKey: ""
# Optional: when set, DATABASE_URL_READ_REPLICA will be sourced from this
# secret key instead of db.readReplicaUrl. Prefer this over the plain
# value: read-replica URLs typically embed credentials, and a value
# written to db.readReplicaUrl ends up visible in the rendered pod spec
# and the Helm release secret.
readReplicaUrlKey: ""
# Optional read-replica routing. When set, the proxy sends read-only
# queries (find_*, count, group_by, query_raw/_first) to this URL while
# writes continue to go to db.url. Useful for Aurora-style clusters with
# separate reader/writer endpoints. Leave empty to keep single-DB behavior.
# When IAM_TOKEN_DB_AUTH is enabled, the reader URL is auto-refreshed
# alongside the writer (host/port/user/db are parsed from this URL once
# at startup; only the IAM token rotates).
#
# If the URL embeds credentials, prefer db.secret.readReplicaUrlKey over
# this field — the plain value is rendered into the pod spec and the
# Helm release secret. This field is intended for credential-less URLs
# only (e.g. when IAM_TOKEN_DB_AUTH supplies the token at runtime).
readReplicaUrl: ""
# Use the Stackgres Helm chart to deploy an instance of a Stackgres cluster.
# The Stackgres Operator must already be installed within the target

View file

@ -1,56 +0,0 @@
apiVersion: apps/v1
kind: Deployment
metadata:
name: litellm-deployment
spec:
replicas: 3
selector:
matchLabels:
app: litellm
template:
metadata:
labels:
app: litellm
spec:
containers:
- name: litellm-container
image: ghcr.io/berriai/litellm:main-latest
imagePullPolicy: Always
env:
- name: AZURE_API_KEY
value: "d6f****"
- name: AZURE_API_BASE
value: "https://openai"
- name: LITELLM_MASTER_KEY
value: "sk-1234"
- name: DATABASE_URL
value: "postgresql://ishaan*********"
args:
- "--config"
- "/app/proxy_config.yaml" # Update the path to mount the config file
volumeMounts: # Define volume mount for proxy_config.yaml
- name: config-volume
mountPath: /app
readOnly: true
livenessProbe:
httpGet:
path: /health/liveliness
port: 4000
initialDelaySeconds: 120
periodSeconds: 15
successThreshold: 1
failureThreshold: 3
timeoutSeconds: 10
readinessProbe:
httpGet:
path: /health/readiness
port: 4000
initialDelaySeconds: 120
periodSeconds: 15
successThreshold: 1
failureThreshold: 3
timeoutSeconds: 10
volumes: # Define volume to mount proxy_config.yaml
- name: config-volume
configMap:
name: litellm-config

View file

@ -1,12 +0,0 @@
apiVersion: v1
kind: Service
metadata:
name: litellm-service
spec:
selector:
app: litellm
ports:
- protocol: TCP
port: 4000
targetPort: 4000
type: LoadBalancer

View file

@ -1,13 +0,0 @@
model_list:
- model_name: fake-openai-endpoint
litellm_params:
model: openai/fake-model
api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/
general_settings:
master_key: sk-1234
litellm_settings:
drop_params: True
telemetry: False

View file

@ -16,6 +16,11 @@ services:
- "4000:4000" # Map the container port to the host, change the host port if necessary
environment:
DATABASE_URL: "postgresql://llmproxy:dbpassword9090@db:5432/litellm"
# Optional: route read-only queries (find_*, count, group_by, query_raw/_first)
# to a separate reader endpoint, e.g. an Aurora reader. Leave unset for
# single-DB deployments. With IAM_TOKEN_DB_AUTH enabled, the reader URL
# is auto-refreshed alongside the writer.
# DATABASE_URL_READ_REPLICA: "postgresql://llmproxy:dbpassword9090@db-reader:5432/litellm"
STORE_MODEL_IN_DB: "True" # allows adding models to proxy via UI
env_file:
- .env # Load local .env file

View file

@ -1,68 +0,0 @@
# Base image for building
ARG LITELLM_BUILD_IMAGE=python:3.11-alpine@sha256:f07e2ace46f560f09a6eeec7b4913b80ee99546e749ef82342a419a326620856
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=python:3.11-alpine@sha256:f07e2ace46f560f09a6eeec7b4913b80ee99546e749ef82342a419a326620856
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin
FROM $LITELLM_BUILD_IMAGE AS builder
WORKDIR /app
COPY --from=uvbin /uv /usr/local/bin/uv
COPY --from=uvbin /uvx /usr/local/bin/uvx
RUN apk add --no-cache gcc python3-dev musl-dev nodejs npm libsndfile
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
UV_PROJECT_ENVIRONMENT=/app/.venv \
UV_LINK_MODE=copy \
XDG_CACHE_HOME=/app/.cache \
PATH="/app/.venv/bin:${PATH}"
# Copy dependency metadata first for layer caching
COPY pyproject.toml uv.lock ./
COPY enterprise/pyproject.toml enterprise/
COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/
# Install third-party dependencies (cached unless pyproject.toml/uv.lock change)
RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--python python3
# Copy full source tree
COPY . .
# Install project and workspace packages (fast - deps already cached)
RUN uv sync --frozen --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--python python3
RUN prisma generate --schema=./schema.prisma
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
FROM $LITELLM_RUNTIME_IMAGE AS runtime
RUN apk upgrade --no-cache && apk add --no-cache libsndfile nodejs npm
WORKDIR /app
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
XDG_CACHE_HOME=/app/.cache \
PATH="/app/.venv/bin:${PATH}"
COPY --from=builder /app /app
EXPOSE 4000/tcp
ENTRYPOINT ["docker/prod_entrypoint.sh"]
CMD ["--port", "4000"]

View file

@ -1,86 +0,0 @@
# Use the provided base image
# NOTE: This is a dev/branch-specific tag. Update digest when the base image is rebuilt.
FROM ghcr.io/berriai/litellm:litellm_fwd_server_root_path-dev
# Set the working directory to /app
WORKDIR /app
# Install Node.js and npm (adjust version as needed)
RUN apt-get update && apt-get upgrade -y \
libxml2 \
libexpat1 \
openssl \
libssl3 \
git \
libkrb5-3 \
libglib2.0-0 \
wget \
libaom3 \
libxslt1.1 \
libgnutls30 \
libc6 && \
apt-get install -y --no-install-recommends nodejs npm && \
npm install -g npm@11.12.1 tar@7.5.11 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \
GLOBAL="$(npm root -g)" && \
find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "glob" -path "*/node_modules/glob" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \
done && \
find /usr/local/lib /usr/lib -path "*/node_modules/npm/package.json" -exec \
sed -i 's/"tar": "\^7\.5\.[0-9]*"/"tar": "^7.5.10"/g; s/"minimatch": "\^10\.[0-9.]*"/"minimatch": "^10.2.4"/g' {} + 2>/dev/null && \
npm cache clean --force && \
apt-get purge -y npm
# Copy the UI source into the container
COPY ./ui/litellm-dashboard /app/ui/litellm-dashboard
# Set an environment variable for UI_BASE_PATH
# This can be overridden at build time
# set UI_BASE_PATH to "<your server root path>/ui"
ENV UI_BASE_PATH="/prod/ui"
# Build the UI with the specified UI_BASE_PATH
WORKDIR /app/ui/litellm-dashboard
RUN npm ci
RUN UI_BASE_PATH=$UI_BASE_PATH npm run build
# Create the destination directory
RUN mkdir -p /app/litellm/proxy/_experimental/out
# Move the built files to the appropriate location
# Assuming the build output is in ./out directory
RUN rm -rf /app/litellm/proxy/_experimental/out/* && \
mv ./out/* /app/litellm/proxy/_experimental/out/
# Switch back to the main app directory
WORKDIR /app
# Make sure your docker/entrypoint.sh is executable
# Convert Windows line endings to Unix for entrypoint scripts
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh
RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
# Run as non-root user
RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser \
&& chown -R appuser:appuser /app
USER appuser
# Expose the necessary port
EXPOSE 4000/tcp
HEALTHCHECK --interval=30s --timeout=5s --start-period=10s --retries=3 \
CMD ["python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:4000/health')"]
# Override the CMD instruction with your desired command and arguments
CMD ["--port", "4000", "--config", "config.yaml", "--detailed_debug"]

View file

@ -1,121 +0,0 @@
# Base image for building
ARG LITELLM_BUILD_IMAGE=python:3.13-slim@sha256:739e7213785e88c0f702dcdc12c0973afcbd606dbf021a589cab77d6b00b579d
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=python:3.13-slim@sha256:739e7213785e88c0f702dcdc12c0973afcbd606dbf021a589cab77d6b00b579d
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin
FROM $LITELLM_BUILD_IMAGE AS builder
WORKDIR /app
USER root
COPY --from=uvbin /uv /usr/local/bin/uv
COPY --from=uvbin /uvx /usr/local/bin/uvx
RUN apt-get update && apt-get install -y --no-install-recommends \
gcc \
g++ \
python3-dev \
libssl-dev \
pkg-config \
nodejs \
npm \
&& rm -rf /var/lib/apt/lists/*
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
UV_PROJECT_ENVIRONMENT=/app/.venv \
UV_LINK_MODE=copy \
XDG_CACHE_HOME=/app/.cache \
PATH="/app/.venv/bin:${PATH}"
# Copy dependency metadata first for layer caching
COPY pyproject.toml uv.lock ./
COPY enterprise/pyproject.toml enterprise/
COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/
# Install third-party dependencies (cached unless pyproject.toml/uv.lock change)
RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--python python
# Copy full source tree
COPY . .
# Build Admin UI before final sync
RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Install project and workspace packages (fast - deps already cached)
RUN uv sync --frozen --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--python python
RUN prisma generate --schema=./schema.prisma
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN apt-get update && apt-get upgrade -y \
libxml2 \
libexpat1 \
openssl \
libssl3 \
git \
libkrb5-3 \
libglib2.0-0 \
wget \
libaom3 \
libxslt1.1 \
libgnutls30 \
libc6 \
&& apt-get install -y --no-install-recommends \
libssl3 \
libatomic1 \
nodejs \
npm \
&& rm -rf /var/lib/apt/lists/* \
&& npm install -g npm@11.12.1 tar@7.5.11 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 \
&& GLOBAL="$(npm root -g)" \
&& find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \
done \
&& find "$GLOBAL/npm" -type d -name "glob" -path "*/node_modules/glob" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \
done \
&& find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \
done \
&& find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \
done \
&& find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \
done \
&& find /usr/local/lib /usr/lib -path "*/node_modules/npm/package.json" -exec \
sed -i 's/"tar": "\^7\.5\.[0-9]*"/"tar": "^7.5.10"/g; s/"minimatch": "\^10\.[0-9.]*"/"minimatch": "^10.2.4"/g' {} + 2>/dev/null \
&& npm cache clean --force \
&& apt-get purge -y npm
WORKDIR /app
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
XDG_CACHE_HOME=/app/.cache \
PATH="/app/.venv/bin:${PATH}"
COPY --from=builder /app /app
EXPOSE 4000/tcp
ENTRYPOINT ["docker/prod_entrypoint.sh"]
CMD ["--port", "4000"]

View file

@ -1,30 +0,0 @@
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin
FROM python:3.13-slim@sha256:739e7213785e88c0f702dcdc12c0973afcbd606dbf021a589cab77d6b00b579d
WORKDIR /app
# Copy the uv binary and the health check script.
COPY --from=uvbin /uv /usr/local/bin/uv
COPY pyproject.toml uv.lock /app/
COPY scripts/health_check/health_check_client.py /app/health_check_client.py
# Resolve and install the health-check dependencies from the project lockfile
# so the runtime image stays self-contained and reproducible.
RUN uv export --frozen --no-default-groups --only-group healthcheck --no-emit-project --no-hashes --output-file /tmp/health-check-requirements.txt \
&& uv pip install --system -r /tmp/health-check-requirements.txt \
&& rm /tmp/health-check-requirements.txt \
&& rm /app/pyproject.toml /app/uv.lock \
&& chmod +x /app/health_check_client.py
# Run as non-root user
RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser
USER appuser
# Health check
HEALTHCHECK --interval=30s --timeout=5s --start-period=5s --retries=3 \
CMD ["python", "/app/health_check_client.py", "--help"]
# Set entrypoint
ENTRYPOINT ["python", "/app/health_check_client.py"]

View file

@ -1,108 +0,0 @@
apiVersion: v1
entries:
litellm-helm:
- apiVersion: v2
appVersion: v1.43.18
created: "2024-08-19T23:58:25.331689+08:00"
dependencies:
- condition: db.deployStandalone
name: postgresql
repository: oci://registry-1.docker.io/bitnamicharts
version: '>=13.3.0'
- condition: redis.enabled
name: redis
repository: oci://registry-1.docker.io/bitnamicharts
version: '>=18.0.0'
description: Call all LLM APIs using the OpenAI format
digest: 0411df3dc42868be8af3ad3e00cb252790e6bd7ad15f5b77f1ca5214573a8531
name: litellm-helm
type: application
urls:
- https://berriai.github.io/litellm/litellm-helm-0.2.3.tgz
version: 0.2.3
postgresql:
- annotations:
category: Database
images: |
- name: os-shell
image: docker.io/bitnami/os-shell:12-debian-12-r16
- name: postgres-exporter
image: docker.io/bitnami/postgres-exporter:0.15.0-debian-12-r14
- name: postgresql
image: docker.io/bitnami/postgresql:16.2.0-debian-12-r6
licenses: Apache-2.0
apiVersion: v2
appVersion: 16.2.0
created: "2024-08-19T23:58:25.335716+08:00"
dependencies:
- name: common
repository: oci://registry-1.docker.io/bitnamicharts
tags:
- bitnami-common
version: 2.x.x
description: PostgreSQL (Postgres) is an open source object-relational database
known for reliability and data integrity. ACID-compliant, it supports foreign
keys, joins, views, triggers and stored procedures.
digest: 3c8125526b06833df32e2f626db34aeaedb29d38f03d15349db6604027d4a167
home: https://bitnami.com
icon: https://bitnami.com/assets/stacks/postgresql/img/postgresql-stack-220x234.png
keywords:
- postgresql
- postgres
- database
- sql
- replication
- cluster
maintainers:
- name: VMware, Inc.
url: https://github.com/bitnami/charts
name: postgresql
sources:
- https://github.com/bitnami/charts/tree/main/bitnami/postgresql
urls:
- https://berriai.github.io/litellm/charts/postgresql-14.3.1.tgz
version: 14.3.1
redis:
- annotations:
category: Database
images: |
- name: kubectl
image: docker.io/bitnami/kubectl:1.29.2-debian-12-r3
- name: os-shell
image: docker.io/bitnami/os-shell:12-debian-12-r16
- name: redis
image: docker.io/bitnami/redis:7.2.4-debian-12-r9
- name: redis-exporter
image: docker.io/bitnami/redis-exporter:1.58.0-debian-12-r4
- name: redis-sentinel
image: docker.io/bitnami/redis-sentinel:7.2.4-debian-12-r7
licenses: Apache-2.0
apiVersion: v2
appVersion: 7.2.4
created: "2024-08-19T23:58:25.339392+08:00"
dependencies:
- name: common
repository: oci://registry-1.docker.io/bitnamicharts
tags:
- bitnami-common
version: 2.x.x
description: Redis(R) is an open source, advanced key-value store. It is often
referred to as a data structure server since keys can contain strings, hashes,
lists, sets and sorted sets.
digest: b2fa1835f673a18002ca864c54fadac3c33789b26f6c5e58e2851b0b14a8f984
home: https://bitnami.com
icon: https://bitnami.com/assets/stacks/redis/img/redis-stack-220x234.png
keywords:
- redis
- keyvalue
- database
maintainers:
- name: VMware, Inc.
url: https://github.com/bitnami/charts
name: redis
sources:
- https://github.com/bitnami/charts/tree/main/bitnami/redis
urls:
- https://berriai.github.io/litellm/charts/redis-18.19.1.tgz
version: 18.19.1
generated: "2024-08-19T23:58:25.322532+08:00"

View file

@ -1,5 +0,0 @@
# Supply-chain hardening
# Packages needing lifecycle scripts: npm rebuild <pkg>
ignore-scripts=true
# Protects local npm install only — npm ci (used in CI) ignores this
min-release-age=3

View file

@ -1,8 +0,0 @@
```
npm install
npm run dev
```
```
npm run deploy
```

File diff suppressed because it is too large Load diff

View file

@ -1,14 +0,0 @@
{
"scripts": {
"dev": "wrangler dev src/index.ts",
"deploy": "wrangler deploy --minify src/index.ts"
},
"dependencies": {
"hono": "4.12.16",
"openai": "4.29.2"
},
"devDependencies": {
"@cloudflare/workers-types": "4.20260501.1",
"wrangler": "4.87.0"
}
}

View file

@ -1,59 +0,0 @@
import { Hono } from 'hono'
import { Context } from 'hono';
import { bearerAuth } from 'hono/bearer-auth'
import OpenAI from "openai";
const openai = new OpenAI({
apiKey: "sk-1234",
baseURL: "https://openai-endpoint.ishaanjaffer0324.workers.dev"
});
async function call_proxy() {
const completion = await openai.chat.completions.create({
messages: [{ role: "system", content: "You are a helpful assistant." }],
model: "gpt-3.5-turbo",
});
return completion
}
const app = new Hono()
// Middleware for API Key Authentication
const apiKeyAuth = async (c: Context, next: Function) => {
const apiKey = c.req.header('Authorization');
if (!apiKey || apiKey !== 'Bearer sk-1234') {
return c.text('Unauthorized', 401);
}
await next();
};
app.use('/*', apiKeyAuth)
app.get('/', (c) => {
return c.text('Hello Hono!')
})
// Handler for chat completions
const chatCompletionHandler = async (c: Context) => {
// Assuming your logic for handling chat completion goes here
// For demonstration, just returning a simple JSON response
const response = await call_proxy()
return c.json(response);
};
// Register the above handler for different POST routes with the apiKeyAuth middleware
app.post('/v1/chat/completions', chatCompletionHandler);
app.post('/chat/completions', chatCompletionHandler);
// Example showing how you might handle dynamic segments within the URL
// Here, using ':model*' to capture the rest of the path as a parameter 'model'
app.post('/openai/deployments/:model*/chat/completions', chatCompletionHandler);
export default app

View file

@ -1,17 +0,0 @@
{
"compilerOptions": {
"target": "ESNext",
"module": "ESNext",
"moduleResolution": "Bundler",
"strict": true,
"lib": [
"ESNext"
],
"types": [
"@cloudflare/workers-types"
],
"jsx": "react-jsx",
"jsxImportSource": "hono/jsx",
"skipLibCheck": true
},
}

View file

@ -1,18 +0,0 @@
name = "my-app"
compatibility_date = "2023-12-01"
# [vars]
# MY_VAR = "my-variable"
# [[kv_namespaces]]
# binding = "MY_KV_NAMESPACE"
# id = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
# [[r2_buckets]]
# binding = "MY_BUCKET"
# bucket_name = "my-bucket"
# [[d1_databases]]
# binding = "DB"
# database_name = "my-database"
# database_id = ""

View file

@ -1,5 +0,0 @@
# Supply-chain hardening
# Packages needing lifecycle scripts: npm rebuild <pkg>
ignore-scripts=true
# Protects local npm install only — npm ci (used in CI) ignores this
min-release-age=3

View file

@ -1,26 +0,0 @@
# Use the specific Node.js v20.11.0 image
FROM node:20.18.1-alpine3.20
# Set the working directory inside the container
WORKDIR /app
# Copy package.json and package-lock.json to the working directory
COPY ./litellm-js/spend-logs/package*.json ./
# Install dependencies
RUN npm ci
# Install Prisma globally
RUN npm install -g prisma
# Copy the rest of the application code
COPY ./litellm-js/spend-logs .
# Generate Prisma client
RUN npx prisma generate
# Expose the port that the Node.js server will run on
EXPOSE 3000
# Command to run the Node.js app with npm run dev
CMD ["npm", "run", "dev"]

View file

@ -1,8 +0,0 @@
```
npm install
npm run dev
```
```
open http://localhost:3000
```

View file

@ -1,597 +0,0 @@
{
"name": "spend-logs",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"dependencies": {
"@hono/node-server": "1.19.13",
"hono": "4.12.16"
},
"devDependencies": {
"@types/node": "20.19.25",
"tsx": "4.20.6"
}
},
"node_modules/@esbuild/aix-ppc64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/aix-ppc64/-/aix-ppc64-0.25.12.tgz",
"integrity": "sha512-Hhmwd6CInZ3dwpuGTF8fJG6yoWmsToE+vYgD4nytZVxcu1ulHpUQRAB1UJ8+N1Am3Mz4+xOByoQoSZf4D+CpkA==",
"cpu": [
"ppc64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"aix"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/android-arm": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/android-arm/-/android-arm-0.25.12.tgz",
"integrity": "sha512-VJ+sKvNA/GE7Ccacc9Cha7bpS8nyzVv0jdVgwNDaR4gDMC/2TTRc33Ip8qrNYUcpkOHUT5OZ0bUcNNVZQ9RLlg==",
"cpu": [
"arm"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"android"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/android-arm64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/android-arm64/-/android-arm64-0.25.12.tgz",
"integrity": "sha512-6AAmLG7zwD1Z159jCKPvAxZd4y/VTO0VkprYy+3N2FtJ8+BQWFXU+OxARIwA46c5tdD9SsKGZ/1ocqBS/gAKHg==",
"cpu": [
"arm64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"android"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/android-x64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/android-x64/-/android-x64-0.25.12.tgz",
"integrity": "sha512-5jbb+2hhDHx5phYR2By8GTWEzn6I9UqR11Kwf22iKbNpYrsmRB18aX/9ivc5cabcUiAT/wM+YIZ6SG9QO6a8kg==",
"cpu": [
"x64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"android"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/darwin-arm64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/darwin-arm64/-/darwin-arm64-0.25.12.tgz",
"integrity": "sha512-N3zl+lxHCifgIlcMUP5016ESkeQjLj/959RxxNYIthIg+CQHInujFuXeWbWMgnTo4cp5XVHqFPmpyu9J65C1Yg==",
"cpu": [
"arm64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"darwin"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/darwin-x64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/darwin-x64/-/darwin-x64-0.25.12.tgz",
"integrity": "sha512-HQ9ka4Kx21qHXwtlTUVbKJOAnmG1ipXhdWTmNXiPzPfWKpXqASVcWdnf2bnL73wgjNrFXAa3yYvBSd9pzfEIpA==",
"cpu": [
"x64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"darwin"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/freebsd-arm64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/freebsd-arm64/-/freebsd-arm64-0.25.12.tgz",
"integrity": "sha512-gA0Bx759+7Jve03K1S0vkOu5Lg/85dou3EseOGUes8flVOGxbhDDh/iZaoek11Y8mtyKPGF3vP8XhnkDEAmzeg==",
"cpu": [
"arm64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"freebsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/freebsd-x64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/freebsd-x64/-/freebsd-x64-0.25.12.tgz",
"integrity": "sha512-TGbO26Yw2xsHzxtbVFGEXBFH0FRAP7gtcPE7P5yP7wGy7cXK2oO7RyOhL5NLiqTlBh47XhmIUXuGciXEqYFfBQ==",
"cpu": [
"x64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"freebsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-arm": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/linux-arm/-/linux-arm-0.25.12.tgz",
"integrity": "sha512-lPDGyC1JPDou8kGcywY0YILzWlhhnRjdof3UlcoqYmS9El818LLfJJc3PXXgZHrHCAKs/Z2SeZtDJr5MrkxtOw==",
"cpu": [
"arm"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-arm64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/linux-arm64/-/linux-arm64-0.25.12.tgz",
"integrity": "sha512-8bwX7a8FghIgrupcxb4aUmYDLp8pX06rGh5HqDT7bB+8Rdells6mHvrFHHW2JAOPZUbnjUpKTLg6ECyzvas2AQ==",
"cpu": [
"arm64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-ia32": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/linux-ia32/-/linux-ia32-0.25.12.tgz",
"integrity": "sha512-0y9KrdVnbMM2/vG8KfU0byhUN+EFCny9+8g202gYqSSVMonbsCfLjUO+rCci7pM0WBEtz+oK/PIwHkzxkyharA==",
"cpu": [
"ia32"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-loong64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/linux-loong64/-/linux-loong64-0.25.12.tgz",
"integrity": "sha512-h///Lr5a9rib/v1GGqXVGzjL4TMvVTv+s1DPoxQdz7l/AYv6LDSxdIwzxkrPW438oUXiDtwM10o9PmwS/6Z0Ng==",
"cpu": [
"loong64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-mips64el": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/linux-mips64el/-/linux-mips64el-0.25.12.tgz",
"integrity": "sha512-iyRrM1Pzy9GFMDLsXn1iHUm18nhKnNMWscjmp4+hpafcZjrr2WbT//d20xaGljXDBYHqRcl8HnxbX6uaA/eGVw==",
"cpu": [
"mips64el"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-ppc64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/linux-ppc64/-/linux-ppc64-0.25.12.tgz",
"integrity": "sha512-9meM/lRXxMi5PSUqEXRCtVjEZBGwB7P/D4yT8UG/mwIdze2aV4Vo6U5gD3+RsoHXKkHCfSxZKzmDssVlRj1QQA==",
"cpu": [
"ppc64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-riscv64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/linux-riscv64/-/linux-riscv64-0.25.12.tgz",
"integrity": "sha512-Zr7KR4hgKUpWAwb1f3o5ygT04MzqVrGEGXGLnj15YQDJErYu/BGg+wmFlIDOdJp0PmB0lLvxFIOXZgFRrdjR0w==",
"cpu": [
"riscv64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-s390x": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/linux-s390x/-/linux-s390x-0.25.12.tgz",
"integrity": "sha512-MsKncOcgTNvdtiISc/jZs/Zf8d0cl/t3gYWX8J9ubBnVOwlk65UIEEvgBORTiljloIWnBzLs4qhzPkJcitIzIg==",
"cpu": [
"s390x"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-x64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/linux-x64/-/linux-x64-0.25.12.tgz",
"integrity": "sha512-uqZMTLr/zR/ed4jIGnwSLkaHmPjOjJvnm6TVVitAa08SLS9Z0VM8wIRx7gWbJB5/J54YuIMInDquWyYvQLZkgw==",
"cpu": [
"x64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/netbsd-arm64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/netbsd-arm64/-/netbsd-arm64-0.25.12.tgz",
"integrity": "sha512-xXwcTq4GhRM7J9A8Gv5boanHhRa/Q9KLVmcyXHCTaM4wKfIpWkdXiMog/KsnxzJ0A1+nD+zoecuzqPmCRyBGjg==",
"cpu": [
"arm64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"netbsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/netbsd-x64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/netbsd-x64/-/netbsd-x64-0.25.12.tgz",
"integrity": "sha512-Ld5pTlzPy3YwGec4OuHh1aCVCRvOXdH8DgRjfDy/oumVovmuSzWfnSJg+VtakB9Cm0gxNO9BzWkj6mtO1FMXkQ==",
"cpu": [
"x64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"netbsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/openbsd-arm64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/openbsd-arm64/-/openbsd-arm64-0.25.12.tgz",
"integrity": "sha512-fF96T6KsBo/pkQI950FARU9apGNTSlZGsv1jZBAlcLL1MLjLNIWPBkj5NlSz8aAzYKg+eNqknrUJ24QBybeR5A==",
"cpu": [
"arm64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"openbsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/openbsd-x64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/openbsd-x64/-/openbsd-x64-0.25.12.tgz",
"integrity": "sha512-MZyXUkZHjQxUvzK7rN8DJ3SRmrVrke8ZyRusHlP+kuwqTcfWLyqMOE3sScPPyeIXN/mDJIfGXvcMqCgYKekoQw==",
"cpu": [
"x64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"openbsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/openharmony-arm64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/openharmony-arm64/-/openharmony-arm64-0.25.12.tgz",
"integrity": "sha512-rm0YWsqUSRrjncSXGA7Zv78Nbnw4XL6/dzr20cyrQf7ZmRcsovpcRBdhD43Nuk3y7XIoW2OxMVvwuRvk9XdASg==",
"cpu": [
"arm64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"openharmony"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/sunos-x64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/sunos-x64/-/sunos-x64-0.25.12.tgz",
"integrity": "sha512-3wGSCDyuTHQUzt0nV7bocDy72r2lI33QL3gkDNGkod22EsYl04sMf0qLb8luNKTOmgF/eDEDP5BFNwoBKH441w==",
"cpu": [
"x64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"sunos"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/win32-arm64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/win32-arm64/-/win32-arm64-0.25.12.tgz",
"integrity": "sha512-rMmLrur64A7+DKlnSuwqUdRKyd3UE7oPJZmnljqEptesKM8wx9J8gx5u0+9Pq0fQQW8vqeKebwNXdfOyP+8Bsg==",
"cpu": [
"arm64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/win32-ia32": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/win32-ia32/-/win32-ia32-0.25.12.tgz",
"integrity": "sha512-HkqnmmBoCbCwxUKKNPBixiWDGCpQGVsrQfJoVGYLPT41XWF8lHuE5N6WhVia2n4o5QK5M4tYr21827fNhi4byQ==",
"cpu": [
"ia32"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/win32-x64": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/@esbuild/win32-x64/-/win32-x64-0.25.12.tgz",
"integrity": "sha512-alJC0uCZpTFrSL0CCDjcgleBXPnCrEAhTBILpeAp7M/OFgoqtAetfBzX0xM00MUsVVPpVjlPuMbREqnZCXaTnA==",
"cpu": [
"x64"
],
"dev": true,
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@hono/node-server": {
"version": "1.19.13",
"resolved": "https://registry.npmjs.org/@hono/node-server/-/node-server-1.19.13.tgz",
"integrity": "sha512-TsQLe4i2gvoTtrHje625ngThGBySOgSK3Xo2XRYOdqGN1teR8+I7vchQC46uLJi8OF62YTYA3AhSpumtkhsaKQ==",
"license": "MIT",
"engines": {
"node": ">=18.14.1"
},
"peerDependencies": {
"hono": "^4"
}
},
"node_modules/@types/node": {
"version": "20.19.25",
"resolved": "https://registry.npmjs.org/@types/node/-/node-20.19.25.tgz",
"integrity": "sha512-ZsJzA5thDQMSQO788d7IocwwQbI8B5OPzmqNvpf3NY/+MHDAS759Wo0gd2WQeXYt5AAAQjzcrTVC6SKCuYgoCQ==",
"dev": true,
"license": "MIT",
"dependencies": {
"undici-types": "~6.21.0"
}
},
"node_modules/esbuild": {
"version": "0.25.12",
"resolved": "https://registry.npmjs.org/esbuild/-/esbuild-0.25.12.tgz",
"integrity": "sha512-bbPBYYrtZbkt6Os6FiTLCTFxvq4tt3JKall1vRwshA3fdVztsLAatFaZobhkBC8/BrPetoa0oksYoKXoG4ryJg==",
"dev": true,
"hasInstallScript": true,
"license": "MIT",
"bin": {
"esbuild": "bin/esbuild"
},
"engines": {
"node": ">=18"
},
"optionalDependencies": {
"@esbuild/aix-ppc64": "0.25.12",
"@esbuild/android-arm": "0.25.12",
"@esbuild/android-arm64": "0.25.12",
"@esbuild/android-x64": "0.25.12",
"@esbuild/darwin-arm64": "0.25.12",
"@esbuild/darwin-x64": "0.25.12",
"@esbuild/freebsd-arm64": "0.25.12",
"@esbuild/freebsd-x64": "0.25.12",
"@esbuild/linux-arm": "0.25.12",
"@esbuild/linux-arm64": "0.25.12",
"@esbuild/linux-ia32": "0.25.12",
"@esbuild/linux-loong64": "0.25.12",
"@esbuild/linux-mips64el": "0.25.12",
"@esbuild/linux-ppc64": "0.25.12",
"@esbuild/linux-riscv64": "0.25.12",
"@esbuild/linux-s390x": "0.25.12",
"@esbuild/linux-x64": "0.25.12",
"@esbuild/netbsd-arm64": "0.25.12",
"@esbuild/netbsd-x64": "0.25.12",
"@esbuild/openbsd-arm64": "0.25.12",
"@esbuild/openbsd-x64": "0.25.12",
"@esbuild/openharmony-arm64": "0.25.12",
"@esbuild/sunos-x64": "0.25.12",
"@esbuild/win32-arm64": "0.25.12",
"@esbuild/win32-ia32": "0.25.12",
"@esbuild/win32-x64": "0.25.12"
}
},
"node_modules/fsevents": {
"version": "2.3.3",
"resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz",
"integrity": "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==",
"dev": true,
"hasInstallScript": true,
"license": "MIT",
"optional": true,
"os": [
"darwin"
],
"engines": {
"node": "^8.16.0 || ^10.6.0 || >=11.0.0"
}
},
"node_modules/get-tsconfig": {
"version": "4.14.0",
"resolved": "https://registry.npmjs.org/get-tsconfig/-/get-tsconfig-4.14.0.tgz",
"integrity": "sha512-yTb+8DXzDREzgvYmh6s9vHsSVCHeC0G3PI5bEXNBHtmshPnO+S5O7qgLEOn0I5QvMy6kpZN8K1NKGyilLb93wA==",
"dev": true,
"license": "MIT",
"dependencies": {
"resolve-pkg-maps": "^1.0.0"
},
"funding": {
"url": "https://github.com/privatenumber/get-tsconfig?sponsor=1"
}
},
"node_modules/hono": {
"version": "4.12.16",
"resolved": "https://registry.npmjs.org/hono/-/hono-4.12.16.tgz",
"integrity": "sha512-jN0ZewiNAWSe5khM3EyCmBb250+b40wWbwNILNfEvq84VREWwOIkuUsFONk/3i3nqkz7Oe1PcpM2mwQEK2L9Kg==",
"license": "MIT",
"engines": {
"node": ">=16.9.0"
}
},
"node_modules/resolve-pkg-maps": {
"version": "1.0.0",
"resolved": "https://registry.npmjs.org/resolve-pkg-maps/-/resolve-pkg-maps-1.0.0.tgz",
"integrity": "sha512-seS2Tj26TBVOC2NIc2rOe2y2ZO7efxITtLZcGSOnHHNOQ7CkiUBfw0Iw2ck6xkIhPwLhKNLS8BO+hEpngQlqzw==",
"dev": true,
"license": "MIT",
"funding": {
"url": "https://github.com/privatenumber/resolve-pkg-maps?sponsor=1"
}
},
"node_modules/tsx": {
"version": "4.20.6",
"resolved": "https://registry.npmjs.org/tsx/-/tsx-4.20.6.tgz",
"integrity": "sha512-ytQKuwgmrrkDTFP4LjR0ToE2nqgy886GpvRSpU0JAnrdBYppuY5rLkRUYPU1yCryb24SsKBTL/hlDQAEFVwtZg==",
"dev": true,
"license": "MIT",
"dependencies": {
"esbuild": "~0.25.0",
"get-tsconfig": "^4.7.5"
},
"bin": {
"tsx": "dist/cli.mjs"
},
"engines": {
"node": ">=18.0.0"
},
"optionalDependencies": {
"fsevents": "~2.3.3"
}
},
"node_modules/undici-types": {
"version": "6.21.0",
"resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.21.0.tgz",
"integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==",
"dev": true,
"license": "MIT"
}
}
}

View file

@ -1,13 +0,0 @@
{
"scripts": {
"dev": "tsx watch src/index.ts"
},
"dependencies": {
"@hono/node-server": "1.19.13",
"hono": "4.12.16"
},
"devDependencies": {
"@types/node": "20.19.25",
"tsx": "4.20.6"
}
}

View file

@ -1,29 +0,0 @@
generator client {
provider = "prisma-client-js"
}
datasource client {
provider = "postgresql"
url = env("DATABASE_URL")
}
model LiteLLM_SpendLogs {
request_id String @id
call_type String
api_key String @default("")
spend Float @default(0.0)
total_tokens Int @default(0)
prompt_tokens Int @default(0)
completion_tokens Int @default(0)
startTime DateTime
endTime DateTime
model String @default("")
api_base String @default("")
user String @default("")
metadata Json @default("{}")
cache_hit String @default("")
cache_key String @default("")
request_tags Json @default("[]")
team_id String?
end_user String?
}

View file

@ -1,32 +0,0 @@
export type LiteLLM_IncrementSpend = {
key_transactions: Array<LiteLLM_IncrementObject>, // [{"key": spend},..]
user_transactions: Array<LiteLLM_IncrementObject>,
team_transactions: Array<LiteLLM_IncrementObject>,
spend_logs_transactions: Array<LiteLLM_SpendLogs>
}
export type LiteLLM_IncrementObject = {
key: string,
spend: number
}
export type LiteLLM_SpendLogs = {
request_id: string; // @id means it's a unique identifier
call_type: string;
api_key: string; // @default("") means it defaults to an empty string if not provided
spend: number; // Float in Prisma corresponds to number in TypeScript
total_tokens: number; // Int in Prisma corresponds to number in TypeScript
prompt_tokens: number;
completion_tokens: number;
startTime: Date; // DateTime in Prisma corresponds to Date in TypeScript
endTime: Date;
model: string; // @default("") means it defaults to an empty string if not provided
api_base: string;
user: string;
metadata: any; // Json type in Prisma is represented by any in TypeScript; could also use a more specific type if the structure of JSON is known
cache_hit: string;
cache_key: string;
request_tags: any; // Similarly, this could be an array or a more specific type depending on the expected structure
team_id?: string | null; // ? indicates it's optional and can be undefined, but could also be null if not provided
end_user?: string | null;
};

View file

@ -1,84 +0,0 @@
import { serve } from '@hono/node-server'
import { Hono } from 'hono'
import { PrismaClient } from '@prisma/client'
import {LiteLLM_SpendLogs, LiteLLM_IncrementSpend, LiteLLM_IncrementObject} from './_types'
const app = new Hono()
const prisma = new PrismaClient()
// In-memory storage for logs
let spend_logs: LiteLLM_SpendLogs[] = [];
const key_logs: LiteLLM_IncrementObject[] = [];
const user_logs: LiteLLM_IncrementObject[] = [];
const transaction_logs: LiteLLM_IncrementObject[] = [];
app.get('/', (c) => {
return c.text('Hello Hono!')
})
const MIN_LOGS = 1; // Minimum number of logs needed to initiate a flush
const FLUSH_INTERVAL = 5000; // Time in ms to wait before trying to flush again
const BATCH_SIZE = 100; // Preferred size of each batch to write to the database
const MAX_LOGS_PER_INTERVAL = 1000; // Maximum number of logs to flush in a single interval
const flushLogsToDb = async () => {
if (spend_logs.length >= MIN_LOGS) {
// Limit the logs to process in this interval to MAX_LOGS_PER_INTERVAL or less
const logsToProcess = spend_logs.slice(0, MAX_LOGS_PER_INTERVAL);
for (let i = 0; i < logsToProcess.length; i += BATCH_SIZE) {
// Create subarray for current batch, ensuring it doesn't exceed the BATCH_SIZE
const batch = logsToProcess.slice(i, i + BATCH_SIZE);
// Convert datetime strings to Date objects
const batchWithDates = batch.map(entry => ({
...entry,
startTime: new Date(entry.startTime),
endTime: new Date(entry.endTime),
// Repeat for any other DateTime fields you may have
}));
await prisma.liteLLM_SpendLogs.createMany({
data: batchWithDates,
});
console.log(`Flushed ${batch.length} logs to the DB.`);
}
// Remove the processed logs from spend_logs
spend_logs = spend_logs.slice(logsToProcess.length);
console.log(`${logsToProcess.length} logs processed. Remaining in queue: ${spend_logs.length}`);
} else {
// This will ensure it doesn't falsely claim "No logs to flush." when it's merely below the MIN_LOGS threshold.
if(spend_logs.length > 0) {
console.log(`Accumulating logs. Currently at ${spend_logs.length}, waiting for at least ${MIN_LOGS}.`);
} else {
console.log("No logs to flush.");
}
}
};
// Setup interval for attempting to flush the logs
setInterval(flushLogsToDb, FLUSH_INTERVAL);
// Route to receive log messages
app.post('/spend/update', async (c) => {
const incomingLogs = await c.req.json<LiteLLM_SpendLogs[]>();
spend_logs.push(...incomingLogs);
console.log(`Received and stored ${incomingLogs.length} logs. Total logs in memory: ${spend_logs.length}`);
return c.json({ message: `Successfully stored ${incomingLogs.length} logs` });
});
const port = 3000
console.log(`Server is running on port ${port}`)
serve({
fetch: app.fetch,
port
})

View file

@ -1,13 +0,0 @@
{
"compilerOptions": {
"target": "ESNext",
"module": "ESNext",
"moduleResolution": "Bundler",
"strict": true,
"types": [
"node"
],
"jsx": "react-jsx",
"jsxImportSource": "hono/jsx",
}
}

View file

@ -206,6 +206,7 @@ add_user_information_to_llm_headers: Optional[bool] = (
)
store_audit_logs = False # Enterprise feature, allow users to see audit logs
skip_system_message_in_guardrail: bool = False
skip_tool_message_in_guardrail: bool = False
### end of callbacks #############
email: Optional[str] = (
@ -1425,6 +1426,12 @@ if TYPE_CHECKING:
)
from .llms.datarobot.chat.transformation import DataRobotConfig as DataRobotConfig
from .llms.anthropic.chat.transformation import AnthropicConfig as AnthropicConfig
from .llms.bedrock.claude_platform.transformation import (
BedrockClaudePlatformConfig as BedrockClaudePlatformConfig,
)
from .llms.bedrock.claude_platform.messages_transformation import (
BedrockClaudePlatformMessagesConfig as BedrockClaudePlatformMessagesConfig,
)
from .llms.anthropic.completion.transformation import (
AnthropicTextConfig as AnthropicTextConfig,
)

View file

@ -131,6 +131,7 @@ LLM_CONFIG_NAMES = (
"OpenrouterConfig",
"DataRobotConfig",
"AnthropicConfig",
"BedrockClaudePlatformConfig",
"AnthropicTextConfig",
"GroqSTTConfig",
"TritonConfig",
@ -170,6 +171,7 @@ LLM_CONFIG_NAMES = (
"SagemakerNovaConfig",
"CohereChatConfig",
"AnthropicMessagesConfig",
"BedrockClaudePlatformMessagesConfig",
"AmazonAnthropicClaudeMessagesConfig",
"AmazonMantleMessagesConfig",
"TogetherAIConfig",
@ -610,6 +612,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
"OpenrouterConfig": (".llms.openrouter.chat.transformation", "OpenrouterConfig"),
"DataRobotConfig": (".llms.datarobot.chat.transformation", "DataRobotConfig"),
"AnthropicConfig": (".llms.anthropic.chat.transformation", "AnthropicConfig"),
"BedrockClaudePlatformConfig": (
".llms.bedrock.claude_platform.transformation",
"BedrockClaudePlatformConfig",
),
"AnthropicTextConfig": (
".llms.anthropic.completion.transformation",
"AnthropicTextConfig",
@ -712,6 +718,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
".llms.anthropic.experimental_pass_through.messages.transformation",
"AnthropicMessagesConfig",
),
"BedrockClaudePlatformMessagesConfig": (
".llms.bedrock.claude_platform.messages_transformation",
"BedrockClaudePlatformMessagesConfig",
),
"AmazonAnthropicClaudeMessagesConfig": (
".llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation",
"AmazonAnthropicClaudeMessagesConfig",

View file

@ -19,6 +19,7 @@ import redis.asyncio as async_redis # type: ignore
from litellm import get_secret, get_secret_str
from litellm._redis_credential_provider import (
AzureADCredentialProvider,
GCPIAMCredentialProvider,
_generate_gcp_iam_access_token,
)
@ -27,6 +28,8 @@ from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from ._logging import verbose_logger
AZURE_REDIS_SCOPE = "https://redis.azure.com/.default"
def _get_redis_kwargs():
arg_spec = inspect.getfullargspec(redis.Redis)
@ -43,6 +46,10 @@ def _get_redis_kwargs():
"redis_connect_func",
"gcp_service_account",
"gcp_ssl_ca_certs",
"azure_redis_ad_token",
"azure_client_id",
"azure_tenant_id",
"azure_client_secret",
]
available_args = [x for x in arg_spec.args if x not in exclude_args] + include_args
@ -89,6 +96,10 @@ def _get_redis_cluster_kwargs(client=None):
) # Needed for sync clusters and IAM detection
available_args.append("gcp_service_account")
available_args.append("gcp_ssl_ca_certs")
available_args.append("azure_redis_ad_token")
available_args.append("azure_client_id")
available_args.append("azure_tenant_id")
available_args.append("azure_client_secret")
available_args.append("max_connections")
return available_args
@ -155,6 +166,125 @@ def create_gcp_iam_redis_connect_func(
return iam_connect
def _build_azure_credential(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
):
"""
Build a long-lived Azure credential object.
Azure SDK credentials cache tokens internally and handle expiry/refresh
transparently, so this should be called once and the result reused.
"""
try:
from azure.identity import (
ClientSecretCredential,
DefaultAzureCredential,
ManagedIdentityCredential,
)
except ImportError:
raise ImportError(
"azure-identity is required for Azure AD Redis authentication. "
"Install it with: pip install azure-identity"
)
_client_id = azure_client_id or os.environ.get("AZURE_CLIENT_ID")
_tenant_id = azure_tenant_id or os.environ.get("AZURE_TENANT_ID")
_client_secret = azure_client_secret or os.environ.get("AZURE_CLIENT_SECRET")
if _client_id and _tenant_id and _client_secret:
return ClientSecretCredential(
client_id=_client_id,
tenant_id=_tenant_id,
client_secret=_client_secret,
)
elif _client_id:
return ManagedIdentityCredential(client_id=_client_id)
else:
return DefaultAzureCredential()
def _generate_azure_ad_redis_token(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
) -> str:
"""
One-shot helper that builds a credential and fetches a single Azure AD
access token for Redis. Each call rebuilds the credential and performs a
network round-trip, so it should not be used in steady-state Redis flows
— the sync (``create_azure_ad_redis_connect_func``) and async paths
(``AzureADCredentialProvider``) keep the credential alive across
connections so the Azure SDK's internal cache + silent refresh apply.
"""
credential = _build_azure_credential(
azure_client_id=azure_client_id,
azure_tenant_id=azure_tenant_id,
azure_client_secret=azure_client_secret,
)
token = credential.get_token(AZURE_REDIS_SCOPE)
return token.token
def create_azure_ad_redis_connect_func(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
) -> Callable:
"""
Creates a custom Redis connection function for Azure AD authentication.
Used for sync Redis clients. The credential is created once (captured by the
closure) and reused across connections — the Azure SDK handles token caching
and silent renewal internally. Only ``get_token`` is called per connection.
"""
credential = _build_azure_credential(
azure_client_id=azure_client_id,
azure_tenant_id=azure_tenant_id,
azure_client_secret=azure_client_secret,
)
def ad_connect(self):
"""Initialize the connection and authenticate using Azure AD"""
from redis.exceptions import (
AuthenticationError,
AuthenticationWrongNumberOfArgsError,
)
from redis.utils import str_if_bytes
self._parser.on_connect(self)
access_token = credential.get_token(AZURE_REDIS_SCOPE).token
# Only include username when explicitly set — sending AUTH "" <token>
# is invalid for most ACL-configured Azure Redis instances.
username = os.environ.get("REDIS_USERNAME", "")
if username:
auth_args = (username, access_token)
else:
auth_args = (access_token,)
self.send_command("AUTH", *auth_args, check_health=False)
try:
auth_response = self.read_response()
except AuthenticationWrongNumberOfArgsError:
# Fallback: try with just the token (Redis < 6 / no ACL)
self.send_command("AUTH", access_token, check_health=False)
auth_response = self.read_response()
if str_if_bytes(auth_response) != "OK":
raise AuthenticationError("Azure AD authentication failed for Redis")
# Attach the live credential object so async paths can wrap it in
# AzureADCredentialProvider for refresh-aware token retrieval. The raw
# client_id/tenant_id/secret are intentionally NOT exposed here — the
# credential closure already holds them.
ad_connect._azure_credential = credential # type: ignore[attr-defined]
return ad_connect
def get_redis_url_from_environment():
if "REDIS_URL" in os.environ:
return os.environ["REDIS_URL"]
@ -179,7 +309,7 @@ def get_redis_url_from_environment():
return f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}"
def _get_redis_client_logic(**env_overrides):
def _get_redis_client_logic(**env_overrides): # noqa: PLR0915
"""
Common functionality across sync + async redis client implementations
"""
@ -253,6 +383,52 @@ def _get_redis_client_logic(**env_overrides):
if _gcp_ssl_ca_certs and redis_kwargs.get("ssl", False):
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
# Handle Azure AD authentication (after GCP IAM block)
_azure_redis_ad_token = redis_kwargs.get("azure_redis_ad_token") or get_secret(
"REDIS_AZURE_AD_TOKEN"
)
_azure_ad_enabled = (
_azure_redis_ad_token is not None
and str(_azure_redis_ad_token).lower() == "true"
)
if _azure_ad_enabled and _gcp_service_account is not None:
verbose_logger.warning(
"Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. "
"Using GCP IAM. Remove one to avoid misconfiguration."
)
if _azure_ad_enabled and _gcp_service_account is None:
_azure_client_id = redis_kwargs.get("azure_client_id") or get_secret_str(
"AZURE_CLIENT_ID"
)
_azure_tenant_id = redis_kwargs.get("azure_tenant_id") or get_secret_str(
"AZURE_TENANT_ID"
)
_azure_client_secret = redis_kwargs.get(
"azure_client_secret"
) or get_secret_str("AZURE_CLIENT_SECRET")
verbose_logger.debug("Setting up Azure AD authentication for Redis.")
redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func(
azure_client_id=_azure_client_id,
azure_tenant_id=_azure_tenant_id,
azure_client_secret=_azure_client_secret,
)
# Marker for async paths to detect Azure AD auth. The live credential
# object is attached separately as `_azure_credential` by
# `create_azure_ad_redis_connect_func`; the raw client_id/tenant_id/secret
# are intentionally NOT exposed on the function to avoid leaking
# credentials via inspection or logging.
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True # type: ignore[attr-defined]
# Always remove Azure-specific kwargs that shouldn't be passed to Redis client
redis_kwargs.pop("azure_redis_ad_token", None)
redis_kwargs.pop("azure_client_id", None)
redis_kwargs.pop("azure_tenant_id", None)
redis_kwargs.pop("azure_client_secret", None)
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
# Only strip host/port/db/password when not routing to a cluster.
# When startup_nodes is also present the cluster path takes priority and
@ -373,7 +549,7 @@ def get_redis_client(**env_overrides):
return redis.Redis(**redis_kwargs)
def get_redis_async_client(
def get_redis_async_client( # noqa: PLR0915
connection_pool: Optional[async_redis.BlockingConnectionPool] = None,
**env_overrides,
) -> Union[async_redis.Redis, async_redis.RedisCluster]:
@ -398,6 +574,14 @@ def get_redis_async_client(
cluster_kwargs["credential_provider"] = GCPIAMCredentialProvider(
redis_connect_func._gcp_service_account
)
# Handle Azure AD authentication for async clusters via CredentialProvider
# so the credential's internal cache + silent refresh runs per connection
# (mirrors GCP IAM above; avoids static-token-baked-in-pool expiry).
elif redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
cluster_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
new_startup_nodes: List[ClusterNode] = []
@ -431,6 +615,22 @@ def get_redis_async_client(
# Check for Redis Sentinel
if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs:
return _init_async_redis_sentinel(redis_kwargs)
# Wrap GCP / Azure AD auth in a CredentialProvider for the standard async
# Redis client. The async client doesn't support redis_connect_func, but it
# does honour credential_provider — which is called per connection, so the
# underlying SDK can refresh tokens silently before they expire.
redis_connect_func = redis_kwargs.pop("redis_connect_func", None)
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(
redis_connect_func._gcp_service_account
)
_pretty_print_redis_config(redis_kwargs=redis_kwargs)
if connection_pool is not None:
@ -464,6 +664,21 @@ def get_redis_connection_pool(
redis_kwargs["max_connections"],
)
return async_redis.BlockingConnectionPool.from_url(**pool_kwargs)
# Wrap GCP / Azure AD auth in a CredentialProvider so pool-managed
# connections re-fetch tokens via the SDK's internal cache + silent refresh
# rather than reusing a single token captured at pool creation.
redis_connect_func = redis_kwargs.pop("redis_connect_func", None)
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(
redis_connect_func._gcp_service_account
)
connection_class = async_redis.Connection
if "ssl" in redis_kwargs:
connection_class = async_redis.SSLConnection

View file

@ -1,10 +1,13 @@
import asyncio
import threading
import time
from typing import Dict, Tuple
from typing import Any, Dict, Optional, Tuple, Union
from redis.credentials import CredentialProvider # type: ignore[attr-defined]
# Azure AD scope for Redis Cache for Azure.
AZURE_REDIS_SCOPE = "https://redis.azure.com/.default"
# GCP IAM tokens are valid for 1 hour. Cache for 55 minutes to refresh before expiry.
_GCP_IAM_TOKEN_TTL_SECONDS = 3300
@ -101,3 +104,33 @@ class GCPIAMCredentialProvider(CredentialProvider):
_get_cached_gcp_iam_token, self._gcp_service_account
)
return (token,)
class AzureADCredentialProvider(CredentialProvider):
"""
redis.credentials.CredentialProvider implementation that supplies Azure AD
tokens for Redis authentication.
Wraps an azure-identity credential object so the Azure SDK's internal token
cache and silent refresh are honoured on every Redis connection. This avoids
the static-token-baked-in-pool issue where pool-managed connections would
fail authentication after the initial token expired (~1 hour TTL).
"""
def __init__(self, credential: Any, username: Optional[str] = None) -> None:
self._credential = credential
self._username = username
def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:
token = self._credential.get_token(AZURE_REDIS_SCOPE).token
if self._username:
return (self._username, token)
return (token,)
async def get_credentials_async(self) -> Union[Tuple[str], Tuple[str, str]]:
token_obj = await asyncio.to_thread(
self._credential.get_token, AZURE_REDIS_SCOPE
)
if self._username:
return (self._username, token_obj.token)
return (token_obj.token,)

View file

@ -617,24 +617,35 @@ def retrieve_batch(
_is_async = kwargs.pop("aretrieve_batch", False) is True
client = kwargs.get("client", None)
# Check if this is an async invoke ARN (different from regular batch ARN)
# Async invoke ARNs have format: arn:aws(-[^:]+)?:bedrock:[a-z0-9-]{1,20}:[0-9]{12}:async-invoke/[a-z0-9]{12}
if (
batch_id.startswith("arn:aws")
and ":bedrock:" in batch_id
and ":async-invoke/" in batch_id
):
# Handle async invoke status check
# Remove aws_region_name from kwargs to avoid duplicate parameter
async_kwargs = kwargs.copy()
async_kwargs.pop("aws_region_name", None)
# Bedrock has two distinct ARN families that need different APIs:
# * async-invoke ARNs (Twelve Labs Marengo embeddings) -> bedrock-runtime data plane
# * model-invocation-job ARNs (CreateModelInvocationJob batch) -> bedrock control plane
# They live on different AWS service endpoints and can't share a handler.
# ARN shapes:
# arn:aws(-[^:]+)?:bedrock:<region>:<account>:async-invoke/<id>
# arn:aws(-[^:]+)?:bedrock:<region>:<account>:model-invocation-job/<id>
if batch_id.startswith("arn:aws") and ":bedrock:" in batch_id:
if ":async-invoke/" in batch_id:
# Remove aws_region_name from kwargs to avoid duplicate parameter
async_kwargs = kwargs.copy()
async_kwargs.pop("aws_region_name", None)
return BedrockBatchesHandler._handle_async_invoke_status(
batch_id=batch_id,
aws_region_name=kwargs.get("aws_region_name", "us-east-1"),
logging_obj=litellm_logging_obj,
**async_kwargs,
)
return BedrockBatchesHandler._handle_async_invoke_status(
batch_id=batch_id,
aws_region_name=kwargs.get("aws_region_name", "us-east-1"),
logging_obj=litellm_logging_obj,
**async_kwargs,
)
if ":model-invocation-job/" in batch_id:
mij_kwargs = kwargs.copy()
mij_kwargs.pop("aws_region_name", None)
return BedrockBatchesHandler._handle_model_invocation_job_status(
batch_id=batch_id,
aws_region_name=kwargs.get("aws_region_name"),
logging_obj=litellm_logging_obj,
**mij_kwargs,
)
# Try to use provider config first (for providers like bedrock)
model: Optional[str] = kwargs.get("model", None)

View file

@ -119,6 +119,20 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
def __init__(self):
pass
def _normalize_tool_choice_for_responses_api(self, tool_choice: Any) -> Any:
"""Chat tool_choice uses function.name; Responses API expects top-level name."""
if not isinstance(tool_choice, dict) or tool_choice.get("type") != "function":
return tool_choice
if isinstance(tool_choice.get("name"), str) and tool_choice.get("name"):
# Return only Responses shape so stray chat ``function`` key is not sent upstream.
return {"type": "function", "name": tool_choice["name"]}
fn = tool_choice.get("function")
if isinstance(fn, dict):
fn_name = fn.get("name")
if isinstance(fn_name, str) and fn_name:
return {"type": "function", "name": fn_name}
return tool_choice
def _handle_raw_dict_response_item(
self, item: Dict[str, Any], index: int
) -> Tuple[Optional[Any], int]:
@ -309,6 +323,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
text_format = self._transform_response_format_to_text_format(value)
if text_format:
responses_api_request["text"] = text_format # type: ignore
elif key == "tool_choice":
responses_api_request["tool_choice"] = ( # type: ignore[assignment]
self._normalize_tool_choice_for_responses_api(value)
)
elif key in ResponsesAPIOptionalRequestParams.__annotations__.keys():
responses_api_request[key] = value # type: ignore
elif key == "previous_response_id":

View file

@ -57,6 +57,17 @@ LITELLM_PROXY_REQUEST_SPAN_NAME = "Received Proxy Server Request"
RAW_REQUEST_SPAN_NAME = "raw_gen_ai_request"
LITELLM_REQUEST_SPAN_NAME = "litellm_request"
CAPTURE_MODE_NO_CONTENT = "NO_CONTENT"
CAPTURE_MODE_SPAN_ONLY = "SPAN_ONLY"
CAPTURE_MODE_EVENT_ONLY = "EVENT_ONLY"
CAPTURE_MODE_SPAN_AND_EVENT = "SPAN_AND_EVENT"
_VALID_CAPTURE_MODES = {
CAPTURE_MODE_NO_CONTENT,
CAPTURE_MODE_SPAN_ONLY,
CAPTURE_MODE_EVENT_ONLY,
CAPTURE_MODE_SPAN_AND_EVENT,
}
@dataclass
class OpenTelemetryConfig:
@ -71,6 +82,9 @@ class OpenTelemetryConfig:
ignore_context_propagation: Optional[bool] = None
# When True, create a private TracerProvider instead of reusing or setting the global one.
skip_set_global: bool = False
# Programmatic override for OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT.
# One of NO_CONTENT, SPAN_ONLY, EVENT_ONLY, SPAN_AND_EVENT (or "true" as legacy alias).
capture_message_content: Optional[str] = None
def __post_init__(self) -> None:
# If endpoint is specified but exporter is still the default "console",
@ -182,6 +196,9 @@ class OpenTelemetry(CustomLogger):
super().__init__(**kwargs)
self._init_metrics(meter_provider)
self._init_logs(logger_provider)
# Sample env-var / config / message_logging at init so subsequent
# _capture_in_span / _capture_in_event calls are deterministic.
self._capture_mode_cached = self._compute_capture_mode_from_init_state()
self._init_otel_logger_on_litellm_proxy()
@staticmethod
@ -306,6 +323,62 @@ class OpenTelemetry(CustomLogger):
hasattr(self, "callback_name") and self.callback_name == "langfuse_otel"
)
def _compute_capture_mode_from_init_state(self) -> Optional[str]:
"""Sample explicit settings at init. Returns the resolved mode or
None if nothing explicit is set (in which case the legacy
``self.message_logging`` flag is consulted dynamically per request).
``"true"``/``"1"`` map to ``EVENT_ONLY`` per the contrib convention.
``"false"``/``"0"`` map to ``NO_CONTENT``.
Unknown values are ignored.
"""
explicit = self.config.capture_message_content or os.getenv(
"OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"
)
if not explicit:
return None
normalized = explicit.upper()
if normalized in ("TRUE", "1"):
return CAPTURE_MODE_EVENT_ONLY
if normalized in ("FALSE", "0"):
return CAPTURE_MODE_NO_CONTENT
if normalized in _VALID_CAPTURE_MODES:
return normalized
return None
def _resolve_capture_mode(self) -> str:
"""Return the active capture mode for this request.
Precedence:
1. ``litellm.turn_off_message_logging=True`` forces ``NO_CONTENT``
(kill-switch checked dynamically).
2. Explicit setting sampled at init from
``OpenTelemetryConfig.capture_message_content`` or
``OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT``.
3. Legacy ``self.message_logging`` (checked dynamically).
"""
if litellm.turn_off_message_logging:
return CAPTURE_MODE_NO_CONTENT
if self._capture_mode_cached is not None:
return self._capture_mode_cached
return (
CAPTURE_MODE_SPAN_AND_EVENT
if self.message_logging
else CAPTURE_MODE_NO_CONTENT
)
def _capture_in_span(self) -> bool:
return self._resolve_capture_mode() in (
CAPTURE_MODE_SPAN_ONLY,
CAPTURE_MODE_SPAN_AND_EVENT,
)
def _capture_in_event(self) -> bool:
return self._resolve_capture_mode() in (
CAPTURE_MODE_EVENT_ONLY,
CAPTURE_MODE_SPAN_AND_EVENT,
)
def _init_tracing(self, tracer_provider):
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
@ -825,8 +898,7 @@ class OpenTelemetry(CustomLogger):
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode
# only log raw LLM request/response if message_logging is on and not globally turned off
if litellm.turn_off_message_logging or not self.message_logging:
if not self._capture_in_span():
return
litellm_params = kwargs.get("litellm_params", {})
@ -1117,9 +1189,14 @@ class OpenTelemetry(CustomLogger):
}
if role == "tool" and msg.get("id"):
attrs["id"] = msg["id"]
if self.message_logging and msg.get("content"):
capture_event_content = self._capture_in_event()
if capture_event_content and msg.get("content"):
attrs["gen_ai.prompt"] = msg["content"]
body = msg.copy()
if not capture_event_content:
body.pop("content", None)
log_record = SdkLogRecord(
timestamp=self._to_ns(datetime.now()),
trace_id=parent_ctx.trace_id,
@ -1127,7 +1204,7 @@ class OpenTelemetry(CustomLogger):
trace_flags=parent_ctx.trace_flags,
severity_number=SeverityNumber.INFO,
severity_text="INFO",
body=msg.copy(),
body=body,
attributes=attrs,
)
otel_logger.emit(log_record)
@ -1141,14 +1218,15 @@ class OpenTelemetry(CustomLogger):
"finish_reason": choice.get("finish_reason"),
}
body_msg = choice.get("message", {})
if self.message_logging and body_msg.get("content"):
capture_event_content = self._capture_in_event()
if capture_event_content and body_msg.get("content"):
attrs["message.content"] = body_msg["content"]
body = {
"index": idx,
"finish_reason": choice.get("finish_reason"),
"message": {"role": body_msg.get("role", "assistant")},
}
if self.message_logging and body_msg.get("content"):
if capture_event_content and body_msg.get("content"):
body["message"]["content"] = body_msg["content"]
log_record = SdkLogRecord(
@ -1674,9 +1752,7 @@ class OpenTelemetry(CustomLogger):
########## LLM Request Medssages / tools / content Attributes ###########
#########################################################################
if litellm.turn_off_message_logging is True:
return
if self.message_logging is not True:
if not self._capture_in_span():
return
if optional_params.get("tools"):
@ -1695,17 +1771,41 @@ class OpenTelemetry(CustomLogger):
value=safe_dumps(transformed_messages),
)
if kwargs.get("system_instructions"):
transformed_system_instructions = (
self._transform_messages_to_otel_semantic_conventions(
kwargs.get("system_instructions")
# Coalesce the different kwarg names that carry the system
# prompt depending on the call path:
# - "system_instructions" — Vertex AI Gemini chat-completion
# - "instructions" — OpenAI Responses API
# - "system" — Anthropic Messages API
# Use `is not None` rather than truthiness to avoid falsy
# values (e.g. []) falling through to the wrong kwarg.
system_instructions = (
kwargs.get("system_instructions")
if kwargs.get("system_instructions") is not None
else (
kwargs.get("instructions")
if kwargs.get("instructions") is not None
else kwargs.get("system")
)
)
if system_instructions:
if isinstance(system_instructions, str):
# Plain text system prompt — no transformation needed
self.safe_set_attribute(
span=span,
key=SpanAttributes.GEN_AI_SYSTEM_INSTRUCTIONS.value,
value=system_instructions,
)
else:
transformed_system_instructions = (
self._transform_messages_to_otel_semantic_conventions(
system_instructions
)
)
self.safe_set_attribute(
span=span,
key=SpanAttributes.GEN_AI_SYSTEM_INSTRUCTIONS.value,
value=safe_dumps(transformed_system_instructions),
)
)
self.safe_set_attribute(
span=span,
key=SpanAttributes.GEN_AI_SYSTEM_INSTRUCTIONS.value,
value=safe_dumps(transformed_system_instructions),
)
self.safe_set_attribute(
span=span,
@ -1764,6 +1864,57 @@ class OpenTelemetry(CustomLogger):
value=value,
)
elif response_obj.get("output"):
# Responses API: ResponsesAPIResponse has an "output"
# list instead of "choices". Each item with
# type="message" contains a "content" list of
# OutputText objects (type="output_text").
output_items = response_obj.get("output")
output_messages = self._transform_responses_api_output_to_otel(
output_items
)
if output_messages:
self.safe_set_attribute(
span=span,
key=SpanAttributes.GEN_AI_OUTPUT_MESSAGES.value,
value=safe_dumps(output_messages),
)
# Emit per-tool-call span attributes (parity with
# the choices branch that calls _tool_calls_kv_pair).
# Convert Responses API function_call items to the
# ChatCompletionMessageToolCall format expected by
# _tool_calls_kv_pair.
tool_calls = []
for out_item in output_items:
item_d = self._to_dict(out_item)
if item_d and item_d.get("type") == "function_call":
tool_calls.append(
{
"function": {
"name": item_d.get("name", ""),
"arguments": item_d.get("arguments", ""),
}
}
)
if tool_calls:
kv_pairs = OpenTelemetry._tool_calls_kv_pair(tool_calls) # type: ignore
for key, value in kv_pairs.items():
self.safe_set_attribute(
span=span,
key=key,
value=value,
)
# Extract finish reason from ResponsesAPIResponse.status
status = response_obj.get("status")
if status:
self.safe_set_attribute(
span=span,
key=SpanAttributes.GEN_AI_RESPONSE_FINISH_REASONS.value,
value=safe_dumps([status]),
)
except Exception as e:
self.handle_callback_failure(
callback_name=self.callback_name or "opentelemetry"
@ -1859,6 +2010,78 @@ class OpenTelemetry(CustomLogger):
transformed.append(transformed_msg)
return transformed
@staticmethod
def _to_dict(obj) -> Optional[dict]:
"""Normalize an object to a plain dict.
Handles three forms that appear in practice:
1. Plain ``dict`` — returned as-is.
2. LiteLLM's ``BaseLiteLLMOpenAIResponseObject`` — exposes a
``.get()`` method that delegates to ``__dict__``.
3. Raw Pydantic v2 models from the ``openai`` SDK (e.g.
``ResponseOutputMessage``, ``ResponseOutputText``) — these do
**not** have ``.get()`` but do have ``.model_dump()``.
Returns ``None`` for anything else so callers can skip it.
"""
if isinstance(obj, dict):
return obj
if hasattr(obj, "get"):
# BaseLiteLLMOpenAIResponseObject duck-type
return obj # type: ignore[return-value]
if hasattr(obj, "model_dump"):
# Raw Pydantic v2 model (e.g. openai SDK types)
return obj.model_dump() # type: ignore[union-attr]
return None
def _transform_responses_api_output_to_otel(self, output: List) -> List[dict]:
"""
Transform Responses API output items into OTEL GenAI 1.38 format.
The Responses API returns output as a list of items, each with a
``type`` field. Message items (``type="message"``) contain a
``content`` list of ``OutputText`` objects with ``type="output_text"``
and ``text`` fields.
Items may be plain dicts, LiteLLM wrapper objects (with ``.get()``),
or raw Pydantic v2 models from the ``openai`` SDK (with
``.model_dump()``). We normalize each item to a dict via
``_to_dict`` before processing.
This method converts them to the same ``{"role": ..., "parts": [...]}``
format used by ``_transform_choices_to_otel_semantic_conventions``.
"""
transformed = []
for raw_item in output:
item = self._to_dict(raw_item)
if item is None:
continue
if item.get("type") == "message":
role = item.get("role", "assistant")
parts = []
for raw_content in item.get("content", []):
content = self._to_dict(raw_content)
if content is None:
continue
if content.get("type") == "output_text":
text = content.get("text", "")
if text:
parts.append({"type": "text", "content": text})
if parts:
transformed.append({"role": role, "parts": parts})
elif item.get("type") == "function_call":
# Surface tool calls from Responses API output
part: dict = {
"type": "tool_call",
"name": item.get("name", ""),
"arguments": item.get("arguments", ""),
}
if item.get("call_id"):
part["id"] = item["call_id"]
transformed.append({"role": "assistant", "parts": [part]})
return transformed
def set_raw_request_attributes(self, span: Span, kwargs, response_obj):
try:
# Only set provider-specific raw payload attributes on this span.

View file

@ -1212,7 +1212,7 @@ class Logging(LiteLLMLoggingBaseClass):
# Log the exact result from the LLM API, for streaming - log the type of response received
litellm.error_logs["POST_CALL"] = locals()
if isinstance(original_response, dict):
original_response = json.dumps(original_response)
original_response = json.dumps(original_response, default=str)
try:
self.model_call_details["input"] = input
self.model_call_details["api_key"] = api_key

View file

@ -23,7 +23,9 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
openai_messages_without_system,
openai_messages_without_tool,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@ -108,6 +110,7 @@ class AnthropicMessagesHandler(BaseTranslation):
return data
skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply)
skip_tool = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
chat_completion_compatible_request = self._translate_to_openai(data)
@ -117,6 +120,8 @@ class AnthropicMessagesHandler(BaseTranslation):
)
if skip_system:
structured_messages = openai_messages_without_system(structured_messages)
if skip_tool:
structured_messages = openai_messages_without_tool(structured_messages)
texts_to_check: List[str] = []
images_to_check: List[str] = []
@ -134,6 +139,7 @@ class AnthropicMessagesHandler(BaseTranslation):
images_to_check=images_to_check,
task_mappings=task_mappings,
skip_system_message=skip_system,
skip_tool_message=skip_tool,
)
# Step 2: Apply guardrail to all texts in batch
@ -198,13 +204,17 @@ class AnthropicMessagesHandler(BaseTranslation):
images_to_check: List[str],
task_mappings: List[Tuple[int, Optional[int]]],
skip_system_message: bool = False,
skip_tool_message: bool = False,
) -> None:
"""
Extract text content and images from a message.
Override this method to customize text/image extraction logic.
"""
if skip_system_message and str(message.get("role") or "").lower() == "system":
role = str(message.get("role") or "").lower()
if skip_system_message and role == "system":
return
if skip_tool_message and role == "tool":
return
content = message.get("content", None)

View file

@ -1809,9 +1809,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
Translate messages to anthropic format.
"""
## VALIDATE REQUEST
"""
Anthropic doesn't support tool calling without `tools=` param specified.
"""
"""Anthropic requires ``tools`` when messages include tool blocks; LiteLLM injects a dummy tool if omitted (no ``modify_params`` needed)."""
from litellm.litellm_core_utils.prompt_templates.factory import (
anthropic_messages_pt,
)
@ -1821,16 +1819,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
and messages is not None
and has_tool_call_blocks(messages)
):
if litellm.modify_params:
optional_params["tools"], _ = self._map_tools(
add_dummy_tool(custom_llm_provider="anthropic")
)
else:
raise litellm.UnsupportedParamsError(
message="Anthropic doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
model="",
llm_provider="anthropic",
)
optional_params["tools"], _ = self._map_tools(
add_dummy_tool(custom_llm_provider="anthropic")
)
# Drop thinking param if thinking is enabled but thinking_blocks are missing
# This prevents the error: "Expected thinking or redacted_thinking, but found tool_use"

View file

@ -14,7 +14,22 @@ def effective_skip_system_message_for_guardrail(guardrail_to_apply: Any) -> bool
return bool(getattr(litellm, "skip_system_message_in_guardrail", False))
def effective_skip_tool_message_for_guardrail(guardrail_to_apply: Any) -> bool:
per = getattr(guardrail_to_apply, "skip_tool_message_in_guardrail", None)
if per is not None:
return bool(per)
import litellm
return bool(getattr(litellm, "skip_tool_message_in_guardrail", False))
def openai_messages_without_system(
messages: List[AllMessageValues],
) -> List[AllMessageValues]:
return [m for m in messages if str((m or {}).get("role") or "").lower() != "system"]
def openai_messages_without_tool(
messages: List[AllMessageValues],
) -> List[AllMessageValues]:
return [m for m in messages if str((m or {}).get("role") or "").lower() != "tool"]

View file

@ -1428,7 +1428,13 @@ class BaseAWSLLM:
def _sign_request(
self,
service_name: Literal["bedrock", "sagemaker", "bedrock-agentcore", "s3vectors"],
service_name: Literal[
"bedrock",
"sagemaker",
"bedrock-agentcore",
"s3vectors",
"aws-external-anthropic",
],
headers: dict,
optional_params: dict,
request_data: dict,

View file

@ -1,8 +1,79 @@
from datetime import datetime
from typing import Any, Optional, cast
from openai.types.batch import BatchRequestCounts
from openai.types.batch import Metadata as OpenAIBatchMetadata
from litellm.types.utils import LiteLLMBatch
# AWS Bedrock model-invocation-job statuses → OpenAI Batch statuses.
# Mirrors the mapping used by `BedrockBatchesConfig.transform_create_batch_response`
# so create / retrieve return consistent statuses.
_BEDROCK_MIJ_STATUS_TO_OPENAI = {
"Submitted": "validating",
"Validating": "validating",
"Scheduled": "validating",
"InProgress": "in_progress",
"Stopping": "cancelling",
"Stopped": "cancelled",
"Completed": "completed",
"PartiallyCompleted": "completed",
"Failed": "failed",
"Expired": "expired",
}
def _extract_region_from_bedrock_arn(arn: str) -> Optional[str]:
"""ARN shape: ``arn:aws:bedrock:<region>:<account>:<type>/<id>``"""
try:
parts = arn.split(":")
if len(parts) >= 4 and parts[2] == "bedrock":
return parts[3] or None
except Exception:
pass
return None
def _extract_job_id_from_arn(arn: str) -> Optional[str]:
"""``arn:aws:bedrock:<region>:<acct>:model-invocation-job/<job-id>`` -> ``<job-id>``."""
if ":model-invocation-job/" not in arn:
return None
return arn.rsplit("/", 1)[-1] or None
def _predict_output_file_uri(
output_prefix: str, input_uri: str, job_id: Optional[str]
) -> Optional[str]:
"""
Compute the deterministic per-job result file URI Bedrock writes to.
Bedrock lays results out as::
<output_prefix>/<job-id>/<basename(input_uri)>.out
We compute it client-side so OpenAI-style ``client.files.content(output_file_id)``
works without an extra S3 ``ListObjectsV2`` round-trip. Returns ``None`` if we
don't have enough info; callers should fall back to the bare prefix.
"""
if not output_prefix or not input_uri or not job_id:
return None
if not output_prefix.endswith("/"):
output_prefix = output_prefix + "/"
input_basename = input_uri.rsplit("/", 1)[-1]
if not input_basename:
return None
return f"{output_prefix}{job_id}/{input_basename}.out"
def _to_epoch(value: Any) -> Optional[int]:
if value is None:
return None
if isinstance(value, (int, float)):
return int(value)
if isinstance(value, datetime):
return int(value.timestamp())
return None
class BedrockBatchesHandler:
"""
@ -97,3 +168,173 @@ class BedrockBatchesHandler:
with concurrent.futures.ThreadPoolExecutor() as executor:
future = executor.submit(run_in_thread)
return future.result()
@staticmethod
def _handle_model_invocation_job_status(
batch_id: str,
aws_region_name: Optional[str] = None,
logging_obj=None,
**kwargs,
) -> "LiteLLMBatch":
"""
Handle ``GetModelInvocationJob`` status check for AWS Bedrock bulk batch
inference jobs (the ARN type returned by ``CreateModelInvocationJob``).
``CreateModelInvocationJob`` lives on the Bedrock **control plane**
(``bedrock.<region>.amazonaws.com``), distinct from the data-plane
``bedrock-runtime`` endpoint that serves Twelve Labs async-invoke ARNs.
The two ARN families therefore can't share a handler — see
``litellm/batches/main.py`` for the dispatch.
Args:
batch_id: A ``arn:aws:bedrock:<region>:<acct>:model-invocation-job/<id>``
ARN (or just the trailing job id; both are accepted by
``GetModelInvocationJob``).
aws_region_name: Region for the boto3 ``bedrock`` client. If omitted,
we fall back to parsing the region out of ``batch_id`` itself.
logging_obj: Optional litellm logging object.
**kwargs: Optional AWS credential overrides
(``aws_access_key_id``, ``aws_secret_access_key``,
``aws_session_token``, ``aws_profile_name``,
``aws_role_name``, ``aws_session_name``,
``aws_web_identity_token``, ``aws_sts_endpoint``,
``aws_external_id``). Unknown keys are ignored.
Returns:
``LiteLLMBatch`` shaped like an OpenAI Batch resource. Note that
``request_counts`` is always ``(0, 0, 0)`` because
``GetModelInvocationJob`` does not surface per-record counts;
callers that need accurate counts should parse
``manifest.json.out`` from the output S3 prefix.
"""
try:
import boto3
except ImportError as exc:
raise ImportError(
"Missing boto3 to call bedrock. Run 'pip install boto3'."
) from exc
# Resolve region: explicit > parsed-from-ARN > us-east-1 (boto3 default).
region = (
aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1"
)
# Resolve credentials through the same path the rest of the bedrock
# provider uses, so model_list / env / role-assumption configs are
# honored. We instantiate BedrockBatchesConfig (which extends
# BaseAWSLLM) lazily to avoid a circular import at module load.
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
creds = BedrockBatchesConfig().get_credentials(
aws_access_key_id=kwargs.get("aws_access_key_id"),
aws_secret_access_key=kwargs.get("aws_secret_access_key"),
aws_session_token=kwargs.get("aws_session_token"),
aws_region_name=region,
aws_session_name=kwargs.get("aws_session_name"),
aws_profile_name=kwargs.get("aws_profile_name"),
aws_role_name=kwargs.get("aws_role_name"),
aws_web_identity_token=kwargs.get("aws_web_identity_token"),
aws_sts_endpoint=kwargs.get("aws_sts_endpoint"),
aws_external_id=kwargs.get("aws_external_id"),
)
client = boto3.client(
"bedrock",
region_name=region,
aws_access_key_id=creds.access_key,
aws_secret_access_key=creds.secret_key,
aws_session_token=creds.token,
)
if logging_obj is not None:
# Use the bare job id in the logged URL so we don't double up the
# `model-invocation-job/` segment when `batch_id` is a full ARN.
# `GetModelInvocationJob` accepts either form, but only the bare id
# produces a sensible-looking URL in logs.
url_path_id = _extract_job_id_from_arn(batch_id) or batch_id
logging_obj.pre_call(
input=batch_id,
api_key="",
additional_args={
"complete_input_dict": {"jobIdentifier": batch_id},
"api_base": (
f"https://bedrock.{region}.amazonaws.com/"
f"model-invocation-job/{url_path_id}"
),
},
)
response = client.get_model_invocation_job(jobIdentifier=batch_id)
if logging_obj is not None:
logging_obj.post_call(
input=batch_id,
api_key="",
original_response=response,
additional_args={"complete_input_dict": {"jobIdentifier": batch_id}},
)
bedrock_status = str(response.get("status", ""))
openai_status = cast(
Any,
_BEDROCK_MIJ_STATUS_TO_OPENAI.get(bedrock_status, "in_progress"),
)
input_uri = (
response.get("inputDataConfig", {})
.get("s3InputDataConfig", {})
.get("s3Uri", "")
)
output_prefix = (
response.get("outputDataConfig", {})
.get("s3OutputDataConfig", {})
.get("s3Uri", "")
)
# Bedrock returns the output *prefix* the user supplied at job creation.
# Actual results land at <prefix>/<job-id>/<basename(input)>.out — we
# surface that single-file URI as `output_file_id` so the OpenAI-style
# download flow works without an extra S3 listing call. We deliberately
# do NOT fall back to the bare prefix when prediction fails: a prefix
# is not a downloadable object, so handing it back as `output_file_id`
# would reproduce the very NoSuchKey bug this handler exists to fix.
# The bare prefix is preserved in metadata for callers that want the
# `manifest.json.out` or want to do their own listing.
job_arn = response.get("jobArn", batch_id)
job_id = _extract_job_id_from_arn(job_arn)
output_file_uri = _predict_output_file_uri(output_prefix, input_uri, job_id)
completed_at = _to_epoch(response.get("endTime"))
# Note: metadata uses "" (not None) for unknown URIs to satisfy the
# OpenAI Batch metadata schema, which is `dict[str, str]`. The
# `output_file_id` field on the LiteLLMBatch itself does carry None
# correctly (see below), so callers should branch on that, not on
# `metadata["output_file_uri"]`.
openai_batch_metadata: OpenAIBatchMetadata = {
"model_arn": response.get("modelId", ""),
"job_arn": job_arn,
"job_name": response.get("jobName", ""),
"failure_message": response.get("message") or "",
"input_s3_uri": input_uri,
"output_s3_uri": output_prefix,
"output_file_uri": output_file_uri or "",
}
return LiteLLMBatch(
id=job_arn,
object="batch",
status=openai_status,
created_at=_to_epoch(response.get("submitTime")) or 0,
in_progress_at=_to_epoch(response.get("lastModifiedTime")),
completed_at=completed_at if openai_status == "completed" else None,
failed_at=completed_at if openai_status == "failed" else None,
cancelled_at=completed_at if openai_status == "cancelled" else None,
expired_at=completed_at if openai_status == "expired" else None,
request_counts=BatchRequestCounts(total=0, completed=0, failed=0),
metadata=openai_batch_metadata,
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id=input_uri,
output_file_id=output_file_uri if openai_status == "completed" else None,
)

View file

@ -0,0 +1,8 @@
from .transformation import (
BedrockClaudePlatformConfig,
)
from .messages_transformation import (
BedrockClaudePlatformMessagesConfig,
)
__all__ = ["BedrockClaudePlatformConfig", "BedrockClaudePlatformMessagesConfig"]

View file

@ -0,0 +1,107 @@
from typing import Literal, Optional, Tuple
import litellm
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.secret_managers.main import get_secret_str
CLAUDE_PLATFORM_SERVICE_NAME: Literal["aws-external-anthropic"] = (
"aws-external-anthropic"
)
CLAUDE_PLATFORM_BEDROCK_ROUTE = "claude_platform/"
def strip_claude_platform_route(model: str) -> str:
if model.startswith(CLAUDE_PLATFORM_BEDROCK_ROUTE):
return model.replace(CLAUDE_PLATFORM_BEDROCK_ROUTE, "", 1)
return model
class BedrockClaudePlatformMixin(BaseAWSLLM):
@staticmethod
def _get_workspace_id(optional_params: dict, litellm_params: dict) -> Optional[str]:
workspace_id = (
optional_params.get("workspace_id")
or litellm_params.get("workspace_id")
or optional_params.get("aws_workspace_id")
or litellm_params.get("aws_workspace_id")
or optional_params.get("anthropic-workspace-id")
or litellm_params.get("anthropic-workspace-id")
)
if workspace_id is None:
workspace_id = optional_params.get(
"anthropic_workspace_id"
) or litellm_params.get("anthropic_workspace_id")
if workspace_id is not None:
return str(workspace_id)
return get_secret_str("ANTHROPIC_AWS_WORKSPACE_ID") or get_secret_str(
"ANTHROPIC_WORKSPACE_ID"
)
def _get_required_aws_region_name(self, optional_params: dict) -> str:
aws_region_name = (
optional_params.get("aws_region_name")
or get_secret_str("AWS_REGION_NAME")
or get_secret_str("AWS_REGION")
or get_secret_str("AWS_DEFAULT_REGION")
)
if aws_region_name is None:
raise litellm.AuthenticationError(
message=(
"Missing AWS region for Claude Platform on AWS. Pass "
"`aws_region_name` or set a standard AWS region environment value."
),
llm_provider="bedrock",
model="",
)
self._validate_aws_region_name(str(aws_region_name))
return str(aws_region_name)
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
api_base = (
api_base
or litellm.api_base
or get_secret_str("ANTHROPIC_AWS_BASE_URL")
or get_secret_str("ANTHROPIC_AWS_API_BASE")
)
if api_base is None:
aws_region_name = self._get_required_aws_region_name(optional_params)
api_base = (
f"https://{CLAUDE_PLATFORM_SERVICE_NAME}.{aws_region_name}.api.aws"
)
if not api_base.endswith("/v1/messages"):
api_base = f"{api_base.rstrip('/')}/v1/messages"
return api_base
def sign_request(
self,
headers: dict,
optional_params: dict,
request_data: dict,
api_base: str,
api_key: Optional[str] = None,
model: Optional[str] = None,
stream: Optional[bool] = None,
fake_stream: Optional[bool] = None,
) -> Tuple[dict, Optional[bytes]]:
if api_key or get_secret_str("ANTHROPIC_AWS_API_KEY"):
return headers, None
return self._sign_request(
service_name=CLAUDE_PLATFORM_SERVICE_NAME,
headers=headers,
optional_params=optional_params,
request_data=request_data,
api_base=api_base,
model=model,
stream=stream,
fake_stream=fake_stream,
)

View file

@ -0,0 +1,71 @@
from typing import Any, Dict, List, Optional, Tuple
import litellm
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
DEFAULT_ANTHROPIC_API_VERSION,
AnthropicMessagesConfig,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.router import GenericLiteLLMParams
from .common_utils import BedrockClaudePlatformMixin, strip_claude_platform_route
class BedrockClaudePlatformMessagesConfig(
BedrockClaudePlatformMixin, AnthropicMessagesConfig
):
def validate_anthropic_messages_environment(
self,
headers: dict,
model: str,
messages: List[Any],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> Tuple[dict, Optional[str]]:
workspace_id = self._get_workspace_id(optional_params, litellm_params)
if workspace_id is None:
raise litellm.AuthenticationError(
message=(
"Missing workspace ID for Claude Platform on AWS. Pass "
"`workspace_id` or configure the provider workspace setting."
),
llm_provider="bedrock",
model=model,
)
resolved_api_key = api_key or get_secret_str("ANTHROPIC_AWS_API_KEY")
headers = {
**headers,
"anthropic-version": headers.get(
"anthropic-version", DEFAULT_ANTHROPIC_API_VERSION
),
"content-type": headers.get("content-type", "application/json"),
"anthropic-workspace-id": workspace_id,
}
if resolved_api_key and "x-api-key" not in headers:
headers["x-api-key"] = resolved_api_key
headers = self._update_headers_with_anthropic_beta(
headers=headers,
optional_params=optional_params,
)
return headers, api_base
def transform_anthropic_messages_request(
self,
model: str,
messages: List[Dict],
anthropic_messages_optional_request_params: Dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> Dict:
return super().transform_anthropic_messages_request(
model=strip_claude_platform_route(model),
messages=messages,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
litellm_params=litellm_params,
headers=headers,
)

View file

@ -0,0 +1,94 @@
from typing import Any, Dict, List, Optional
import litellm
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from .common_utils import BedrockClaudePlatformMixin
class BedrockClaudePlatformConfig(BedrockClaudePlatformMixin, AnthropicConfig):
"""
Bedrock Claude Platform uses Anthropic's Messages API with AWS gateway auth.
"""
@property
def custom_llm_provider(self) -> Optional[str]:
return "bedrock"
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> Dict:
workspace_id = self._get_workspace_id(optional_params, litellm_params)
if workspace_id is None:
raise litellm.AuthenticationError(
message=(
"Missing workspace ID for Claude Platform on AWS. Pass "
"`workspace_id` or configure the provider workspace setting."
),
llm_provider="bedrock",
model=model,
)
api_key = api_key or get_secret_str("ANTHROPIC_AWS_API_KEY")
anthropic_headers = self.get_anthropic_headers(
api_key=api_key,
auth_token=None,
computer_tool_used=self.is_computer_tool_used(
tools=optional_params.get("tools")
),
prompt_caching_set=self.is_cache_control_set(messages=messages),
pdf_used=self.is_pdf_used(messages=messages),
file_id_used=self.is_file_id_used(messages=messages),
mcp_server_used=self.is_mcp_server_used(
mcp_servers=optional_params.get("mcp_servers")
),
web_search_tool_used=self.is_web_search_tool_used(
tools=optional_params.get("tools")
),
tool_search_used=self.is_tool_search_used(
tools=optional_params.get("tools")
),
programmatic_tool_calling_used=self.is_programmatic_tool_calling_used(
tools=optional_params.get("tools")
),
input_examples_used=self.is_input_examples_used(
tools=optional_params.get("tools")
),
effort_used=self.is_effort_used(
optional_params=optional_params, model=model
),
user_anthropic_beta_headers=self._get_user_anthropic_beta_headers(
anthropic_beta_header=headers.get("anthropic-beta")
),
code_execution_tool_used=self.is_code_execution_tool_used(
tools=optional_params.get("tools")
),
container_with_skills_used=self.is_container_with_skills_used(
optional_params=optional_params
),
)
anthropic_headers["anthropic-workspace-id"] = workspace_id
return {**headers, **anthropic_headers}
def get_model_response_iterator(
self,
streaming_response: Any,
sync_stream: bool,
json_mode: Optional[bool] = False,
) -> Any:
from litellm.llms.anthropic.chat.handler import ModelResponseIterator
return ModelResponseIterator(
streaming_response=streaming_response,
sync_stream=sync_stream,
json_mode=bool(json_mode),
)

View file

@ -692,6 +692,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
) -> Literal[
"converse",
"invoke",
"claude_platform",
"converse_like",
"agent",
"agentcore",
@ -706,6 +707,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
str,
Literal[
"invoke",
"claude_platform",
"converse_like",
"converse",
"agent",
@ -716,6 +718,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
],
] = {
"invoke/": "invoke",
"claude_platform/": "claude_platform",
"converse_like/": "converse_like",
"converse/": "converse",
"agent/": "agent",
@ -753,6 +756,36 @@ class BedrockModelInfo(BaseLLMModelInfo):
"""
return "converse/" in model
@staticmethod
def _explicit_claude_platform_route(model: str) -> bool:
"""
Check if the model is an explicit Claude Platform on AWS route.
"""
return "claude_platform/" in model
@staticmethod
def get_claude_platform_model(model: str) -> str:
"""
Strip the Claude Platform route prefix from a Bedrock model name.
"""
return model.replace("claude_platform/", "", 1)
@staticmethod
def map_claude_platform_auth_params(
passed_params: dict, optional_params: dict
) -> dict:
"""
Map Claude Platform route auth params that are not OpenAI request params.
"""
for key in (
"workspace_id",
"aws_workspace_id",
"anthropic_workspace_id",
):
if key in passed_params:
optional_params[key] = passed_params[key]
return optional_params
@staticmethod
def _explicit_invoke_route(model: str) -> bool:
"""
@ -815,6 +848,12 @@ class BedrockModelInfo(BaseLLMModelInfo):
All other routes should return None since they will go through litellm.completion
"""
#########################################################
# Claude Platform route uses Anthropic Messages API via the AWS gateway.
#########################################################
if BedrockModelInfo._explicit_claude_platform_route(model):
return litellm.BedrockClaudePlatformMessagesConfig()
#########################################################
# Converse routes should go through litellm.completion()
if BedrockModelInfo._explicit_converse_route(model):
@ -860,7 +899,9 @@ def get_bedrock_chat_config(model: str):
base_model = BedrockModelInfo.get_base_model(model)
# Handle explicit routes first
if bedrock_route == "converse" or bedrock_route == "converse_like":
if bedrock_route == "claude_platform":
return litellm.BedrockClaudePlatformConfig()
elif bedrock_route == "converse" or bedrock_route == "converse_like":
return litellm.AmazonConverseConfig()
elif bedrock_route == "openai":
return litellm.AmazonBedrockOpenAIConfig()

View file

@ -408,6 +408,47 @@ class AmazonAnthropicClaudeMessagesConfig(
if self._supports_tool_search_on_bedrock(model):
beta_set.add("tool-search-tool-2025-10-19")
@staticmethod
def _filter_context_management_for_bedrock_invoke(
anthropic_messages_request: Dict,
beta_set: set,
) -> None:
"""
Bedrock InvokeModel accepts ``context_management`` only when it carries
``compact_20260112`` edits paired with the ``compact-2026-01-12``
anthropic-beta header. Other edit types (notably ``clear_thinking_20251015``,
which Claude Code sends on every request) are LiteLLM-internal and would
cause Bedrock to 400 with ``"context_management: Extra inputs are not
permitted"``.
Filter the edits list to the supported subset, add the beta header when
compact edits remain, and drop ``context_management`` entirely when no
supported edits are left so the safety-net allowlist can pass it through.
Ref: https://github.com/BerriAI/litellm/issues/27532
"""
cm = anthropic_messages_request.get("context_management")
if not isinstance(cm, dict):
return
edits = cm.get("edits")
if not isinstance(edits, list):
anthropic_messages_request.pop("context_management", None)
return
compact_edits = [
e
for e in edits
if isinstance(e, dict) and e.get("type") == "compact_20260112"
]
if compact_edits:
beta_set.add("compact-2026-01-12")
anthropic_messages_request["context_management"] = {
**cm,
"edits": compact_edits,
}
else:
anthropic_messages_request.pop("context_management", None)
def _convert_output_format_to_inline_schema(
self,
output_format: Dict,
@ -551,6 +592,11 @@ class AmazonAnthropicClaudeMessagesConfig(
if injected_thinking_for_clear_thinking:
beta_set.add("interleaved-thinking-2025-05-14")
self._filter_context_management_for_bedrock_invoke(
anthropic_messages_request=anthropic_messages_request,
beta_set=beta_set,
)
self._get_tool_search_beta_header_for_bedrock(
model=model,
tool_search_used=tool_search_used,
@ -597,8 +643,9 @@ class AmazonAnthropicClaudeMessagesConfig(
anthropic_messages_request.pop("output_config", None)
# 7. Final safety net: filter top-level fields to the Bedrock Invoke allowlist.
# Catches Anthropic-only extensions (context_management, output_config, speed,
# mcp_servers, ...) and any future additions Claude Code may start sending.
# Catches Anthropic-only extensions (output_config, speed, mcp_servers, ...)
# and any future additions Claude Code may start sending. ``context_management``
# has already been pre-filtered to its Bedrock-supported subset above.
allowed = self.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS
stripped = sorted(k for k in anthropic_messages_request if k not in allowed)
if stripped:

View file

@ -3,7 +3,10 @@
from typing import Optional, Union
import litellm
from litellm.utils import _is_explicitly_disabled_factory, _supports_factory
from litellm.utils import (
_is_explicitly_disabled_factory,
_supports_factory,
)
from .gpt_transformation import OpenAIGPTConfig

View file

@ -21,7 +21,9 @@ from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
openai_messages_without_system,
openai_messages_without_tool,
)
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
@ -73,6 +75,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
return data
skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply)
skip_tool = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
texts_to_check: List[str] = []
images_to_check: List[str] = []
@ -91,6 +94,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
text_task_mappings=text_task_mappings,
tool_call_task_mappings=tool_call_task_mappings,
skip_system_message=skip_system,
skip_tool_message=skip_tool,
)
# Step 2: Apply guardrail to all texts and tool calls in batch
@ -102,11 +106,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
inputs["tool_calls"] = tool_calls_to_check # type: ignore
structured_messages = self.get_structured_messages(data)
if structured_messages:
inputs["structured_messages"] = (
openai_messages_without_system(structured_messages)
if skip_system
else structured_messages
)
if skip_system:
structured_messages = openai_messages_without_system(
structured_messages
)
if skip_tool:
structured_messages = openai_messages_without_tool(
structured_messages
)
inputs["structured_messages"] = structured_messages
# Pass tools (function definitions) to the guardrail
tools = data.get("tools")
if tools:
@ -176,13 +184,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
text_task_mappings: List[Tuple[int, Optional[int]]],
tool_call_task_mappings: List[Tuple[int, int]],
skip_system_message: bool = False,
skip_tool_message: bool = False,
) -> None:
"""
Extract text content, images, and tool calls from a message.
Override this method to customize text/image/tool call extraction logic.
"""
if skip_system_message and str(message.get("role") or "").lower() == "system":
role = str(message.get("role") or "").lower()
if skip_system_message and role == "system":
return
if skip_tool_message and role == "tool":
return
content = message.get("content", None)

View file

@ -156,5 +156,17 @@ class OVHCloudAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
text = response_json.get("text") or response_json.get("transcript") or ""
response = TranscriptionResponse(text=text)
# OVHCloud field migration (deadline: 2026-05-11):
# `duration` is replaced by `seconds` in STT responses.
# Prefer `seconds`, fall back to `duration`, normalize to `duration`
# so downstream consumers see a consistent key.
duration = (
response_json["seconds"]
if "seconds" in response_json and response_json["seconds"] is not None
else response_json.get("duration")
)
if duration is not None:
response_json["duration"] = duration
response._hidden_params = response_json
return response

View file

@ -13,6 +13,7 @@ from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.llms.ovhcloud.utils import OVHCloudException
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.openai import AllMessageValues
@ -98,10 +99,16 @@ class OVHCloudChatCompletionStreamingHandler(BaseModelResponseIterator):
new_choices = []
for choice in chunk["choices"]:
if "delta" in choice and "reasoning" in choice["delta"]:
choice["delta"]["reasoning_content"] = choice["delta"].get(
"reasoning"
)
if "delta" in choice:
delta = choice["delta"]
# OVHCloud field migration (deadline: 2026-05-11):
# `reasoning_content` is replaced by `reasoning`.
# Normalise to `reasoning_content` so downstream consumers
# see a consistent key during the transition window.
reasoning_new = delta.get("reasoning")
reasoning_legacy = delta.get("reasoning_content")
if reasoning_new is not None and reasoning_legacy is None:
delta["reasoning_content"] = reasoning_new
new_choices.append(choice)
return ModelResponseStream(

View file

@ -1,3 +1,5 @@
# LiteLLM main module: public completion, embedding, streaming, and moderation entrypoints.
#
# +-----------------------------------------------+
# | |
# | Give Feedback / Get Help |
@ -59,7 +61,13 @@ import litellm
from litellm import client
# Other utils are imported directly to avoid circular imports
from litellm.utils import exception_type, get_litellm_params, get_optional_params
from litellm.utils import (
exception_type,
get_litellm_params,
get_optional_params,
peek_reasoning_summary_aliases,
strip_reasoning_summary_aliases_from_optional_params,
)
# Logging is imported lazily when needed to avoid loading litellm_logging at import time
if TYPE_CHECKING:
@ -946,6 +954,7 @@ def responses_api_bridge_check(
web_search_options: Optional[OpenAIWebSearchOptions] = None,
tools: Optional[List[Any]] = None,
reasoning_effort: Optional[Any] = None,
reasoning_summary: Optional[Any] = None,
) -> Tuple[dict, str]:
model_info: Dict[str, Any] = {}
@ -982,14 +991,23 @@ def responses_api_bridge_check(
mode = "responses"
model_info["mode"] = mode
# OpenAI/Azure gpt-5.4+ chat-completions calls with both tools + reasoning_effort
# must be bridged to Responses API.
# OpenAI/Azure GPT-5 chat-completions that need Responses-only fields (e.g.
# ``reasoningSummary`` in ``extra_body``) must be bridged; Chat Completions rejects
# those keys.
#
# - gpt-5.4+: tools + reasoning_effort (original) or any reasoning-summary alias.
# - Older GPT-5 names (e.g. ``gpt-5``, ``gpt-5.1``): bridge only when a reasoning
# summary alias is present with ``reasoning_effort`` (tools alone stay on chat).
if (
custom_llm_provider in ("openai", "azure")
and OpenAIGPT5Config.is_model_gpt_5_4_plus_model(model)
and tools
and reasoning_effort is not None
and model_info.get("mode") != "responses"
and OpenAIGPT5Config.is_model_gpt_5_model(model)
and not OpenAIGPT5Config.is_model_gpt_5_search_model(model)
and reasoning_effort is not None
and (
reasoning_summary is not None
or (OpenAIGPT5Config.is_model_gpt_5_4_plus_model(model) and tools)
)
):
model_info["mode"] = "responses"
model = model.replace("responses/", "")
@ -1510,7 +1528,11 @@ def completion( # type: ignore # noqa: PLR0915
"logit_bias": logit_bias,
"user": user,
# params to identify the model
"model": model,
"model": (
model_info.get("base_model")
if isinstance(model_info, dict) and model_info.get("base_model")
else model
),
"custom_llm_provider": custom_llm_provider,
"response_format": response_format,
"seed": seed,
@ -1634,8 +1656,10 @@ def completion( # type: ignore # noqa: PLR0915
## RESPONSES API BRIDGE LOGIC ## - check if model has 'mode: responses' in litellm.model_cost map
# Only run the second bridge check if the first one didn't already
# detect responses mode (e.g. via the "responses/" prefix). The second
# check handles cases like gpt-5.4+ with tools+reasoning_effort that
# the first (early) check doesn't cover.
# check handles cases like gpt-5.4+ with tools+reasoning_effort or
# reasoningSummary/reasoning_summary without tools (AI SDK) that the first
# (early) check doesn't cover.
_reasoning_summary_for_bridge = peek_reasoning_summary_aliases(optional_params)
if responses_api_model_info.get("mode") != "responses":
responses_api_model_info, model = responses_api_bridge_check(
model=model,
@ -1643,14 +1667,29 @@ def completion( # type: ignore # noqa: PLR0915
web_search_options=web_search_options,
tools=tools,
reasoning_effort=reasoning_effort,
reasoning_summary=_reasoning_summary_for_bridge,
)
if responses_api_model_info.get("mode") == "responses":
from litellm.completion_extras import responses_api_bridge
optional_params, rs_val = (
strip_reasoning_summary_aliases_from_optional_params(optional_params)
)
if isinstance(reasoning_effort, dict) and "summary" in reasoning_effort:
optional_params = dict(optional_params)
optional_params["reasoning_effort"] = reasoning_effort
elif rs_val is not None:
eff = optional_params.get("reasoning_effort", reasoning_effort)
if isinstance(eff, dict):
optional_params["reasoning_effort"] = {**eff, "summary": rs_val}
elif eff is not None:
optional_params["reasoning_effort"] = {
"effort": eff,
"summary": rs_val,
}
else:
optional_params["reasoning_effort"] = {"summary": rs_val}
return responses_api_bridge.completion(
model=model,
@ -1669,6 +1708,16 @@ def completion( # type: ignore # noqa: PLR0915
encoding=_get_encoding(),
stream=stream,
)
elif (
custom_llm_provider == "openai"
and OpenAIGPT5Config.is_model_gpt_5_model(model)
) or (
custom_llm_provider == "azure"
and litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model)
):
optional_params, _ = strip_reasoning_summary_aliases_from_optional_params(
optional_params
)
if custom_llm_provider == "azure":
# azure configs
@ -3813,7 +3862,33 @@ def completion( # type: ignore # noqa: PLR0915
)
bedrock_route = BedrockModelInfo.get_bedrock_route(model)
if bedrock_route == "converse":
if bedrock_route == "claude_platform":
provider_config = ProviderConfigManager.get_provider_chat_config(
model=model,
provider=LlmProviders.BEDROCK,
)
model = BedrockModelInfo.get_claude_platform_model(model)
response = base_llm_http_handler.completion(
model=model,
stream=stream,
messages=messages,
acompletion=acompletion,
api_base=api_base,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
shared_session=shared_session,
custom_llm_provider="bedrock",
timeout=timeout,
headers=headers,
encoding=_get_encoding(),
api_key=api_key,
logging_obj=logging,
client=client,
provider_config=provider_config,
)
return response
elif bedrock_route == "converse":
model = model.replace("converse/", "")
response = bedrock_converse_chat_completion.completion(
model=model,

View file

@ -21104,6 +21104,38 @@
"supports_system_messages": true,
"supports_tool_choice": true
},
"gpt-realtime-2": {
"cache_creation_input_audio_token_cost": 4e-07,
"cache_read_input_token_cost": 4e-07,
"input_cost_per_audio_token": 3.2e-05,
"input_cost_per_image": 5e-06,
"input_cost_per_token": 4e-06,
"litellm_provider": "openai",
"max_input_tokens": 32000,
"max_output_tokens": 4096,
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_audio_token": 6.4e-05,
"output_cost_per_token": 1.6e-05,
"supported_endpoints": [
"/v1/realtime"
],
"supported_modalities": [
"text",
"image",
"audio"
],
"supported_output_modalities": [
"text",
"audio"
],
"supports_audio_input": true,
"supports_audio_output": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_system_messages": true,
"supports_tool_choice": true
},
"gpt-realtime-mini": {
"cache_creation_input_audio_token_cost": 3e-07,
"cache_read_input_audio_token_cost": 3e-07,
@ -27187,6 +27219,20 @@
"supports_reasoning": true,
"supports_tool_choice": true
},
"openrouter/qwen/qwen3.6-plus": {
"input_cost_per_token": 3.25e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1000000,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 1.95e-06,
"source": "https://openrouter.ai/qwen/qwen3.6-plus",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"openrouter/qwen/qwen3.5-35b-a3b": {
"input_cost_per_token": 2.5e-07,
"litellm_provider": "openrouter",

View file

@ -12,7 +12,8 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.proxy._experimental.mcp_server.oauth_utils import (
TOKEN_NO_CACHE_HEADERS,
validate_loopback_redirect_uri,
get_request_base_url,
validate_trusted_redirect_uri,
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
@ -29,51 +30,6 @@ router = APIRouter(
)
def get_request_base_url(request: Request) -> str:
"""
Get the base URL for the request, considering X-Forwarded-* headers.
X-Forwarded-Proto / X-Forwarded-Host / X-Forwarded-Port are only honoured
when the request comes from a configured trusted proxy
(``use_x_forwarded_for`` enabled AND caller in ``mcp_trusted_proxy_ranges``).
Otherwise the request's literal ``base_url`` is returned, so an
untrusted caller cannot poison OAuth-discovery / redirect_uri values
by injecting headers.
Args:
request: FastAPI Request object
Returns:
The reconstructed base URL (e.g., "https://proxy.example.com")
"""
base_url = str(request.base_url).rstrip("/")
parsed = urlparse(base_url)
if not IPAddressUtils.is_request_from_trusted_proxy(request):
return base_url
x_forwarded_proto = request.headers.get("X-Forwarded-Proto")
x_forwarded_host = request.headers.get("X-Forwarded-Host")
x_forwarded_port = request.headers.get("X-Forwarded-Port")
scheme = x_forwarded_proto if x_forwarded_proto else parsed.scheme
if x_forwarded_host:
# X-Forwarded-Host may already include port (e.g., "example.com:8080")
if ":" in x_forwarded_host and not x_forwarded_host.startswith("["):
netloc = x_forwarded_host
elif x_forwarded_port:
netloc = f"{x_forwarded_host}:{x_forwarded_port}"
else:
netloc = x_forwarded_host
else:
netloc = parsed.netloc
if x_forwarded_port and ":" not in netloc:
netloc = f"{netloc}:{x_forwarded_port}"
return urlunparse((scheme, netloc, parsed.path, "", "", ""))
def encode_state_with_base_url(
base_url: str,
original_state: str,
@ -127,12 +83,14 @@ def decode_state_hash(encrypted_state: str) -> dict:
return state_data
def _get_validated_client_redirect_uri(state_data: Dict[str, Any]) -> str:
"""Return a loopback client redirect URI from OAuth state."""
def _get_validated_client_redirect_uri(
request: Request, state_data: Dict[str, Any]
) -> str:
"""Return a trusted (same-origin or loopback) client redirect URI from OAuth state."""
redirect_uri = state_data.get("client_redirect_uri") or state_data.get("base_url")
if not redirect_uri or not isinstance(redirect_uri, str):
raise HTTPException(status_code=400, detail="Invalid redirect URI")
validate_loopback_redirect_uri(redirect_uri)
validate_trusted_redirect_uri(request, redirect_uri)
return redirect_uri
@ -338,12 +296,12 @@ async def authorize_with_server(
status_code=400, detail="MCP server authorization url is not set"
)
# Loopback-only redirect_uri. The URI is encrypted into the OAuth
# state and decoded on /callback to redirect the user back; a non-
# loopback URI would be an open-redirect + code-theft primitive
# (VERIA-57 root cause B). MCP clients are native apps — loopback is
# the spec-compliant callback pattern.
validate_loopback_redirect_uri(redirect_uri)
# Loopback OR same-origin redirect_uri. The URI is encrypted into the
# OAuth state and decoded on /callback to redirect the user back;
# restricting to trusted origins blocks the open-redirect +
# code-theft primitive (VERIA-57 root cause B). Loopback supports
# native MCP clients; same-origin supports the proxy's own UI callback.
validate_trusted_redirect_uri(request, redirect_uri)
parsed = urlparse(redirect_uri)
base_url = urlunparse(parsed._replace(query=""))
request_base_url = get_request_base_url(request)
@ -660,17 +618,18 @@ async def token_endpoint(
@router.get("/callback")
async def callback(code: str, state: str):
async def callback(request: Request, code: str, state: str):
try:
state_data = decode_state_hash(state)
original_state = state_data["original_state"]
# Re-validate loopback at the sink. /authorize rejects non-loopback
# Re-validate at the sink. /authorize rejects untrusted
# redirect_uri before encoding into state, but encrypted states
# minted before that check was added have no expiry and remain
# valid indefinitely. Validating here blocks the open-redirect +
# code-theft primitive even for pre-fix states.
redirect_uri = _get_validated_client_redirect_uri(state_data)
# valid indefinitely. Validating here (same-origin OR loopback)
# blocks the open-redirect + code-theft primitive even for pre-fix
# states while allowing the UI's same-origin callback to work.
redirect_uri = _get_validated_client_redirect_uri(request, state_data)
params = {"code": code, "state": original_state}
complete_returned_url = _append_query_params(redirect_uri, params)

View file

@ -506,7 +506,8 @@ class MCPServerManager:
# Add any static headers from server config.
#
# Note: `extra_headers` on MCPServer is a List[str] of header names to forward
# from the client request (not available in this OpenAPI tool generation step).
# from each client MCP request; values are applied at call time via
# `_request_extra_headers` in server.py (not baked in here).
# `static_headers` is a dict of concrete headers to always send.
headers = (
merge_mcp_headers(
@ -598,16 +599,57 @@ class MCPServerManager:
)
raise e
def _cleanup_server_tool_routing_artifacts(self, server: MCPServer) -> None:
"""Drop OpenAPI global tools and name-mapping rows owned by ``server``.
When a server leaves ``self.registry`` (eviction, ``remove_server``, etc.),
OpenAPI tools remain in ``global_mcp_tool_registry`` and
``tool_name_to_mcp_server_name_mapping`` unless removed here. Stale
mappings make ``_get_mcp_server_from_tool_name`` resolve to a prefix that
no longer exists in the live registry.
"""
from litellm.proxy._experimental.mcp_server.tool_registry import (
global_mcp_tool_registry,
)
prefix_root = normalize_server_name(get_server_prefix(server))
if server.spec_path and prefix_root:
openapi_key_prefix = prefix_root + MCP_TOOL_PREFIX_SEPARATOR
global_mcp_tool_registry.unregister_tools_with_prefix(openapi_key_prefix)
owned_raw: Set[str] = set()
for p in iter_known_server_prefixes(server):
if p:
owned_raw.add(p)
if server.name:
owned_raw.add(server.name)
owned_normalized = {normalize_server_name(x) for x in owned_raw}
stale_mapping_keys: List[str] = []
for tool_name, mapped_server in list(
self.tool_name_to_mcp_server_name_mapping.items()
):
if mapped_server in owned_raw:
stale_mapping_keys.append(tool_name)
elif normalize_server_name(str(mapped_server)) in owned_normalized:
stale_mapping_keys.append(tool_name)
for key in stale_mapping_keys:
del self.tool_name_to_mcp_server_name_mapping[key]
def remove_server(self, mcp_server: LiteLLM_MCPServerTable):
"""
Remove a server from the registry
"""
if mcp_server.server_name in self.get_registry():
del self.registry[mcp_server.server_name]
verbose_logger.debug(f"Removed MCP Server: {mcp_server.server_name}")
elif mcp_server.server_id in self.get_registry():
del self.registry[mcp_server.server_id]
verbose_logger.debug(f"Removed MCP Server: {mcp_server.server_id}")
evicted: Optional[MCPServer] = self.registry.pop(mcp_server.server_id, None)
if evicted is None and mcp_server.server_name:
evicted = self.registry.pop(mcp_server.server_name, None)
if evicted is not None:
verbose_logger.debug(
"Removed MCP Server: %s", mcp_server.server_id or mcp_server.server_name
)
self._cleanup_server_tool_routing_artifacts(evicted)
else:
verbose_logger.warning(
f"Server ID {mcp_server.server_id} not found in registry"
@ -805,6 +847,13 @@ class MCPServerManager:
self.initialize_tool_name_to_mcp_server_name_mapping()
async def add_server(self, mcp_server: LiteLLM_MCPServerTable):
# The runtime registry is the allowlist for tool calls and health
# probes (which spawn the underlying transport, including stdio
# subprocesses). Match the eligibility set used by the bulk DB
# filter in reload_servers_from_database() — NULL is legacy and
# "approved" is a legacy alias for "active".
if mcp_server.approval_status not in (None, "active", "approved"):
return
try:
if mcp_server.server_id not in self.registry:
new_server = await self.build_mcp_server_from_table(mcp_server)
@ -818,6 +867,16 @@ class MCPServerManager:
raise e
async def update_server(self, mcp_server: LiteLLM_MCPServerTable):
# If a previously-active server has been moved out of the active
# state, evict any stale registry entry so subsequent tool calls and
# health probes can't reach it.
if mcp_server.approval_status not in (None, "active", "approved"):
evicted = self.registry.pop(mcp_server.server_id, None)
if evicted is None and mcp_server.server_name:
evicted = self.registry.pop(mcp_server.server_name, None)
if evicted is not None:
self._cleanup_server_tool_routing_artifacts(evicted)
return
try:
if mcp_server.server_id in self.registry:
new_server = await self.build_mcp_server_from_table(mcp_server)

View file

@ -2,15 +2,63 @@
(BYOK + discoverable / pass-through OAuth proxy)."""
from ipaddress import ip_address
from urllib.parse import urlparse
from urllib.parse import urlparse, urlunparse
from fastapi import HTTPException
from fastapi import HTTPException, Request
from litellm._logging import verbose_logger
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
# RFC 6749 §5.1 / OAuth 2.1 draft-15 §4.1.3: token-endpoint responses
# must not be cached — both success and error bodies may reveal secrets.
TOKEN_NO_CACHE_HEADERS = {"Cache-Control": "no-store", "Pragma": "no-cache"}
def get_request_base_url(request: Request) -> str:
"""
Get the base URL for the request, considering X-Forwarded-* headers.
X-Forwarded-Proto / X-Forwarded-Host / X-Forwarded-Port are only honoured
when the request comes from a configured trusted proxy
(``use_x_forwarded_for`` enabled AND caller in ``mcp_trusted_proxy_ranges``).
Otherwise the request's literal ``base_url`` is returned, so an
untrusted caller cannot poison OAuth-discovery / redirect_uri values
by injecting headers.
Args:
request: FastAPI Request object
Returns:
The reconstructed base URL (e.g., "https://proxy.example.com")
"""
base_url = str(request.base_url).rstrip("/")
parsed = urlparse(base_url)
if not IPAddressUtils.is_request_from_trusted_proxy(request):
return base_url
x_forwarded_proto = request.headers.get("X-Forwarded-Proto")
x_forwarded_host = request.headers.get("X-Forwarded-Host")
x_forwarded_port = request.headers.get("X-Forwarded-Port")
scheme = x_forwarded_proto if x_forwarded_proto else parsed.scheme
if x_forwarded_host:
# X-Forwarded-Host may already include port (e.g., "example.com:8080")
if ":" in x_forwarded_host and not x_forwarded_host.startswith("["):
netloc = x_forwarded_host
elif x_forwarded_port:
netloc = f"{x_forwarded_host}:{x_forwarded_port}"
else:
netloc = x_forwarded_host
else:
netloc = parsed.netloc
if x_forwarded_port and ":" not in netloc:
netloc = f"{netloc}:{x_forwarded_port}"
return urlunparse((scheme, netloc, parsed.path, "", "", ""))
def validate_loopback_redirect_uri(redirect_uri: str) -> None:
"""Require a loopback ``redirect_uri`` (OAuth 2.1 §4.1.2.1 + RFC 8252
§7.3 native-app pattern). MCP clients are native apps that listen on
@ -46,3 +94,60 @@ def validate_loopback_redirect_uri(redirect_uri: str) -> None:
# don't let it bubble up as a 500.
pass
raise HTTPException(status_code=400, detail="invalid_request")
def validate_trusted_redirect_uri(request: Request, redirect_uri: str) -> None:
"""Accept same-origin (proxy's own origin) OR loopback ``redirect_uri``.
Same-origin is required for the LiteLLM UI's OAuth flow: the UI
redirects to ``<proxy>/ui/mcp/oauth/callback`` which is not loopback
but is on the proxy's own trusted HTTPS origin. An attacker cannot
host content on the proxy's own origin without already owning the
proxy, so the open-redirect / code-theft primitive that motivated
:func:`validate_loopback_redirect_uri` does not apply here.
Loopback continues to be accepted for native MCP clients (per
OAuth 2.1 §4.1.2.1 + RFC 8252 §7.3).
Use this in the discoverable OAuth proxy endpoints that serve both
native clients and the proxy's own UI. BYOK endpoints that only
support native clients should keep
:func:`validate_loopback_redirect_uri`.
"""
try:
parsed = urlparse(redirect_uri)
except ValueError:
raise HTTPException(status_code=400, detail="invalid_request")
if parsed.scheme not in ("http", "https"):
raise HTTPException(status_code=400, detail="invalid_request")
if parsed.fragment:
raise HTTPException(status_code=400, detail="invalid_request")
# Same-origin: scheme + netloc (host[:port]) must match the proxy's
# own base URL at this request (honouring trusted X-Forwarded-*).
try:
proxy_base = urlparse(get_request_base_url(request))
if (
parsed.netloc
and parsed.scheme == proxy_base.scheme
and parsed.netloc.lower() == proxy_base.netloc.lower()
):
return
except Exception as exc:
# If we can't determine the proxy's origin, fall through to
# loopback. Log so the failure is diagnosable in production.
verbose_logger.warning(
"validate_trusted_redirect_uri: could not determine proxy origin, "
"falling back to loopback-only check. error=%s",
exc,
)
host = (parsed.hostname or "").lower()
if host == "localhost":
return
try:
if ip_address(host).is_loopback:
return
except ValueError:
pass
raise HTTPException(status_code=400, detail="invalid_request")

View file

@ -55,6 +55,13 @@ _request_auth_header: contextvars.ContextVar[Optional[str]] = contextvars.Contex
"_request_auth_header", default=None
)
# Per-request extra headers forwarded from the client request.
# Populated from MCPServer.extra_headers names matched against raw request
# headers in server.py before dispatching to a local/OpenAPI tool handler.
_request_extra_headers: contextvars.ContextVar[Optional[Dict[str, str]]] = (
contextvars.ContextVar("_request_extra_headers", default=None)
)
def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str:
"""Ensure path params cannot introduce directory traversal."""
@ -297,6 +304,46 @@ def build_input_schema(operation: Dict[str, Any]) -> Dict[str, Any]:
}
def _merge_openapi_tool_request_headers(
static_headers: Dict[str, str]
) -> Dict[str, str]:
"""Merge static closure headers with per-request ContextVar overrides.
Precedence (highest to lowest):
1. ``_request_auth_header`` — BYOK override of ``Authorization``
2. ``static_headers`` — operator-configured headers baked into the
tool closure at registration time
3. ``_request_extra_headers`` — per-request headers forwarded from
the MCP caller (allowlisted by ``MCPServer.extra_headers``)
This matches the existing MCP invariant in
:func:`litellm.proxy._experimental.mcp_server.utils.merge_mcp_headers`
and the managed MCP path, where ``static_headers`` always wins over
caller-forwarded headers. Keeping the same precedence here prevents an
authenticated caller from overriding an operator-configured value
(e.g. a tenant id or upstream API key) by sending the same header name.
Header names are compared case-insensitively so different casing cannot
bypass the precedence rules.
"""
request_extra = _request_extra_headers.get() or {}
static = static_headers or {}
static_lower_names = {k.lower() for k in static}
effective_headers: Dict[str, str] = {
k: v for k, v in request_extra.items() if k.lower() not in static_lower_names
}
effective_headers.update(static)
override_auth = _request_auth_header.get()
if override_auth:
for existing in [k for k in effective_headers if k.lower() == "authorization"]:
del effective_headers[existing]
effective_headers["Authorization"] = override_auth
return effective_headers
def create_tool_function(
path: str,
method: str,
@ -334,14 +381,7 @@ def create_tool_function(
The function safely handles parameter names that aren't valid Python identifiers
by using **kwargs instead of named parameters.
"""
# Allow per-request auth override (e.g. BYOK credential set via ContextVar).
# The ContextVar holds the full Authorization header value, including the
# correct prefix (Bearer / ApiKey / Basic) formatted by the caller in
# server.py based on the server's configured auth_type.
effective_headers = dict(headers)
override_auth = _request_auth_header.get()
if override_auth:
effective_headers["Authorization"] = override_auth
effective_headers = _merge_openapi_tool_request_headers(headers)
# Build URL from base_url and path
url = base_url + path

View file

@ -158,6 +158,7 @@ if MCP_AVAILABLE:
)
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
_request_auth_header,
_request_extra_headers,
)
from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport
from litellm.proxy._experimental.mcp_server.tool_registry import (
@ -2195,11 +2196,40 @@ if MCP_AVAILABLE:
auth_header_value = f"Basic {mcp_auth_header}"
else:
auth_header_value = f"Bearer {mcp_auth_header}"
# Forward named client headers to OpenAPI tool upstream requests.
# MCPServer.extra_headers lists header names to copy from raw_headers.
# OAuth2 M2M: never take Authorization from the caller (matches
# _prepare_mcp_server_headers for managed MCP).
forwarded_headers: Optional[Dict[str, str]] = None
if mcp_server and mcp_server.extra_headers and raw_headers:
normalized_raw = {
str(k).lower(): v
for k, v in raw_headers.items()
if isinstance(k, str)
}
skip_caller_authorization = bool(mcp_server.has_client_credentials)
for header_name in mcp_server.extra_headers:
if not isinstance(header_name, str):
continue
if (
skip_caller_authorization
and header_name.lower() == "authorization"
):
continue
value = normalized_raw.get(header_name.lower())
if value is not None:
if forwarded_headers is None:
forwarded_headers = {}
forwarded_headers[header_name] = value
_auth_token = _request_auth_header.set(auth_header_value)
_extra_token = _request_extra_headers.set(forwarded_headers)
try:
local_content = await _handle_local_mcp_tool(name, arguments)
finally:
_request_auth_header.reset(_auth_token)
_request_extra_headers.reset(_extra_token)
response = CallToolResult(content=cast(Any, local_content), isError=False)
# Try managed MCP server tool (pass the full prefixed name)

View file

@ -59,6 +59,22 @@ class MCPToolRegistry:
]
return list(self.tools.values())
def unregister_tools_with_prefix(self, prefix: str) -> int:
"""Remove tools whose registered name starts with ``prefix``.
Used when an OpenAPI-backed MCP server leaves the runtime registry so
stale tool handlers cannot be invoked after eviction.
"""
if not prefix:
return 0
removed = 0
for name in list(self.tools.keys()):
if name.startswith(prefix):
del self.tools[name]
removed += 1
verbose_logger.debug("Unregistered MCP tool %s", name)
return removed
def convert_tools_to_mcp_sdk_tool_type(
self, tools: List[MCPTool]
) -> List["MCPToolSDKTool"]:

View file

@ -239,6 +239,7 @@ class KeyManagementRoutes(str, enum.Enum):
KEY_BLOCK = "/key/block"
KEY_UNBLOCK = "/key/unblock"
KEY_BULK_UPDATE = "/key/bulk_update"
TEAM_KEY_BULK_UPDATE = "/team/key/bulk_update"
KEY_RESET_SPEND = "/key/{key_id}/reset_spend"
# info and health routes
@ -353,8 +354,10 @@ class LiteLLMRoutes(enum.Enum):
# realtime
"/realtime",
"/v1/realtime",
"/openai/v1/realtime",
"/realtime?{model}",
"/v1/realtime?{model}",
"/openai/v1/realtime?{model}",
# responses API
"/responses",
"/v1/responses",
@ -538,6 +541,7 @@ class LiteLLMRoutes(enum.Enum):
KeyManagementRoutes.KEY_BLOCK.value,
KeyManagementRoutes.KEY_UNBLOCK.value,
KeyManagementRoutes.KEY_BULK_UPDATE.value,
KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value,
KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value,
KeyManagementRoutes.SPEND_LOGS.value,
KeyManagementRoutes.KEY_RESET_SPEND.value,
@ -673,6 +677,10 @@ class LiteLLMRoutes(enum.Enum):
"/global/activity",
"/global/activity/model",
"/global/activity/cache_hits",
# Tag usage endpoints scope internal users to tags produced by
# their own keys in tag_management_endpoints.py.
"/tag/daily/activity",
"/tag/list",
"/v1/models/{model_id}",
"/models/{model_id}",
"/guardrails/list",
@ -685,7 +693,16 @@ class LiteLLMRoutes(enum.Enum):
+ compliance_check_routes
)
internal_user_view_only_routes = spend_tracking_routes
internal_user_view_only_routes = (
spend_tracking_routes
+ compliance_check_routes
+ [
# Tag usage endpoints scope internal viewers to tags produced by
# their own keys in tag_management_endpoints.py.
"/tag/daily/activity",
"/tag/list",
]
)
self_managed_routes = [
"/team/member_add",
@ -707,6 +724,8 @@ class LiteLLMRoutes(enum.Enum):
# Project read routes - endpoint scopes results to caller's teams (non-admin)
"/project/list",
"/project/info",
# Endpoint enforces proxy-admin vs team-admin model access itself.
"/health/test_connection",
# Invitation routes - org/team admins checked in endpoint via _user_has_admin_privileges
"/invitation/new",
"/invitation/delete",
@ -4342,10 +4361,16 @@ class JWTRoutingOverride(BaseModel):
A rule matches when all provided selectors match token claims.
If matched, request is routed to the configured auth path.
Wildcard selectors use shell-style patterns (* and ?) and are matched with
case-sensitive semantics; use the same casing your IdP emits in JWT claims.
Space-delimited tokenization applies only to the ``scope`` claim (OAuth/OIDC
scope strings), not to ``iss``, ``aud``, or ``client_id``.
"""
iss: Union[str, List[str]]
client_id: Optional[Union[str, List[str]]] = None
scope: Optional[Union[str, List[str]]] = None
aud: Optional[Union[str, List[str]]] = None
path: Literal["oauth2"] = "oauth2"

View file

@ -2849,7 +2849,7 @@ def _can_object_call_model(
object_type=object_type
),
param="model",
code=status.HTTP_401_UNAUTHORIZED,
code=status.HTTP_403_FORBIDDEN,
)
@ -3082,7 +3082,7 @@ async def can_user_call_model(
message=f"User not allowed to access model. No default model access, only team models allowed. Tried to access {model}",
type=ProxyErrorTypes.key_model_access_denied,
param="model",
code=status.HTTP_401_UNAUTHORIZED,
code=status.HTTP_403_FORBIDDEN,
)
return _can_object_call_model(
@ -3625,7 +3625,7 @@ async def _check_team_member_model_access(
message=f"Team member not allowed to access model. User={valid_token.user_id}, Team={team_object.team_id}, Model={model}. Allowed member models = {member_allowed_models}",
type=ProxyErrorTypes.team_model_access_denied,
param="model",
code=status.HTTP_401_UNAUTHORIZED,
code=status.HTTP_403_FORBIDDEN,
)

View file

@ -123,7 +123,7 @@ class UserAPIKeyAuthExceptionHandler:
message=e.message,
type=ProxyErrorTypes.budget_exceeded,
param=None,
code=400,
code=getattr(e, "status_code", status.HTTP_429_TOO_MANY_REQUESTS),
)
if isinstance(e, HTTPException):
raise ProxyException(

View file

@ -224,6 +224,41 @@ class JWTHandler:
return []
def get_all_jwt_team_ids(self, token: dict) -> List[str]:
"""
Return team IDs from both the plural ``team_ids_jwt_field`` and the
singular ``team_id_jwt_field`` claim (string or list of strings), as a
deduplicated list preserving plural-first order.
Membership-reconciliation paths (SSO callback, JWT-bearer sync) need
to consider both claim shapes. Reading only the plural field — as
callers historically did — silently dropped users whose IdP populates
the singular field, which is what Okta and Auth0 default to when a
user has a single primary team.
This intentionally does NOT consult ``team_id_default``: that fallback
is a property of how the JWT-bearer auth flow resolves a single
request-bound team, not of the token's claims. Callers that want the
default-team behavior should still go through ``get_team_id``.
"""
team_ids: List[str] = list(self.get_team_ids_from_jwt(token))
if self.litellm_jwtauth.team_id_jwt_field is not None:
singular = get_nested_value(
data=token,
key_path=self.litellm_jwtauth.team_id_jwt_field,
default=None,
)
if isinstance(singular, list):
for item in singular:
if item is None:
continue
sid = str(item)
if sid and sid not in team_ids:
team_ids.append(sid)
elif singular and str(singular) not in team_ids:
team_ids.append(str(singular))
return team_ids
def get_end_user_id(
self, token: dict, default_value: Optional[str]
) -> Optional[str]:

View file

@ -50,6 +50,7 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES = frozenset(
KeyManagementRoutes.KEY_BLOCK.value,
KeyManagementRoutes.KEY_UNBLOCK.value,
KeyManagementRoutes.KEY_BULK_UPDATE.value,
KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value,
]
)
@ -671,6 +672,7 @@ class RouteChecks:
"/key/service-account/generate",
"/key/block",
"/key/unblock",
"/team/key/bulk_update",
]
)

View file

@ -8,6 +8,7 @@ Returns a UserAPIKeyAuth object if the API key is valid
"""
import asyncio
import fnmatch
import re
import secrets
from datetime import datetime, timezone
@ -183,22 +184,54 @@ def _get_bearer_token_or_received_api_key(api_key: str) -> str:
def _routing_selector_matches_claim(
selector_value: Optional[Any], claim_value: Optional[Any]
selector_value: Optional[Any],
claim_value: Optional[Any],
*,
split_space_delimited: bool = False,
) -> bool:
if selector_value is None:
return True
selector_list = (
selector_list: List[str] = (
[str(v) for v in selector_value]
if isinstance(selector_value, list)
else [str(selector_value)]
)
if claim_value is None:
return False
if isinstance(claim_value, list):
claim_list = [str(v) for v in claim_value]
return any(v in claim_list for v in selector_list)
elif (
split_space_delimited
and isinstance(claim_value, str)
and " " in claim_value.strip()
):
# OAuth/OIDC often sends scope as a single space-delimited string. Only split
# for the scope selector: iss/aud/client_id must stay exact full-string match
# on unverified claims (see routing override security review). The elif guard
# (`" " in claim_value.strip()`) ensures at least two non-empty tokens survive.
claim_list = [v for v in claim_value.strip().split(" ") if v]
else:
claim_list = [str(claim_value)]
return str(claim_value) in selector_list if claim_value is not None else False
def _selector_matches_claim(selector: str, claim: str) -> bool:
# NOTE: wildcard matching is case-sensitive (fnmatch.fnmatchcase).
if "*" in selector or "?" in selector:
# Without scope splitting, do not let `*` span whitespace: a malformed
# iss like "trusted.example.com evil.com" must not match "trusted.*".
# Scope uses split_space_delimited so each claim token is checked separately.
if not split_space_delimited and any(ch.isspace() for ch in claim):
return False
return fnmatch.fnmatchcase(claim, selector)
return selector == claim
return any(
_selector_matches_claim(selector=s, claim=c)
for s in selector_list
for c in claim_list
)
def _matches_routing_override(
@ -209,6 +242,11 @@ def _matches_routing_override(
and _routing_selector_matches_claim(
override.client_id, token_claims.get("client_id")
)
and _routing_selector_matches_claim(
override.scope,
token_claims.get("scope"),
split_space_delimited=True,
)
and _routing_selector_matches_claim(override.aud, token_claims.get("aud"))
)
@ -1107,7 +1145,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
raise ProxyException(
message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}",
type=ProxyErrorTypes.expired_key,
code=400,
code=status.HTTP_401_UNAUTHORIZED,
param=abbreviate_api_key(api_key=api_key),
)
valid_token = update_valid_token_with_end_user_params(
@ -1432,7 +1470,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
raise ProxyException(
message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}",
type=ProxyErrorTypes.expired_key,
code=400,
code=status.HTTP_401_UNAUTHORIZED,
param=abbreviate_api_key(api_key=api_key),
)
@ -2417,7 +2455,7 @@ async def _run_post_custom_auth_checks(
raise ProxyException(
message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}",
type=ProxyErrorTypes.expired_key,
code=400,
code=status.HTTP_401_UNAUTHORIZED,
param=(
abbreviate_api_key(api_key=valid_token.token)
if valid_token.token

View file

@ -14,7 +14,7 @@ jwt_display_template = """
padding: 20px;
display: flex;
justify-content: center;
align-items: center;
align-items: flex-start;
min-height: 100vh;
color: #333;
}
@ -27,18 +27,18 @@ jwt_display_template = """
width: 800px;
max-width: 100%;
}
.logo-container {
text-align: center;
margin-bottom: 30px;
}
.logo {
font-size: 24px;
font-weight: 600;
color: #1e293b;
}
h2 {
margin: 0 0 10px;
color: #1e293b;
@ -46,7 +46,14 @@ jwt_display_template = """
font-weight: 600;
text-align: center;
}
h3 {
margin: 0 0 12px;
color: #1e293b;
font-size: 18px;
font-weight: 600;
}
.subtitle {
color: #64748b;
margin: 0 0 20px;
@ -58,15 +65,15 @@ jwt_display_template = """
background-color: #f1f5f9;
border-radius: 6px;
padding: 20px;
margin-bottom: 30px;
margin-bottom: 20px;
border-left: 4px solid #2563eb;
}
.success-box {
background-color: #f0fdf4;
border-radius: 6px;
padding: 20px;
margin-bottom: 30px;
margin-bottom: 20px;
border-left: 4px solid #16a34a;
}
@ -78,7 +85,7 @@ jwt_display_template = """
font-weight: 600;
font-size: 16px;
}
.success-header {
display: flex;
align-items: center;
@ -87,46 +94,53 @@ jwt_display_template = """
font-weight: 600;
font-size: 16px;
}
.info-header svg, .success-header svg {
margin-right: 8px;
}
.data-container {
margin-top: 20px;
}
.data-row {
display: flex;
border-bottom: 1px solid #e2e8f0;
padding: 12px 0;
}
.data-row:last-child {
border-bottom: none;
}
.data-label {
font-weight: 500;
color: #334155;
width: 180px;
width: 220px;
flex-shrink: 0;
}
.data-value {
color: #475569;
word-break: break-all;
}
.empty-note {
color: #64748b;
font-style: italic;
margin: 0;
font-size: 14px;
}
.jwt-container {
background-color: #f8fafc;
border-radius: 6px;
padding: 15px;
margin-top: 20px;
margin-top: 12px;
overflow-x: auto;
border: 1px solid #e2e8f0;
}
.jwt-text {
font-family: monospace;
white-space: pre-wrap;
@ -134,7 +148,7 @@ jwt_display_template = """
margin: 0;
color: #334155;
}
.back-button {
display: inline-block;
background-color: #6466E9;
@ -146,18 +160,18 @@ jwt_display_template = """
margin-top: 20px;
text-align: center;
}
.back-button:hover {
background-color: #4138C2;
text-decoration: none;
}
.buttons {
display: flex;
gap: 10px;
margin-top: 20px;
margin-top: 12px;
}
.copy-button {
background-color: #e2e8f0;
color: #334155;
@ -169,11 +183,11 @@ jwt_display_template = """
display: flex;
align-items: center;
}
.copy-button:hover {
background-color: #cbd5e1;
}
.copy-button svg {
margin-right: 6px;
}
@ -188,7 +202,7 @@ jwt_display_template = """
</div>
<h2>SSO Debug Information</h2>
<p class="subtitle">Results from the SSO authentication process.</p>
<div class="success-box">
<div class="success-header">
<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
@ -199,11 +213,7 @@ jwt_display_template = """
</div>
<p>The SSO authentication completed successfully. Below is the information returned by the provider.</p>
</div>
<div class="data-container" id="userData">
<!-- Data will be inserted here by JavaScript -->
</div>
<div class="info-box">
<div class="info-header">
<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
@ -211,22 +221,62 @@ jwt_display_template = """
<line x1="12" y1="16" x2="12" y2="12"></line>
<line x1="12" y1="8" x2="12.01" y2="8"></line>
</svg>
JSON Representation
Parsed by Proxy
</div>
<p class="empty-note">Fields the proxy extracted into its internal user model.</p>
<div class="data-container" id="parsedByProxy">
<!-- Populated by JavaScript -->
</div>
</div>
<div class="info-box">
<div class="info-header">
<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
<circle cx="12" cy="12" r="10"></circle>
<line x1="12" y1="16" x2="12" y2="12"></line>
<line x1="12" y1="8" x2="12.01" y2="8"></line>
</svg>
Raw Claims (userinfo)
</div>
<p class="empty-note">Complete set of claims returned by the IdP's userinfo endpoint.</p>
<div class="jwt-container">
<pre class="jwt-text" id="jsonData">Loading...</pre>
<pre class="jwt-text" id="rawClaims">Loading...</pre>
</div>
<div class="buttons">
<button class="copy-button" onclick="copyToClipboard('jsonData')">
<button class="copy-button" onclick="copyToClipboard('rawClaims')">
<svg xmlns="http://www.w3.org/2000/svg" width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
<rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect>
<path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path>
</svg>
Copy to Clipboard
Copy
</button>
</div>
</div>
<div class="info-box">
<div class="info-header">
<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
<circle cx="12" cy="12" r="10"></circle>
<line x1="12" y1="16" x2="12" y2="12"></line>
<line x1="12" y1="8" x2="12.01" y2="8"></line>
</svg>
Access Token Claims
</div>
<p class="empty-note">Decoded payload of the access token JWT (when the IdP issues one).</p>
<div class="jwt-container">
<pre class="jwt-text" id="accessTokenClaims">Loading...</pre>
</div>
<div class="buttons">
<button class="copy-button" onclick="copyToClipboard('accessTokenClaims')">
<svg xmlns="http://www.w3.org/2000/svg" width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
<rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect>
<path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path>
</svg>
Copy
</button>
</div>
</div>
<a href="/sso/debug/login" class="back-button">
Try Another SSO Login
</a>
@ -234,39 +284,58 @@ jwt_display_template = """
<script>
// This will be populated with the actual data from the server
const userData = SSO_DATA;
function renderUserData() {
const container = document.getElementById('userData');
const jsonDisplay = document.getElementById('jsonData');
// Format JSON with indentation for display
jsonDisplay.textContent = JSON.stringify(userData, null, 2);
// Clear container
const ssoData = SSO_DATA;
function renderParsed(container, parsed) {
container.innerHTML = '';
// Add each key-value pair to the UI
for (const [key, value] of Object.entries(userData)) {
if (typeof value !== 'object' || value === null) {
const row = document.createElement('div');
row.className = 'data-row';
const label = document.createElement('div');
label.className = 'data-label';
label.textContent = key;
const dataValue = document.createElement('div');
dataValue.className = 'data-value';
dataValue.textContent = value !== null ? value : 'null';
row.appendChild(label);
row.appendChild(dataValue);
container.appendChild(row);
const entries = Object.entries(parsed || {});
if (entries.length === 0) {
const note = document.createElement('p');
note.className = 'empty-note';
note.textContent = 'No fields available.';
container.appendChild(note);
return;
}
for (const [key, value] of entries) {
const row = document.createElement('div');
row.className = 'data-row';
const label = document.createElement('div');
label.className = 'data-label';
label.textContent = key;
const dataValue = document.createElement('div');
dataValue.className = 'data-value';
if (value === null || value === undefined) {
dataValue.textContent = 'null';
} else if (typeof value === 'object') {
dataValue.textContent = JSON.stringify(value);
} else {
dataValue.textContent = String(value);
}
row.appendChild(label);
row.appendChild(dataValue);
container.appendChild(row);
}
}
function renderJson(elementId, value) {
const el = document.getElementById(elementId);
const obj = value || {};
if (Object.keys(obj).length === 0) {
el.textContent = '(empty — provider returned no claims for this section)';
} else {
el.textContent = JSON.stringify(obj, null, 2);
}
}
function renderUserData() {
renderParsed(document.getElementById('parsedByProxy'), ssoData.parsed_by_proxy);
renderJson('rawClaims', ssoData.raw_claims);
renderJson('accessTokenClaims', ssoData.access_token_claims);
}
function copyToClipboard(elementId) {
const text = document.getElementById(elementId).textContent;
navigator.clipboard.writeText(text).then(() => {
@ -275,7 +344,7 @@ jwt_display_template = """
console.error('Could not copy text: ', err);
});
}
// Render the data when the page loads
document.addEventListener('DOMContentLoaded', renderUserData);
</script>

View file

@ -2,7 +2,7 @@ import asyncio
import json
import time
from datetime import datetime, timezone
from typing import Any, List, Literal, Optional, Union
from typing import Any, Callable, List, Literal, Optional, Union
import litellm
from litellm._logging import verbose_proxy_logger
@ -83,93 +83,146 @@ class ResetBudgetJob:
"Failed to reset spend counter %s: %s", counter_key, e
)
@staticmethod
async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None:
"""Drop a stale management-cache entry so the next read fetches from DB.
Tags and end-users are not reseeded by ``SpendCounterReseed.from_db``;
for those, when the spend counter expires the budget check falls back
to ``cached_obj.spend``. Keys, orgs, and team memberships are reseeded
from the DB, but auth still may consult ``user_api_key_cache`` objects
whose ``.spend`` field can lag a cross-pod DB reset. Deleting the cache
entry forces the next auth-time fetch to reload the zeroed row from
Postgres.
"""
try:
from litellm.proxy.proxy_server import user_api_key_cache
await user_api_key_cache.async_delete_cache(key=cache_key)
except Exception as e:
verbose_proxy_logger.warning(
"Failed to invalidate user_api_key_cache entry %s: %s",
cache_key,
e,
)
async def _cascade_reset_spend_for_budget_link(
self,
budgets_to_reset: List[LiteLLM_BudgetTableFull],
table: Any,
counter_key_fn: Callable[[Any], str],
log_subject: str,
extra_where: Optional[dict] = None,
cache_key_fn: Optional[Callable[[Any], Union[str, List[str]]]] = None,
):
"""
Generic cascade: zero spend on rows whose budget_id is in the reset set.
``cache_key_fn`` is optional: when provided, after the DB update each
matching row's entry or entries in ``user_api_key_cache`` are dropped so
cached spend cannot stay pinned above the zeroed DB row after a reset.
"""
budget_ids = [b.budget_id for b in budgets_to_reset if b.budget_id is not None]
if not budget_ids:
return
where: dict = {"budget_id": {"in": budget_ids}}
if extra_where:
where.update(extra_where)
try:
rows = await table.find_many(where=where)
except Exception as e:
rows = []
verbose_proxy_logger.warning(
"Failed to fetch %s for counter invalidation: %s", log_subject, e
)
update_result = await table.update_many(where=where, data={"spend": 0})
for row in rows:
await self._invalidate_spend_counter(counter_key_fn(row))
if cache_key_fn is not None:
cache_keys = cache_key_fn(row)
if isinstance(cache_keys, str):
cache_keys = [cache_keys]
for cache_key in cache_keys:
await self._invalidate_user_api_key_cache_entry(cache_key)
return update_result
async def reset_budget_for_litellm_team_members(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
"""
Resets the budget for all LiteLLM Team Members if their budget has expired
"""
budget_ids = [
budget.budget_id
for budget in budgets_to_reset
if budget.budget_id is not None
]
try:
memberships = await self.prisma_client.db.litellm_teammembership.find_many(
where={"budget_id": {"in": budget_ids}}
)
except Exception as e:
memberships = []
verbose_proxy_logger.warning(
"Failed to fetch team memberships for counter invalidation: %s", e
)
update_result = await self.prisma_client.db.litellm_teammembership.update_many(
where={"budget_id": {"in": budget_ids}},
data={
"spend": 0,
},
return await self._cascade_reset_spend_for_budget_link(
budgets_to_reset=budgets_to_reset,
table=self.prisma_client.db.litellm_teammembership,
counter_key_fn=lambda m: f"spend:team_member:{m.user_id}:{m.team_id}",
log_subject="team memberships",
cache_key_fn=lambda m: f"{m.team_id}_{m.user_id}",
)
for m in memberships:
await self._invalidate_spend_counter(
f"spend:team_member:{m.user_id}:{m.team_id}"
)
return update_result
async def reset_budget_for_keys_linked_to_budgets(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
"""
Resets the spend for keys linked to budget tiers that are being reset.
This handles keys that have budget_id but no budget_duration set on the key
itself. Keys with budget_id rely on their linked budget tier's reset schedule
rather than having their own budget_duration.
Keys that have their own budget_duration are already handled by
reset_budget_for_litellm_keys() and are excluded here to avoid
double-resetting.
Excludes keys with their own budget_duration; those are reset by
reset_budget_for_litellm_keys() to avoid double-resetting.
"""
budget_ids = [
budget.budget_id
for budget in budgets_to_reset
if budget.budget_id is not None
]
if not budget_ids:
return
where_clause: dict = {
"budget_id": {"in": budget_ids},
"budget_duration": None, # only keys without their own reset schedule
"spend": {"gt": 0}, # only reset keys that have accumulated spend
}
try:
keys = await self.prisma_client.db.litellm_verificationtoken.find_many(
where=where_clause
)
except Exception as e:
keys = []
verbose_proxy_logger.warning(
"Failed to fetch keys for counter invalidation: %s", e
)
update_result = (
await self.prisma_client.db.litellm_verificationtoken.update_many(
where=where_clause,
data={
"spend": 0,
},
)
return await self._cascade_reset_spend_for_budget_link(
budgets_to_reset=budgets_to_reset,
table=self.prisma_client.db.litellm_verificationtoken,
counter_key_fn=lambda k: f"spend:key:{k.token}",
log_subject="keys",
extra_where={"budget_duration": None, "spend": {"gt": 0}},
cache_key_fn=lambda k: k.token,
)
for k in keys:
await self._invalidate_spend_counter(f"spend:key:{k.token}")
async def reset_budget_for_orgs_linked_to_budgets(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
"""
Resets the spend for orgs linked to budget tiers that are being reset.
"""
return await self._cascade_reset_spend_for_budget_link(
budgets_to_reset=budgets_to_reset,
table=self.prisma_client.db.litellm_organizationtable,
counter_key_fn=lambda o: f"spend:org:{o.organization_id}",
log_subject="orgs",
extra_where={"spend": {"gt": 0}},
cache_key_fn=lambda o: [
f"org_id:{o.organization_id}",
f"org_id:{o.organization_id}:with_budget",
],
)
return update_result
async def reset_budget_for_tags_linked_to_budgets(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
"""
Resets the spend for tags linked to budget tiers that are being reset.
Also drops each tag's ``user_api_key_cache`` entry so the next
``_tag_max_budget_check`` reloads the zeroed row from the DB.
``SpendCounterReseed.from_db`` intentionally returns ``None`` for
tags, so the budget check falls back to the cached
``LiteLLM_TagTable.spend`` once the spend counter expires; without
this invalidation, that stale ``.spend`` keeps the tag over-budget
indefinitely.
"""
return await self._cascade_reset_spend_for_budget_link(
budgets_to_reset=budgets_to_reset,
table=self.prisma_client.db.litellm_tagtable,
counter_key_fn=lambda t: f"spend:tag:{t.tag_name}",
log_subject="tags",
extra_where={"spend": {"gt": 0}},
cache_key_fn=lambda t: f"tag:{t.tag_name}",
)
async def reset_budget_for_litellm_budget_table(self):
"""
@ -237,6 +290,14 @@ class ResetBudgetJob:
budgets_to_reset=budgets_to_reset
)
await self.reset_budget_for_orgs_linked_to_budgets(
budgets_to_reset=budgets_to_reset
)
await self.reset_budget_for_tags_linked_to_budgets(
budgets_to_reset=budgets_to_reset
)
if endusers_to_reset is not None and len(endusers_to_reset) > 0:
for enduser in endusers_to_reset:
try:

View file

@ -1131,10 +1131,12 @@ class DBSpendUpdateWriter:
timeout=timedelta(seconds=60)
) as transaction:
async with transaction.batch_() as batcher:
for (
user_id,
response_cost,
) in user_list_transactions.items():
# Sort by ID for consistent lock ordering across pods to prevent deadlocks.
# batch_() issues statements sequentially within the tx, so iteration
# order = lock acquisition order.
for user_id, response_cost in sorted(
user_list_transactions.items()
):
batcher.litellm_usertable.update_many(
where={"user_id": user_id},
data={"spend": {"increment": response_cost}},
@ -1186,10 +1188,10 @@ class DBSpendUpdateWriter:
timeout=timedelta(seconds=60)
) as transaction:
async with transaction.batch_() as batcher:
for (
token,
response_cost,
) in key_list_transactions.items():
# Sort by token for consistent lock ordering across pods to prevent deadlocks.
for token, response_cost in sorted(
key_list_transactions.items()
):
batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists
where={"token": token},
data={
@ -1230,10 +1232,10 @@ class DBSpendUpdateWriter:
timeout=timedelta(seconds=60)
) as transaction:
async with transaction.batch_() as batcher:
for (
team_id,
response_cost,
) in team_list_transactions.items():
# Sort by team_id for consistent lock ordering across pods to prevent deadlocks.
for team_id, response_cost in sorted(
team_list_transactions.items()
):
verbose_proxy_logger.debug(
"Updating spend for team id={} by {}".format(
team_id, response_cost
@ -1288,10 +1290,11 @@ class DBSpendUpdateWriter:
timeout=timedelta(seconds=60)
) as transaction:
async with transaction.batch_() as batcher:
for (
key,
response_cost,
) in team_member_list_transactions.items():
# Sort by composite key for consistent lock ordering across pods to prevent deadlocks.
# Key format "team_id::<v>::user_id::<v>" makes the string sort equivalent to sorting by (team_id, user_id).
for key, response_cost in sorted(
team_member_list_transactions.items()
):
# key is "team_id::<value>::user_id::<value>"
team_id = key.split("::")[1]
user_id = key.split("::")[3]
@ -1348,10 +1351,10 @@ class DBSpendUpdateWriter:
timeout=timedelta(seconds=60)
) as transaction:
async with transaction.batch_() as batcher:
for (
org_id,
response_cost,
) in org_list_transactions.items():
# Sort by org_id for consistent lock ordering across pods to prevent deadlocks.
for org_id, response_cost in sorted(
org_list_transactions.items()
):
batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists
where={"organization_id": org_id},
data={"spend": {"increment": response_cost}},
@ -1439,7 +1442,10 @@ class DBSpendUpdateWriter:
timeout=timedelta(seconds=60)
) as transaction:
async with transaction.batch_() as batcher:
for entity_id, response_cost in transactions.items():
# Sort by entity_id for consistent lock ordering across pods to prevent deadlocks.
for entity_id, response_cost in sorted(
transactions.items()
):
verbose_proxy_logger.debug(
f"Updating spend for {entity_name} {where_field}={entity_id} by {response_cost}"
)

View file

@ -10,13 +10,64 @@ import subprocess
import time
import urllib
import urllib.parse
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Any, Optional, Union
from typing import Any, Dict, Optional, Union
from litellm._logging import verbose_proxy_logger
from litellm.secret_managers.main import str_to_bool
@dataclass(frozen=True)
class IAMEndpoint:
"""Static parts of an RDS IAM-authenticated Postgres connection.
The IAM token rotates every ~15 minutes; everything else (host, port, user,
database name, schema) stays fixed. We capture the static fields once so
refresh just regenerates the token and reassembles the URL.
"""
host: str
port: str
user: str
name: str
schema: Optional[str] = None
def build_url(self, token: str) -> str:
url = f"postgresql://{self.user}:{token}@{self.host}:{self.port}/{self.name}"
if self.schema:
url += f"?schema={self.schema}"
return url
def parse_iam_endpoint_from_url(url: str) -> IAMEndpoint:
"""Parse an IAMEndpoint from a Postgres URL.
Used so a reader URL can drive its own IAM refresh without requiring
callers to set parallel DATABASE_HOST_READ_REPLICA / etc. env vars.
"""
parsed = urllib.parse.urlparse(url)
if not parsed.hostname or not parsed.username:
raise ValueError("Cannot parse IAM endpoint from URL: missing host or username")
name = (parsed.path or "/").lstrip("/")
if not name:
raise ValueError("Cannot parse IAM endpoint from URL: missing database name")
port = str(parsed.port) if parsed.port else "5432"
schema: Optional[str] = None
if parsed.query:
qs = urllib.parse.parse_qs(parsed.query)
schema_vals = qs.get("schema")
if schema_vals:
schema = schema_vals[0]
return IAMEndpoint(
host=parsed.hostname,
port=port,
user=parsed.username,
name=name,
schema=schema,
)
class PrismaWrapper:
"""
Wrapper around Prisma client that handles RDS IAM token authentication.
@ -37,10 +88,33 @@ class PrismaWrapper:
# Fallback refresh interval if token parsing fails (10 minutes)
FALLBACK_REFRESH_INTERVAL_SECONDS = 600
def __init__(self, original_prisma: Any, iam_token_db_auth: bool):
def __init__(
self,
original_prisma: Any,
iam_token_db_auth: bool,
*,
db_url_env_var: str = "DATABASE_URL",
iam_endpoint: Optional[IAMEndpoint] = None,
recreate_uses_datasource: bool = False,
log_prefix: str = "",
):
self._original_prisma = original_prisma
self.iam_token_db_auth = iam_token_db_auth
# Per-connection knobs so the same wrapper can be used for the writer
# (defaults: DATABASE_URL env, IAM endpoint from DATABASE_HOST/etc.,
# recreate via env reload) or for a reader (DATABASE_URL_READ_REPLICA
# env, IAM endpoint parsed from that URL, recreate via datasource
# override since Prisma only auto-reads DATABASE_URL).
self._db_url_env_var = db_url_env_var
self._iam_endpoint = iam_endpoint
self._recreate_uses_datasource = recreate_uses_datasource
# Tag every log line emitted by this wrapper instance so writer and
# reader can be told apart in interleaved output (e.g. "[writer] RDS
# IAM token refresh scheduled in 720 seconds"). Empty string (default)
# keeps backward-compatible logs for the single-DB case.
self._log_prefix = f"{log_prefix} " if log_prefix else ""
# Background token refresh task management
self._token_refresh_task: Optional[asyncio.Task] = None
self._reconnection_lock = asyncio.Lock()
@ -157,7 +231,7 @@ class PrismaWrapper:
Returns 0 if token should be refreshed immediately.
Returns FALLBACK_REFRESH_INTERVAL_SECONDS if parsing fails.
"""
db_url = os.getenv("DATABASE_URL")
db_url = os.getenv(self._db_url_env_var)
token = self._extract_token_from_db_url(db_url)
expiration_time = self._parse_token_expiration(token)
@ -199,12 +273,30 @@ class PrismaWrapper:
return datetime.utcnow() > expiration_time
def get_rds_iam_token(self) -> Optional[str]:
"""Generate a new RDS IAM token and update DATABASE_URL."""
if self.iam_token_db_auth:
from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token
"""Generate a new RDS IAM token and update the configured DB URL env var.
When the wrapper was constructed with an explicit `iam_endpoint`
(typical for a reader wrapper whose host/port/user came from a parsed
URL), use that. Otherwise fall back to the legacy DATABASE_HOST/PORT/
USER/NAME/SCHEMA env vars (writer behavior).
"""
if not self.iam_token_db_auth:
return None
from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token
if self._iam_endpoint is not None:
endpoint = self._iam_endpoint
token = generate_iam_auth_token(
db_host=endpoint.host, db_port=endpoint.port, db_user=endpoint.user
)
_db_url = endpoint.build_url(token)
else:
db_host = os.getenv("DATABASE_HOST")
db_port = os.getenv("DATABASE_PORT")
# Default to the Postgres standard port; passing None to
# `generate_iam_auth_token` makes botocore embed the literal
# string "None" in the presigned URL, which then fails to parse.
db_port = os.getenv("DATABASE_PORT", "5432")
db_user = os.getenv("DATABASE_USER")
db_name = os.getenv("DATABASE_NAME")
db_schema = os.getenv("DATABASE_SCHEMA")
@ -217,9 +309,8 @@ class PrismaWrapper:
if db_schema:
_db_url += f"?schema={db_schema}"
os.environ["DATABASE_URL"] = _db_url
return _db_url
return None
os.environ[self._db_url_env_var] = _db_url
return _db_url
async def recreate_prisma_client(
self, new_db_url: str, http_client: Optional[Any] = None
@ -231,6 +322,11 @@ class PrismaWrapper:
synchronous `subprocess.Popen.wait()` that can freeze the asyncio event
loop for 30-120+ seconds when the engine is stuck on TCP close,
breaking `/health/liveliness` and causing Kubernetes pod restarts.
The writer wrapper relies on Prisma re-reading `DATABASE_URL` from env;
the reader wrapper opts into `recreate_uses_datasource=True` so the
new URL is passed explicitly via `datasource={"url": ...}` (Prisma
does not auto-read alternate env vars like DATABASE_URL_READ_REPLICA).
"""
from prisma import Prisma # type: ignore
@ -238,10 +334,12 @@ class PrismaWrapper:
if old_engine_pid > 0:
await self._kill_engine_process(old_engine_pid)
kwargs: Dict[str, Any] = {}
if http_client is not None:
self._original_prisma = Prisma(http=http_client)
else:
self._original_prisma = Prisma()
kwargs["http"] = http_client
if self._recreate_uses_datasource:
kwargs["datasource"] = {"url": new_db_url}
self._original_prisma = Prisma(**kwargs)
await self._original_prisma.connect()
@ -265,7 +363,8 @@ class PrismaWrapper:
self._token_refresh_task = asyncio.create_task(self._token_refresh_loop())
verbose_proxy_logger.info(
"Started RDS IAM token proactive refresh background task"
"%sStarted RDS IAM token proactive refresh background task",
self._log_prefix,
)
async def stop_token_refresh_task(self) -> None:
@ -283,7 +382,9 @@ class PrismaWrapper:
except asyncio.CancelledError:
pass
self._token_refresh_task = None
verbose_proxy_logger.info("Stopped RDS IAM token refresh background task")
verbose_proxy_logger.info(
"%sStopped RDS IAM token refresh background task", self._log_prefix
)
async def _token_refresh_loop(self) -> None:
"""
@ -294,7 +395,7 @@ class PrismaWrapper:
This is more efficient than polling, requiring only 1 wake-up per token cycle.
"""
verbose_proxy_logger.info(
f"RDS IAM token refresh loop started. "
f"{self._log_prefix}RDS IAM token refresh loop started. "
f"Tokens will be refreshed {self.TOKEN_REFRESH_BUFFER_SECONDS}s before expiration."
)
@ -305,21 +406,25 @@ class PrismaWrapper:
if sleep_seconds > 0:
verbose_proxy_logger.info(
f"RDS IAM token refresh scheduled in {sleep_seconds:.0f} seconds "
f"({sleep_seconds / 60:.1f} minutes)"
f"{self._log_prefix}RDS IAM token refresh scheduled in "
f"{sleep_seconds:.0f} seconds ({sleep_seconds / 60:.1f} minutes)"
)
await asyncio.sleep(sleep_seconds)
# Refresh the token
verbose_proxy_logger.info("Proactively refreshing RDS IAM token...")
verbose_proxy_logger.info(
"%sProactively refreshing RDS IAM token...", self._log_prefix
)
await self._safe_refresh_token()
except asyncio.CancelledError:
verbose_proxy_logger.info("RDS IAM token refresh loop cancelled")
verbose_proxy_logger.info(
"%sRDS IAM token refresh loop cancelled", self._log_prefix
)
break
except Exception as e:
verbose_proxy_logger.error(
f"Error in RDS IAM token refresh loop: {e}. "
f"{self._log_prefix}Error in RDS IAM token refresh loop: {e}. "
f"Retrying in {self.FALLBACK_REFRESH_INTERVAL_SECONDS}s..."
)
# On error, wait before retrying to avoid tight error loops
@ -341,65 +446,75 @@ class PrismaWrapper:
await self.recreate_prisma_client(new_db_url)
self._last_refresh_time = datetime.utcnow()
verbose_proxy_logger.info(
"RDS IAM token refreshed successfully. New token valid for ~15 minutes."
"%sRDS IAM token refreshed successfully. New token valid for ~15 minutes.",
self._log_prefix,
)
else:
verbose_proxy_logger.error(
"Failed to generate new RDS IAM token during proactive refresh"
"%sFailed to generate new RDS IAM token during proactive refresh",
self._log_prefix,
)
def __getattr__(self, name: str):
"""
Proxy attribute access to the underlying Prisma client.
If IAM token auth is enabled and the token is expired, this method
provides a synchronous fallback to refresh the token. However, this
should rarely be needed since the background task proactively refreshes
tokens before they expire.
If IAM token auth is enabled and the token is found expired here, the
proactive refresh task has missed its window. Behavior depends on
whether we're called from inside a running event loop:
FIXED: Now properly waits for reconnection to complete before returning,
instead of the previous fire-and-forget pattern that caused the bug.
- Inside the loop (typical: from a coroutine): schedule a refresh as a
background task and return the (stale) attribute. The caller's await
will likely fail with a connection error and be retried by upper
layers (`call_with_db_reconnect_retry`); by that time the refresh
has either completed or escalated to the proactive loop's error
path. We CANNOT block here — `run_coroutine_threadsafe(...)` +
`future.result()` from inside the same loop deadlocks the loop
(loop thread is blocked, scheduled coroutine never runs, 30s timeout).
- No running loop (sync caller, mostly tests): run the refresh in a
fresh loop and re-fetch the attribute.
"""
original_attr = getattr(self._original_prisma, name)
if self.iam_token_db_auth:
db_url = os.getenv("DATABASE_URL")
db_url = os.getenv(self._db_url_env_var)
# Check if token is expired (should be rare if background task is running)
if self.is_token_expired(db_url):
verbose_proxy_logger.warning(
"RDS IAM token expired in __getattr__ - proactive refresh may have failed. "
"Triggering synchronous fallback refresh..."
)
try:
running_loop = asyncio.get_running_loop()
except RuntimeError:
running_loop = None
new_db_url = self.get_rds_iam_token()
if new_db_url:
loop = asyncio.get_event_loop()
if loop.is_running():
# FIXED: Actually wait for the reconnection to complete!
# The previous code used fire-and-forget which caused the bug.
future = asyncio.run_coroutine_threadsafe(
self.recreate_prisma_client(new_db_url), loop
)
try:
# Wait up to 30 seconds for reconnection
future.result(timeout=30)
verbose_proxy_logger.info(
"Synchronous token refresh completed successfully"
)
except Exception as e:
verbose_proxy_logger.error(
f"Failed to refresh token synchronously: {e}"
)
raise
else:
asyncio.run(self.recreate_prisma_client(new_db_url))
# Get the NEW attribute after reconnection
original_attr = getattr(self._original_prisma, name)
if running_loop is not None:
verbose_proxy_logger.warning(
"%sRDS IAM token expired in __getattr__ — proactive refresh "
"may have failed. Scheduling async refresh; the current "
"request may fail and be retried with the fresh token.",
self._log_prefix,
)
# Non-blocking: schedule the locked refresh on the
# running loop. The reconnection lock inside
# `_safe_refresh_token` coalesces concurrent triggers.
running_loop.create_task(self._safe_refresh_token())
else:
raise ValueError("Failed to get RDS IAM token")
verbose_proxy_logger.warning(
"%sRDS IAM token expired in __getattr__ — proactive refresh "
"may have failed. Triggering synchronous fallback refresh...",
self._log_prefix,
)
new_db_url = self.get_rds_iam_token()
if new_db_url:
asyncio.run(self.recreate_prisma_client(new_db_url))
# Re-fetch attribute against the recreated Prisma instance.
original_attr = getattr(self._original_prisma, name)
verbose_proxy_logger.info(
"%sSynchronous token refresh completed successfully",
self._log_prefix,
)
else:
raise ValueError("Failed to get RDS IAM token")
return original_attr

View file

@ -0,0 +1,213 @@
"""
RoutingPrismaWrapper: routes Prisma reads to a read-replica client and writes
to a writer client. Used when DATABASE_URL_READ_REPLICA is configured;
otherwise PrismaClient uses the writer-only PrismaWrapper directly.
"""
import os
from typing import Any, Callable, Optional
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.prisma_client import PrismaWrapper
# Per-model action methods that read from the database. These are routed to
# the read replica when one is configured.
_MODEL_READ_METHODS = frozenset(
{
"find_first",
"find_first_or_raise",
"find_many",
"find_unique",
"find_unique_or_raise",
"count",
"group_by",
"query_first",
"query_raw",
}
)
# Top-level Prisma client methods that read from the database.
_TOP_LEVEL_READ_METHODS = frozenset({"query_first", "query_raw"})
class _RoutedActions:
"""Per-model accessor that sends reads to the reader and writes to the writer.
`should_use_reader` is consulted on every read dispatch so a mid-call flip
of the routing wrapper's reader-availability flag (e.g. after the reader
fails a recreate) is observed without re-fetching the actions accessor.
"""
__slots__ = ("_writer_actions", "_reader_actions", "_should_use_reader")
def __init__(
self,
writer_actions: Any,
reader_actions: Any,
should_use_reader: Callable[[], bool],
):
self._writer_actions = writer_actions
self._reader_actions = reader_actions
self._should_use_reader = should_use_reader
def __getattr__(self, name: str) -> Any:
if name in _MODEL_READ_METHODS and self._should_use_reader():
return getattr(self._reader_actions, name)
return getattr(self._writer_actions, name)
class RoutingPrismaWrapper:
"""
Routes Prisma operations between a writer and a reader Prisma client.
Reads (find_*, count, group_by, query_raw, query_first) go to the reader;
everything else (writes, transactions, raw execute) goes to the writer.
Lifecycle methods (connect, disconnect, IAM token refresh) act on both
clients so callers do not need to know about the split. When
IAM_TOKEN_DB_AUTH is enabled, both writer and reader refresh their tokens
independently on their own ~12-minute cadence.
Reader degradation: a reader-side failure (failed connect, failed
recreate) is non-fatal — the wrapper sets `_reader_unavailable=True`, logs
a warning, and routes subsequent reads to the writer. The next successful
`connect()` or `recreate_prisma_client()` clears the flag. This keeps the
proxy serving traffic during transient reader outages instead of failing
startup or returning errors for read-heavy endpoints.
"""
def __init__(self, writer: PrismaWrapper, reader: PrismaWrapper):
self._writer = writer
self._reader = reader
# When True, reads fall back to the writer. Flipped on by reader
# connect/recreate failures and flipped off on the next reader recovery.
self._reader_unavailable: bool = False
@property
def writer(self) -> PrismaWrapper:
return self._writer
@property
def reader(self) -> PrismaWrapper:
return self._reader
@property
def reader_unavailable(self) -> bool:
return self._reader_unavailable
def _should_use_reader(self) -> bool:
return not self._reader_unavailable
async def connect(self, *args: Any, **kwargs: Any) -> None:
await self._writer.connect(*args, **kwargs)
verbose_proxy_logger.info("[writer] DB connected")
try:
await self._reader.connect(*args, **kwargs)
self._reader_unavailable = False
verbose_proxy_logger.info("[reader] DB connected")
except Exception as e:
# Degrade gracefully: the proxy keeps serving traffic with reads
# routed to the writer until the reader endpoint is reachable.
# Aborting startup here would tie proxy availability to an
# opt-in, best-effort reader endpoint.
self._reader_unavailable = True
verbose_proxy_logger.warning(
"Failed to connect to read replica DB: %s. "
"Falling back to the writer for reads until the reader is reachable.",
e,
)
async def disconnect(self, *args: Any, **kwargs: Any) -> None:
first_error: Optional[BaseException] = None
for client in (self._writer, self._reader):
try:
await client.disconnect(*args, **kwargs)
except Exception as e:
if first_error is None:
first_error = e
verbose_proxy_logger.warning("Error disconnecting Prisma client: %s", e)
if first_error is not None:
raise first_error
def is_connected(self) -> bool:
# Reflects writer health only. The reader is best-effort; its
# availability is tracked via `_reader_unavailable` and a degraded
# reader must NOT cause a writer reconnect (would loop indefinitely
# since recreate_prisma_client only fixes writer-side problems).
return bool(self._writer.is_connected())
async def start_token_refresh_task(self) -> None:
await self._writer.start_token_refresh_task()
await self._reader.start_token_refresh_task()
async def stop_token_refresh_task(self) -> None:
await self._writer.stop_token_refresh_task()
await self._reader.stop_token_refresh_task()
async def recreate_prisma_client(
self, new_db_url: str, http_client: Optional[Any] = None
) -> None:
"""Recreate both writer and reader Prisma clients.
The writer reconnect path in PrismaClient calls
`self.db.recreate_prisma_client(...)`. Without this method, a DB-wide
connectivity event would only re-create the writer; the reader engine
would stay broken and every routed read would fail. We always recreate
the writer first (its URL is the one passed in), then best-effort
recreate the reader. A reader failure flips `_reader_unavailable=True`
so reads transparently fall through to the writer.
"""
await self._writer.recreate_prisma_client(new_db_url, http_client=http_client)
try:
await self._recreate_reader(http_client=http_client)
self._reader_unavailable = False
except Exception as e:
self._reader_unavailable = True
verbose_proxy_logger.warning(
"Failed to recreate reader Prisma client: %s. "
"Reads will fall back to the writer until the reader recovers.",
e,
)
async def _recreate_reader(self, http_client: Optional[Any] = None) -> None:
"""Resolve the reader URL and recreate its Prisma client.
IAM-enabled readers regenerate their token (host/port/user came from
the parsed reader URL at construction time). Non-IAM readers reuse
the URL stored in `DATABASE_URL_READ_REPLICA`.
"""
if self._reader.iam_token_db_auth:
new_reader_url = self._reader.get_rds_iam_token()
if not new_reader_url:
raise RuntimeError(
"Failed to generate fresh IAM token for read replica"
)
await self._reader.recreate_prisma_client(
new_reader_url, http_client=http_client
)
return
reader_url = os.getenv("DATABASE_URL_READ_REPLICA", "")
if not reader_url:
raise RuntimeError(
"DATABASE_URL_READ_REPLICA not set; cannot recreate read replica client"
)
await self._reader.recreate_prisma_client(reader_url, http_client=http_client)
def __getattr__(self, name: str) -> Any:
if name in _TOP_LEVEL_READ_METHODS:
target = self._writer if self._reader_unavailable else self._reader
return getattr(target, name)
writer_attr = getattr(self._writer, name)
# Per-model action accessors are non-callable instances that expose
# both `find_many` and `create`. Methods like execute_raw / batch_ /
# tx are callables and stay on the writer untouched.
if (
not callable(writer_attr)
and hasattr(writer_attr, "find_many")
and hasattr(writer_attr, "create")
):
try:
reader_attr = getattr(self._reader, name)
except AttributeError:
return writer_attr
return _RoutedActions(writer_attr, reader_attr, self._should_use_reader)
return writer_attr

View file

@ -242,6 +242,10 @@ async def list_guardrails_v2(
gid = guardrail.get("guardrail_id")
if gid in seen_guardrail_ids:
continue
# Skip stale DB-backed entries — the DB row was deleted (likely by
# another pod) and reconciliation hasn't fired yet on this pod.
if gid is not None and IN_MEMORY_GUARDRAIL_HANDLER.get_source(gid) == "db":
continue
if not is_admin:
g_team_id = guardrail.get("team_id")
if g_team_id is not None and g_team_id not in caller_team_ids:
@ -360,7 +364,7 @@ async def create_guardrail(
try:
IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail(
guardrail=cast(Guardrail, result)
guardrail=cast(Guardrail, result), source="db"
)
verbose_proxy_logger.info(
f"Immediate sync: Successfully initialized guardrail '{guardrail_name}' (ID: {guardrail_id})"
@ -1017,7 +1021,7 @@ async def approve_guardrail_submission(
}
try:
IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail(
guardrail=cast(Guardrail, guardrail_dict)
guardrail=cast(Guardrail, guardrail_dict), source="db"
)
verbose_proxy_logger.info(
"Approved guardrail %s (ID: %s) and initialized in memory",
@ -1295,10 +1299,18 @@ async def get_guardrail_info(guardrail_id: str):
guardrail_id=guardrail_id, prisma_client=prisma_client
)
if result is None:
result = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_id(
in_memory = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_id(
guardrail_id=guardrail_id
)
guardrail_definition_location = GUARDRAIL_DEFINITION_LOCATION.CONFIG
# Only return config-loaded entries here. A DB-backed entry that's
# missing from the DB is stale (deleted on another pod, awaiting
# reconciliation on this one) and must surface as 404.
if (
in_memory is not None
and IN_MEMORY_GUARDRAIL_HANDLER.get_source(guardrail_id) == "config"
):
result = in_memory
guardrail_definition_location = GUARDRAIL_DEFINITION_LOCATION.CONFIG
if result is None:
raise HTTPException(

View file

@ -3,7 +3,7 @@
import importlib
import os
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Type, cast
from typing import Any, Dict, List, Literal, Optional, Set, Type, cast
import litellm
from litellm import Router
@ -403,11 +403,19 @@ class InMemoryGuardrailHandler:
Guardrail id to CustomGuardrail object mapping
"""
self._sources: Dict[str, Literal["db", "config"]] = {}
"""
Guardrail id to provenance marker. "db" entries are reconciled against
the DB on each polling tick; "config" entries are owned by proxy_config.yaml
and never deleted by reconciliation.
"""
def initialize_guardrail(
self,
guardrail: Guardrail,
config_file_path: Optional[str] = None,
llm_router: Optional["Router"] = None,
source: Literal["db", "config"] = "config",
) -> Optional[Guardrail]:
"""
Initialize a guardrail from a dictionary and add it to the litellm callback manager
@ -420,6 +428,10 @@ class InMemoryGuardrailHandler:
verbose_proxy_logger.debug(
"guardrail_id already exists in IN_MEMORY_GUARDRAILS"
)
# Honor the caller's source even on the early-return path so a
# racing polling tick or a hot-reload of config can correct an
# entry's provenance.
self._sources[guardrail_id] = source
return self.IN_MEMORY_GUARDRAILS[guardrail_id]
custom_guardrail_callback: Optional[CustomGuardrail] = None
@ -482,6 +494,11 @@ class InMemoryGuardrailHandler:
"skip_system_message_in_guardrail",
getattr(litellm_params, "skip_system_message_in_guardrail", None),
)
setattr(
custom_guardrail_callback,
"skip_tool_message_in_guardrail",
getattr(litellm_params, "skip_tool_message_in_guardrail", None),
)
parsed_guardrail = Guardrail(
guardrail_id=guardrail.get("guardrail_id"),
@ -492,6 +509,7 @@ class InMemoryGuardrailHandler:
# store references to the guardrail in memory
self.IN_MEMORY_GUARDRAILS[guardrail_id] = parsed_guardrail
self.guardrail_id_to_custom_guardrail[guardrail_id] = custom_guardrail_callback
self._sources[guardrail_id] = source
return parsed_guardrail
@ -552,7 +570,10 @@ class InMemoryGuardrailHandler:
return _guardrail_callback
def update_in_memory_guardrail(
self, guardrail_id: str, guardrail: Guardrail
self,
guardrail_id: str,
guardrail: Guardrail,
source: Literal["db", "config"] = "db",
) -> None:
"""
Update a guardrail in memory
@ -561,6 +582,7 @@ class InMemoryGuardrailHandler:
- updates the guardrail params in litellm.callback_manager
"""
self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail
self._sources[guardrail_id] = source
custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.get(
guardrail_id
@ -579,6 +601,7 @@ class InMemoryGuardrailHandler:
"""
# Remove from in-memory storage
self.IN_MEMORY_GUARDRAILS.pop(guardrail_id, None)
self._sources.pop(guardrail_id, None)
# Remove the callback from litellm.callbacks
custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.pop(
@ -603,6 +626,34 @@ class InMemoryGuardrailHandler:
"""
return self.IN_MEMORY_GUARDRAILS.get(guardrail_id)
def get_source(self, guardrail_id: str) -> Optional[Literal["db", "config"]]:
"""
Return the provenance of an in-memory guardrail.
"""
return self._sources.get(guardrail_id)
def reconcile_db_guardrails(self, db_guardrail_ids: Set[str]) -> List[str]:
"""
Drop in-memory entries that originated from the DB but are no longer
present in db_guardrail_ids. Config-loaded guardrails are never touched.
Called by the periodic DB polling tick so that a guardrail deleted
on another pod is eventually purged from this pod's memory + callbacks.
"""
stale_ids = [
guardrail_id
for guardrail_id, source in self._sources.items()
if source == "db" and guardrail_id not in db_guardrail_ids
]
for guardrail_id in stale_ids:
verbose_proxy_logger.info(
"Reconcile: removing stale DB-backed guardrail '%s' from memory "
"(deleted in DB by another pod)",
guardrail_id,
)
self.delete_in_memory_guardrail(guardrail_id)
return stale_ids
def _has_guardrail_params_changed(
self, guardrail_id: str, new_guardrail: Guardrail
) -> bool:
@ -656,7 +707,10 @@ class InMemoryGuardrailHandler:
return len(changed_fields) > 0
def reinitialize_guardrail(
self, guardrail: Guardrail, config_file_path: Optional[str] = None
self,
guardrail: Guardrail,
config_file_path: Optional[str] = None,
source: Literal["db", "config"] = "config",
) -> Optional[Guardrail]:
"""
Force re-initialization of a guardrail even if it exists in memory.
@ -675,7 +729,7 @@ class InMemoryGuardrailHandler:
# Initialize fresh (will add new callback to litellm.callbacks)
return self.initialize_guardrail(
guardrail=guardrail, config_file_path=config_file_path
guardrail=guardrail, config_file_path=config_file_path, source=source
)
def sync_guardrail_from_db(
@ -696,9 +750,15 @@ class InMemoryGuardrailHandler:
f"Guardrail '{guardrail_name}' (ID: {guardrail_id}) params changed, re-initializing..."
)
return self.reinitialize_guardrail(
guardrail=guardrail, config_file_path=config_file_path
guardrail=guardrail,
config_file_path=config_file_path,
source="db",
)
# Params unchanged but the entry is still DB-backed; make sure the
# source marker reflects that even if it was previously set differently
# (e.g. a config entry whose UUID later collided with a DB row).
self._sources[guardrail_id] = "db"
return self.IN_MEMORY_GUARDRAILS.get(guardrail_id)

View file

@ -30,6 +30,7 @@ def init_guardrails_v2(
guardrail=cast(Guardrail, guardrail),
config_file_path=config_file_path,
llm_router=llm_router,
source="config",
)
if initialized_guardrail:
guardrail_list.append(initialized_guardrail)

View file

@ -253,27 +253,63 @@ class SharedHealthCheckManager:
# Always release the lock
await self.release_health_check_lock()
else:
# Lock not acquired, wait briefly and try to get cached results
# If Redis is not configured, skip polling — there is no cache
# to wait for.
if self.redis_cache is None:
return await perform_health_check(
model_list=model_list,
details=details,
max_concurrency=max_concurrency,
)
# Lock not acquired — poll for cached results until the lock
# holder finishes or the lock expires, rather than falling back
# to a redundant local health check after only 2 seconds.
verbose_proxy_logger.debug(
"Pod %s waiting for other pod to complete health check", self.pod_id
)
# Wait a bit for the other pod to complete
await asyncio.sleep(2)
poll_interval = 5 # seconds between cache checks
max_wait = self.lock_ttl # wait at most as long as the lock can live
elapsed = 0
# Try to get cached results again
cached_results = await self.get_cached_health_check_results()
if cached_results is not None:
return (
cached_results.get("healthy_endpoints", []),
cached_results.get("unhealthy_endpoints", []),
{},
)
while elapsed < max_wait:
await asyncio.sleep(poll_interval)
elapsed += poll_interval
# Still no cache, fall back to local health check
cached_results = await self.get_cached_health_check_results()
if cached_results is not None:
verbose_proxy_logger.info(
"Pod %s using cached health check results after waiting %ds",
self.pod_id,
elapsed,
)
return (
cached_results.get("healthy_endpoints", []),
cached_results.get("unhealthy_endpoints", []),
{},
)
# Check if the lock is still held — if it was released without
# caching (e.g. the holder crashed), stop waiting early.
try:
lock_key = self.get_health_check_lock_key()
current_owner = await self.redis_cache.async_get_cache(lock_key)
if current_owner is None:
verbose_proxy_logger.debug(
"Pod %s detected lock released without cache, stopping wait",
self.pod_id,
)
break
except Exception:
# Redis hiccup — continue polling rather than crashing out
pass
# Exhausted wait — fall back to local health check
verbose_proxy_logger.warning(
"Pod %s falling back to local health check (no cache available)",
"Pod %s falling back to local health check after waiting %ds (no cache available)",
self.pod_id,
elapsed,
)
return await perform_health_check(

View file

@ -1742,29 +1742,54 @@ async def test_model_connection(
# Look up model configuration from router if model name is provided
# This gets the litellm_params from proxy config (with resolved env vars)
config_litellm_params: dict = {}
if model_name and llm_router is not None:
if llm_router is not None:
# Prefer disambiguation by deployment id (`model_info.id`) when
# the caller supplies it. This is required when multiple
# deployments share a `model_name` (e.g. wildcard `openai/*`
# with multiple `api_base` values for failover): the UI's
# "Test Connection" button targets a specific row, and that
# row's id is the only thing that uniquely identifies which
# deployment to probe. Without this, all duplicates collapse
# onto `deployments[0]`.
request_model_info = model_info or {}
request_model_id = request_model_info.get("id")
try:
# First try to find by proxy model_name (e.g., "gpt-4o")
deployments = llm_router.get_model_list(model_name=model_name)
# If not found, try to find by litellm model name (e.g., "azure/gpt-4o")
if not deployments or len(deployments) == 0:
all_deployments = llm_router.get_model_list(model_name=None)
if all_deployments:
for deployment in all_deployments:
if (
deployment.get("litellm_params", {}).get("model")
== model_name
):
deployments = [deployment]
break
if deployments and len(deployments) > 0:
# Use the first deployment's litellm_params as base config
# These already have resolved environment variables from proxy config
config_litellm_params = dict(
deployments[0].get("litellm_params", {})
deployment_by_id = None
if request_model_id:
deployment_by_id = llm_router.get_deployment(
model_id=request_model_id
)
if deployment_by_id is not None:
config_litellm_params = deployment_by_id.litellm_params.model_dump(
exclude_none=True
)
elif model_name:
# Fall back to model_name lookup for callers (e.g. the
# "Add Model" wizard, or curl) that don't supply an id.
# First try to find by proxy model_name (e.g., "gpt-4o")
deployments = llm_router.get_model_list(model_name=model_name)
# If not found, try to find by litellm model name
# (e.g., "azure/gpt-4o")
if not deployments or len(deployments) == 0:
all_deployments = llm_router.get_model_list(model_name=None)
if all_deployments:
for deployment in all_deployments:
if (
deployment.get("litellm_params", {}).get("model")
== model_name
):
deployments = [deployment]
break
if deployments and len(deployments) > 0:
# Use the first deployment's litellm_params as base
# config. These already have resolved environment
# variables from proxy config.
config_litellm_params = dict(
deployments[0].get("litellm_params", {})
)
except Exception as e:
verbose_proxy_logger.debug(
f"Could not find model {model_name} in router: {e}. "

View file

@ -319,6 +319,9 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
response_cost=response_cost,
)
if self.dual_cache.redis_cache is not None:
await self._push_in_memory_increments_to_redis()
verbose_proxy_logger.debug(
"current state of in memory cache %s",
json.dumps(

View file

@ -1,5 +1,6 @@
import asyncio
import copy
import json
import re
import time
from collections import OrderedDict
@ -794,8 +795,17 @@ class LiteLLMProxyRequestSetup:
)
)
for k, v in litellm_logging_metadata_headers.items():
if v is not None:
if v is None:
continue
# httpx requires header values to be str or bytes; coerce numbers/bools
# to str and JSON-encode dict/list (e.g. user_api_key_spend is float,
# user_api_key_auth_metadata is dict). See #27458.
if isinstance(v, (dict, list)):
returned_headers["x-litellm-{}".format(k)] = json.dumps(v)
elif isinstance(v, (str, bytes)):
returned_headers["x-litellm-{}".format(k)] = v
else:
returned_headers["x-litellm-{}".format(k)] = str(v)
return returned_headers
@ -1731,6 +1741,7 @@ async def add_litellm_data_to_request( # noqa: PLR0915
data=data,
user_api_key_dict=user_api_key_dict,
pre_alias_model_name=_pre_alias_model,
llm_router=llm_router,
)
## ENFORCED PARAMS CHECK
@ -1864,6 +1875,7 @@ def _apply_credential_overrides_from_model_config(
data: dict,
user_api_key_dict: UserAPIKeyAuth,
pre_alias_model_name: Optional[str] = None,
llm_router: Optional[Router] = None,
) -> None:
"""
Walk the model_config precedence chain in team/project metadata.
@ -1899,10 +1911,19 @@ def _apply_credential_overrides_from_model_config(
if not project_model_config and not team_model_config:
return
# Extract provider hint from model name (e.g. "azure/gpt-4" -> "azure")
# Extract provider hint from model name (e.g. "azure/gpt-4" -> "azure").
# When the user-facing name has no provider prefix, fall back to the
# deployment's litellm_params so multi-provider defaultconfig entries
# don't silently match the first dict key (#27516).
provider: Optional[str] = None
if "/" in model_name:
provider = model_name.split("/", 1)[0]
elif llm_router is not None:
provider = _resolve_provider_from_deployment(
llm_router=llm_router,
model_name=model_name,
pre_alias_model_name=pre_alias_model_name,
)
credential_name = _resolve_credential_from_model_config(
model_name=model_name,
@ -1938,6 +1959,48 @@ def _apply_credential_overrides_from_model_config(
)
def _resolve_provider_from_deployment(
llm_router: Router,
model_name: str,
pre_alias_model_name: Optional[str] = None,
) -> Optional[str]:
"""
Resolve a provider hint from the deployment's litellm_params when the
user-facing model name has no provider prefix.
Tries the post-alias name first (the resolved model group), then the
pre-alias name. Returns None if no deployment is found or the deployment
has no usable provider info.
"""
candidates = [model_name]
if pre_alias_model_name and pre_alias_model_name != model_name:
candidates.append(pre_alias_model_name)
for name in candidates:
try:
deployment = llm_router.get_deployment_by_model_group_name(
model_group_name=name
)
except Exception:
deployment = None
if deployment is None:
continue
litellm_params = getattr(deployment, "litellm_params", None)
if litellm_params is None:
continue
custom_provider = getattr(litellm_params, "custom_llm_provider", None)
if custom_provider:
return custom_provider
deployment_model = getattr(litellm_params, "model", "") or ""
if "/" in deployment_model:
return deployment_model.split("/", 1)[0]
return None
def _resolve_credential_from_model_config(
model_name: str,
project_model_config: Optional[dict],

View file

@ -88,8 +88,8 @@ from litellm.router import Router
from litellm.secret_managers.main import get_secret
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateKeyRequest,
BulkUpdateKeyRequestItem,
BulkUpdateKeyResponse,
BulkUpdateTeamKeysRequest,
FailedKeyUpdate,
SuccessfulKeyUpdate,
)
@ -1881,7 +1881,7 @@ async def _get_and_validate_existing_key(
async def _process_single_key_update(
key_update_item: BulkUpdateKeyRequestItem,
update_key_request: UpdateKeyRequest,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str],
prisma_client: Optional[PrismaClient],
@ -1889,6 +1889,7 @@ async def _process_single_key_update(
proxy_logging_obj: Any,
llm_router: Optional[Router],
user_custom_key_update: Optional[Callable] = None,
existing_key_row: Optional[LiteLLM_VerificationToken] = None,
) -> Dict[str, Any]:
"""
Process a single key update with all validations and checks.
@ -1897,13 +1898,14 @@ async def _process_single_key_update(
including validation, permission checks, team checks, and database updates.
Args:
key_update_item: The key update request item
update_key_request: Fully-constructed UpdateKeyRequest for the target key
user_api_key_dict: The authenticated user's API key info
litellm_changed_by: Optional header for tracking who made the change
prisma_client: Prisma client instance
user_api_key_cache: User API key cache
proxy_logging_obj: Proxy logging object
llm_router: LLM router instance
existing_key_row: Optional pre-fetched key row to avoid redundant lookups
Returns:
Dict containing the updated key information
@ -1912,13 +1914,14 @@ async def _process_single_key_update(
HTTPException: For various validation and permission errors
"""
# Validate max_budget
_validate_max_budget(key_update_item.max_budget)
_validate_max_budget(update_key_request.max_budget)
# Get and validate existing key
existing_key_row = await _get_and_validate_existing_key(
token=key_update_item.key,
prisma_client=prisma_client,
)
if existing_key_row is None:
existing_key_row = await _get_and_validate_existing_key(
token=update_key_request.key,
prisma_client=prisma_client,
)
# Check team member permissions
if prisma_client is not None:
@ -1930,15 +1933,6 @@ async def _process_single_key_update(
user_api_key_cache=user_api_key_cache,
)
# Create UpdateKeyRequest from BulkUpdateKeyRequestItem
update_key_request = UpdateKeyRequest(
key=key_update_item.key,
budget_id=key_update_item.budget_id,
max_budget=key_update_item.max_budget,
team_id=key_update_item.team_id,
tags=key_update_item.tags,
)
# Custom key update hook
if user_custom_key_update is not None:
if inspect.iscoroutinefunction(user_custom_key_update):
@ -2003,12 +1997,12 @@ async def _process_single_key_update(
detail={"error": "Database not connected"},
)
_data = {**non_default_values, "token": key_update_item.key}
response = await prisma_client.update_data(token=key_update_item.key, data=_data)
_data = {**non_default_values, "token": update_key_request.key}
response = await prisma_client.update_data(token=update_key_request.key, data=_data)
# Delete cache
await _delete_cache_key_object(
hashed_token=_hash_token_if_needed(key_update_item.key),
hashed_token=_hash_token_if_needed(update_key_request.key),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
@ -2598,9 +2592,15 @@ async def bulk_update_keys(
for key_update_item in data.keys:
try:
# Process single key update using reusable function
update_key_request = UpdateKeyRequest(
key=key_update_item.key,
budget_id=key_update_item.budget_id,
max_budget=key_update_item.max_budget,
team_id=key_update_item.team_id,
tags=key_update_item.tags,
)
updated_key_info = await _process_single_key_update(
key_update_item=key_update_item,
update_key_request=update_key_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
prisma_client=prisma_client,
@ -2665,6 +2665,223 @@ async def bulk_update_keys(
)
def _build_failed_team_key_update(
token: str,
exception: Exception,
existing_key_row: Optional[LiteLLM_VerificationToken],
) -> FailedKeyUpdate:
"""Normalize an exception from the per-key update loop into a FailedKeyUpdate."""
if isinstance(exception, HTTPException):
detail = exception.detail
if isinstance(detail, dict):
error_message = detail.get("error", str(exception))
else:
error_message = str(detail)
elif isinstance(exception, ProxyException):
error_message = exception.message
else:
error_message = str(exception)
key_info: Optional[Dict[str, Any]] = None
if existing_key_row is not None:
if hasattr(existing_key_row, "model_dump"):
key_info = existing_key_row.model_dump()
elif hasattr(existing_key_row, "dict"):
key_info = existing_key_row.dict()
if key_info:
key_info.pop("token", None)
return FailedKeyUpdate(key=token, key_info=key_info, failed_reason=error_message)
@router.post(
"/team/key/bulk_update",
tags=["key management"],
dependencies=[Depends(user_api_key_auth)],
response_model=BulkUpdateKeyResponse,
)
@management_endpoint_wrapper
async def bulk_update_team_keys(
data: BulkUpdateTeamKeysRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
litellm_changed_by: Optional[str] = Header(
None,
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
),
):
"""
Apply one update payload to many keys inside a single team.
Pass `team_id` plus either `key_ids` or `all_keys_in_team=True`. The
`update_fields` payload is broadcast to every selected key. Per-key
failures are returned in `failed_updates` rather than aborting the batch.
Callable by proxy admins, or by team admins with `KEY_UPDATE` permission.
"""
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
user_custom_key_update,
)
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": "Database not connected"},
)
if not data.team_id:
raise HTTPException(
status_code=400,
detail={"error": "team_id is required"},
)
MAX_BATCH_SIZE = 500
if data.key_ids is not None and len(data.key_ids) > MAX_BATCH_SIZE:
raise HTTPException(
status_code=400,
detail={
"error": f"Maximum {MAX_BATCH_SIZE} keys can be updated at once. Found {len(data.key_ids)} key_ids."
},
)
if data.all_keys_in_team:
# "all" excludes blocked/expired — bulk refresh shouldn't revive a key an admin disabled.
# `blocked` is Boolean? with no default; `/key/generate` writes NULL. Prisma's `NOT`
# excludes NULLs, so explicitly OR `false` with `null` to include them.
now = datetime.now(timezone.utc)
existing_keys = await prisma_client.db.litellm_verificationtoken.find_many(
where={
"team_id": data.team_id,
"AND": [
{"OR": [{"blocked": False}, {"blocked": None}]},
{"OR": [{"expires": None}, {"expires": {"gt": now}}]},
],
},
order={"token": "asc"},
take=MAX_BATCH_SIZE + 1,
)
if len(existing_keys) > MAX_BATCH_SIZE:
raise HTTPException(
status_code=400,
detail={
"error": f"Team {data.team_id} has more than {MAX_BATCH_SIZE} keys. Use `key_ids` to update in batches of {MAX_BATCH_SIZE}."
},
)
requested_tokens = [row.token for row in existing_keys]
else:
if data.key_ids is None or len(data.key_ids) == 0:
raise HTTPException(
status_code=400,
detail={
"error": "key_ids must be provided when all_keys_in_team is False"
},
)
# Dedupe by hashed form — duplicates collapse to one update.
requested_tokens = []
hashed_key_ids = []
seen_hashes = set()
for k in data.key_ids:
h = _hash_token_if_needed(k)
if h in seen_hashes:
continue
seen_hashes.add(h)
requested_tokens.append(k)
hashed_key_ids.append(h)
existing_keys = await prisma_client.db.litellm_verificationtoken.find_many(
where={"team_id": data.team_id, "token": {"in": hashed_key_ids}}
)
# Anchor membership check on data.team_id (not existing_keys[0]); empty result must still gate non-admins.
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
auth_anchor = (
existing_keys[0]
if existing_keys
else LiteLLM_VerificationToken(
token="__team_scope_auth_check__",
team_id=data.team_id,
models=[],
)
)
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
user_api_key_dict=user_api_key_dict,
route=KeyManagementRoutes.KEY_UPDATE,
prisma_client=prisma_client,
existing_key_row=auth_anchor,
user_api_key_cache=user_api_key_cache,
)
# Block metadata.allowed_passthrough_routes for non-admins — the runtime
# route checker reads it from key/team metadata to grant passthrough.
_check_passthrough_routes_caller_permission(
data=data.update_fields, user_api_key_dict=user_api_key_dict
)
if not requested_tokens:
raise HTTPException(
status_code=404,
detail={"error": f"No keys found for team {data.team_id}"},
)
existing_by_token = {row.token: row for row in existing_keys}
update_field_dict = data.update_fields.model_dump(exclude_unset=True)
successful_updates: List[SuccessfulKeyUpdate] = []
failed_updates: List[FailedKeyUpdate] = []
for token in requested_tokens:
db_token = _hash_token_if_needed(token)
try:
if db_token not in existing_by_token:
raise HTTPException(
status_code=404,
detail={"error": f"Key not found in team {data.team_id}"},
)
# team_id from validated scope, never user payload — drives _check_team_key_limits.
update_key_request = UpdateKeyRequest(
key=token,
team_id=data.team_id,
**update_field_dict,
)
updated_key_info = await _process_single_key_update(
update_key_request=update_key_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
user_custom_key_update=user_custom_key_update,
existing_key_row=existing_by_token[db_token],
)
successful_updates.append(
SuccessfulKeyUpdate(key=token, key_info=updated_key_info)
)
except Exception as e:
# Log the hashed prefix — `token` may be a raw sk-... and ERROR logs persist.
verbose_proxy_logger.exception(
f"Failed to update key {db_token[:12]}... in team {data.team_id}: {e}"
)
failed_updates.append(
_build_failed_team_key_update(
token=token,
exception=e,
existing_key_row=existing_by_token.get(db_token),
)
)
return BulkUpdateKeyResponse(
total_requested=len(requested_tokens),
successful_updates=successful_updates,
failed_updates=failed_updates,
)
async def validate_key_team_change(
key: LiteLLM_VerificationToken,
team: LiteLLM_TeamTable,

View file

@ -142,6 +142,7 @@ if MCP_AVAILABLE:
MCPOAuthUserCredentialRequest,
MCPOAuthUserCredentialStatus,
MCPSubmissionsSummary,
MCPTransport,
MCPUserCredentialListItem,
MCPUserCredentialRequest,
MCPUserCredentialResponse,
@ -1070,6 +1071,24 @@ if MCP_AVAILABLE:
},
)
# stdio servers spawn a local subprocess on the proxy host with the
# configured command + args, so accepting them from non-admin callers
# would let a team member propose a server config that an admin could
# rubber-stamp into local code execution. Restrict stdio submission to
# the admin POST /v1/mcp/server path or to config.yaml.
if payload.transport == MCPTransport.stdio:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": (
"stdio MCP servers cannot be submitted via the user "
"registration workflow. Ask a proxy admin to add this "
"server via POST /v1/mcp/server or to declare it in "
"config.yaml."
)
},
)
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to your proxy"
)

View file

@ -12,12 +12,13 @@ All /tag management endpoints
import asyncio
import json
from typing import TYPE_CHECKING, Dict, List, Optional
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from fastapi import APIRouter, Depends, HTTPException
from fastapi import APIRouter, Depends, HTTPException, Query
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
@ -39,6 +40,72 @@ if TYPE_CHECKING:
router = APIRouter()
async def _get_internal_user_api_keys(
prisma_client,
user_api_key_dict: UserAPIKeyAuth,
) -> List[str]:
user_role = user_api_key_dict.user_role
if user_role is None or not user_role.is_internal_user_role:
return []
user_api_keys = set()
if user_api_key_dict.api_key:
user_api_keys.add(user_api_key_dict.api_key)
user_id = user_api_key_dict.user_id
if user_id is None:
return sorted(user_api_keys)
key_records = await prisma_client.db.litellm_verificationtoken.find_many(
where={"user_id": user_id},
select={"token": True},
)
user_api_keys.update(
key_record.token
for key_record in key_records
if getattr(key_record, "token", None)
)
return sorted(user_api_keys)
async def _get_tag_list_scope(
prisma_client,
user_api_key_dict: UserAPIKeyAuth,
) -> Optional[Dict[str, dict]]:
user_role = user_api_key_dict.user_role
if user_api_key_has_admin_view(user_api_key_dict) or (
user_role is None or not user_role.is_internal_user_role
):
return None
scoped_api_keys = await _get_internal_user_api_keys(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
)
return {"api_key": {"in": scoped_api_keys}}
async def _get_tag_daily_activity_api_key_filter(
prisma_client,
user_api_key_dict: UserAPIKeyAuth,
requested_api_key: Optional[str],
) -> Optional[Union[str, List[str]]]:
user_role = user_api_key_dict.user_role
if user_api_key_has_admin_view(user_api_key_dict) or (
user_role is None or not user_role.is_internal_user_role
):
return requested_api_key
scoped_api_keys = await _get_internal_user_api_keys(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
)
if requested_api_key is not None:
return requested_api_key if requested_api_key in scoped_api_keys else []
return scoped_api_keys
async def _get_model_names(prisma_client, model_ids: list) -> Dict[str, str]:
"""Helper function to get model names from model IDs"""
try:
@ -395,6 +462,32 @@ async def info_tag(
raise HTTPException(status_code=500, detail=str(e))
def _validate_tag_list_date_range(
start_date: Optional[str], end_date: Optional[str]
) -> None:
"""Require both dates together, and enforce YYYY-MM-DD format with start <= end."""
if (start_date is None) != (end_date is None):
raise HTTPException(
status_code=400,
detail="start_date and end_date must be provided together",
)
if start_date is None:
return
try:
start = datetime.strptime(start_date, "%Y-%m-%d")
end = datetime.strptime(end_date, "%Y-%m-%d") # type: ignore[arg-type]
except ValueError as e:
raise HTTPException(
status_code=400,
detail=f"Invalid date format, expected YYYY-MM-DD: {e}",
)
if start > end:
raise HTTPException(
status_code=400,
detail="start_date must be on or before end_date",
)
@router.get(
"/tag/list",
tags=["tag management"],
@ -402,6 +495,18 @@ async def info_tag(
)
async def list_tags(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
start_date: Optional[str] = Query(
None,
description=(
"Optional start date (YYYY-MM-DD). When provided together with "
"end_date, dynamic tags are limited to those active in the window. "
"Stored tags are always returned."
),
),
end_date: Optional[str] = Query(
None,
description="Optional end date (YYYY-MM-DD). Must be given with start_date.",
),
):
"""
List all available tags with their budget information.
@ -411,10 +516,44 @@ async def list_tags(
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
_validate_tag_list_date_range(start_date, end_date)
try:
tag_scope = await _get_tag_list_scope(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
)
## QUERY DYNAMIC TAGS ##
# Use group_by instead of find_many(distinct=["tag"]).
# Prisma's distinct fetches all columns for all rows and deduplicates
# in application code, which is extremely slow on large tables.
# See: https://www.prisma.io/docs/orm/prisma-client/queries/aggregation-grouping-summarizing#distinct-under-the-hood
dynamic_tag_where: Dict[str, Any] = {"tag": {"not": None}}
if tag_scope:
dynamic_tag_where = {**dynamic_tag_where, **tag_scope}
if start_date is not None and end_date is not None:
dynamic_tag_where["date"] = {"gte": start_date, "lte": end_date}
dynamic_tag_rows = await prisma_client.db.litellm_dailytagspend.group_by(
by=["tag"],
where=dynamic_tag_where,
min={"created_at": True},
max={"updated_at": True},
)
used_tag_names = [row["tag"] for row in dynamic_tag_rows if row["tag"]]
if tag_scope is not None and not used_tag_names:
return []
stored_tag_where = (
{"tag_name": {"in": used_tag_names}} if tag_scope is not None else None
)
## QUERY STORED TAGS ##
tag_records = await prisma_client.db.litellm_tagtable.find_many(
include={"litellm_budget_table": True}
where=stored_tag_where,
include={"litellm_budget_table": True},
)
stored_tag_names = set()
@ -448,18 +587,6 @@ async def list_tags(
list_of_tags.append(tag_dict)
## QUERY DYNAMIC TAGS ##
# Use group_by instead of find_many(distinct=["tag"]).
# Prisma's distinct fetches all columns for all rows and deduplicates
# in application code, which is extremely slow on large tables.
# See: https://www.prisma.io/docs/orm/prisma-client/queries/aggregation-grouping-summarizing#distinct-under-the-hood
dynamic_tag_rows = await prisma_client.db.litellm_dailytagspend.group_by(
by=["tag"],
where={"tag": {"not": None}},
min={"created_at": True},
max={"updated_at": True},
)
dynamic_tag_config = [
{
"name": row["tag"],
@ -527,6 +654,7 @@ async def get_tag_daily_activity(
api_key: Optional[str] = None,
page: int = 1,
page_size: int = 10,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get daily activity for specific tags or all tags.
@ -545,8 +673,18 @@ async def get_tag_daily_activity(
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
# Convert comma-separated tags string to list if provided
tag_list = tags.split(",") if tags else None
scoped_api_key_filter = await _get_tag_daily_activity_api_key_filter(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
requested_api_key=api_key,
)
if scoped_api_key_filter == []:
return SpendAnalyticsPaginatedResponse(results=[])
return await get_daily_activity(
prisma_client=prisma_client,
@ -557,7 +695,7 @@ async def get_tag_daily_activity(
start_date=start_date,
end_date=end_date,
model=model,
api_key=api_key,
api_key=scoped_api_key_filter,
page=page,
page_size=page_size,
# metadata_metrics_func=None because litellm_dailytagspend rows are

View file

@ -3805,6 +3805,7 @@ async def _build_team_list_where_conditions(
organization_id: Optional[str],
user_id: Optional[str],
use_deleted_table: bool,
search: Optional[str] = None,
org_admin_org_ids: Optional[List[str]] = None,
user_api_key_cache: Optional[Any] = None,
proxy_logging_obj: Optional[Any] = None,
@ -3826,6 +3827,12 @@ async def _build_team_list_where_conditions(
"mode": "insensitive", # Case-insensitive search
}
if search:
where_conditions["OR"] = [
{"team_id": search},
{"team_alias": {"contains": search, "mode": "insensitive"}},
]
if organization_id:
where_conditions["organization_id"] = organization_id
elif org_admin_org_ids is not None:
@ -4019,6 +4026,10 @@ async def list_team_v2(
default=None,
description="Only return teams which this 'team_alias' belongs to. Supports partial matching.",
),
search: Optional[str] = fastapi.Query(
default=None,
description="Combined search: matches teams whose 'team_id' equals the value OR whose 'team_alias' contains it (case-insensitive).",
),
page: int = fastapi.Query(
default=1, description="Page number for pagination", ge=1
),
@ -4104,6 +4115,7 @@ async def list_team_v2(
organization_id=organization_id,
user_id=user_id,
use_deleted_table=use_deleted_table,
search=search,
org_admin_org_ids=org_admin_org_ids,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,

View file

@ -740,7 +740,7 @@ def generic_response_convertor(
all_teams = []
if sso_jwt_handler is not None:
team_ids = sso_jwt_handler.get_team_ids_from_jwt(cast(dict, response))
team_ids = sso_jwt_handler.get_all_jwt_team_ids(cast(dict, response))
all_teams.extend(team_ids)
if team_mappings is not None and team_mappings.team_ids_jwt_field is not None:
@ -755,7 +755,7 @@ def generic_response_convertor(
f"Loaded team_ids from DB team_mappings.team_ids_jwt_field='{team_mappings.team_ids_jwt_field}': {team_ids_from_db_mapping}"
)
else:
team_ids = jwt_handler.get_team_ids_from_jwt(cast(dict, response))
team_ids = jwt_handler.get_all_jwt_team_ids(cast(dict, response))
all_teams.extend(team_ids)
# Determine user role based on role_mappings if available
@ -4078,6 +4078,8 @@ async def debug_sso_callback(request: Request):
redirect_url += "/sso/debug/callback"
result = None
received_response: Optional[dict] = None
access_token_payload: Optional[dict] = None
if google_client_id is not None:
result = await GoogleSSOHandler.get_google_callback_response(
request=request,
@ -4094,12 +4096,14 @@ async def debug_sso_callback(request: Request):
)
elif generic_client_id is not None:
result, _, _ = await get_generic_sso_response(
request=request,
jwt_handler=jwt_handler,
generic_client_id=generic_client_id,
redirect_url=redirect_url,
sso_jwt_handler=sso_jwt_handler,
result, received_response, access_token_payload = (
await get_generic_sso_response(
request=request,
jwt_handler=jwt_handler,
generic_client_id=generic_client_id,
redirect_url=redirect_url,
sso_jwt_handler=sso_jwt_handler,
)
)
# If result is None, return a basic error message
@ -4128,10 +4132,32 @@ async def debug_sso_callback(request: Request):
except Exception as e:
filtered_result[key] = f"Complex value (not displayable): {str(e)}"
# Defense-in-depth: ensure no bearer tokens leak into the rendered HTML even if
# a non-conforming IdP places them in its userinfo response.
safe_raw_claims = {
k: v
for k, v in (received_response or {}).items()
if k not in _OAUTH_TOKEN_FIELDS
}
safe_access_token_claims = {
k: v
for k, v in (access_token_payload or {}).items()
if k not in _OAUTH_TOKEN_FIELDS
}
sso_payload = {
"parsed_by_proxy": filtered_result,
"raw_claims": safe_raw_claims,
"access_token_claims": safe_access_token_claims,
}
# Replace the placeholder in the template with the actual data
sso_payload_json = json.dumps(sso_payload, indent=2, default=str).replace(
"</", "<\\/"
)
html_content = jwt_display_template.replace(
"const userData = SSO_DATA;",
f"const userData = {json.dumps(filtered_result, indent=2)};",
"const ssoData = SSO_DATA;",
f"const ssoData = {sso_payload_json};",
)
return HTMLResponse(content=html_content)

View file

@ -79,7 +79,10 @@ class PrometheusAuthMiddleware:
# Send 401 response directly via ASGI protocol
error_message = getattr(e, "message", str(e))
body = json.dumps(
f"Unauthorized access to metrics endpoint: {error_message}"
f"Unauthorized access to metrics endpoint: {error_message} "
f"To allow unauthenticated access, set "
f"`litellm_settings.require_auth_for_metrics_endpoint: false` "
f"in your proxy_config.yaml."
).encode("utf-8")
await send(
{

View file

@ -805,7 +805,8 @@ def run_server( # noqa: PLR0915
)
db_connection_pool_limit = 100
db_connection_timeout = 60
# Starts optional due to config fallback checks; guaranteed non-None before use.
db_connection_timeout: Optional[Union[int, float]] = 60
general_settings = {}
### GET DB TOKEN FOR IAM AUTH ###
@ -813,7 +814,12 @@ def run_server( # noqa: PLR0915
from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token
db_host = os.getenv("DATABASE_HOST")
db_port = os.getenv("DATABASE_PORT")
# Default to the Postgres standard port. Without a default,
# `db_port=None` flows into `boto.generate_db_auth_token(Port=None)`
# and botocore stringifies it to `"None"` while building the
# presigned URL, which then blows up with `ValueError: Port could
# not be cast to integer value as 'None'` during signing.
db_port = os.getenv("DATABASE_PORT", "5432")
db_user = os.getenv("DATABASE_USER")
db_name = os.getenv("DATABASE_NAME")
db_schema = os.getenv("DATABASE_SCHEMA")
@ -909,10 +915,15 @@ def run_server( # noqa: PLR0915
"database_connection_pool_limit",
LiteLLMDatabaseConnectionPool.database_connection_pool_limit.value,
)
db_connection_timeout = general_settings.get(
"database_connection_pool_timeout",
LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value,
)
db_connection_timeout = general_settings.get("database_connection_timeout")
if db_connection_timeout is None:
db_connection_timeout = general_settings.get(
"database_connection_pool_timeout"
)
if db_connection_timeout is None:
db_connection_timeout = (
LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value
)
if database_url and database_url.startswith("os.environ/"):
original_dir = os.getcwd()
# set the working directory to where this script is

View file

@ -211,6 +211,7 @@ from litellm import Router
from litellm._logging import verbose_proxy_logger, verbose_router_logger
from litellm.caching.caching import DualCache, RedisCache
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.constants import (
_REALTIME_BODY_CACHE_SIZE,
@ -1060,6 +1061,52 @@ vertex_live_passthrough_vertex_base = VertexBase()
from fastapi.routing import APIWebSocketRoute
def _inject_websocket_stubs_into_openapi_schema(
openapi_schema: dict, websocket_routes: list
) -> dict:
"""
Add a synthetic GET stub for each WebSocket route so it appears in Swagger UI.
Merges into any existing path entry rather than replacing it — a WebSocket route
that shares its path with an HTTP route must not erase the HTTP operation. If
a "get" operation is already documented on the path, the WebSocket stub is
skipped to preserve the real GET.
"""
for route in websocket_routes:
base_path = route.path.split("{")[0].rstrip("?")
parameters = []
try:
if hasattr(route, "dependant") and route.dependant is not None:
# Handle both FastAPI <0.120 and >=0.120
query_params = getattr(route.dependant, "query_params", [])
if query_params:
for param in query_params:
parameters.append(
{
"name": param.name,
"in": "query",
"required": param.required,
"schema": {"type": "string"},
}
)
except (AttributeError, TypeError):
pass
path_entry = openapi_schema["paths"].setdefault(base_path, {})
if "get" not in path_entry:
path_entry["get"] = {
"summary": f"WebSocket: {route.name or base_path}",
"description": "WebSocket connection endpoint",
"operationId": f"websocket_{route.name or base_path.replace('/', '_')}",
"parameters": parameters,
"responses": {"101": {"description": "WebSocket Protocol Switched"}},
"tags": ["WebSocket"],
}
return openapi_schema
def get_openapi_schema():
if app.openapi_schema:
return app.openapi_schema
@ -1082,43 +1129,11 @@ def get_openapi_schema():
route for route in app.routes if isinstance(route, APIWebSocketRoute)
]
# Add each WebSocket route to the schema
for route in websocket_routes:
# Get the base path without query parameters
base_path = route.path.split("{")[0].rstrip("?")
# Extract parameters from the route
parameters = []
try:
if hasattr(route, "dependant") and route.dependant is not None:
# Handle both FastAPI <0.120 and >=0.120
query_params = getattr(route.dependant, "query_params", [])
if query_params:
for param in query_params:
parameters.append(
{
"name": param.name,
"in": "query",
"required": param.required,
"schema": {
"type": "string"
}, # You can make this more specific if needed
}
)
except (AttributeError, TypeError):
# If we can't access query_params, continue without them
pass
openapi_schema["paths"][base_path] = {
"get": {
"summary": f"WebSocket: {route.name or base_path}",
"description": "WebSocket connection endpoint",
"operationId": f"websocket_{route.name or base_path.replace('/', '_')}",
"parameters": parameters,
"responses": {"101": {"description": "WebSocket Protocol Switched"}},
"tags": ["WebSocket"],
}
}
# Add a synthetic GET stub for each so they render in Swagger UI,
# without clobbering existing HTTP operations on the same path.
openapi_schema = _inject_websocket_stubs_into_openapi_schema(
openapi_schema, websocket_routes
)
# Add LLM API request schema bodies for documentation
from litellm.proxy.common_utils.custom_openapi_spec import CustomOpenAPISpec
@ -5936,10 +5951,20 @@ class ProxyConfig:
verbose_proxy_logger.debug(
"guardrails from the DB %s", str(guardrails_in_db)
)
db_guardrail_ids: set = set()
for guardrail in guardrails_in_db:
guardrail_id = guardrail.get("guardrail_id")
if guardrail_id:
db_guardrail_ids.add(guardrail_id)
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
guardrail=cast(Guardrail, guardrail),
)
# Drop in-memory DB-backed entries whose row was deleted on another
# pod. Config-loaded entries are never touched.
IN_MEMORY_GUARDRAIL_HANDLER.reconcile_db_guardrails(
db_guardrail_ids=db_guardrail_ids
)
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - {}".format(
@ -6750,27 +6775,64 @@ class ProxyStartupEvent:
"budget_duration not set on Proxy. budget_duration is required to use max_budget."
)
# add proxy budget to db in the user table
asyncio.create_task(
generate_key_helper_fn( # type: ignore
request_type="user",
table_name="user",
user_id=litellm_proxy_budget_name,
duration=None,
models=[],
aliases={},
config={},
spend=0,
max_budget=litellm.max_budget,
budget_duration=litellm.budget_duration,
query_type="update_data",
update_key_values={
"max_budget": litellm.max_budget,
"budget_duration": litellm.budget_duration,
},
)
cls._upsert_proxy_budget_with_reset_at_backfill(litellm_proxy_budget_name)
)
@classmethod
async def _upsert_proxy_budget_with_reset_at_backfill(
cls, litellm_proxy_budget_name: str
) -> None:
"""
Upsert the proxy admin user row with the configured max_budget /
budget_duration, then backfill budget_reset_at if currently NULL.
The backfill uses `WHERE budget_reset_at IS NULL` so it only fires
when the row pre-existed without a reset schedule (e.g. row created
via a different path before the proxy budget was configured). On
subsequent restarts it no-ops, so an active reset window is never
slid forward.
"""
await generate_key_helper_fn( # type: ignore
request_type="user",
table_name="user",
user_id=litellm_proxy_budget_name,
duration=None,
models=[],
aliases={},
config={},
spend=0,
max_budget=litellm.max_budget,
budget_duration=litellm.budget_duration,
query_type="update_data",
update_key_values={
"max_budget": litellm.max_budget,
"budget_duration": litellm.budget_duration,
},
)
# Without this, the upsert leaves budget_reset_at=NULL on rows that
# took the UPDATE path, and reset_budget_for_litellm_users never
# matches them (NULL < now() is unknown in SQL) — so the proxy-wide
# spend cap blocks forever once it's hit.
if prisma_client is not None and litellm.budget_duration is not None:
try:
await prisma_client.db.litellm_usertable.update_many(
where={
"user_id": litellm_proxy_budget_name,
"budget_reset_at": None,
},
data={
"budget_reset_at": get_budget_reset_time(
budget_duration=litellm.budget_duration
)
},
)
except Exception as e:
verbose_proxy_logger.warning(
"Failed to backfill budget_reset_at on proxy admin row: %s", e
)
@classmethod
async def _warm_global_spend_cache(
cls,
@ -8828,6 +8890,7 @@ def _realtime_query_params_template(
return tuple(params)
@app.websocket("/openai/v1/realtime")
@app.websocket("/v1/realtime")
@app.websocket("/realtime")
async def realtime_websocket_endpoint(

View file

@ -95,11 +95,9 @@ async def reserve_budget_for_request(
route=route,
llm_router=llm_router,
)
if reservation_cost is None:
reservation_cost = await _get_smallest_remaining_budget(
counters=counters,
current_spend_by_counter_key=current_spend_by_counter_key,
)
# estimate_request_max_cost still returns None when the model is unknown
# to the cost map (no token-priced cost fields, e.g. image/audio routes).
# In that case we fall back to read-time enforcement only.
if reservation_cost is None or reservation_cost <= 0:
return None
@ -553,32 +551,6 @@ def _coerce_window(window: Any) -> dict:
return {}
async def _get_smallest_remaining_budget(
counters: List[_BudgetCounter],
current_spend_by_counter_key: Dict[str, float],
) -> Optional[float]:
remaining_budget: Optional[float] = None
for counter in counters:
current_spend = await _get_current_counter_value(counter=counter)
current_spend_by_counter_key[counter.counter_key] = current_spend
remaining = counter.max_budget - current_spend
if remaining <= 0:
raise litellm.BudgetExceededError(
current_cost=current_spend,
max_budget=counter.max_budget,
message=(
"Budget has been exceeded! "
f"{counter.entity_type}={counter.entity_id} "
f"Current cost: {current_spend}, "
f"Max budget: {counter.max_budget}"
),
)
remaining_budget = (
remaining if remaining_budget is None else min(remaining_budget, remaining)
)
return remaining_budget
async def _reserve_counter(
counter: _BudgetCounter,
reservation_cost: float,
@ -855,6 +827,13 @@ def _estimate_request_max_cost_for_model(
if model_info is None:
return None
image_cost = _estimate_image_generation_cost(
request_body=request_body,
model_info=model_info,
)
if image_cost is not None:
return image_cost
input_cost_per_token = _to_float(model_info.get("input_cost_per_token"))
output_cost_per_token = _to_float(model_info.get("output_cost_per_token"))
input_tokens = _estimate_input_tokens(
@ -886,6 +865,44 @@ def _estimate_request_max_cost_for_model(
return cost
def _estimate_image_generation_cost(
request_body: dict,
model_info: Dict[str, Any],
) -> Optional[float]:
"""
Reserve `n × per-image cost` for image-generation requests so concurrent
requests against a depleted budget cannot all slip past the admission gate
onto the provider. Token-based pricing (e.g. gpt-image-1) is handled by
the chat-route token path; per-pixel and size/quality-tiered pricing
(DALL-E 2 size variants, premium tiers) are not handled here and fall
through to read-time enforcement.
The "output" vs "input" cost-per-image naming is inconsistent across
providers — OpenAI's dall-e-3 entry uses ``input_cost_per_image`` while
aiml/dall-e-3 uses ``output_cost_per_image`` — so both are summed.
"""
# Gate strictly on `mode`. Several chat and embedding models carry
# ``input_cost_per_image`` / ``output_cost_per_image`` to price multimodal
# *vision input* (e.g. ``gemini-3.1-pro-preview``, ``azure/gpt-realtime-*``,
# ``amazon.titan-embed-image-v1``). Falling back to "treat as image-gen if
# an image cost field is present" would short-circuit the token-priced
# path for those models and reserve a fraction of a cent instead of the
# true per-token cost. All real image-generation entries in
# ``model_prices_and_context_window.json`` carry ``mode: image_generation``
# or ``mode: image_edit``, so the field-presence fallback is unnecessary.
if model_info.get("mode") not in ("image_generation", "image_edit"):
return None
output_cost_per_image = _to_float(model_info.get("output_cost_per_image"))
input_cost_per_image = _to_float(model_info.get("input_cost_per_image"))
cost_per_image = (output_cost_per_image or 0.0) + (input_cost_per_image or 0.0)
if cost_per_image <= 0:
return None
n = _to_int(request_body.get("n")) or 1
return cost_per_image * max(n, 1)
def _get_model_cost_info(
model: str,
llm_router: Optional[Router],
@ -946,6 +963,9 @@ def _estimate_input_tokens(
return None
DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK = 16384
def _estimate_output_tokens(
request_body: dict,
route: str,
@ -954,15 +974,27 @@ def _estimate_output_tokens(
if _is_input_only_route(route=route):
return 0
requested: Optional[int] = None
for key in ("max_completion_tokens", "max_tokens", "max_output_tokens"):
max_tokens = _to_int(request_body.get(key))
if max_tokens is not None:
return max_tokens
requested = _to_int(request_body.get(key))
if requested is not None:
break
# If the caller did not cap output tokens, avoid reserving a model's
# theoretical maximum context. The caller can still admit one request by
# reserving the smallest remaining budget in reserve_budget_for_request().
return None
# Clamp at min(requested-or-default, model_max-or-default). Two purposes:
# (1) Without an explicit cap we still need a finite reservation so the
# atomic admission counter actually bounds concurrent in-flight cost
# (mirrors parallel_request_limiter_v3's DEFAULT_MAX_TOKENS_ESTIMATE).
# (2) An adversarial caller cannot send max_tokens=999999999 to inflate
# the reservation up to remaining team headroom and pin the counter
# at the cap — the model can only physically emit max_output_tokens
# anyway, so reserving more is both wasteful and a DoS surface.
model_ceiling = (
_to_int(model_info.get("max_output_tokens"))
or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK
)
if requested is None:
requested = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK
return min(requested, model_ceiling)
def _count_text_tokens(model: str, text: Any) -> int:

View file

@ -113,7 +113,11 @@ from litellm.proxy.db.exception_handler import (
call_with_db_reconnect_retry,
)
from litellm.proxy.db.log_db_metrics import log_db_metrics
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.db.prisma_client import (
PrismaWrapper,
parse_iam_endpoint_from_url,
)
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -2569,24 +2573,101 @@ class PrismaClient:
raise Exception(
"Unable to find Prisma binaries. Please run 'prisma generate' first."
)
iam_flag = (
self.iam_token_db_auth if self.iam_token_db_auth is not None else False
)
# When read-replica routing is on, tag log lines with [writer]/[reader]
# so the two wrappers' interleaved IAM refresh logs can be told apart.
# Single-DB deployments get an empty prefix (logs unchanged).
read_replica_url = os.getenv("DATABASE_URL_READ_REPLICA")
writer_log_prefix = "[writer]" if read_replica_url else ""
if http_client is not None:
self.db = PrismaWrapper(
writer_wrapper = PrismaWrapper(
original_prisma=Prisma(http=http_client),
iam_token_db_auth=(
self.iam_token_db_auth
if self.iam_token_db_auth is not None
else False
),
iam_token_db_auth=iam_flag,
log_prefix=writer_log_prefix,
)
else:
self.db = PrismaWrapper(
writer_wrapper = PrismaWrapper(
original_prisma=Prisma(),
iam_token_db_auth=(
self.iam_token_db_auth
if self.iam_token_db_auth is not None
else False
),
) # Client to connect to Prisma db
iam_token_db_auth=iam_flag,
log_prefix=writer_log_prefix,
)
# Optional read-replica routing. When DATABASE_URL_READ_REPLICA is set,
# reads (find_*, count, group_by, query_raw/_first) are routed to the
# reader endpoint and writes stay on the writer. Falls back to the
# writer-only wrapper when the env var is unset, preserving existing
# single-DB deployments.
self.db: Union[PrismaWrapper, RoutingPrismaWrapper]
if read_replica_url:
try:
# If IAM auth is enabled, the reader refreshes its own token on
# the same cadence as the writer. We parse the static endpoint
# pieces (host/port/user/db) once from the reader URL — only
# the IAM token rotates after that.
reader_iam_endpoint = (
parse_iam_endpoint_from_url(read_replica_url) if iam_flag else None
)
# Mint a fresh IAM token for the reader BEFORE constructing the
# Prisma client. Mirrors what `proxy_cli.py` already does for
# the writer (proxy_cli.py:812-832) — without this, the reader
# Prisma is built with whatever placeholder URL the user
# supplied (no real token), and the first query falls through
# to the synchronous fallback path in
# `PrismaWrapper.__getattr__`, which deadlocks the event loop
# and times out after 30s.
if iam_flag and reader_iam_endpoint is not None:
from litellm.proxy.auth.rds_iam_token import (
generate_iam_auth_token,
)
reader_token = generate_iam_auth_token(
db_host=reader_iam_endpoint.host,
db_port=reader_iam_endpoint.port,
db_user=reader_iam_endpoint.user,
)
read_replica_url = reader_iam_endpoint.build_url(reader_token)
os.environ["DATABASE_URL_READ_REPLICA"] = read_replica_url
reader_kwargs: Dict[str, Any] = {
"datasource": {"url": read_replica_url}
}
if http_client is not None:
reader_prisma = Prisma(http=http_client, **reader_kwargs)
else:
reader_prisma = Prisma(**reader_kwargs)
reader_wrapper = PrismaWrapper(
original_prisma=reader_prisma,
iam_token_db_auth=iam_flag,
db_url_env_var="DATABASE_URL_READ_REPLICA",
iam_endpoint=reader_iam_endpoint,
recreate_uses_datasource=True,
log_prefix="[reader]",
)
self.db = RoutingPrismaWrapper(
writer=writer_wrapper, reader=reader_wrapper
)
verbose_proxy_logger.info(
"PrismaClient: read-replica routing enabled via DATABASE_URL_READ_REPLICA"
+ (" (with IAM token auto-refresh)" if iam_flag else "")
)
except Exception as e:
# Reader is opt-in; never let its construction fail proxy
# startup. Mirrors the runtime contract from
# `RoutingPrismaWrapper.connect`: reader-side failures are
# logged and we keep serving traffic via the writer alone.
# This recovers from transient AWS STS hiccups during the
# reader IAM token mint, malformed DATABASE_URL_READ_REPLICA,
# and Prisma construction errors. Operator restart is required
# to retry read-routing once the underlying issue is resolved.
verbose_proxy_logger.warning(
"Failed to initialize read replica Prisma client: %s. "
"Falling back to writer-only mode (no read routing) until proxy restart.",
e,
)
self.db = writer_wrapper
else:
self.db = writer_wrapper # Client to connect to Prisma db
self._db_reconnect_lock = asyncio.Lock()
self._db_health_watchdog_task: Optional[asyncio.Task] = None
self._db_last_reconnect_attempt_ts: float = 0.0
@ -2624,6 +2705,13 @@ class PrismaClient:
self._engine_wait_thread: Optional[threading.Thread] = None
verbose_proxy_logger.debug("Success - Created Prisma Client")
@property
def writer_db(self) -> PrismaWrapper:
"""Underlying writer Prisma wrapper, regardless of read-replica routing."""
if isinstance(self.db, RoutingPrismaWrapper):
return self.db.writer
return self.db
def get_request_status(
self, payload: Union[dict, SpendLogsPayload]
) -> Literal["success", "failure"]:
@ -4272,7 +4360,10 @@ class PrismaClient:
self._cleanup_engine_watcher()
await self.db.recreate_prisma_client(db_url)
await self._start_engine_watcher()
await self.db.query_raw("SELECT 1")
# Smoke-test the writer specifically; query_raw on the routing
# wrapper sends to the reader, which would not validate the
# newly-recreated writer engine.
await self.writer_db.query_raw("SELECT 1")
await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout)
@ -4886,10 +4977,10 @@ class ProxyUpdateSpend:
timeout=timedelta(seconds=60)
) as transaction:
async with transaction.batch_() as batcher:
for (
end_user_id,
response_cost,
) in end_user_list_transactions.items():
# Sort by end_user_id for consistent lock ordering across pods to prevent deadlocks.
for end_user_id, response_cost in sorted(
end_user_list_transactions.items()
):
if litellm.max_end_user_budget is not None:
pass
batcher.litellm_endusertable.upsert(

View file

@ -1673,7 +1673,7 @@ class Router:
for cb in self.optional_callbacks
)
if not already_registered:
ec_callback = EncryptedContentAffinityCheck()
ec_callback = EncryptedContentAffinityCheck(router=self)
self.optional_callbacks.append(ec_callback)
litellm.logging_callback_manager.add_litellm_callback(ec_callback)
@ -7076,11 +7076,11 @@ class Router:
_shared_model_info = {
k: v for k, v in _model_info.items() if k not in _custom_pricing_fields
}
litellm.register_model(
model_cost={
_model_name: _shared_model_info,
}
)
_backend_alias_cost = {_model_name: _shared_model_info}
if "responses/" in _model_name:
_stripped_model_name = _model_name.replace("responses/", "")
_backend_alias_cost[_stripped_model_name] = _shared_model_info
litellm.register_model(model_cost=_backend_alias_cost)
## Check if LLM Deployment is allowed for this deployment
if (
@ -7752,6 +7752,12 @@ class Router:
# initialize client
self._add_deployment(deployment=deployment)
_model_info_dict: dict = deployment.model_info.model_dump(exclude_none=True)
for field in CustomPricingLiteLLMParams.model_fields.keys():
field_value = deployment.litellm_params.get(field)
if field_value is not None:
_model_info_dict[field] = field_value
# Register custom pricing in litellm.model_cost.
# Mirrors _create_deployment() logic to ensure dynamically-added deployments
# (e.g., loaded from DB) also have their custom pricing registered.
@ -7759,13 +7765,31 @@ class Router:
# zero-cost models, causing budget checks to block free models.
_model_id = deployment.model_info.id
if _model_id is not None:
_model_info_dict: dict = deployment.model_info.model_dump(exclude_none=True)
for field in CustomPricingLiteLLMParams.model_fields.keys():
field_value = deployment.litellm_params.get(field)
if field_value is not None:
_model_info_dict[field] = field_value
litellm.register_model(model_cost={_model_id: _model_info_dict})
## REGISTER MODEL INFO IN LITELLM MODEL COST MAP
## OLD MODEL REGISTRATION ## Kept to prevent breaking changes
_model_name = deployment.litellm_params.model
if deployment.litellm_params.custom_llm_provider is not None:
_model_name = (
deployment.litellm_params.custom_llm_provider + "/" + _model_name
)
# For the shared backend key, strip custom pricing fields so that
# one deployment's pricing overrides don't pollute another
# deployment sharing the same backend model name.
# Each deployment's full pricing is already stored under its
# unique model_id above (when present).
_custom_pricing_fields = CustomPricingLiteLLMParams.model_fields.keys()
_shared_model_info = {
k: v for k, v in _model_info_dict.items() if k not in _custom_pricing_fields
}
_backend_alias_cost = {_model_name: _shared_model_info}
if "responses/" in _model_name:
_stripped_model_name = _model_name.replace("responses/", "")
_backend_alias_cost[_stripped_model_name] = _shared_model_info
litellm.register_model(model_cost=_backend_alias_cost)
# add to model names
self._add_model_to_list_and_index_map(
model=_deployment, model_id=deployment.model_info.id

View file

@ -36,13 +36,16 @@ Safe to enable globally:
- No cache required.
"""
from typing import Any, List, Optional, cast
from typing import TYPE_CHECKING, Any, List, Optional, cast
from litellm._logging import verbose_router_logger
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import AllMessageValues
if TYPE_CHECKING:
from litellm.router import Router
class EncryptedContentAffinityCheck(CustomLogger):
"""
@ -55,8 +58,9 @@ class EncryptedContentAffinityCheck(CustomLogger):
Wired via ``Router(optional_pre_call_checks=["encrypted_content_affinity"])``.
"""
def __init__(self) -> None:
def __init__(self, router: Optional["Router"] = None) -> None:
super().__init__()
self.router = router
# ------------------------------------------------------------------
# Helpers
@ -119,6 +123,58 @@ class EncryptedContentAffinityCheck(CustomLogger):
return deployment
return None
@staticmethod
def _encryption_boundary_key(
litellm_params: Any,
) -> Optional[tuple]:
"""
``(api_base, api_key)`` pair identifying an Azure resource. Two
deployments sharing both are interchangeable for ``encrypted_content``
follow-ups; Azure rejects content produced by any other resource.
Accepts any object exposing dict-style ``.get(key, default)``: plain
dicts (the common case in ``healthy_deployments``) as well as
``LiteLLM_Params``-style Pydantic instances, which define a custom
``.get()``. A stricter ``isinstance(dict)`` guard would silently drop
the latter from boundary matching and fall back to the full pool —
i.e. trigger the exact ``invalid_encrypted_content`` failure this
check exists to prevent.
"""
getter = getattr(litellm_params, "get", None)
if not callable(getter):
return None
api_base = getter("api_base")
api_key = getter("api_key")
if not api_base or not api_key:
return None
return (api_base, api_key)
def _find_deployments_on_same_encryption_boundary(
self,
healthy_deployments: List[dict],
model_id: str,
) -> List[dict]:
"""
Deployments in ``healthy_deployments`` sharing the originating
deployment's ``(api_base, api_key)``. Returns ``[]`` if router isn't
wired in, the originating deployment was removed, or no boundary match.
"""
if self.router is None:
return []
originating = self.router.get_deployment(model_id=model_id)
if originating is None:
return []
boundary = self._encryption_boundary_key(
originating.litellm_params.model_dump(exclude_none=True)
)
if boundary is None:
return []
return [
d
for d in healthy_deployments
if self._encryption_boundary_key(d.get("litellm_params", {})) == boundary
]
# ------------------------------------------------------------------
# Request routing (pre-call filter)
# ------------------------------------------------------------------
@ -172,8 +228,25 @@ class EncryptedContentAffinityCheck(CustomLogger):
request_kwargs["_encrypted_content_affinity_pinned"] = True
return [deployment]
# Follow-up switched model_name (LIT-2531): pin by Azure resource instead.
boundary_matches = self._find_deployments_on_same_encryption_boundary(
healthy_deployments=typed_healthy_deployments,
model_id=model_id,
)
if boundary_matches:
verbose_router_logger.debug(
"EncryptedContentAffinityCheck: model_id=%s not in healthy_deployments; "
"pinning to %d deployment(s) on same encryption boundary",
model_id,
len(boundary_matches),
)
request_kwargs["_encrypted_content_affinity_pinned"] = True
return boundary_matches
verbose_router_logger.error(
"EncryptedContentAffinityCheck: decoded deployment=%s not found in healthy_deployments",
"EncryptedContentAffinityCheck: decoded deployment=%s not found in "
"healthy_deployments and no boundary match available; falling back to "
"full deployment pool (encrypted_content may be rejected upstream)",
model_id,
)
return typed_healthy_deployments

View file

@ -633,6 +633,16 @@ class BaseLitellmParams(
),
)
skip_tool_message_in_guardrail: Optional[bool] = Field(
default=None,
description=(
"When True, unified guardrails skip tool-role messages when building "
"evaluation inputs (texts and structured_messages). When False, tool "
"messages are included even if litellm_settings sets a global skip. When "
"None, use the global litellm.skip_tool_message_in_guardrail setting."
),
)
# Lakera specific params
category_thresholds: Optional[LakeraCategoryThresholds] = Field(
default=None,

View file

@ -1042,3 +1042,10 @@ class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False):
thinking: dict
metadata: dict
output_config: dict
# `context_management` is allowed for Bedrock InvokeModel only when it
# carries `compact_20260112` edits paired with the `compact-2026-01-12`
# anthropic-beta header. The Invoke transformation filters edits to the
# supported subset and strips the field entirely when nothing remains, so
# other edit types (e.g. `clear_thinking_20251015`) never reach Bedrock.
context_management: dict

View file

@ -1,6 +1,7 @@
from typing import Any, Dict, List, Optional
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel
from pydantic import BaseModel, ConfigDict, model_validator
class BulkUpdateKeyRequestItem(BaseModel):
@ -40,3 +41,78 @@ class BulkUpdateKeyResponse(BaseModel):
total_requested: int
successful_updates: List[SuccessfulKeyUpdate]
failed_updates: List[FailedKeyUpdate]
class KeyUpdateFields(BaseModel):
"""Allowlist of bulk-broadcastable fields for /team/key/bulk_update; `extra="forbid"` blocks RBAC/ownership/scope mutations even by team admins."""
model_config = ConfigDict(extra="forbid", protected_namespaces=())
# Budgets
max_budget: Optional[float] = None
budget_id: Optional[str] = None
budget_duration: Optional[str] = None
budget_limits: Optional[List[Any]] = None
model_max_budget: Optional[Dict[str, Any]] = None
# Rate limits
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
model_tpm_limit: Optional[Dict[str, Any]] = None
model_rpm_limit: Optional[Dict[str, Any]] = None
max_parallel_requests: Optional[int] = None
rpm_limit_type: Optional[
Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]
] = None
tpm_limit_type: Optional[
Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]
] = None
# Temporary budget grants (auto-expire). `spend` deliberately omitted — bulk-zeroing it bypasses budget enforcement; admin-only via /key/update.
temp_budget_increase: Optional[float] = None
temp_budget_expiry: Optional[datetime] = None
# Expiry
duration: Optional[str] = None
# Operational metadata
tags: Optional[List[str]] = None
metadata: Optional[Dict[str, Any]] = None
@model_validator(mode="after")
def validate_temp_budget(self) -> "KeyUpdateFields":
if self.temp_budget_increase is not None or self.temp_budget_expiry is not None:
if self.temp_budget_increase is None or self.temp_budget_expiry is None:
raise ValueError(
"temp_budget_increase and temp_budget_expiry must be set together"
)
return self
@model_validator(mode="after")
def require_at_least_one_field(self) -> "KeyUpdateFields":
# Reject empty payload — would iterate every key with no-op writes.
if not self.model_fields_set:
raise ValueError("update_fields must specify at least one field to update.")
return self
class BulkUpdateTeamKeysRequest(BaseModel):
"""Apply one update payload to many keys inside a team; provide either `key_ids` or `all_keys_in_team=True`."""
team_id: str
key_ids: Optional[List[str]] = None
all_keys_in_team: bool = False
update_fields: KeyUpdateFields
@model_validator(mode="after")
def validate_selection(self) -> "BulkUpdateTeamKeysRequest":
has_key_ids = self.key_ids is not None and len(self.key_ids) > 0
if has_key_ids and self.all_keys_in_team:
raise ValueError(
"Provide either `key_ids` or `all_keys_in_team=True`, not both."
)
if not has_key_ids and not self.all_keys_in_team:
raise ValueError(
"Must provide either `key_ids` (non-empty) or `all_keys_in_team=True`."
)
return self

View file

@ -825,6 +825,7 @@ API_ROUTE_TO_CALL_TYPES = {
# Realtime API
"/realtime": [CallTypes.arealtime],
"/v1/realtime": [CallTypes.arealtime],
"/openai/v1/realtime": [CallTypes.arealtime],
# Provider-specific routes
"/anthropic/v1/messages": [CallTypes.anthropic_messages],
# Google GenAI routes

View file

@ -1,3 +1,5 @@
"""Utility helpers for LiteLLM core request handling and provider support."""
# from __future__ import annotations must be the first non-comment statement
from __future__ import annotations
@ -4405,6 +4407,10 @@ def get_optional_params( # noqa: PLR0915
else False
),
)
if bedrock_route == "claude_platform":
optional_params = BedrockModelInfo.map_claude_platform_auth_params(
passed_params=passed_params, optional_params=optional_params
)
elif custom_llm_provider == "cloudflare":
optional_params = litellm.CloudflareChatConfig().map_openai_params(
model=model,
@ -9490,6 +9496,49 @@ def get_non_default_completion_params(kwargs: dict) -> dict:
return non_default_params
def peek_reasoning_summary_aliases(optional_params: dict) -> Optional[Any]:
"""Read AI-SDK-style reasoning summary from optional_params or nested extra_body.
Uses key membership (not ``or`` chains) so falsy values like ``""`` are not skipped.
"""
if "reasoningSummary" in optional_params:
return optional_params["reasoningSummary"]
if "reasoning_summary" in optional_params:
return optional_params["reasoning_summary"]
extra_body = optional_params.get("extra_body")
if isinstance(extra_body, dict):
if "reasoningSummary" in extra_body:
return extra_body["reasoningSummary"]
if "reasoning_summary" in extra_body:
return extra_body["reasoning_summary"]
return None
def strip_reasoning_summary_aliases_from_optional_params(
optional_params: dict,
) -> Tuple[dict, Optional[Any]]:
"""Copy optional_params; remove reasoningSummary aliases from top-level and extra_body."""
op = dict(optional_params)
rs_val = op.pop("reasoningSummary", None)
snake_rs_val = op.pop("reasoning_summary", None)
if rs_val is None:
rs_val = snake_rs_val
eb = op.get("extra_body")
if isinstance(eb, dict):
eb = dict(eb)
eb_rs_val = eb.pop("reasoningSummary", None)
eb_snake_rs_val = eb.pop("reasoning_summary", None)
if rs_val is None:
rs_val = eb_rs_val
if rs_val is None:
rs_val = eb_snake_rs_val
if eb:
op["extra_body"] = eb
else:
op.pop("extra_body", None)
return op, rs_val
def get_non_default_transcription_params(kwargs: dict) -> dict:
from litellm.constants import OPENAI_TRANSCRIPTION_PARAMS

View file

@ -21109,6 +21109,38 @@
"supports_system_messages": true,
"supports_tool_choice": true
},
"gpt-realtime-2": {
"cache_creation_input_audio_token_cost": 4e-07,
"cache_read_input_token_cost": 4e-07,
"input_cost_per_audio_token": 3.2e-05,
"input_cost_per_image": 5e-06,
"input_cost_per_token": 4e-06,
"litellm_provider": "openai",
"max_input_tokens": 32000,
"max_output_tokens": 4096,
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_audio_token": 6.4e-05,
"output_cost_per_token": 1.6e-05,
"supported_endpoints": [
"/v1/realtime"
],
"supported_modalities": [
"text",
"image",
"audio"
],
"supported_output_modalities": [
"text",
"audio"
],
"supports_audio_input": true,
"supports_audio_output": true,
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_system_messages": true,
"supports_tool_choice": true
},
"gpt-realtime-mini": {
"cache_creation_input_audio_token_cost": 3e-07,
"cache_read_input_audio_token_cost": 3e-07,
@ -27192,6 +27224,20 @@
"supports_reasoning": true,
"supports_tool_choice": true
},
"openrouter/qwen/qwen3.6-plus": {
"input_cost_per_token": 3.25e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1000000,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 1.95e-06,
"source": "https://openrouter.ai/qwen/qwen3.6-plus",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"openrouter/qwen/qwen3.5-35b-a3b": {
"input_cost_per_token": 2.5e-07,
"litellm_provider": "openrouter",

View file

@ -22,7 +22,7 @@ dependencies = [
"importlib-metadata>=8.0.0,<9.0",
"tokenizers>=0.21.0,<1.0",
"click>=8.0.0,<9.0",
"jinja2>=3.1.0,<4.0",
"jinja2>=3.1.6,<4.0",
"aiohttp>=3.10,<4.0",
"pydantic>=2.10.0,<3.0.0",
"jsonschema>=4.0.0,<5.0",
@ -275,6 +275,33 @@ filterwarnings = [
"ignore::DeprecationWarning:pytest_asyncio.plugin",
]
[tool.mutmut]
# Mutation-testing scope. Driven by the manually-triggered workflow at
# .github/workflows/mutation-test.yml. mutmut is not part of the project's
# default install; it is pulled in via `uv run --with mutmut==<version>` in CI.
# `also_copy = ["litellm/"]` is required because mutmut runs in a `mutants/`
# sandbox and the test conftest imports from across the litellm package.
paths_to_mutate = [
"litellm/proxy/management_endpoints/",
]
tests_dir = [
"tests/test_litellm/proxy/management_endpoints/",
]
also_copy = [
"litellm/",
]
# Disable rerun/parallel plugins for mutation runs:
# - pytest-retry triggers an `INTERNALERROR: no option named 'filtered_exceptions'`
# when invoked via mutmut's in-process `pytest.main()` call.
# - rerunning a "failed" test on a mutant would mask which mutants are killed
# vs. survive, so reruns are wrong for mutation testing regardless.
# - xdist is unnecessary inside mutmut (mutmut handles its own parallelism).
pytest_add_cli_args = [
"-p", "no:retry",
"-p", "no:rerunfailures",
"-p", "no:xdist",
]
[tool.coverage.run]
source = ["litellm"]
relative_files = true

423
scripts/mutation_report.py Normal file
View file

@ -0,0 +1,423 @@
#!/usr/bin/env python3
"""Generate an agent-actionable mutation testing report.
Reads the mutmut sandbox state at `mutants/` and produces a single
`mutation-report.md` grouped by function. For each function with surviving
mutants, the report embeds the original function source (via AST), the
unified diff for each surviving mutation (via `mutmut show`), and the
existing test file(s) — followed by an ACH-style instruction asking the
reader to write tests that kill the survivors.
Run after `mutmut run` and `mutmut export-cicd-stats`. Expects mutmut to be
invokable as `uv run --no-sync --with mutmut==<version> mutmut <subcommand>`.
"""
from __future__ import annotations
import ast
import json
import re
import subprocess
import sys
import tomllib
from collections import defaultdict
from difflib import SequenceMatcher
from pathlib import Path
from textwrap import dedent
ROOT = Path(__file__).resolve().parent.parent
MUTMUT_INVOCATION = ["uv", "run", "--no-sync", "--with", "mutmut==3.5.0", "mutmut"]
def load_mutmut_config() -> dict:
with open(ROOT / "pyproject.toml", "rb") as f:
return tomllib.load(f)["tool"]["mutmut"]
def get_survivors() -> list[str]:
proc = subprocess.run(
[*MUTMUT_INVOCATION, "results"], capture_output=True, text=True, check=False
)
survivors = []
for line in proc.stdout.splitlines():
m = re.match(r"\s*(\S+):\s*survived\s*$", line)
if m:
survivors.append(m.group(1))
return survivors
def get_mutmut_show(mutant_name: str) -> str:
proc = subprocess.run(
[*MUTMUT_INVOCATION, "show", mutant_name],
capture_output=True,
text=True,
check=False,
)
return proc.stdout.strip() or "(mutmut show produced no output)"
def parse_mutant_name(name: str) -> tuple[str, str, str]:
"""Parse `<dotted.module>.x_<function>__mutmut_<N>` -> (module, function, N).
mutmut prefixes mutated functions with `x_` (single underscore). For a
function named `foo`, mutants are `x_foo__mutmut_N`. For a function named
`_foo` (leading underscore), the mutant becomes `x__foo__mutmut_N` — so
the regex matches a single underscore after `x` and captures everything
(including any leading underscores) up to `__mutmut_<N>`.
"""
m = re.match(r"^(.+)\.x_(.+)__mutmut_(\d+)$", name)
if not m:
return name, name, "?"
return m.group(1), m.group(2), m.group(3)
def function_anchor(module_path: str, function_name: str) -> str:
return re.sub(r"[^a-z0-9_-]+", "-", f"{module_path}-{function_name}".lower()).strip(
"-"
)
def module_to_file(module_path: str) -> Path | None:
candidate = ROOT / Path(*module_path.split(".")).with_suffix(".py")
return candidate if candidate.exists() else None
def find_function_in_file(
file_path: Path, function_name: str
) -> tuple[int, int, str, list[int]] | None:
"""Find a top-level or nested function by name; returns the first match.
Returns ``(start_line, end_line, source, all_match_lines)`` or ``None``.
``all_match_lines`` is the start line of every function (any nesting
level) in the file with this name. When ``len(all_match_lines) > 1`` the
file defines the same name in multiple places (e.g., a module-level
helper and a class method) — mutmut's mutant identifier does not carry
class context, so we can't determine which definition was mutated.
Callers surface a disambiguation note in that case.
"""
src = file_path.read_text()
tree = ast.parse(src)
matches = [
node
for node in ast.walk(tree)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
and node.name == function_name
]
if not matches:
return None
first = matches[0]
lines = src.splitlines()
return (
first.lineno,
first.end_lineno,
"\n".join(lines[first.lineno - 1 : first.end_lineno]),
[m.lineno for m in matches],
)
def collect_test_files(tests_dir: list[str]) -> list[Path]:
found: list[Path] = []
for entry in tests_dir:
p = ROOT / entry
if p.is_file():
found.append(p)
elif p.is_dir():
found.extend(sorted(p.rglob("test_*.py")))
return found
def _indent_of(line: str) -> str:
return line[: len(line) - len(line.lstrip())]
def render_meta_style_mutant(
module_path: str, function_name: str, mutant_num: str
) -> str | None:
"""Render the mutated function with `# MUTANT START`/`# MUTANT END` delimiters.
Reads `mutants/<module>.py` (the trampoline file mutmut emits), finds
`x_<func>__mutmut_orig` and `x_<func>__mutmut_<N>`, and renders the
mutated version with the lines that differ from `__mutmut_orig` wrapped
in `# MUTANT START`/`# MUTANT END` comments — the format from Meta's
ACH paper (arXiv 2501.12862, Table 1).
The function header is rewritten to use the original function name so
the agent sees the source as it would appear in the file (rather than
mutmut's internal `x_*__mutmut_<N>` name).
Returns None if the trampoline file or either function cannot be found
(the caller falls back to the unified diff).
"""
trampoline = ROOT / "mutants" / Path(*module_path.split(".")).with_suffix(".py")
if not trampoline.exists():
return None
src = trampoline.read_text()
try:
tree = ast.parse(src)
except SyntaxError:
return None
file_lines = src.splitlines()
orig_def = f"x_{function_name}__mutmut_orig"
mutant_def = f"x_{function_name}__mutmut_{mutant_num}"
orig_node = mutated_node = None
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
if node.name == orig_def:
orig_node = node
elif node.name == mutant_def:
mutated_node = node
if orig_node is None or mutated_node is None:
return None
orig_lines = file_lines[orig_node.lineno - 1 : orig_node.end_lineno]
mutated_lines = file_lines[mutated_node.lineno - 1 : mutated_node.end_lineno]
if not orig_lines or not mutated_lines:
return None
# Rewrite the def line to use the original (non-trampolined) function name
# so the agent sees the function as it appears in the source file.
orig_lines[0] = orig_lines[0].replace(orig_def, function_name, 1)
mutated_lines[0] = mutated_lines[0].replace(mutant_def, function_name, 1)
matcher = SequenceMatcher(a=orig_lines, b=mutated_lines)
out: list[str] = []
in_diff = False
for op, i1, i2, j1, j2 in matcher.get_opcodes():
if op == "equal":
if in_diff:
# Close the block at the indent of the line just inside it.
indent = _indent_of(out[-1]) if out else ""
out.append(f"{indent}# MUTANT END")
in_diff = False
out.extend(mutated_lines[j1:j2])
else:
if not in_diff:
# Open the block at the indent of the first differing line.
if j1 < len(mutated_lines):
indent = _indent_of(mutated_lines[j1])
elif i1 < len(orig_lines):
indent = _indent_of(orig_lines[i1])
else:
indent = ""
out.append(f"{indent}# MUTANT START")
in_diff = True
if op == "delete":
# Mutation removed lines — surface what was deleted as a
# comment so the agent can see the intent of the change.
for deleted in orig_lines[i1:i2]:
indent = _indent_of(deleted)
out.append(f"{indent}# (deleted by mutation): {deleted.lstrip()}")
else:
# replace / insert: take from mutated_lines
out.extend(mutated_lines[j1:j2])
if in_diff:
indent = _indent_of(out[-1]) if out else ""
out.append(f"{indent}# MUTANT END")
return "\n".join(out)
def render(config: dict, survivors: list[str], stats: dict | None) -> str:
by_function: dict[tuple[str, str], list[tuple[str, str]]] = defaultdict(list)
for survivor in survivors:
module_path, function_name, mutant_num = parse_mutant_name(survivor)
by_function[(module_path, function_name)].append((survivor, mutant_num))
out: list[str] = []
out.append("# Mutation Test Report")
out.append("")
out.append("## Summary")
out.append("")
if stats:
total = stats.get("total", 0) or sum(
stats.get(k, 0)
for k in (
"killed",
"survived",
"no_tests",
"skipped",
"suspicious",
"timeout",
"segfault",
)
)
killed = stats.get("killed", 0)
survived = stats.get("survived", 0)
score = (killed / total * 100) if total else 0.0
out.append(f"- Total mutants: **{total}**")
out.append(f"- Killed: **{killed}**")
out.append(f"- Survived: **{survived}**")
out.append(f"- Mutation score: **{score:.1f}%**")
for k in ("no_tests", "skipped", "suspicious", "timeout", "segfault"):
v = stats.get(k, 0)
if v:
out.append(f"- {k.replace('_', ' ').title()}: {v}")
else:
out.append(f"- Survivors found: **{len(survivors)}**")
out.append("- (mutmut-cicd-stats.json not available — full counts unavailable)")
out.append("")
if not survivors:
out.append("**No surviving mutants — the test suite caught every mutation.**")
out.append("")
return "\n".join(out)
out.append("## Surviving mutants by function")
out.append("")
for (module_path, function_name), items in by_function.items():
anchor = function_anchor(module_path, function_name)
out.append(
f"- [`{function_name}`](#{anchor}) — {len(items)} mutant"
f"{'s' if len(items) != 1 else ''} ({module_path})"
)
out.append("")
for (module_path, function_name), items in by_function.items():
anchor = function_anchor(module_path, function_name)
out.append(f'<a id="{anchor}"></a>')
out.append(f"## `{module_path}.{function_name}`")
out.append("")
out.append(f"**Module:** `{module_path}`")
file_path = module_to_file(module_path)
if file_path is None:
out.append("")
out.append(f"_(could not locate source file for module `{module_path}`)_")
out.append("")
else:
rel = file_path.relative_to(ROOT)
out.append(f"**File:** `{rel}`")
out.append("")
found = find_function_in_file(file_path, function_name)
if found:
start, end, fn_src, all_lines = found
out.append(f"### Original function (lines {start}-{end})")
out.append("")
if len(all_lines) > 1:
line_list = ", ".join(str(line) for line in all_lines)
out.append(
f"> **Note:** {len(all_lines)} functions named "
f"`{function_name}` are defined in this file at lines "
f"{line_list}. Showing the first match. mutmut's "
f"mutant identifier does not carry class context, so "
f"the body below may not correspond to the function "
f"that was actually mutated — verify manually before "
f"writing the killing test."
)
out.append("")
out.append("```python")
out.append(fn_src)
out.append("```")
out.append("")
else:
out.append(f"_(could not locate `{function_name}` in {rel} via AST)_")
out.append("")
out.append(f"### Surviving mutations ({len(items)})")
out.append("")
for i, (mutant_name, mutant_num) in enumerate(items, 1):
out.append(f"#### Mutation {i} of {len(items)} — `{mutant_name}`")
out.append("")
meta_style = render_meta_style_mutant(
module_path, function_name, mutant_num
)
if meta_style is not None:
out.append(
"Mutated function (the bug is delimited by "
"`# MUTANT START` / `# MUTANT END`):"
)
out.append("")
out.append("```python")
out.append(meta_style)
out.append("```")
out.append("")
out.append("<details><summary>Unified diff (`mutmut show`)</summary>")
out.append("")
out.append("```diff")
out.append(get_mutmut_show(mutant_name))
out.append("```")
out.append("")
out.append("</details>")
out.append("")
else:
# Fallback: trampoline file or function lookup failed.
out.append("```diff")
out.append(get_mutmut_show(mutant_name))
out.append("```")
out.append("")
test_files = collect_test_files(config.get("tests_dir", []))
if test_files:
out.append("## Existing tests")
out.append("")
out.append(
"These are the test files that mutmut considered when classifying the "
"mutants above. New tests should be added here, matching existing "
"conventions, fixtures, and naming."
)
out.append("")
for tf in test_files:
rel = tf.relative_to(ROOT)
out.append(f"### `{rel}`")
out.append("")
out.append("```python")
out.append(tf.read_text())
out.append("```")
out.append("")
out.append("## Task")
out.append("")
out.append(
dedent(
"""\
For each surviving mutant listed above, write a new test in the
existing test file (matching its conventions, fixtures, and naming
style) that:
- **Fails** when the mutated version of the function is in place.
- **Passes** when the original (correct) version is in place.
Aim for one test per surviving mutant. If multiple mutants in the
same function can be killed by a single test, that is fine — note
which mutant numbers in the test name or docstring.
Do not modify the source file. Only add tests.
"""
).strip()
)
out.append("")
return "\n".join(out)
def main() -> int:
config = load_mutmut_config()
stats_file = ROOT / "mutants" / "mutmut-cicd-stats.json"
stats: dict | None = None
if stats_file.exists():
try:
stats = json.loads(stats_file.read_text())
except json.JSONDecodeError as exc:
print(f"warning: could not parse {stats_file}: {exc}", file=sys.stderr)
survivors = get_survivors()
report = render(config, survivors, stats)
out_path = ROOT / "mutation-report.md"
out_path.write_text(report)
print(
f"Wrote {out_path} ({len(survivors)} survivor"
f"{'s' if len(survivors) != 1 else ''}, {len(report)} chars)"
)
return 0
if __name__ == "__main__":
raise SystemExit(main())

Some files were not shown because too many files have changed in this diff Show more