tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-01 04:22:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'litellm_internal_staging' into litellm_agent_oss_staging_05_06_2026

Browse Source
This commit is contained in:
Sameer Kankute
2026-05-11 09:07:08 +05:30
committed by GitHub
parent a466c2eb98 0af33fbe70
commit ce17e9490f
170 changed files with 11361 additions and 5562 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/test-unit-proxy-db.yml
+1
View File
@@ -141,6 +141,7 @@ jobs:
tests/proxy_unit_tests/test_server_root_path.py
tests/proxy_unit_tests/test_proxy_pass_user_config.py
tests/proxy_unit_tests/test_proxy_token_counter.py
tests/proxy_unit_tests/test_request_size_limit_middleware.py
workers: 4
dist: loadscope
timeout: 15
.gitignore
+2 -1
View File
@@ -100,4 +100,5 @@ STABILIZATION_TODO.md
**/playwright-report
**/*.storageState.json
**/coverage
test-config
test-config
.vscode
AGENTS.md
+19 -2
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.
CLAUDE.md
+1 -1
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).
Dockerfile
+10 -10
View File
@@ -1,8 +1,8 @@
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin
@@ -68,8 +68,8 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile supervisor && \
npm install -g npm@11.12.1 tar@7.5.11 glob@13.0.6 @isaacs/brace-expansion@5.0.1 brace-expansion@5.0.5 minimatch@10.2.4 diff@8.0.3 picomatch@4.0.4 && \
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile && \
npm install -g npm@11.14.0 tar@7.5.11 glob@13.0.6 @isaacs/brace-expansion@5.0.1 brace-expansion@5.0.5 minimatch@10.2.4 diff@8.0.3 picomatch@4.0.4 && \
GLOBAL="$(npm root -g)" && \
for pkg in tar glob @isaacs/brace-expansion brace-expansion minimatch diff picomatch; do \
name="${pkg##*/}"; \
@@ -85,17 +85,17 @@ ENV PATH="/app/.venv/bin:${PATH}"
COPY --from=builder /app /app
# Prisma binaries live in $HOME/.cache (default prisma-python location),
# which is /root/.cache here. Copy them from the builder so they survive
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
COPY --from=builder /root/.cache /root/.cache
# which is /root/.cache here. Copy only the Prisma subdirs — copying the
# whole /root/.cache drags in the uv build cache (~660 MB, includes a
# setuptools wheel that surfaces as a CVE finding even though it's not
# on the runtime sys.path).
COPY --from=builder /root/.cache/prisma /root/.cache/prisma
COPY --from=builder /root/.cache/prisma-python /root/.cache/prisma-python
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete
EXPOSE 4000/tcp
COPY docker/supervisord.conf /etc/supervisord.conf
ENTRYPOINT ["docker/prod_entrypoint.sh"]
CMD ["--port", "4000"]
codecov.yaml
+6 -1
View File
@@ -1,3 +1,8 @@
codecov:
require_ci_to_pass: false # post coverage status even if CI has unrelated failures
notify:
wait_for_ci: false # post as soon as expected uploads arrive, don't wait on CI
component_management:
individual_components:
- component_id: "Router"
@@ -28,7 +33,7 @@ coverage:
project:
default:
target: auto
threshold: 1% # at maximum allow project coverage to drop by 1%
threshold: 0% # do not allow project coverage to drop
patch:
default:
target: auto
deploy/Dockerfile.ghcr_base
-18
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"]
deploy/charts/litellm-helm/templates/deployment.yaml
+13 -14
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:
@@ -129,12 +139,6 @@ spec:
value: {{ $val | quote }}
{{- end }}
{{- end }}
{{- if .Values.separateHealthApp }}
- name: SEPARATE_HEALTH_APP
value: "1"
- name: SEPARATE_HEALTH_PORT
value: {{ .Values.separateHealthPort | default "8081" | quote }}
{{- end }}
{{- with .Values.extraEnvVars }}
{{- toYaml . | nindent 12 }}
{{- end }}
@@ -175,15 +179,10 @@ spec:
- name: http
containerPort: {{ .Values.service.port }}
protocol: TCP
{{- if .Values.separateHealthApp }}
- name: health
containerPort: {{ .Values.separateHealthPort | default 8081 }}
protocol: TCP
{{- end }}
livenessProbe:
httpGet:
path: {{ .Values.livenessProbe.path | quote }}
port: {{ if .Values.separateHealthApp }}"health"{{ else }}"http"{{ end }}
port: "http"
initialDelaySeconds: {{ .Values.livenessProbe.initialDelaySeconds }}
periodSeconds: {{ .Values.livenessProbe.periodSeconds }}
timeoutSeconds: {{ .Values.livenessProbe.timeoutSeconds }}
@@ -192,7 +191,7 @@ spec:
readinessProbe:
httpGet:
path: {{ .Values.readinessProbe.path | quote }}
port: {{ if .Values.separateHealthApp }}"health"{{ else }}"http"{{ end }}
port: "http"
initialDelaySeconds: {{ .Values.readinessProbe.initialDelaySeconds }}
periodSeconds: {{ .Values.readinessProbe.periodSeconds }}
timeoutSeconds: {{ .Values.readinessProbe.timeoutSeconds }}
@@ -201,7 +200,7 @@ spec:
startupProbe:
httpGet:
path: {{ .Values.startupProbe.path | quote }}
port: {{ if .Values.separateHealthApp }}"health"{{ else }}"http"{{ end }}
port: "http"
initialDelaySeconds: {{ .Values.startupProbe.initialDelaySeconds }}
periodSeconds: {{ .Values.startupProbe.periodSeconds }}
timeoutSeconds: {{ .Values.startupProbe.timeoutSeconds }}
deploy/charts/litellm-helm/values.yaml
+20 -6
View File
@@ -88,12 +88,6 @@ service:
# optionally specify loadBalancerClass
# loadBalancerClass: tailscale
# Separate health app configuration
# When enabled, health checks will use a separate port and the application
# will receive SEPARATE_HEALTH_APP=1 and SEPARATE_HEALTH_PORT from environment variables
separateHealthApp: false
separateHealthPort: 8081
# Probes for LiteLLM gateway container
livenessProbe:
path: /health/liveliness
@@ -258,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
deploy/kubernetes/kub.yaml
-56
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
deploy/kubernetes/service.yaml
-12
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
dev_config.yaml
-13
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
docker-compose.yml
+5
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
docker/Dockerfile.alpine
-68
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"]
docker/Dockerfile.custom_ui
-86
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"]
docker/Dockerfile.database
+1 -3
View File
@@ -66,7 +66,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile supervisor && \
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile && \
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 \
@@ -102,7 +102,5 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
EXPOSE 4000/tcp
COPY docker/supervisord.conf /etc/supervisord.conf
ENTRYPOINT ["docker/prod_entrypoint.sh"]
CMD ["--port", "4000"]
docker/Dockerfile.dev
-121
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"]
docker/Dockerfile.health_check
-30
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"]
docker/Dockerfile.non_root
+1 -2
View File
@@ -103,13 +103,12 @@ RUN for i in 1 2 3; do \
apk upgrade --no-cache && break || sleep 5; \
done && \
for i in 1 2 3; do \
apk add --no-cache python3 bash openssl tzdata supervisor libsndfile nodejs && break || sleep 5; \
apk add --no-cache python3 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
done
COPY --from=builder /app /app
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets
COPY --from=builder /app/docker/supervisord.conf /etc/supervisord.conf
ENV PATH="/app/.venv/bin:${PATH}" \
PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
docker/prod_entrypoint.sh
+1 -7
View File
@@ -1,14 +1,8 @@
#!/bin/sh
if [ "$SEPARATE_HEALTH_APP" = "1" ]; then
export LITELLM_ARGS="$@"
export SUPERVISORD_STOPWAITSECS="${SUPERVISORD_STOPWAITSECS:-3600}"
exec supervisord -c /etc/supervisord.conf
fi
if [ "$USE_DDTRACE" = "true" ]; then
export DD_TRACE_OPENAI_ENABLED="False"
exec ddtrace-run litellm "$@"
else
exec litellm "$@"
fi
fi
docker/supervisord.conf
-46
View File
@@ -1,46 +0,0 @@
[supervisord]
nodaemon=true
loglevel=info
logfile=/tmp/supervisord.log
pidfile=/tmp/supervisord.pid
[group:litellm]
programs=main,health
[program:main]
command=sh -c 'if [ "$USE_DDTRACE" = "true" ]; then export DD_TRACE_OPENAI_ENABLED="False"; exec ddtrace-run python -m litellm.proxy.proxy_cli --host 0.0.0.0 --port=4000 $LITELLM_ARGS; else exec python -m litellm.proxy.proxy_cli --host 0.0.0.0 --port=4000 $LITELLM_ARGS; fi'
autostart=true
autorestart=true
startretries=3
priority=1
exitcodes=0
stopasgroup=true
killasgroup=true
stopwaitsecs=%(ENV_SUPERVISORD_STOPWAITSECS)s
stdout_logfile=/dev/stdout
stderr_logfile=/dev/stderr
stdout_logfile_maxbytes = 0
stderr_logfile_maxbytes = 0
environment=PYTHONUNBUFFERED=true
[program:health]
command=sh -c '[ "$SEPARATE_HEALTH_APP" = "1" ] && exec uvicorn litellm.proxy.health_endpoints.health_app_factory:build_health_app --factory --host 0.0.0.0 --port=${SEPARATE_HEALTH_PORT:-4001} || exit 0'
autostart=true
autorestart=true
startretries=3
priority=2
exitcodes=0
stopasgroup=true
killasgroup=true
stopwaitsecs=%(ENV_SUPERVISORD_STOPWAITSECS)s
stdout_logfile=/dev/stdout
stderr_logfile=/dev/stderr
stdout_logfile_maxbytes = 0
stderr_logfile_maxbytes = 0
environment=PYTHONUNBUFFERED=true
[eventlistener:process_monitor]
command=python -c "from supervisor import childutils; import os, signal; [os.kill(os.getppid(), signal.SIGTERM) for h,p in iter(lambda: childutils.listener.wait(), None) if h['eventname'] in ['PROCESS_STATE_FATAL', 'PROCESS_STATE_EXITED'] and dict([x.split(':') for x in p.split(' ')])['processname'] in ['main', 'health'] or childutils.listener.ok()]"
events=PROCESS_STATE_EXITED,PROCESS_STATE_FATAL
autostart=true
autorestart=true
index.yaml
-108
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"
litellm-js/proxy/.npmrc
-5
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
litellm-js/proxy/README.md
-8
View File
@@ -1,8 +0,0 @@
```
npm install
npm run dev
```
```
npm run deploy
```
litellm-js/proxy/package-lock.json
Generated
-2054
View File
File diff suppressed because it is too large Load Diff
litellm-js/proxy/package.json
-14
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"
}
}
litellm-js/proxy/src/index.ts
-59
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
litellm-js/proxy/tsconfig.json
-17
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
},
}
litellm-js/proxy/wrangler.toml
-18
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 = ""
litellm-js/spend-logs/.npmrc
-5
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
litellm-js/spend-logs/Dockerfile
-26
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"]
litellm-js/spend-logs/README.md
-8
View File
@@ -1,8 +0,0 @@
```
npm install
npm run dev
```
```
open http://localhost:3000
```
litellm-js/spend-logs/package-lock.json
Generated
-597
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.13.0",
"resolved": "https://registry.npmjs.org/get-tsconfig/-/get-tsconfig-4.13.0.tgz",
"integrity": "sha512-1VKTZJCwBrvbd+Wn3AOgQP/2Av+TfTCOlE4AcRJE72W1ksZXbAx8PPBR9RzgTeSPzlPMHrbANMH3LbltH73wxQ==",
"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"
}
}
}
litellm-js/spend-logs/package.json
-13
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"
}
}
litellm-js/spend-logs/schema.prisma
-29
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?
}
litellm-js/spend-logs/src/_types.ts
-32
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;
};
litellm-js/spend-logs/src/index.ts
-84
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
})
litellm-js/spend-logs/tsconfig.json
-13
View File
@@ -1,13 +0,0 @@
{
"compilerOptions": {
"target": "ESNext",
"module": "ESNext",
"moduleResolution": "Bundler",
"strict": true,
"types": [
"node"
],
"jsx": "react-jsx",
"jsxImportSource": "hono/jsx",
}
}
litellm/__init__.py
+5
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] = (
@@ -388,6 +389,7 @@ anthropic_beta_headers_url: str = os.getenv(
suppress_debug_info = False
dynamodb_table_name: Optional[str] = None
s3_callback_params: Optional[Dict] = None
s3_audit_callback_params: Optional[Dict] = None
datadog_llm_observability_params: Optional[Union[DatadogLLMObsInitParams, Dict]] = None
datadog_params: Optional[Union[DatadogInitParams, Dict]] = None
aws_sqs_callback_params: Optional[Dict] = None
@@ -414,6 +416,9 @@ custom_prometheus_metadata_labels: List[str] = []
custom_prometheus_tags: List[str] = []
prometheus_metrics_config: Optional[List] = None
prometheus_emit_stream_label: bool = False
prometheus_end_user_metrics_max_series_per_metric: Optional[int] = 10000
prometheus_end_user_metrics_ttl_seconds: Optional[float] = 3600.0
prometheus_end_user_metrics_cleanup_interval_seconds: Optional[float] = 60.0
disable_add_prefix_to_prompt: bool = (
False # used by anthropic, to disable adding prefix to prompt
)
litellm/constants.py
+11
View File
@@ -161,6 +161,11 @@ MCP_STDIO_ALLOWED_COMMANDS: frozenset = frozenset(
| (set(_MCP_STDIO_EXTRA_COMMANDS.split(",")) - {""})
)
# MCP OAuth2 Token Exchange (OBO) Defaults
MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE = int(
os.getenv("MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE", "500")
)
LITELLM_UI_ALLOW_HEADERS = [
"x-litellm-semantic-filter",
"x-litellm-semantic-filter-tools",
@@ -1457,6 +1462,12 @@ KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job"
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME = "litellm_expired_ui_session_key_cleanup_job"
SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500))
SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000))
SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(
os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)
)
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float(
os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5)
)
SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
SPEND_LOG_QUEUE_POLL_INTERVAL = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0))
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = int(
litellm/experimental_mcp_client/client.py
+2
View File
@@ -366,6 +366,8 @@ class MCPClient:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.token:
headers["Authorization"] = f"token {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.oauth2_token_exchange:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
elif isinstance(self._mcp_auth_value, dict):
headers.update(self._mcp_auth_value)
# Note: aws_sigv4 auth is not handled here — SigV4 requires per-request
litellm/integrations/azure_sentinel/azure_sentinel.py
+111 -28
View File
@@ -14,16 +14,18 @@ For batching specific details see CustomBatchLogger class
import asyncio
import os
import time
import traceback
from typing import List, Optional
from typing import List, Optional, Union
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.utils import StandardLoggingPayload
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
class AzureSentinelLogger(CustomBatchLogger):
@@ -39,6 +41,7 @@ class AzureSentinelLogger(CustomBatchLogger):
tenant_id: Optional[str] = None,
client_id: Optional[str] = None,
client_secret: Optional[str] = None,
audit_stream_name: Optional[str] = None,
**kwargs,
):
"""
@@ -57,57 +60,77 @@ class AzureSentinelLogger(CustomBatchLogger):
If not provided, will use AZURE_SENTINEL_CLIENT_ID or AZURE_CLIENT_ID env var.
client_secret (str, optional): Azure Client Secret for OAuth2 authentication.
If not provided, will use AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET env var.
audit_stream_name (str, optional): Stream name from DCR for audit logs.
If not provided, audit logs use the standard stream name.
"""
self.async_httpx_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
self.dcr_immutable_id = dcr_immutable_id or os.getenv(
resolved_dcr_immutable_id = dcr_immutable_id or os.getenv(
"AZURE_SENTINEL_DCR_IMMUTABLE_ID"
)
self.stream_name = stream_name or os.getenv(
"AZURE_SENTINEL_STREAM_NAME", "Custom-LiteLLM"
resolved_stream_name = (
stream_name or os.getenv("AZURE_SENTINEL_STREAM_NAME") or "Custom-LiteLLM"
)
self.endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT")
self.tenant_id = (
resolved_audit_stream_name = audit_stream_name or resolved_stream_name
resolved_endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT")
resolved_tenant_id = (
tenant_id
or os.getenv("AZURE_SENTINEL_TENANT_ID")
or os.getenv("AZURE_TENANT_ID")
)
self.client_id = (
resolved_client_id = (
client_id
or os.getenv("AZURE_SENTINEL_CLIENT_ID")
or os.getenv("AZURE_CLIENT_ID")
)
self.client_secret = (
resolved_client_secret = (
client_secret
or os.getenv("AZURE_SENTINEL_CLIENT_SECRET")
or os.getenv("AZURE_CLIENT_SECRET")
)
if not self.dcr_immutable_id:
if not resolved_dcr_immutable_id:
raise ValueError(
"AZURE_SENTINEL_DCR_IMMUTABLE_ID is required. Set it as an environment variable or pass dcr_immutable_id parameter."
)
if not self.endpoint:
if not resolved_endpoint:
raise ValueError(
"AZURE_SENTINEL_ENDPOINT is required. Set it as an environment variable or pass endpoint parameter."
)
if not self.tenant_id:
if not resolved_tenant_id:
raise ValueError(
"AZURE_SENTINEL_TENANT_ID or AZURE_TENANT_ID is required. Set it as an environment variable or pass tenant_id parameter."
)
if not self.client_id:
if not resolved_client_id:
raise ValueError(
"AZURE_SENTINEL_CLIENT_ID or AZURE_CLIENT_ID is required. Set it as an environment variable or pass client_id parameter."
)
if not self.client_secret:
if not resolved_client_secret:
raise ValueError(
"AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET is required. Set it as an environment variable or pass client_secret parameter."
)
self.dcr_immutable_id = resolved_dcr_immutable_id
self.stream_name = resolved_stream_name
self.audit_stream_name = resolved_audit_stream_name
self.endpoint = resolved_endpoint
self.tenant_id = resolved_tenant_id
self.client_id = resolved_client_id
self.client_secret = resolved_client_secret
# Build API endpoint: {Endpoint}/dataCollectionRules/{DCR Immutable ID}/streams/{Stream Name}?api-version=2023-01-01
self.api_endpoint = f"{self.endpoint.rstrip('/')}/dataCollectionRules/{self.dcr_immutable_id}/streams/{self.stream_name}?api-version=2023-01-01"
self.api_endpoint = self._build_api_endpoint(
endpoint=resolved_endpoint,
dcr_immutable_id=resolved_dcr_immutable_id,
stream_name=resolved_stream_name,
)
self.audit_api_endpoint = self._build_api_endpoint(
endpoint=resolved_endpoint,
dcr_immutable_id=resolved_dcr_immutable_id,
stream_name=resolved_audit_stream_name,
)
# OAuth2 scope for Azure Monitor
self.oauth_scope = "https://monitor.azure.com/.default"
@@ -118,6 +141,13 @@ class AzureSentinelLogger(CustomBatchLogger):
super().__init__(**kwargs, flush_lock=self.flush_lock)
asyncio.create_task(self.periodic_flush())
self.log_queue: List[StandardLoggingPayload] = []
self.audit_log_queue: List[StandardAuditLogPayload] = []
@staticmethod
def _build_api_endpoint(
endpoint: str, dcr_immutable_id: str, stream_name: str
) -> str:
return f"{endpoint.rstrip('/')}/dataCollectionRules/{dcr_immutable_id}/streams/{stream_name}?api-version=2023-01-01"
async def _get_oauth_token(self) -> str:
"""
@@ -126,9 +156,6 @@ class AzureSentinelLogger(CustomBatchLogger):
Returns:
Bearer token string
"""
# Check if we have a valid cached token
import time
if (
self.oauth_token
and self.oauth_token_expires_at
@@ -170,9 +197,6 @@ class AzureSentinelLogger(CustomBatchLogger):
if not self.oauth_token:
raise Exception("OAuth2 token response did not contain access_token")
# Cache token expiry time
import time
self.oauth_token_expires_at = time.time() + expires_in
return self.oauth_token
@@ -246,6 +270,34 @@ class AzureSentinelLogger(CustomBatchLogger):
)
pass
async def async_log_audit_log_event(
self, audit_log: StandardAuditLogPayload
) -> None:
"""
Async log LiteLLM audit log events to Azure Sentinel.
Audit logs are queued separately from standard LLM logs so mixed callback
usage never sends schema-mismatched records in the same ingestion batch.
"""
try:
verbose_logger.debug(
"Azure Sentinel: Logging audit event id=%s action=%s table=%s",
audit_log.get("id"),
audit_log.get("action"),
audit_log.get("table_name"),
)
self.audit_log_queue.append(audit_log)
if len(self.audit_log_queue) >= self.batch_size:
await self.async_send_audit_batch()
except Exception as e:
verbose_logger.exception(
f"Azure Sentinel Audit Log Layer Error - {str(e)}\n{traceback.format_exc()}"
)
pass
async def async_send_batch(self):
"""
Sends the batch of logs to Azure Monitor Logs Ingestion API
@@ -253,22 +305,42 @@ class AzureSentinelLogger(CustomBatchLogger):
Raises:
Raises a NON Blocking verbose_logger.exception if an error occurs
"""
await self._async_send_batch_to_api(
log_queue=self.log_queue,
api_endpoint=self.api_endpoint,
log_type="logs",
)
async def async_send_audit_batch(self):
"""
Sends the batch of audit logs to Azure Monitor Logs Ingestion API
"""
await self._async_send_batch_to_api(
log_queue=self.audit_log_queue,
api_endpoint=self.audit_api_endpoint,
log_type="audit logs",
)
async def _async_send_batch_to_api(
self,
log_queue: List[Union[StandardLoggingPayload, StandardAuditLogPayload]],
api_endpoint: str,
log_type: str,
) -> None:
try:
if not self.log_queue:
if not log_queue:
return
verbose_logger.debug(
"Azure Sentinel - about to flush %s events", len(self.log_queue)
"Azure Sentinel - about to flush %s %s", len(log_queue), log_type
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
# Get OAuth2 token
bearer_token = await self._get_oauth_token()
# Convert log queue to JSON array format expected by Logs Ingestion API
# Each log entry should be a JSON object in the array
body = safe_dumps(self.log_queue)
body = safe_dumps(log_queue)
# Set headers for Logs Ingestion API
headers = {
@@ -278,7 +350,7 @@ class AzureSentinelLogger(CustomBatchLogger):
# Send the request
response = await self.async_httpx_client.post(
url=self.api_endpoint, data=body.encode("utf-8"), headers=headers
url=api_endpoint, data=body.encode("utf-8"), headers=headers
)
if response.status_code not in [200, 204]:
@@ -301,4 +373,15 @@ class AzureSentinelLogger(CustomBatchLogger):
f"Azure Sentinel Error sending batch API - {str(e)}\n{traceback.format_exc()}"
)
finally:
self.log_queue.clear()
log_queue.clear()
async def flush_queue(self):
if self.flush_lock is None:
return
async with self.flush_lock:
if self.log_queue:
await self.async_send_batch()
if self.audit_log_queue:
await self.async_send_audit_batch()
self.last_flush_time = time.time()
litellm/integrations/opentelemetry.py
+85 -9
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"):
litellm/integrations/prometheus.py
+148 -42
View File
@@ -25,6 +25,9 @@ from typing import (
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import (
BoundedPrometheusSeriesTracker,
)
from litellm.integrations.prometheus_helpers import (
PrometheusLabelFactoryContext,
_get_cached_end_user_id_for_cost_tracking,
@@ -81,6 +84,7 @@ class PrometheusLogger(CustomLogger):
if _custom_buckets is not None
else LATENCY_BUCKETS
)
self._bounded_prometheus_series_tracker = BoundedPrometheusSeriesTracker()
# Create metric factory functions
self._counter_factory = self._create_metric_factory(Counter)
@@ -984,6 +988,40 @@ class PrometheusLogger(CustomLogger):
return filtered_labels
def _track_end_user_metric_series(
self,
metric: Any,
metric_name: DEFINED_PROMETHEUS_METRICS,
labels: Dict[str, Optional[str]],
) -> None:
"""
Cap the cardinality of metrics that include the ``end_user`` label.
Called *after* ``metric.labels(...).inc()/observe()`` so the emission is
recorded in prometheus-client's child map before any eviction runs.
Series that get evicted before the next scrape lose updates accrued
since the last scrape — this is inherent to any cardinality cap.
"""
labelnames = self.get_labels_for_metric(metric_name)
if UserAPIKeyLabelNames.END_USER.value not in labelnames:
return
if labels.get(UserAPIKeyLabelNames.END_USER.value) is None:
return
max_series = litellm.prometheus_end_user_metrics_max_series_per_metric
ttl_seconds = litellm.prometheus_end_user_metrics_ttl_seconds
if max_series is None and ttl_seconds is None:
return
self._bounded_prometheus_series_tracker.track_series(
metric=metric,
metric_name=metric_name,
label_values=tuple(labels.get(label) for label in labelnames),
max_series=max_series,
ttl_seconds=ttl_seconds,
cleanup_interval_seconds=litellm.prometheus_end_user_metrics_cleanup_interval_seconds,
)
def _inc_labeled_counter(
self,
counter: Any,
@@ -998,6 +1036,7 @@ class PrometheusLogger(CustomLogger):
label_context=label_context,
)
counter.labels(**_labels).inc(amount)
self._track_end_user_metric_series(counter, metric_name, _labels)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
# Define prometheus client
@@ -1047,21 +1086,9 @@ class PrometheusLogger(CustomLogger):
output_tokens = standard_logging_payload["completion_tokens"]
tokens_used = standard_logging_payload["total_tokens"]
response_cost = standard_logging_payload["response_cost"]
_requester_metadata: Optional[dict] = standard_logging_payload["metadata"].get(
"requester_metadata"
combined_metadata = _get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload=standard_logging_payload
)
user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[
"metadata"
].get("user_api_key_auth_metadata")
spend_logs_metadata: Optional[dict] = standard_logging_payload["metadata"].get(
"spend_logs_metadata"
)
combined_metadata: Dict[str, Any] = {
**(_requester_metadata if _requester_metadata else {}),
**(user_api_key_auth_metadata if user_api_key_auth_metadata else {}),
**(spend_logs_metadata if spend_logs_metadata else {}),
}
if standard_logging_payload is not None and isinstance(
standard_logging_payload, dict
):
@@ -1416,26 +1443,46 @@ class PrometheusLogger(CustomLogger):
)
remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}"
remaining_requests = (
metadata.get(remaining_requests_variable_name, sys.maxsize) or sys.maxsize
remaining_requests = metadata.get(remaining_requests_variable_name)
if remaining_requests is None:
remaining_requests = sys.maxsize
remaining_tokens = metadata.get(remaining_tokens_variable_name)
if remaining_tokens is None:
remaining_tokens = sys.maxsize
enum_values = UserAPIKeyLabelValues(
hashed_api_key=user_api_key,
api_key_alias=user_api_key_alias,
model=model_group,
model_id=model_id,
custom_metadata_labels=get_custom_labels_from_metadata(
metadata=_get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload=kwargs.get("standard_logging_object")
)
),
)
remaining_tokens = (
metadata.get(remaining_tokens_variable_name, sys.maxsize) or sys.maxsize
label_context = PrometheusLabelFactoryContext(enum_values)
requests_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
"litellm_remaining_api_key_requests_for_model"
),
enum_values=enum_values,
label_context=label_context,
)
self.litellm_remaining_api_key_requests_for_model.labels(**requests_labels).set(
remaining_requests
)
self.litellm_remaining_api_key_requests_for_model.labels(
_sanitize_prometheus_label_value(user_api_key),
_sanitize_prometheus_label_value(user_api_key_alias),
_sanitize_prometheus_label_value(model_group),
_sanitize_prometheus_label_value(model_id),
).set(remaining_requests)
self.litellm_remaining_api_key_tokens_for_model.labels(
_sanitize_prometheus_label_value(user_api_key),
_sanitize_prometheus_label_value(user_api_key_alias),
_sanitize_prometheus_label_value(model_group),
_sanitize_prometheus_label_value(model_id),
).set(remaining_tokens)
tokens_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
"litellm_remaining_api_key_tokens_for_model"
),
enum_values=enum_values,
label_context=label_context,
)
self.litellm_remaining_api_key_tokens_for_model.labels(**tokens_labels).set(
remaining_tokens
)
def _set_latency_metrics(
self,
@@ -1471,6 +1518,11 @@ class PrometheusLogger(CustomLogger):
self.litellm_llm_api_time_to_first_token_metric.labels(
**_ttft_labels
).observe(time_to_first_token_seconds)
self._track_end_user_metric_series(
self.litellm_llm_api_time_to_first_token_metric,
"litellm_llm_api_time_to_first_token_metric",
_ttft_labels,
)
else:
verbose_logger.debug(
"Time to first token metric not emitted, stream option in model_parameters is not True"
@@ -1491,6 +1543,11 @@ class PrometheusLogger(CustomLogger):
self.litellm_llm_api_latency_metric.labels(**_labels).observe(
api_call_total_time_seconds
)
self._track_end_user_metric_series(
self.litellm_llm_api_latency_metric,
"litellm_llm_api_latency_metric",
_labels,
)
# total request latency
total_time_seconds = self._safe_duration_seconds(
@@ -1508,6 +1565,11 @@ class PrometheusLogger(CustomLogger):
self.litellm_request_total_latency_metric.labels(**_labels).observe(
total_time_seconds
)
self._track_end_user_metric_series(
self.litellm_request_total_latency_metric,
"litellm_request_total_latency_metric",
_labels,
)
# request queue time (time from arrival to processing start)
_litellm_params = kwargs.get("litellm_params", {}) or {}
@@ -1525,6 +1587,11 @@ class PrometheusLogger(CustomLogger):
self.litellm_request_queue_time_metric.labels(**_labels).observe(
queue_time_seconds
)
self._track_end_user_metric_series(
self.litellm_request_queue_time_metric,
"litellm_request_queue_time_seconds",
_labels,
)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
verbose_logger.debug(
@@ -1561,18 +1628,27 @@ class PrometheusLogger(CustomLogger):
)
try:
self.litellm_llm_api_failed_requests_metric.labels(
_sanitize_prometheus_label_value(end_user_id),
_sanitize_prometheus_label_value(user_api_key),
_sanitize_prometheus_label_value(user_api_key_alias),
_sanitize_prometheus_label_value(model),
_sanitize_prometheus_label_value(user_api_team),
_sanitize_prometheus_label_value(user_api_team_alias),
_sanitize_prometheus_label_value(user_id),
_sanitize_prometheus_label_value(
standard_logging_payload.get("model_id", "")
enum_values = UserAPIKeyLabelValues(
end_user=end_user_id,
hashed_api_key=user_api_key,
api_key_alias=user_api_key_alias,
model=model,
team=user_api_team,
team_alias=user_api_team_alias,
user=user_id,
model_id=standard_logging_payload.get("model_id", ""),
custom_metadata_labels=get_custom_labels_from_metadata(
metadata=_get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload=standard_logging_payload
)
),
).inc()
)
PrometheusLogger._inc_labeled_counter(
self,
self.litellm_llm_api_failed_requests_metric,
"litellm_llm_api_failed_requests_metric",
enum_values,
)
self.set_llm_deployment_failure_metrics(kwargs)
await self._set_org_budget_metrics_after_api_request(
org_id=user_api_key_org_id,
@@ -3622,6 +3698,36 @@ def get_custom_labels_from_metadata(metadata: dict) -> Dict[str, str]:
return result
def _get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload: Optional[dict],
) -> Dict[str, Any]:
"""
Combine the metadata sources that can supply custom Prometheus labels.
"""
if not isinstance(standard_logging_payload, dict):
return {}
standard_logging_metadata = standard_logging_payload.get("metadata") or {}
if not isinstance(standard_logging_metadata, dict):
return {}
requester_metadata = standard_logging_metadata.get("requester_metadata")
user_api_key_auth_metadata = standard_logging_metadata.get(
"user_api_key_auth_metadata"
)
spend_logs_metadata = standard_logging_metadata.get("spend_logs_metadata")
return {
**(requester_metadata if isinstance(requester_metadata, dict) else {}),
**(
user_api_key_auth_metadata
if isinstance(user_api_key_auth_metadata, dict)
else {}
),
**(spend_logs_metadata if isinstance(spend_logs_metadata, dict) else {}),
}
def _tag_matches_wildcard_configured_pattern(
tags: Sequence[str], configured_tag: str
) -> bool:
litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py
+107
View File
@@ -0,0 +1,107 @@
from __future__ import annotations
import time
from collections import OrderedDict
from threading import RLock
from typing import Any, Dict, Optional
class BoundedPrometheusSeriesTracker:
"""
Tracks Prometheus child series and removes stale/excess labelsets.
The tracker is label-agnostic: callers decide which series should be tracked
and pass the full label tuple used by the Prometheus metric.
"""
def __init__(self) -> None:
self._series: Dict[str, OrderedDict[tuple[Optional[str], ...], float]] = {}
self._last_ttl_cleanup: Dict[str, float] = {}
self.lock = RLock()
def track_series(
self,
metric: Any,
metric_name: str,
label_values: tuple[Optional[str], ...],
max_series: Optional[int],
ttl_seconds: Optional[float],
cleanup_interval_seconds: Optional[float],
) -> None:
if max_series is None and ttl_seconds is None:
return
now = time.monotonic()
with self.lock:
series = self._series.setdefault(metric_name, OrderedDict())
series[label_values] = now
series.move_to_end(label_values)
if ttl_seconds is not None and self._should_run_ttl_cleanup(
metric_name=metric_name,
now=now,
cleanup_interval_seconds=cleanup_interval_seconds,
):
expired_label_values = [
tracked_label_values
for tracked_label_values, last_seen in series.items()
if now - last_seen > ttl_seconds
]
for tracked_label_values in expired_label_values:
self._remove_metric_series(metric, series, tracked_label_values)
# max_series <= 0 is treated as "unlimited" so a misconfigured zero
# value cannot silently drop every emission for this metric.
if max_series is not None and max_series > 0:
while len(series) > max_series:
tracked_label_values = next(iter(series))
if not self._remove_metric_child(metric, tracked_label_values):
break
del series[tracked_label_values]
def _should_run_ttl_cleanup(
self,
metric_name: str,
now: float,
cleanup_interval_seconds: Optional[float],
) -> bool:
if cleanup_interval_seconds is None or cleanup_interval_seconds <= 0:
self._last_ttl_cleanup[metric_name] = now
return True
last_cleanup = self._last_ttl_cleanup.get(metric_name)
if last_cleanup is None or now - last_cleanup >= cleanup_interval_seconds:
self._last_ttl_cleanup[metric_name] = now
return True
return False
def _remove_metric_series(
self,
metric: Any,
series: OrderedDict[tuple[Optional[str], ...], float],
label_values: tuple[Optional[str], ...],
) -> None:
if self._remove_metric_child(metric, label_values):
series.pop(label_values, None)
@staticmethod
def _remove_metric_child(
metric: Any, label_values: tuple[Optional[str], ...]
) -> bool:
"""
Remove the Prometheus child for ``label_values`` and report whether the
tracker should commit the matching state change.
Returns ``True`` when the child is no longer present in Prometheus
(either it was just removed or it was already gone), and ``False`` when
``metric.remove()`` raised an unexpected error and the child likely
still exists.
"""
try:
metric.remove(*label_values)
return True
except KeyError:
return True
except (AttributeError, ValueError):
return False
litellm/integrations/s3_v2.py
+51 -49
View File
@@ -16,6 +16,7 @@ from litellm._logging import print_verbose, verbose_logger
from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS
from litellm.integrations.s3 import get_s3_object_key
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
@@ -53,15 +54,25 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_strip_base64_files: bool = False,
s3_use_key_prefix: bool = False,
s3_use_virtual_hosted_style: bool = False,
s3_callback_params_override: Optional[dict] = None,
**kwargs,
):
try:
verbose_logger.debug(
f"in init s3 logger - s3_callback_params {litellm.s3_callback_params}"
)
_masker = SensitiveDataMasker()
if s3_callback_params_override is not None:
verbose_logger.debug(
f"in init s3 logger (audit override) - "
f"{_masker.mask_dict(dict(s3_callback_params_override))}"
)
else:
verbose_logger.debug(
f"in init s3 logger - s3_callback_params "
f"{_masker.mask_dict(dict(litellm.s3_callback_params or {}))}"
)
# Initialize S3 params first to get the correct s3_verify value
self._init_s3_params(
params_source=s3_callback_params_override,
s3_bucket_name=s3_bucket_name,
s3_region_name=s3_region_name,
s3_api_version=s3_api_version,
@@ -139,94 +150,85 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_strip_base64_files: bool = False,
s3_use_key_prefix: bool = False,
s3_use_virtual_hosted_style: bool = False,
params_source: Optional[dict] = None,
):
"""
Initialize the s3 params for this logging callback
Initialize the s3 params for this logging callback. Reads from
`params_source` if given (e.g. `s3_audit_callback_params` for the
audit-log instance), otherwise falls back to `litellm.s3_callback_params`.
Resolves `os.environ/X` markers into a local dict; never mutates the source.
"""
litellm.s3_callback_params = litellm.s3_callback_params or {}
# read in .env variables - example os.environ/AWS_BUCKET_NAME
for key, value in litellm.s3_callback_params.items():
if isinstance(value, str) and value.startswith("os.environ/"):
litellm.s3_callback_params[key] = litellm.get_secret(value)
if params_source is None:
params_source = litellm.s3_callback_params or {}
params: dict = {
key: (
litellm.get_secret(value)
if isinstance(value, str) and value.startswith("os.environ/")
else value
)
for key, value in params_source.items()
}
self.s3_bucket_name = (
litellm.s3_callback_params.get("s3_bucket_name") or s3_bucket_name
)
self.s3_region_name = (
litellm.s3_callback_params.get("s3_region_name") or s3_region_name
)
self.s3_api_version = (
litellm.s3_callback_params.get("s3_api_version") or s3_api_version
)
self.s3_bucket_name = params.get("s3_bucket_name") or s3_bucket_name
self.s3_region_name = params.get("s3_region_name") or s3_region_name
self.s3_api_version = params.get("s3_api_version") or s3_api_version
self.s3_use_ssl = (
litellm.s3_callback_params.get("s3_use_ssl", True)
if litellm.s3_callback_params.get("s3_use_ssl") is not None
params.get("s3_use_ssl", True)
if params.get("s3_use_ssl") is not None
else s3_use_ssl
)
self.s3_verify = (
litellm.s3_callback_params.get("s3_verify")
if litellm.s3_callback_params.get("s3_verify") is not None
params.get("s3_verify")
if params.get("s3_verify") is not None
else s3_verify
)
self.s3_endpoint_url = (
litellm.s3_callback_params.get("s3_endpoint_url") or s3_endpoint_url
)
self.s3_endpoint_url = params.get("s3_endpoint_url") or s3_endpoint_url
self.s3_aws_access_key_id = (
litellm.s3_callback_params.get("s3_aws_access_key_id")
or s3_aws_access_key_id
params.get("s3_aws_access_key_id") or s3_aws_access_key_id
)
self.s3_aws_secret_access_key = (
litellm.s3_callback_params.get("s3_aws_secret_access_key")
or s3_aws_secret_access_key
params.get("s3_aws_secret_access_key") or s3_aws_secret_access_key
)
self.s3_aws_session_token = (
litellm.s3_callback_params.get("s3_aws_session_token")
or s3_aws_session_token
params.get("s3_aws_session_token") or s3_aws_session_token
)
self.s3_aws_session_name = (
litellm.s3_callback_params.get("s3_aws_session_name") or s3_aws_session_name
params.get("s3_aws_session_name") or s3_aws_session_name
)
self.s3_aws_profile_name = (
litellm.s3_callback_params.get("s3_aws_profile_name") or s3_aws_profile_name
params.get("s3_aws_profile_name") or s3_aws_profile_name
)
self.s3_aws_role_name = (
litellm.s3_callback_params.get("s3_aws_role_name") or s3_aws_role_name
)
self.s3_aws_role_name = params.get("s3_aws_role_name") or s3_aws_role_name
self.s3_aws_web_identity_token = (
litellm.s3_callback_params.get("s3_aws_web_identity_token")
or s3_aws_web_identity_token
params.get("s3_aws_web_identity_token") or s3_aws_web_identity_token
)
self.s3_aws_sts_endpoint = (
litellm.s3_callback_params.get("s3_aws_sts_endpoint") or s3_aws_sts_endpoint
params.get("s3_aws_sts_endpoint") or s3_aws_sts_endpoint
)
self.s3_config = litellm.s3_callback_params.get("s3_config") or s3_config
self.s3_path = litellm.s3_callback_params.get("s3_path") or s3_path
# done reading litellm.s3_callback_params
self.s3_config = params.get("s3_config") or s3_config
self.s3_path = params.get("s3_path") or s3_path
self.s3_use_team_prefix = (
bool(litellm.s3_callback_params.get("s3_use_team_prefix", False))
or s3_use_team_prefix
bool(params.get("s3_use_team_prefix", False)) or s3_use_team_prefix
)
self.s3_use_key_prefix = (
bool(litellm.s3_callback_params.get("s3_use_key_prefix", False))
or s3_use_key_prefix
bool(params.get("s3_use_key_prefix", False)) or s3_use_key_prefix
)
self.s3_strip_base64_files = (
bool(litellm.s3_callback_params.get("s3_strip_base64_files", False))
or s3_strip_base64_files
bool(params.get("s3_strip_base64_files", False)) or s3_strip_base64_files
)
self.s3_use_virtual_hosted_style = (
bool(litellm.s3_callback_params.get("s3_use_virtual_hosted_style", False))
bool(params.get("s3_use_virtual_hosted_style", False))
or s3_use_virtual_hosted_style
)
litellm/litellm_core_utils/custom_logger_registry.py
+2
View File
@@ -14,6 +14,7 @@ from litellm import _custom_logger_compatible_callbacks_literal
from litellm.integrations.agentops import AgentOps
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.integrations.argilla import ArgillaLogger
from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger
from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger
from litellm.integrations.bitbucket import BitBucketPromptManager
from litellm.integrations.braintrust_logging import BraintrustLogger
@@ -73,6 +74,7 @@ class CustomLoggerRegistry:
"opik": OpikLogger,
"argilla": ArgillaLogger,
"opentelemetry": OpenTelemetry,
"azure_sentinel": AzureSentinelLogger,
"azure_storage": AzureBlobStorageLogger,
"humanloop": HumanloopLogger,
# OTEL compatible loggers
litellm/litellm_core_utils/prompt_templates/common_utils.py
+32 -3
View File
@@ -436,12 +436,21 @@ def update_messages_with_model_file_ids(
"""
Updates messages with model file ids.
For managed files (unified file IDs), uses model_file_id_mapping if it
resolves the id, otherwise decodes the base64-encoded unified file ID
and extracts the llm_output_file_id directly. Mirrors the Responses-API
sibling `update_responses_input_with_model_file_ids`.
model_file_id_mapping: Dict[str, Dict[str, str]] = {
"litellm_proxy/file_id": {
"model_id": "provider_file_id"
}
}
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
convert_b64_uid_to_unified_uid,
)
for message in messages:
if message.get("role") == "user":
@@ -450,7 +459,13 @@ def update_messages_with_model_file_ids(
if isinstance(content, str):
continue
for c in content:
if c["type"] == "file":
if not isinstance(c, dict):
# Content list items aren't always dicts. e.g.
# text_completion forwards a token-ids list/list-of-
# lists through this path. Skip non-dict items
# instead of indexing into them.
continue
if c.get("type") == "file":
file_object = cast(ChatCompletionFileObject, c)
file_object_file_field = file_object.get("file")
if not isinstance(file_object_file_field, dict):
@@ -468,9 +483,23 @@ def update_messages_with_model_file_ids(
if file_id:
provider_file_id = (
model_file_id_mapping.get(file_id, {}).get(model_id)
or file_id
if model_file_id_mapping
else None
)
if (
not provider_file_id
and _is_base64_encoded_unified_file_id(file_id)
):
unified_file_id = convert_b64_uid_to_unified_uid(
file_id
)
if "llm_output_file_id," in unified_file_id:
provider_file_id = unified_file_id.split(
"llm_output_file_id,"
)[1].split(";")[0]
file_object_file_field["file_id"] = (
provider_file_id or file_id
)
file_object_file_field["file_id"] = provider_file_id
if format:
file_object_file_field["format"] = format
return messages
litellm/llms/anthropic/chat/guardrail_translation/handler.py
+11 -1
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)
litellm/llms/anthropic/chat/handler.py
+14 -2
View File
@@ -579,6 +579,10 @@ class ModelResponseIterator:
# Accumulate compaction blocks for multi-turn reconstruction
self.compaction_blocks: List[Dict[str, Any]] = []
# Accumulate streamed thinking text so final usage can split reasoning
# tokens from regular output tokens.
self.reasoning_content_chunks: List[str] = []
# Track server tool use inputs and results for code_interpreter_results
self._server_tool_inputs: Dict[str, Any] = {}
self.tool_results: List[Dict[str, Any]] = []
@@ -609,9 +613,14 @@ class ModelResponseIterator:
return False
def _handle_usage(self, anthropic_usage_chunk: Union[dict, UsageDelta]) -> Usage:
reasoning_content = (
"".join(self.reasoning_content_chunks)
if self.reasoning_content_chunks
else None
)
return AnthropicConfig().calculate_usage(
usage_object=cast(dict, anthropic_usage_chunk),
reasoning_content=None,
reasoning_content=reasoning_content,
speed=self.speed,
)
@@ -658,10 +667,13 @@ class ModelResponseIterator:
"thinking" in content_block["delta"]
or "signature" in content_block["delta"]
):
thinking_content = content_block["delta"].get("thinking")
if isinstance(thinking_content, str) and thinking_content:
self.reasoning_content_chunks.append(thinking_content)
thinking_blocks = [
ChatCompletionThinkingBlock(
type="thinking",
thinking=content_block["delta"].get("thinking") or "",
thinking=thinking_content or "",
signature=str(content_block["delta"].get("signature") or ""),
)
]
litellm/llms/anthropic/chat/transformation.py
+12 -3
View File
@@ -2156,8 +2156,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
speed: Optional[str] = None,
) -> Usage:
# NOTE: Sometimes the usage object has None set explicitly for token counts, meaning .get() & key access returns None, and we need to account for this
prompt_tokens = usage_object.get("input_tokens", 0) or 0
completion_tokens = usage_object.get("output_tokens", 0) or 0
raw_prompt_tokens = usage_object.get("input_tokens", 0) or 0
prompt_tokens: int = (
int(raw_prompt_tokens) if isinstance(raw_prompt_tokens, (int, float)) else 0
)
raw_completion_tokens = usage_object.get("output_tokens", 0) or 0
completion_tokens: int = (
int(raw_completion_tokens)
if isinstance(raw_completion_tokens, (int, float))
else 0
)
_usage = usage_object
cache_creation_input_tokens: int = 0
cache_read_input_tokens: int = 0
@@ -2226,11 +2234,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
text_tokens=raw_input_tokens,
)
# Always populate completion_token_details, not just when there's reasoning_content
reasoning_tokens = (
estimated_reasoning_tokens = (
token_counter(text=reasoning_content, count_response_tokens=True)
if reasoning_content
else 0
)
reasoning_tokens = min(estimated_reasoning_tokens, completion_tokens)
completion_token_details = CompletionTokensDetailsWrapper(
reasoning_tokens=reasoning_tokens if reasoning_tokens > 0 else 0,
text_tokens=(
litellm/llms/base_llm/guardrail_translation/utils.py
+15
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"]
litellm/llms/bedrock/chat/invoke_agent/transformation.py
+2 -22
View File
@@ -299,29 +299,9 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
)
def _get_response_stream_shape(self):
"""Get the response stream shape for parsing, reusing existing logic."""
try:
# Try to reuse the cached shape from the existing decoder
from litellm.llms.bedrock.chat.invoke_handler import (
get_response_stream_shape,
)
from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE
return get_response_stream_shape()
except ImportError:
# Fallback: create our own shape
try:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
bedrock_service_dict = loader.load_service_model(
"bedrock-runtime", "service-2"
)
bedrock_service_model = ServiceModel(bedrock_service_dict)
return bedrock_service_model.shape_for("ResponseStream")
except Exception as e:
verbose_logger.warning(f"Could not load response stream shape: {e}")
return None
return BEDROCK_RESPONSE_STREAM_SHAPE
def _extract_response_content(self, events: InvokeAgentEventList) -> str:
"""Extract the final response content from parsed events."""
litellm/llms/bedrock/chat/invoke_handler.py
+17 -17
View File
@@ -67,9 +67,13 @@ from litellm.types.utils import (
from litellm.utils import CustomStreamWrapper, get_secret
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError, ModelResponseIterator, get_bedrock_tool_name
from ..common_utils import (
BEDROCK_RESPONSE_STREAM_SHAPE,
BedrockError,
ModelResponseIterator,
get_bedrock_tool_name,
)
_response_stream_shape_cache = None
bedrock_tool_name_mappings: InMemoryCache = InMemoryCache(
max_size_in_memory=50, default_ttl=600
)
@@ -1391,20 +1395,6 @@ class BedrockLLM(BaseAWSLLM):
return None
def get_response_stream_shape():
global _response_stream_shape_cache
if _response_stream_shape_cache is None:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
bedrock_service_dict = loader.load_service_model("bedrock-runtime", "service-2")
bedrock_service_model = ServiceModel(bedrock_service_dict)
_response_stream_shape_cache = bedrock_service_model.shape_for("ResponseStream")
return _response_stream_shape_cache
class AWSEventStreamDecoder:
def __init__(self, model: str, json_mode: Optional[bool] = False) -> None:
from botocore.parsers import EventStreamJSONParser
@@ -1838,8 +1828,18 @@ class AWSEventStreamDecoder:
yield self._chunk_parser(chunk_data=_data)
def _parse_message_from_event(self, event) -> Optional[str]:
if BEDROCK_RESPONSE_STREAM_SHAPE is None:
raise BedrockError(
status_code=500,
message=(
"Bedrock event-stream shape could not be loaded from botocore. "
"Ensure botocore is correctly installed."
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(response_dict, get_response_stream_shape())
parsed_response = self.parser.parse(
response_dict, BEDROCK_RESPONSE_STREAM_SHAPE
)
if response_dict["status_code"] != 200:
decoded_body = response_dict["body"].decode()
litellm/llms/bedrock/common_utils.py
+39 -19
View File
@@ -14,6 +14,7 @@ if TYPE_CHECKING:
import httpx
import litellm
from litellm import verbose_logger
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
@@ -917,38 +918,57 @@ def get_bedrock_chat_config(model: str):
return litellm.AmazonInvokeConfig()
def _load_bedrock_response_stream_shape():
"""
Load the ResponseStream shape from botocore's bundled bedrock-runtime schema.
Called once at module import time; the result is stored in
``BEDROCK_RESPONSE_STREAM_SHAPE`` and reused for the process lifetime.
Returns ``None`` if botocore is unavailable or the service model cannot be
loaded, so the module still imports cleanly.
"""
try:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
service_dict = loader.load_service_model("bedrock-runtime", "service-2")
return ServiceModel(service_dict).shape_for("ResponseStream")
except Exception as e:
verbose_logger.warning(
"litellm: could not pre-load bedrock-runtime response stream shape "
"— Bedrock event-stream decoding will be unavailable. Error: %s",
e,
)
return None
# Eagerly resolved once per process — avoids per-instance or per-request disk I/O.
BEDROCK_RESPONSE_STREAM_SHAPE = _load_bedrock_response_stream_shape()
class BedrockEventStreamDecoderBase:
"""
Base class for event stream decoding for Bedrock
"""
_response_stream_shape_cache = None
def __init__(self):
from botocore.parsers import EventStreamJSONParser
self.parser = EventStreamJSONParser()
def get_response_stream_shape(self):
if self._response_stream_shape_cache is None:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
bedrock_service_dict = loader.load_service_model(
"bedrock-runtime", "service-2"
)
bedrock_service_model = ServiceModel(bedrock_service_dict)
self._response_stream_shape_cache = bedrock_service_model.shape_for(
"ResponseStream"
)
return self._response_stream_shape_cache
def _parse_message_from_event(self, event) -> Optional[str]:
if BEDROCK_RESPONSE_STREAM_SHAPE is None:
raise BedrockError(
status_code=500,
message=(
"Bedrock event-stream shape could not be loaded from botocore. "
"Ensure botocore is correctly installed."
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(
response_dict, self.get_response_stream_shape()
response_dict, BEDROCK_RESPONSE_STREAM_SHAPE
)
if response_dict["status_code"] != 200:
litellm/llms/custom_httpx/http_handler.py
+47 -6
View File
@@ -1,4 +1,5 @@
import asyncio
import concurrent.futures
import inspect
import os
import socket
@@ -133,6 +134,11 @@ _DEFAULT_TIMEOUT = httpx.Timeout(
timeout=COMPLETION_HTTP_FALLBACK_SECONDS,
connect=HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS,
)
_STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS = 5.0
_STREAMING_ERROR_BODY_READ_EXECUTOR = concurrent.futures.ThreadPoolExecutor(
max_workers=50,
thread_name_prefix="litellm-streaming-error-body-read",
)
def _prepare_request_data_and_content(
@@ -386,17 +392,30 @@ def _safe_get_response_text(response: httpx.Response) -> str:
return ""
async def _safe_aread_response(response: httpx.Response) -> bytes:
async def _safe_aread_response(
response: httpx.Response, timeout: Optional[float] = None
) -> bytes:
"""Safely read async response body, falling back to empty bytes on errors."""
try:
if timeout is not None:
return await asyncio.wait_for(response.aread(), timeout=timeout)
return await response.aread()
except Exception:
return b""
def _safe_read_response(response: httpx.Response) -> bytes:
def _safe_read_response(
response: httpx.Response, timeout: Optional[float] = None
) -> bytes:
"""Safely read sync response body, falling back to empty bytes on errors."""
try:
if timeout is not None:
future = _STREAMING_ERROR_BODY_READ_EXECUTOR.submit(response.read)
try:
return future.result(timeout=timeout)
except Exception:
response.close()
return b""
return response.read()
except Exception:
return b""
@@ -405,8 +424,19 @@ def _safe_read_response(response: httpx.Response) -> bytes:
def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None:
"""Raise a MaskedHTTPStatusError for sync HTTP handlers."""
if stream:
_body = mask_sensitive_info(_safe_read_response(e.response))
raise MaskedHTTPStatusError(e, message=_body, text=_body) from None
try:
_body = mask_sensitive_info(
_safe_read_response(
e.response,
timeout=_STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS,
)
)
raise MaskedHTTPStatusError(e, message=_body, text=_body) from None
finally:
try:
e.response.close()
except Exception:
pass
_text = mask_sensitive_info(_safe_get_response_text(e.response))
raise MaskedHTTPStatusError(e, message=_text, text=_text) from None
@@ -414,8 +444,19 @@ def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None:
async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> None:
"""Raise a MaskedHTTPStatusError for async HTTP handlers."""
if stream:
_body = mask_sensitive_info(await _safe_aread_response(e.response))
raise MaskedHTTPStatusError(e, message=_body, text=_body) from None
try:
_body = mask_sensitive_info(
await _safe_aread_response(
e.response,
timeout=_STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS,
)
)
raise MaskedHTTPStatusError(e, message=_body, text=_body) from None
finally:
try:
await e.response.aclose()
except Exception:
pass
_text = mask_sensitive_info(_safe_get_response_text(e.response))
raise MaskedHTTPStatusError(e, message=_text, text=_text) from None
litellm/llms/openai/chat/guardrail_translation/handler.py
+18 -6
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)
litellm/llms/sagemaker/common_utils.py
+32 -19
View File
@@ -9,7 +9,27 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.utils import GenericStreamingChunk as GChunk
from litellm.types.utils import StreamingChatCompletionChunk
_response_stream_shape_cache = None
def _load_sagemaker_response_stream_shape():
try:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
service_dict = loader.load_service_model("sagemaker-runtime", "service-2")
return ServiceModel(service_dict).shape_for(
"InvokeEndpointWithResponseStreamOutput"
)
except Exception as e:
verbose_logger.warning(
"litellm: could not pre-load sagemaker-runtime response stream shape "
"— SageMaker event-stream decoding will be unavailable. Error: %s",
e,
)
return None
SAGEMAKER_RESPONSE_STREAM_SHAPE = _load_sagemaker_response_stream_shape()
class SagemakerError(BaseLLMException):
@@ -187,8 +207,18 @@ class AWSEventStreamDecoder:
verbose_logger.error(f"Final error parsing accumulated JSON: {e}")
def _parse_message_from_event(self, event) -> Optional[str]:
if SAGEMAKER_RESPONSE_STREAM_SHAPE is None:
raise SagemakerError(
status_code=500,
message=(
"SageMaker event-stream shape could not be loaded from botocore. "
"Ensure botocore is correctly installed."
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(response_dict, get_response_stream_shape())
parsed_response = self.parser.parse(
response_dict, SAGEMAKER_RESPONSE_STREAM_SHAPE
)
if response_dict["status_code"] != 200:
raise ValueError(f"Bad response code, expected 200: {response_dict}")
@@ -204,20 +234,3 @@ class AWSEventStreamDecoder:
return None
return chunk.decode() # type: ignore[no-any-return]
def get_response_stream_shape():
global _response_stream_shape_cache
if _response_stream_shape_cache is None:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
sagemaker_service_dict = loader.load_service_model(
"sagemaker-runtime", "service-2"
)
sagemaker_service_model = ServiceModel(sagemaker_service_dict)
_response_stream_shape_cache = sagemaker_service_model.shape_for(
"InvokeEndpointWithResponseStreamOutput"
)
return _response_stream_shape_cache
litellm/main.py
+8 -8
View File
@@ -1461,14 +1461,14 @@ def completion( # type: ignore # noqa: PLR0915
if eos_token:
custom_prompt_dict[model]["eos_token"] = eos_token
if kwargs.get("model_file_id_mapping"):
messages = update_messages_with_model_file_ids(
messages=messages,
model_id=kwargs.get("model_info", {}).get("id", None),
model_file_id_mapping=cast(
Dict[str, Dict[str, str]], kwargs.get("model_file_id_mapping")
),
)
messages = update_messages_with_model_file_ids(
messages=messages,
model_id=kwargs.get("model_info", {}).get("id", None),
model_file_id_mapping=cast(
Dict[str, Dict[str, str]],
kwargs.get("model_file_id_mapping") or {},
),
)
provider_config: Optional[BaseConfig] = None
if custom_llm_provider is not None and custom_llm_provider in [
litellm/model_prices_and_context_window_backup.json
+69
View File
@@ -27187,6 +27187,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",
@@ -28874,6 +28888,19 @@
"mode": "chat",
"output_cost_per_token": 0.0
},
"sambanova/MiniMax-M2.7": {
"input_cost_per_token": 3e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 204800,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"source": "https://cloud.sambanova.ai/plans/pricing",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"sambanova/DeepSeek-R1": {
"input_cost_per_token": 5e-06,
"litellm_provider": "sambanova",
@@ -34927,6 +34954,48 @@
"supports_vision": true,
"supports_web_search": true
},
"xai/grok-4.3": {
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_above_200k_tokens": 2.5e-06,
"litellm_provider": "xai",
"max_input_tokens": 1000000,
"max_output_tokens": 1000000,
"max_tokens": 1000000,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"output_cost_per_token_above_200k_tokens": 5e-06,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"xai/grok-4.3-latest": {
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_above_200k_tokens": 2.5e-06,
"litellm_provider": "xai",
"max_input_tokens": 1000000,
"max_output_tokens": 1000000,
"max_tokens": 1000000,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"output_cost_per_token_above_200k_tokens": 5e-06,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"xai/grok-beta": {
"input_cost_per_token": 5e-06,
"litellm_provider": "xai",
litellm/proxy/_experimental/mcp_server/auth/token_exchange.py
+196
View File
@@ -0,0 +1,196 @@
"""
OAuth 2.0 Token Exchange (RFC 8693) handler for MCP servers.
Exchanges a user's incoming JWT (subject_token) for a scoped access token
at an IDP's token exchange endpoint. The exchanged token is then used to
authenticate requests to the upstream MCP server.
See: https://datatracker.ietf.org/doc/html/rfc8693
"""
import asyncio
import hashlib
import weakref
from typing import TYPE_CHECKING, Dict, Tuple
import httpx
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import (
MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL,
MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS,
MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE,
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
if TYPE_CHECKING:
from litellm.types.mcp_server.mcp_server_manager import MCPServer
# RFC 8693 grant type constant
TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
DEFAULT_SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token"
class TokenExchangeHandler:
"""Handles OAuth 2.0 Token Exchange (RFC 8693) for MCP servers.
Caches exchanged tokens keyed by ``hash(subject_token + server_id)`` so
repeated calls with the same user token skip the IDP round-trip.
"""
def __init__(self) -> None:
self._cache = InMemoryCache(
max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE,
default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
)
# WeakValueDictionary so locks are GC'd once no coroutine holds a reference,
# preventing unbounded growth with many rotating user tokens.
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = (
weakref.WeakValueDictionary()
)
def _get_lock(self, cache_key: str) -> asyncio.Lock:
lock = self._locks.get(cache_key)
if lock is None:
lock = asyncio.Lock()
self._locks[cache_key] = lock
return lock
@staticmethod
def _cache_key(subject_token: str, server_id: str) -> str:
raw = f"{subject_token}:{server_id}"
return hashlib.sha256(raw.encode()).hexdigest()
async def exchange_token(
self,
subject_token: str,
server: "MCPServer",
) -> str:
"""Exchange *subject_token* for a scoped access token.
Returns the exchanged ``access_token`` string (suitable for a
``Bearer`` header).
Raises ``ValueError`` on configuration or IDP errors.
"""
cache_key = self._cache_key(subject_token, server.server_id)
# Fast path
cached = self._cache.get_cache(cache_key)
if cached is not None:
return cached
# Slow path — one exchange at a time per (user, server) pair
async with self._get_lock(cache_key):
cached = self._cache.get_cache(cache_key)
if cached is not None:
return cached
token, ttl = await self._do_exchange(subject_token, server)
self._cache.set_cache(cache_key, token, ttl=ttl)
return token
async def _do_exchange(
self,
subject_token: str,
server: "MCPServer",
) -> Tuple[str, int]:
"""POST to the token exchange endpoint with RFC 8693 parameters.
Returns ``(access_token, ttl_seconds)``.
"""
endpoint = server.token_exchange_endpoint or server.token_url
if not endpoint:
raise ValueError(
f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange "
f"but no token_exchange_endpoint or token_url configured"
)
if not server.client_id or not server.client_secret:
raise ValueError(
f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange "
f"but missing client_id or client_secret"
)
data: Dict[str, str] = {
"grant_type": TOKEN_EXCHANGE_GRANT_TYPE,
"subject_token": subject_token,
"subject_token_type": server.subject_token_type
or DEFAULT_SUBJECT_TOKEN_TYPE,
"client_id": server.client_id,
"client_secret": server.client_secret,
}
if server.audience:
data["audience"] = server.audience
if server.scopes:
data["scope"] = " ".join(server.scopes)
verbose_logger.debug(
"Exchanging token for MCP server %s at %s (audience=%s)",
server.server_id,
endpoint,
server.audience,
)
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
try:
response = await client.post(endpoint, data=data)
response.raise_for_status()
except httpx.HTTPStatusError as exc:
verbose_logger.debug(
"Token exchange IDP error for MCP server %s (status %d)",
server.server_id,
exc.response.status_code,
)
raise ValueError(
f"Token exchange for MCP server '{server.server_id}' "
f"failed with status {exc.response.status_code}"
) from exc
body = response.json()
if not isinstance(body, dict):
raise ValueError(
f"Token exchange response for MCP server '{server.server_id}' "
f"returned non-object JSON (got {type(body).__name__})"
)
access_token = body.get("access_token")
if not access_token:
raise ValueError(
f"Token exchange response for MCP server '{server.server_id}' "
f"missing 'access_token'"
)
raw_expires_in = body.get("expires_in")
try:
expires_in = (
int(raw_expires_in)
if raw_expires_in is not None
else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
)
except (TypeError, ValueError):
expires_in = MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
ttl = max(
expires_in - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS,
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL,
)
verbose_logger.info(
"Token exchange succeeded for MCP server %s (expires in %ds)",
server.server_id,
expires_in,
)
return access_token, ttl
def invalidate(self, subject_token: str, server_id: str) -> None:
"""Remove a cached exchanged token (e.g. after a 401)."""
cache_key = self._cache_key(subject_token, server_id)
self._cache.delete_cache(cache_key)
# Module-level singleton
mcp_token_exchange_handler = TokenExchangeHandler()
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+132 -45
View File
@@ -411,6 +411,15 @@ class MCPServerManager:
aws_role_name=server_config.get("aws_role_name", None),
aws_session_name=server_config.get("aws_session_name", None),
instructions=server_config.get("instructions", None),
# Token Exchange (OBO) fields
token_exchange_endpoint=server_config.get(
"token_exchange_endpoint", None
),
audience=server_config.get("audience", None),
subject_token_type=server_config.get(
"subject_token_type",
"urn:ietf:params:oauth:token-type:access_token",
),
)
self._assign_unique_short_prefix(new_server)
self.config_mcp_servers[server_id] = new_server
@@ -497,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(
@@ -765,10 +775,23 @@ class MCPServerManager:
aws_role_name=aws_creds.get("aws_role_name"),
aws_session_name=aws_creds.get("aws_session_name"),
instructions=mcp_server.instructions,
# Token Exchange (OBO) fields — read from credentials JSON blob
token_exchange_endpoint=(
credentials_dict.get("token_exchange_endpoint")
if credentials_dict
else None
),
audience=(credentials_dict.get("audience") if credentials_dict else None),
subject_token_type=(
credentials_dict.get("subject_token_type") if credentials_dict else None
)
or "urn:ietf:params:oauth:token-type:access_token",
)
return new_server
async def _maybe_register_openapi_tools(self, server: MCPServer):
async def _maybe_register_openapi_tools(
self, server: MCPServer, *, initialize_mapping: bool = True
):
"""Register OpenAPI tools if the server has a spec_path configured."""
if server.spec_path:
verbose_logger.info(
@@ -779,7 +802,8 @@ class MCPServerManager:
server=server,
base_url=server.url or "",
)
self.initialize_tool_name_to_mcp_server_name_mapping()
if initialize_mapping:
self.initialize_tool_name_to_mcp_server_name_mapping()
async def add_server(self, mcp_server: LiteLLM_MCPServerTable):
try:
@@ -1136,6 +1160,29 @@ class MCPServerManager:
#########################################################
# Methods that call the upstream MCP servers
#########################################################
@staticmethod
def _extract_bearer_token(
oauth2_headers: Optional[Dict[str, str]],
raw_headers: Optional[Dict[str, str]],
) -> Optional[str]:
"""Extract the bare Bearer token from oauth2_headers or raw_headers.
Returns the token string without the ``Bearer `` prefix, or ``None``
if no Authorization header is found.
"""
auth_value: Optional[str] = None
if oauth2_headers and "Authorization" in oauth2_headers:
auth_value = oauth2_headers["Authorization"]
elif raw_headers:
# raw_headers may have lowercase keys depending on the ASGI server
normalized = {k.lower(): v for k, v in raw_headers.items()}
auth_value = normalized.get("authorization")
if auth_value:
if auth_value.startswith("Bearer "):
return auth_value[len("Bearer ") :]
return auth_value
return None
def _build_stdio_env(
self,
server: MCPServer,
@@ -1169,25 +1216,30 @@ class MCPServerManager:
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
extra_headers: Optional[Dict[str, str]] = None,
stdio_env: Optional[Dict[str, str]] = None,
subject_token: Optional[str] = None,
) -> MCPClient:
"""
Create an MCPClient instance for the given server.
Auth resolution (single place for all auth logic):
1. ``mcp_auth_header`` — per-request/per-user override
2. OAuth2 client_credentials token — auto-fetched and cached
3. ``server.authentication_token`` — static token from config/DB
2. OAuth2 Token Exchange (OBO) — exchange user token for scoped token
3. OAuth2 client_credentials token — auto-fetched and cached
4. ``server.authentication_token`` — static token from config/DB
Args:
server: The server configuration.
mcp_auth_header: Optional per-request auth override.
extra_headers: Additional headers to forward.
stdio_env: Environment variables for stdio transport.
subject_token: Optional user JWT for token exchange (OBO) flow.
Returns:
Configured MCP client instance.
"""
auth_value = await resolve_mcp_auth(server, mcp_auth_header)
auth_value = await resolve_mcp_auth(
server, mcp_auth_header, subject_token=subject_token
)
transport = server.transport or MCPTransport.sse
@@ -1978,7 +2030,11 @@ class MCPServerManager:
_SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024
def _assign_unique_short_prefix(self, server: MCPServer) -> None:
def _assign_unique_short_prefix(
self,
server: MCPServer,
registry: Optional[Dict[str, MCPServer]] = None,
) -> None:
"""Resolve and cache a collision-free short tool prefix on ``server``.
Called at registration time for every MCP server entering the
@@ -2002,7 +2058,8 @@ class MCPServerManager:
return
used: Dict[str, str] = {}
for other in self.get_registry().values():
registry_for_collision_check = registry or self.get_registry()
for other in registry_for_collision_check.values():
if other.server_id == server.server_id:
continue
if other.short_prefix:
@@ -2534,9 +2591,12 @@ class MCPServerManager:
if server_auth_header is None:
server_auth_header = mcp_auth_header
# oauth2 headers
# Extract subject token for OAuth2 Token Exchange (OBO) flow
subject_token: Optional[str] = None
extra_headers: Optional[Dict[str, str]] = None
if mcp_server.auth_type == MCPAuth.oauth2:
if mcp_server.auth_type == MCPAuth.oauth2_token_exchange:
subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
elif mcp_server.auth_type == MCPAuth.oauth2:
if mcp_server.has_client_credentials:
# For M2M OAuth servers, Authorization must come from token fetch.
extra_headers = None
@@ -2604,6 +2664,7 @@ class MCPServerManager:
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
)
call_tool_params = MCPCallToolRequestParams(
@@ -2916,46 +2977,72 @@ class MCPServerManager:
# against the *full* set so dedup is deterministic regardless of
# iteration order.
for server in db_mcp_servers:
existing_server = previous_registry.get(server.server_id)
try:
existing_server = previous_registry.get(server.server_id)
if (
existing_server is not None
and existing_server.updated_at is not None
and server.updated_at is not None
and existing_server.updated_at == server.updated_at
):
# Re-use existing server instance to avoid re-running build_mcp_server_from_table()
# which can perform network discovery for OAuth2 servers.
new_registry[server.server_id] = existing_server
continue
if (
existing_server is not None
and existing_server.updated_at is not None
and server.updated_at is not None
and existing_server.updated_at == server.updated_at
):
# Re-use existing server instance to avoid re-running build_mcp_server_from_table()
# which can perform network discovery for OAuth2 servers.
new_registry[server.server_id] = existing_server
continue
_warn_on_server_name_fields(
server_id=server.server_id,
alias=getattr(server, "alias", None),
server_name=getattr(server, "server_name", None),
)
verbose_logger.debug(
f"Building server from DB: {server.server_id} ({server.server_name})"
)
new_server = await self.build_mcp_server_from_table(server)
# Carry the cached short_prefix from the previous registry entry
# (if any) so the prefix is stable across reloads.
if existing_server is not None and existing_server.short_prefix:
new_server.short_prefix = existing_server.short_prefix
new_registry[server.server_id] = new_server
_warn_on_server_name_fields(
server_id=server.server_id,
alias=getattr(server, "alias", None),
server_name=getattr(server, "server_name", None),
)
verbose_logger.debug(
f"Building server from DB: {server.server_id} ({server.server_name})"
)
new_server = await self.build_mcp_server_from_table(server)
# Carry the cached short_prefix from the previous registry entry
# (if any) so the prefix is stable across reloads.
if existing_server is not None and existing_server.short_prefix:
new_server.short_prefix = existing_server.short_prefix
new_registry[server.server_id] = new_server
except Exception as e:
verbose_logger.exception(
"Skipping MCP server %s (%s) during DB reload: %s",
server.server_id,
getattr(server, "alias", None),
e,
)
# Swap in the new registry first so _assign_unique_short_prefix
# sees the complete set when checking for collisions.
self.registry = new_registry
for new_server in new_registry.values():
self._assign_unique_short_prefix(new_server)
# Register OpenAPI tools *after* the final short prefix is assigned
# so the tools are stored in the global registry under the same
# prefix that lookups will use.
await self._maybe_register_openapi_tools(new_server)
# Assign short prefixes against the full candidate set without
# publishing the staged registry to concurrent callers.
registered_registry: Dict[str, MCPServer] = {}
registered_openapi_tools = False
for server_id, new_server in new_registry.items():
try:
self._assign_unique_short_prefix(new_server, registry=new_registry)
# Register OpenAPI tools *after* the final short prefix is assigned
# so the tools are stored in the global registry under the same
# prefix that lookups will use.
await self._maybe_register_openapi_tools(
new_server, initialize_mapping=False
)
registered_registry[server_id] = new_server
if new_server.spec_path:
registered_openapi_tools = True
except Exception as e:
verbose_logger.exception(
"Skipping MCP server %s (%s) during DB reload: %s",
new_server.server_id,
getattr(new_server, "alias", None),
e,
)
self.registry = registered_registry
if registered_openapi_tools:
self.initialize_tool_name_to_mcp_server_name_mapping()
verbose_logger.debug(
"MCP registry refreshed (%s servers in registry)", len(new_registry)
"MCP registry refreshed (%s servers in registry)", len(registered_registry)
)
def get_mcp_servers_from_ids(self, server_ids: List[str]) -> List[MCPServer]:
litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py
+38 -4
View File
@@ -26,6 +26,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
from litellm.proxy._experimental.mcp_server.auth import token_exchange
from litellm.types.llms.custom_http import httpxSpecialProvider
if TYPE_CHECKING:
@@ -50,12 +51,23 @@ class MCPOAuth2TokenCache(InMemoryCache):
def _get_lock(self, server_id: str) -> asyncio.Lock:
return self._locks.setdefault(server_id, asyncio.Lock())
async def async_get_token(self, server: "MCPServer") -> Optional[str]:
@staticmethod
def _has_client_credentials_config(server: "MCPServer") -> bool:
return bool(server.client_id and server.client_secret and server.token_url)
async def async_get_token(
self,
server: "MCPServer",
*,
require_client_credentials_flow: bool = True,
) -> Optional[str]:
"""Return a valid access token, fetching or refreshing as needed.
Returns ``None`` when the server lacks client credentials config.
"""
if not server.has_client_credentials:
if require_client_credentials_flow and not server.has_client_credentials:
return None
if not self._has_client_credentials_config(server):
return None
server_id = server.server_id
@@ -263,16 +275,38 @@ mcp_per_user_token_cache = MCPPerUserTokenCache()
async def resolve_mcp_auth(
server: "MCPServer",
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
subject_token: Optional[str] = None,
) -> Optional[Union[str, Dict[str, str]]]:
"""Resolve the auth value for an MCP server.
Priority:
1. ``mcp_auth_header`` — per-request/per-user override
2. OAuth2 client_credentials token — auto-fetched and cached
3. ``server.authentication_token`` — static token from config/DB
2. OAuth2 Token Exchange (OBO / RFC 8693) — exchange user token for scoped token
3. OAuth2 client_credentials token — auto-fetched and cached
4. ``server.authentication_token`` — static token from config/DB
"""
if mcp_auth_header:
return mcp_auth_header
if server.has_token_exchange_config:
if subject_token:
return await token_exchange.mcp_token_exchange_handler.exchange_token(
subject_token, server
)
# No subject_token — fall back to client_credentials using the same client
# credentials and token_url so M2M scenarios still work.
if server.client_id and server.client_secret and server.token_url:
return await mcp_oauth2_token_cache.async_get_token(
server,
require_client_credentials_flow=False,
)
# OBO configured but no subject_token and missing client credentials — warn
# rather than silently proceeding unauthenticated.
verbose_logger.warning(
"MCP server '%s' is configured for token exchange (OBO) but no subject_token "
"was provided and client credentials (client_id/client_secret/token_url) are "
"incomplete. The request will proceed without authentication.",
server.server_id,
)
if server.has_client_credentials:
return await mcp_oauth2_token_cache.async_get_token(server)
return server.authentication_token
litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py
+48 -8
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
litellm/proxy/_experimental/mcp_server/server.py
+30
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)
litellm/proxy/_types.py
+17
View File
@@ -353,8 +353,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",
@@ -707,6 +709,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",
@@ -3356,6 +3360,19 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
],
)
azure_sentinel: CallbackOnUI = CallbackOnUI(
litellm_callback_name="azure_sentinel",
ui_callback_name="Azure Sentinel",
litellm_callback_params=[
"AZURE_SENTINEL_DCR_IMMUTABLE_ID",
"AZURE_SENTINEL_ENDPOINT",
"AZURE_SENTINEL_TENANT_ID",
"AZURE_SENTINEL_CLIENT_ID",
"AZURE_SENTINEL_CLIENT_SECRET",
"AZURE_SENTINEL_STREAM_NAME",
],
)
openmeter: CallbackOnUI = CallbackOnUI(
litellm_callback_name="openmeter",
ui_callback_name="OpenMeter",
litellm/proxy/auth/auth_checks.py
+4 -4
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(
@@ -3516,7 +3516,6 @@ async def _check_team_member_budget(
if (
team_object is not None
and team_object.team_id is not None
and user_object is not None
and valid_token is not None
and valid_token.user_id is not None
):
@@ -3619,13 +3618,14 @@ async def _check_team_member_model_access(
llm_router=llm_router,
models=member_allowed_models,
object_type="team",
team_id=team_object.team_id,
)
except ProxyException:
raise ProxyException(
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,
)
litellm/proxy/auth/auth_exception_handler.py
+1 -1
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(
litellm/proxy/auth/user_api_key_auth.py
+3 -3
View File
@@ -1107,7 +1107,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 +1432,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 +2417,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
litellm/proxy/common_request_processing.py
+1 -1
View File
@@ -1512,7 +1512,7 @@ class ProxyBaseLLMRequestProcessing:
status_code=result.status_code,
headers=HttpPassThroughEndpointHelpers.get_response_headers(
headers=result.headers,
custom_headers=None,
custom_headers=dict(fastapi_response.headers),
),
)
litellm/proxy/common_utils/html_forms/jwt_display_template.py
+136 -67
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>
litellm/proxy/common_utils/reset_budget_job.py
+82 -70
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,97 @@ class ResetBudgetJob:
"Failed to reset spend counter %s: %s", counter_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,
):
"""
Generic cascade: zero spend on rows whose budget_id is in the reset set.
"""
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))
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",
)
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}},
)
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}},
)
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.
"""
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}},
)
async def reset_budget_for_litellm_budget_table(self):
"""
@@ -237,6 +241,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:
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
+53 -12
View File
@@ -5,8 +5,10 @@ from typing import Optional
from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import (
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS,
SPEND_LOG_CLEANUP_BATCH_SIZE,
SPEND_LOG_CLEANUP_JOB_NAME,
SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES,
SPEND_LOG_RUN_LOOPS,
)
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
@@ -74,6 +76,7 @@ class SpendLogCleanup:
"""
total_deleted = 0
run_count = 0
consecutive_failures = 0
while True:
if run_count > SPEND_LOG_RUN_LOOPS:
verbose_proxy_logger.info(
@@ -82,18 +85,50 @@ class SpendLogCleanup:
break
# Step 1: Find logs and delete them in one go without fetching to application
# Delete in batches, limited by self.batch_size
deleted_result = await prisma_client.db.execute_raw(
"""
DELETE FROM "LiteLLM_SpendLogs"
WHERE "request_id" IN (
SELECT "request_id" FROM "LiteLLM_SpendLogs"
WHERE "startTime" < $1::timestamptz
LIMIT $2
try:
deleted_result = await prisma_client.db.execute_raw(
"""
DELETE FROM "LiteLLM_SpendLogs"
WHERE "request_id" IN (
SELECT "request_id" FROM "LiteLLM_SpendLogs"
WHERE "startTime" < $1::timestamptz
LIMIT $2
)
""",
cutoff_date,
self.batch_size,
)
""",
cutoff_date,
self.batch_size,
)
except Exception as batch_exc:
# A single batch failure (e.g. Prisma/DB timeout) must not abort
# the whole run — subsequent batches may still succeed.
consecutive_failures += 1
verbose_proxy_logger.exception(
"Spend log cleanup batch failed "
"(run_count=%d, consecutive_failures=%d, batch_size=%d, "
"cutoff=%s, total_deleted_so_far=%d): %s: %s",
run_count,
consecutive_failures,
self.batch_size,
cutoff_date.isoformat(),
total_deleted,
type(batch_exc).__name__,
batch_exc,
)
if (
consecutive_failures
>= SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES
):
verbose_proxy_logger.error(
"Aborting spend log cleanup after %d consecutive batch "
"failures; total deleted before abort: %d",
consecutive_failures,
total_deleted,
)
break
await asyncio.sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS)
continue
consecutive_failures = 0
deleted_count = 0
if isinstance(deleted_result, int):
@@ -168,7 +203,13 @@ class SpendLogCleanup:
verbose_proxy_logger.info(f"Deleted {total_deleted} logs")
except Exception as e:
verbose_proxy_logger.error(f"Error during cleanup: {str(e)}")
# .exception() captures the traceback; str(e) alone on a Prisma/DB
# timeout is often empty and gives operators no signal to diagnose.
verbose_proxy_logger.exception(
"Error during spend log cleanup: %s: %s",
type(e).__name__,
e,
)
return # Return after error handling
finally:
# Only release the lock if it was actually acquired
litellm/proxy/db/prisma_client.py
+176 -61
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
litellm/proxy/db/routing_prisma_wrapper.py
+213
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
litellm/proxy/guardrails/guardrail_registry.py
+5
View File
@@ -482,6 +482,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"),
litellm/proxy/health_endpoints/health_app_factory.py
-8
View File
@@ -1,8 +0,0 @@
from fastapi import FastAPI
from litellm.proxy.health_endpoints._health_endpoints import router as health_router
def build_health_app():
health_app = FastAPI(title="LiteLLM Health Endpoints")
health_app.include_router(health_router)
return health_app
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
+5
View File
@@ -498,6 +498,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"error": f"Model capacity reached for {model}. "
f"Priority: {priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Model TPM: {model_group_info.tpm if model_group_info.tpm is not None else 'not configured'}, "
f"Model RPM: {model_group_info.rpm if model_group_info.rpm is not None else 'not configured'}, "
f"Remaining: {status['limit_remaining']}"
},
headers={
@@ -515,8 +517,11 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
status_code=429,
detail={
"error": f"Priority-based rate limit exceeded. "
f"Model: {model}, "
f"Priority: {priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Model TPM: {model_group_info.tpm if model_group_info.tpm is not None else 'not configured'}, "
f"Model RPM: {model_group_info.rpm if model_group_info.rpm is not None else 'not configured'}, "
f"Remaining: {status['limit_remaining']}, "
f"Model saturation: {saturation:.1%}"
},
litellm/proxy/hooks/model_max_budget_limiter.py
+3
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(
litellm/proxy/litellm_pre_call_utils.py
+4 -1
View File
@@ -1667,7 +1667,10 @@ async def add_litellm_data_to_request( # noqa: PLR0915
)
if tags is not None and _admin_allow_client_tags:
data[_metadata_variable_name]["tags"] = tags
data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags(
request_tags=data[_metadata_variable_name].get("tags"),
tags_to_add=tags,
)
elif tags is not None:
verbose_proxy_logger.warning(
"Ignored caller-supplied tags from header/root body: this "
litellm/proxy/management_endpoints/common_daily_activity.py
+254 -21
View File
@@ -1,3 +1,4 @@
import asyncio
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
@@ -543,6 +544,13 @@ def _build_aggregated_sql_query(
where_clause = " AND ".join(sql_conditions)
# Postgres computes every rollup level the response needs — per-date
# totals, per-(date, model), per-(date, model, api_key), per-provider,
# etc. — in a single pass via GROUPING SETS. The GROUPING() bitmask
# encodes which level a row belongs to so Python can dispatch rows
# straight into their buckets without re-summing. The leaf grouping
# is omitted on purpose: nothing in the response shape needs it once
# all the rollups are present.
sql_query = f"""
SELECT
date,
@@ -552,6 +560,9 @@ def _build_aggregated_sql_query(
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
GROUPING(date, api_key, model, model_group,
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
SUM(spend)::float AS spend,
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
@@ -562,32 +573,35 @@ def _build_aggregated_sql_query(
SUM(failed_requests)::bigint AS failed_requests
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY date, api_key, model, model_group, custom_llm_provider,
mcp_namespaced_tool_name, endpoint
ORDER BY date DESC
GROUP BY GROUPING SETS (
(date),
(date, api_key),
(date, model),
(date, model, api_key),
(date, model_group),
(date, model_group, api_key),
(date, custom_llm_provider),
(date, custom_llm_provider, api_key),
(date, mcp_namespaced_tool_name),
(date, mcp_namespaced_tool_name, api_key),
(date, endpoint),
(date, endpoint, api_key),
()
)
"""
return sql_query, sql_params
async def _aggregate_spend_records(
def _aggregate_spend_records_sync(
*,
prisma_client: PrismaClient,
records: List[Any],
api_key_metadata: Dict[str, Dict[str, Any]],
entity_id_field: Optional[str],
entity_metadata_field: Optional[Dict[str, dict]],
) -> Dict[str, Any]:
"""Aggregate rows into DailySpendData list and total metrics."""
api_keys: Set[str] = set()
for record in records:
if record.api_key:
api_keys.add(record.api_key)
api_key_metadata: Dict[str, Dict[str, Any]] = {}
model_metadata: Dict[str, Dict[str, Any]] = {}
provider_metadata: Dict[str, Dict[str, Any]] = {}
if api_keys:
api_key_metadata = await get_api_key_metadata(prisma_client, api_keys)
results: List[DailySpendData] = []
total_metrics = SpendMetrics()
@@ -631,6 +645,228 @@ async def _aggregate_spend_records(
return {"results": results, "totals": total_metrics}
async def _aggregate_spend_records(
*,
prisma_client: PrismaClient,
records: List[Any],
entity_id_field: Optional[str],
entity_metadata_field: Optional[Dict[str, dict]],
) -> Dict[str, Any]:
"""Aggregate rows into DailySpendData list and total metrics.
The per-row loop is offloaded to a worker thread via asyncio.to_thread so
a large result set doesn't peg the event loop.
"""
api_keys: Set[str] = {record.api_key for record in records if record.api_key}
api_key_metadata: Dict[str, Dict[str, Any]] = {}
if api_keys:
api_key_metadata = await get_api_key_metadata(prisma_client, api_keys)
return await asyncio.to_thread(
_aggregate_spend_records_sync,
records=records,
api_key_metadata=api_key_metadata,
entity_id_field=entity_id_field,
entity_metadata_field=entity_metadata_field,
)
# GROUPING() bitmask values for each grouping set emitted by
# _build_aggregated_sql_query. Per Postgres semantics, the rightmost argument
# is the least-significant bit. Argument order:
# date, api_key, model, model_group, custom_llm_provider,
# mcp_namespaced_tool_name, endpoint
# A bit is 1 when the corresponding column is rolled up (i.e. NOT in the
# current grouping set's key), 0 when the column is part of the key.
_GROUP_GRAND_TOTAL = 127 # 0b1111111 — all rolled up
_GROUP_DATE = 63 # 0b0111111 — only date kept
_GROUP_DATE_API_KEY = 31 # 0b0011111
_GROUP_DATE_MODEL = 47 # 0b0101111
_GROUP_DATE_MODEL_API_KEY = 15 # 0b0001111
_GROUP_DATE_MODEL_GROUP = 55 # 0b0110111
_GROUP_DATE_MODEL_GROUP_API_KEY = 23 # 0b0010111
_GROUP_DATE_PROVIDER = 59 # 0b0111011
_GROUP_DATE_PROVIDER_API_KEY = 27 # 0b0011011
_GROUP_DATE_MCP = 61 # 0b0111101
_GROUP_DATE_MCP_API_KEY = 29 # 0b0011101
_GROUP_DATE_ENDPOINT = 62 # 0b0111110
_GROUP_DATE_ENDPOINT_API_KEY = 30 # 0b0011110
def _record_to_spend_metrics(record: Any) -> SpendMetrics:
"""Build a SpendMetrics directly from one already-aggregated rollup row."""
return SpendMetrics(
spend=record.spend,
prompt_tokens=record.prompt_tokens,
completion_tokens=record.completion_tokens,
total_tokens=record.prompt_tokens + record.completion_tokens,
cache_read_input_tokens=record.cache_read_input_tokens,
cache_creation_input_tokens=record.cache_creation_input_tokens,
api_requests=record.api_requests,
successful_requests=record.successful_requests,
failed_requests=record.failed_requests,
)
def _key_metadata(
api_key_metadata: Dict[str, Dict[str, Any]], api_key: str
) -> KeyMetadata:
meta = api_key_metadata.get(api_key, {})
return KeyMetadata(key_alias=meta.get("key_alias"), team_id=meta.get("team_id"))
def _aggregate_grouping_sets_records_sync( # noqa: PLR0915
*,
records: List[Any],
api_key_metadata: Dict[str, Dict[str, Any]],
) -> Dict[str, Any]:
"""Build the response from rollup rows produced by the GROUPING SETS query.
Each row carries a `group_level` bitmask (from Postgres GROUPING()) that
identifies which rollup level it belongs to. We dispatch the row's
pre-aggregated metrics straight into the matching bucket — no per-row
summing in Python and no nested update_metrics calls.
"""
total_metrics = SpendMetrics()
grouped_data: Dict[str, Dict[str, Any]] = {}
def ensure_date(date_str: str) -> Dict[str, Any]:
bucket = grouped_data.get(date_str)
if bucket is None:
bucket = {"metrics": SpendMetrics(), "breakdown": BreakdownMetrics()}
grouped_data[date_str] = bucket
return bucket
def assign_metric_with_metadata(
target: Dict[str, MetricWithMetadata], key: str, metrics: SpendMetrics
) -> None:
existing = target.get(key)
if existing is None:
target[key] = MetricWithMetadata(metrics=metrics, metadata={})
else:
existing.metrics = metrics
def assign_api_key_breakdown(
target: Dict[str, MetricWithMetadata],
parent_key: str,
api_key: str,
metrics: SpendMetrics,
) -> None:
parent = target.get(parent_key)
if parent is None:
parent = MetricWithMetadata(metrics=SpendMetrics(), metadata={})
target[parent_key] = parent
parent.api_key_breakdown[api_key] = KeyMetricWithMetadata(
metrics=metrics, metadata=_key_metadata(api_key_metadata, api_key)
)
for record in records:
level = record.group_level
metrics = _record_to_spend_metrics(record)
if level == _GROUP_GRAND_TOTAL:
total_metrics = metrics
continue
if level == _GROUP_DATE:
ensure_date(record.date)["metrics"] = metrics
continue
breakdown = ensure_date(record.date)["breakdown"]
if level == _GROUP_DATE_API_KEY:
if record.api_key:
breakdown.api_keys[record.api_key] = KeyMetricWithMetadata(
metrics=metrics,
metadata=_key_metadata(api_key_metadata, record.api_key),
)
elif level == _GROUP_DATE_MODEL:
if record.model:
assign_metric_with_metadata(breakdown.models, record.model, metrics)
elif level == _GROUP_DATE_MODEL_API_KEY:
if record.model and record.api_key:
assign_api_key_breakdown(
breakdown.models, record.model, record.api_key, metrics
)
elif level == _GROUP_DATE_MODEL_GROUP:
if record.model_group:
assign_metric_with_metadata(
breakdown.model_groups, record.model_group, metrics
)
elif level == _GROUP_DATE_MODEL_GROUP_API_KEY:
if record.model_group and record.api_key:
assign_api_key_breakdown(
breakdown.model_groups,
record.model_group,
record.api_key,
metrics,
)
elif level == _GROUP_DATE_PROVIDER:
provider = record.custom_llm_provider or "unknown"
assign_metric_with_metadata(breakdown.providers, provider, metrics)
elif level == _GROUP_DATE_PROVIDER_API_KEY:
if record.api_key:
provider = record.custom_llm_provider or "unknown"
assign_api_key_breakdown(
breakdown.providers, provider, record.api_key, metrics
)
elif level == _GROUP_DATE_MCP:
if record.mcp_namespaced_tool_name:
assign_metric_with_metadata(
breakdown.mcp_servers, record.mcp_namespaced_tool_name, metrics
)
elif level == _GROUP_DATE_MCP_API_KEY:
if record.mcp_namespaced_tool_name and record.api_key:
assign_api_key_breakdown(
breakdown.mcp_servers,
record.mcp_namespaced_tool_name,
record.api_key,
metrics,
)
elif level == _GROUP_DATE_ENDPOINT:
if record.endpoint:
assign_metric_with_metadata(
breakdown.endpoints, record.endpoint, metrics
)
elif level == _GROUP_DATE_ENDPOINT_API_KEY:
if record.endpoint and record.api_key:
assign_api_key_breakdown(
breakdown.endpoints, record.endpoint, record.api_key, metrics
)
results = [
DailySpendData(
date=datetime.strptime(date_str, "%Y-%m-%d").date(),
metrics=data["metrics"],
breakdown=data["breakdown"],
)
for date_str, data in grouped_data.items()
]
results.sort(key=lambda x: x.date, reverse=True)
return {"results": results, "totals": total_metrics}
async def _aggregate_grouping_sets_records(
*,
prisma_client: PrismaClient,
records: List[Any],
) -> Dict[str, Any]:
"""Async wrapper: fetch api_key_metadata, then dispatch on a worker thread."""
api_keys: Set[str] = {r.api_key for r in records if r.api_key}
api_key_metadata: Dict[str, Dict[str, Any]] = {}
if api_keys:
api_key_metadata = await get_api_key_metadata(prisma_client, api_keys)
return await asyncio.to_thread(
_aggregate_grouping_sets_records_sync,
records=records,
api_key_metadata=api_key_metadata,
)
async def get_daily_activity(
prisma_client: Optional[PrismaClient],
table_name: str,
@@ -771,21 +1007,18 @@ async def get_daily_activity_aggregated(
timezone_offset_minutes=timezone_offset_minutes,
)
# Execute GROUP BY query — returns pre-aggregated dicts
# Execute GROUPING SETS query — returns one row per rollup level.
rows = await prisma_client.db.query_raw(sql_query, *sql_params)
if rows is None:
rows = []
# Convert dicts to objects for compatibility with _aggregate_spend_records
records = [SimpleNamespace(**row) for row in rows]
# entity_id_field=None skips entity breakdown (entity dimension was
# collapsed by the GROUP BY, so per-entity data is not available)
aggregated = await _aggregate_spend_records(
# The grouping-sets dispatcher places each row directly in its bucket
# using the row's GROUPING() bitmask. No Python-side summing needed.
aggregated = await _aggregate_grouping_sets_records(
prisma_client=prisma_client,
records=records,
entity_id_field=None,
entity_metadata_field=None,
)
return SpendAnalyticsPaginatedResponse(
litellm/proxy/management_endpoints/mcp_management_endpoints.py
+61 -6
View File
@@ -152,8 +152,14 @@ if MCP_AVAILABLE:
UserAPIKeyAuth,
UserMCPManagementMode,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.auth.user_api_key_auth import (
_user_api_key_auth_builder,
user_api_key_auth,
)
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
populate_request_with_path_params,
)
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
from litellm.types.mcp import MCPCredentials
@@ -1492,6 +1498,55 @@ if MCP_AVAILABLE:
return _redact_mcp_credentials(temp_record)
async def _mcp_oauth_user_api_key_auth(request: Request) -> UserAPIKeyAuth:
"""
Auth dependency for MCP OAuth browser-navigation endpoints (/authorize, /token).
Tries the Authorization header first. Falls back to decoding the UI
'token' session cookie (set by SSO login) to extract the API key, which
allows browser-based OAuth redirects to work without an explicit
Authorization header.
"""
import jwt as _jwt
from litellm.proxy.proxy_server import master_key
auth_header = request.headers.get("Authorization", "")
api_key = auth_header # _get_bearer_token will strip "Bearer " prefix
if not api_key:
token_cookie = request.cookies.get("token")
if token_cookie and master_key:
try:
decoded = _jwt.decode(
token_cookie,
master_key,
algorithms=["HS256"],
# UI session cookies may omit exp; don't require it.
options={"verify_exp": False},
)
if decoded.get("login_method") in ("sso", "username_password"):
cookie_key = decoded.get("key", "")
if cookie_key:
api_key = f"Bearer {cookie_key}"
except _jwt.InvalidTokenError:
pass
request_data = await _read_request_body(request=request)
request_data = populate_request_with_path_params(
request_data=request_data, request=request
)
return await _user_api_key_auth_builder(
request=request,
api_key=api_key,
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data=request_data,
)
async def _get_cached_temporary_mcp_server_or_404(
server_id: str,
user_api_key_dict: UserAPIKeyAuth,
@@ -1542,12 +1597,12 @@ if MCP_AVAILABLE:
@router.get(
"/server/oauth/{server_id}/authorize",
include_in_schema=False,
dependencies=[Depends(user_api_key_auth)],
dependencies=[Depends(_mcp_oauth_user_api_key_auth)],
)
async def mcp_authorize(
request: Request,
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_api_key_dict: UserAPIKeyAuth = Depends(_mcp_oauth_user_api_key_auth),
client_id: Optional[str] = None,
redirect_uri: str = Query(...),
state: str = "",
@@ -1587,12 +1642,12 @@ if MCP_AVAILABLE:
@router.post(
"/server/oauth/{server_id}/token",
include_in_schema=False,
dependencies=[Depends(user_api_key_auth)],
dependencies=[Depends(_mcp_oauth_user_api_key_auth)],
)
async def mcp_token(
request: Request,
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_api_key_dict: UserAPIKeyAuth = Depends(_mcp_oauth_user_api_key_auth),
grant_type: str = Form(...),
code: Optional[str] = Form(None),
redirect_uri: Optional[str] = Form(None),
litellm/proxy/management_endpoints/scim/scim_v2.py
+27 -8
View File
@@ -4,6 +4,7 @@
This is an enterprise feature and requires a premium license.
"""
import re
from typing import Any, Dict, List, Optional, Set, Tuple
from fastapi import (
@@ -843,6 +844,18 @@ async def get_service_provider_config(request: Request):
return SCIMServiceProviderConfig(meta=meta)
def _parse_scim_eq_filter(scim_filter: str) -> Optional[Tuple[str, str]]:
"""Parse the SCIM equality filters Okta uses before user lifecycle changes."""
match = re.match(
r"""\s*([\w.]+)\s+eq\s+(['"]?)(.*?)\2\s*$""",
scim_filter,
flags=re.IGNORECASE,
)
if not match:
return None
return match.group(1).lower(), match.group(3)
# User Endpoints
@scim_router.get(
"/Users",
@@ -867,15 +880,21 @@ async def get_users(
try:
prisma_client = await _get_prisma_client_or_raise_exception()
# Parse filter if provided (basic support)
where_conditions = {}
where_conditions: Dict[str, Any] = {}
if filter:
# Very basic filter support - only handling userName eq and emails.value eq
if "userName eq" in filter:
user_id = filter.split("userName eq ")[1].strip("\"'")
where_conditions["user_id"] = user_id
elif "emails.value eq" in filter:
email = filter.split("emails.value eq ")[1].strip("\"'")
where_conditions["user_email"] = email
# Okta locates users by userName before deprovisioning. LiteLLM
# exposes SCIM userName from user_email, while older SCIM-created
# users may still have user_id == userName, so support both.
parsed_filter = _parse_scim_eq_filter(filter)
if parsed_filter:
filter_attribute, filter_value = parsed_filter
if filter_attribute == "username":
where_conditions["OR"] = [
{"user_email": filter_value},
{"user_id": filter_value},
]
elif filter_attribute == "emails.value":
where_conditions["user_email"] = filter_value
# Get users from database
users: List[LiteLLM_UserTable] = (
litellm/proxy/management_endpoints/ui_sso.py
+34 -8
View File
@@ -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)
litellm/proxy/management_helpers/audit_logs.py
+30 -9
View File
@@ -46,25 +46,46 @@ def get_audit_log_changed_by(
def _resolve_audit_log_callback(name: str) -> Optional[CustomLogger]:
"""Resolve a string callback name to a CustomLogger instance, with caching."""
"""Resolve a string callback name to a CustomLogger instance, with caching.
For "s3_v2" with `litellm.s3_audit_callback_params` set, constructs a
dedicated `S3Logger` so audit logs can target a different bucket than the
normal-log singleton served by `_init_custom_logger_compatible_class`.
"""
if name in _audit_log_callback_cache:
return _audit_log_callback_cache[name]
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
)
instance: Optional[CustomLogger]
if (
name == "s3_v2"
and getattr(litellm, "s3_audit_callback_params", None) is not None
):
from litellm.integrations.s3_v2 import S3Logger as S3V2Logger
instance = _init_custom_logger_compatible_class(
logging_integration=name, # type: ignore
internal_usage_cache=None,
llm_router=None,
)
instance = S3V2Logger(
s3_callback_params_override=litellm.s3_audit_callback_params
)
else:
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
)
instance = _init_custom_logger_compatible_class(
logging_integration=name, # type: ignore
internal_usage_cache=None,
llm_router=None,
)
if instance is not None:
_audit_log_callback_cache[name] = instance
return instance
def reset_audit_log_callback_cache() -> None:
"""Clear cached audit-log callback instances. Call on config reload."""
_audit_log_callback_cache.clear()
def _build_audit_log_payload(
request_data: LiteLLM_AuditLogs,
) -> StandardAuditLogPayload:
litellm/proxy/middleware/prometheus_auth_middleware.py
+4 -1
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(
{
litellm/proxy/middleware/request_size_limit_middleware.py
+121
View File
@@ -0,0 +1,121 @@
import json
from typing import Callable, Optional, Union
from starlette.types import ASGIApp, Message, Receive, Scope, Send
MaxRequestSizeGetter = Callable[[], Optional[Union[int, float]]]
RequestSizeLimitEnabledGetter = Callable[[], bool]
class RequestEntityTooLarge(Exception):
pass
class RequestSizeLimitMiddleware:
"""
Reject oversized requests before downstream auth/routes parse the body.
Content-Length can be rejected without reading any body bytes. Requests
without Content-Length are counted as the ASGI stream is consumed, limiting
memory exposure to the configured threshold plus the current chunk.
"""
def __init__(
self,
app: ASGIApp,
get_max_request_size_mb: MaxRequestSizeGetter,
is_request_size_limit_enabled: RequestSizeLimitEnabledGetter,
) -> None:
self.app = app
self.get_max_request_size_mb = get_max_request_size_mb
self.is_request_size_limit_enabled = is_request_size_limit_enabled
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
await self.app(scope, receive, send)
return
max_request_size_mb = self.get_max_request_size_mb()
max_request_size_bytes = _mb_to_bytes(max_request_size_mb)
if max_request_size_bytes is None or not self.is_request_size_limit_enabled():
await self.app(scope, receive, send)
return
content_length = _get_content_length(scope=scope)
if content_length is not None and content_length > max_request_size_bytes:
await _send_request_too_large(
send=send, max_request_size_mb=max_request_size_mb
)
return
received_body_bytes = 0
response_started = False
async def limited_receive() -> Message:
nonlocal received_body_bytes
message = await receive()
if message["type"] != "http.request":
return message
received_body_bytes += len(message.get("body", b""))
if received_body_bytes > max_request_size_bytes:
raise RequestEntityTooLarge
return message
async def tracking_send(message: Message) -> None:
nonlocal response_started
if message["type"] == "http.response.start":
response_started = True
await send(message)
try:
await self.app(scope, limited_receive, tracking_send)
except RequestEntityTooLarge:
if response_started:
raise
await _send_request_too_large(
send=send, max_request_size_mb=max_request_size_mb
)
def _mb_to_bytes(max_request_size_mb: Optional[Union[int, float]]) -> Optional[int]:
if max_request_size_mb is None:
return None
if max_request_size_mb <= 0:
return None
return int(max_request_size_mb * 1024 * 1024)
def _get_content_length(scope: Scope) -> Optional[int]:
headers = dict(scope.get("headers") or [])
raw_content_length = headers.get(b"content-length")
if raw_content_length is None:
return None
try:
return int(raw_content_length)
except ValueError:
return None
async def _send_request_too_large(
send: Send,
max_request_size_mb: Optional[Union[int, float]],
) -> None:
body = json.dumps(
{"error": f"Request size is too large. Max size is {max_request_size_mb} MB"},
separators=(",", ":"),
).encode("utf-8")
await send(
{
"type": "http.response.start",
"status": 413,
"headers": [
(b"content-type", b"application/json"),
(b"content-length", str(len(body)).encode("latin-1")),
],
}
)
await send({"type": "http.response.body", "body": body, "more_body": False})
litellm/proxy/proxy_cli.py
+72 -8
View File
@@ -169,6 +169,66 @@ class ProxyInitializationHelpers:
)
return uvicorn_args
@staticmethod
def _get_reload_options(config_path: Optional[str]) -> dict:
"""Build uvicorn reload kwargs so --reload also reacts to YAML edits."""
options: dict = {"reload": True}
if not config_path:
return options
config_abs = os.path.abspath(config_path)
config_dir = os.path.dirname(config_abs)
cwd = os.path.abspath(os.getcwd())
reload_dirs = [cwd]
if config_dir and config_dir != cwd:
reload_dirs.append(config_dir)
options["reload_dirs"] = reload_dirs
# Must be a basename, not an absolute path: uvicorn's
# resolve_reload_patterns() calls pathlib.Path.glob(), which raises
# NotImplementedError on absolute patterns (uvicorn discussion #2156).
options["reload_includes"] = ["*.py", os.path.basename(config_abs)]
return options
@staticmethod
def _patch_statreload_for_config(config_path: str) -> bool:
"""Make uvicorn's StatReload reloader notice YAML config changes.
Uvicorn uses WatchFilesReload when the optional `watchfiles` package
is installed, otherwise StatReload. StatReload hard-codes `*.py` in
`iter_py_files()` and silently ignores `reload_includes`, so the
kwargs from `_get_reload_options` alone don't trigger reloads on YAML
edits. We monkey-patch `iter_py_files` to also yield the config path.
Idempotent across calls and a no-op for the WatchFilesReload path.
"""
try:
from uvicorn.supervisors.statreload import StatReload
except ImportError: # pragma: no cover - uvicorn is a hard dep
return False
if not config_path:
return False
from pathlib import Path
config_abs = Path(config_path).resolve()
patched_paths = getattr(StatReload, "_litellm_patched_config_paths", None)
if patched_paths is None:
original_iter = StatReload.iter_py_files
patched_paths = set()
def _iter_with_config(self): # type: ignore[no-untyped-def]
yield from original_iter(self)
for path in StatReload._litellm_patched_config_paths:
if path.exists():
yield path
StatReload.iter_py_files = _iter_with_config # type: ignore[assignment]
StatReload._litellm_patched_config_paths = patched_paths # type: ignore[attr-defined]
patched_paths.add(config_abs)
return True
@staticmethod
def _init_hypercorn_server(
app: FastAPI,
@@ -619,7 +679,7 @@ class ProxyInitializationHelpers:
"--reload",
is_flag=True,
default=False,
help="Enable uvicorn hot reload (dev only). Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.",
help="Enable uvicorn hot reload (dev only). Also reloads when the --config YAML file changes. Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.",
)
def run_server( # noqa: PLR0915
host,
@@ -753,7 +813,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")
@@ -990,11 +1055,6 @@ def run_server( # noqa: PLR0915
litellm_settings=litellm_settings if config else None, # type: ignore[possibly-unbound]
)
# --- SEPARATE HEALTH APP LOGIC ---
# To run the health app separately, use:
# uvicorn litellm.proxy.health_app_factory:build_health_app --factory --host 0.0.0.0 --port=4001
# This is compatible with the SEPARATE_HEALTH_APP Docker/supervisord pattern.
# --- END SEPARATE HEALTH APP LOGIC ---
# Skip server startup if requested (after all setup is done)
if skip_server_startup:
print( # noqa
@@ -1028,7 +1088,11 @@ def run_server( # noqa: PLR0915
uvicorn_args["loop"] = loop_type
if reload:
uvicorn_args["reload"] = True
uvicorn_args.update(
ProxyInitializationHelpers._get_reload_options(config)
)
if config:
ProxyInitializationHelpers._patch_statreload_for_config(config)
uvicorn.run(
**uvicorn_args,
litellm/proxy/proxy_server.py
+85 -18
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,
@@ -403,6 +404,9 @@ from litellm.proxy.middleware.in_flight_requests_middleware import (
InFlightRequestsMiddleware,
)
from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware
from litellm.proxy.middleware.request_size_limit_middleware import (
RequestSizeLimitMiddleware,
)
from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router
from litellm.proxy.openai_files_endpoints.files_endpoints import (
router as openai_files_router,
@@ -3834,6 +3838,11 @@ class ProxyConfig:
f"{blue_color_code} Initialized Failure Callbacks - {litellm.failure_callback} {reset_color_code}"
) # noqa
elif key == "audit_log_callbacks":
from litellm.proxy.management_helpers.audit_logs import (
reset_audit_log_callback_cache,
)
reset_audit_log_callback_cache()
litellm.audit_log_callbacks = []
for callback in value:
@@ -3915,6 +3924,21 @@ class ProxyConfig:
f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}"
)
setattr(litellm, key, value)
if key in {"s3_audit_callback_params", "s3_callback_params"}:
from litellm.proxy.management_helpers.audit_logs import (
reset_audit_log_callback_cache,
)
from litellm.litellm_core_utils.litellm_logging import (
_in_memory_loggers,
)
from litellm.integrations.s3_v2 import S3Logger as S3V2Logger
reset_audit_log_callback_cache()
_in_memory_loggers[:] = [
cb
for cb in _in_memory_loggers
if not isinstance(cb, S3V2Logger)
]
## GENERAL SERVER SETTINGS (e.g. master key,..) # do this after initializing litellm, to ensure sentry logging works for proxylogging
general_settings = config.get("general_settings", {})
@@ -6741,27 +6765,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,
@@ -8819,6 +8880,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(
@@ -14895,6 +14957,11 @@ app.include_router(ui_discovery_endpoints_router)
app.include_router(google_router)
attach_lazy_features(app)
app.add_middleware(
RequestSizeLimitMiddleware,
get_max_request_size_mb=lambda: general_settings.get("max_request_size_mb"),
is_request_size_limit_enabled=lambda: premium_user is True,
)
async def _stream_mcp_asgi_response(
litellm/proxy/spend_tracking/budget_reservation.py
+70 -38
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:
litellm/proxy/utils.py
+113 -15
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,
)
@@ -2274,6 +2278,13 @@ class ProxyLogging:
Covers:
1. /chat/completions
"""
from litellm.proxy.proxy_server import llm_router
# Merge model-level guardrails before checking which guardrails to run
request_data = _check_and_merge_model_level_guardrails(
data=request_data, llm_router=llm_router
)
current_response = response
for callback in litellm.callbacks:
@@ -2562,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
@@ -2617,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"]:
@@ -4265,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)
litellm/types/guardrails.py
+10
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,
litellm/types/mcp.py
+18
View File
@@ -37,6 +37,7 @@ class MCPAuth(str, enum.Enum):
oauth2 = "oauth2"
aws_sigv4 = "aws_sigv4"
token = "token"
oauth2_token_exchange = "oauth2_token_exchange"
# MCP Literals
@@ -54,6 +55,7 @@ MCPAuthType = Optional[
MCPAuth.oauth2,
MCPAuth.aws_sigv4,
MCPAuth.token,
MCPAuth.oauth2_token_exchange,
]
]
@@ -117,6 +119,22 @@ class MCPCredentials(TypedDict, total=False):
aws_session_name: Optional[str]
"""Session name for STS AssumeRole (used in CloudTrail). Not a secret — stored unencrypted."""
audience: Optional[str]
"""
Target audience for OAuth 2.0 Token Exchange (RFC 8693)
"""
token_exchange_endpoint: Optional[str]
"""
IDP token endpoint for OAuth 2.0 Token Exchange (RFC 8693)
"""
subject_token_type: Optional[str]
"""
Subject token type for OAuth 2.0 Token Exchange (RFC 8693).
Default: urn:ietf:params:oauth:token-type:access_token
"""
class MCPServerCostInfo(TypedDict, total=False):
default_cost_per_query: Optional[float]
litellm/types/mcp_server/mcp_server_manager.py
+13
View File
@@ -57,6 +57,10 @@ class MCPServer(BaseModel):
aws_service_name: Optional[str] = None # defaults to "bedrock-agentcore"
aws_role_name: Optional[str] = None # IAM role ARN for STS AssumeRole
aws_session_name: Optional[str] = None # session name for CloudTrail auditing
# Token Exchange (OBO) fields — RFC 8693
token_exchange_endpoint: Optional[str] = None
audience: Optional[str] = None
subject_token_type: str = "urn:ietf:params:oauth:token-type:access_token"
# Stdio-specific fields
command: Optional[str] = None
args: Optional[List[str]] = None
@@ -127,3 +131,12 @@ class MCPServer(BaseModel):
return any(h.lower() in auth_header_names for h in self.extra_headers)
return False
@property
def has_token_exchange_config(self) -> bool:
"""True if this server is configured for OAuth2 token exchange (OBO / RFC 8693)."""
return (
self.auth_type == MCPAuth.oauth2_token_exchange
and bool(self.client_id and self.client_secret)
and bool(self.token_exchange_endpoint or self.token_url)
)
litellm/types/utils.py
+1
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
model_prices_and_context_window.json
+69
View File
@@ -27192,6 +27192,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",
@@ -28879,6 +28893,19 @@
"mode": "chat",
"output_cost_per_token": 0.0
},
"sambanova/MiniMax-M2.7": {
"input_cost_per_token": 3e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 204800,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"source": "https://cloud.sambanova.ai/plans/pricing",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"sambanova/DeepSeek-R1": {
"input_cost_per_token": 5e-06,
"litellm_provider": "sambanova",
@@ -34932,6 +34959,48 @@
"supports_vision": true,
"supports_web_search": true
},
"xai/grok-4.3": {
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_above_200k_tokens": 2.5e-06,
"litellm_provider": "xai",
"max_input_tokens": 1000000,
"max_output_tokens": 1000000,
"max_tokens": 1000000,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"output_cost_per_token_above_200k_tokens": 5e-06,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"xai/grok-4.3-latest": {
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_above_200k_tokens": 2.5e-06,
"litellm_provider": "xai",
"max_input_tokens": 1000000,
"max_output_tokens": 1000000,
"max_tokens": 1000000,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"output_cost_per_token_above_200k_tokens": 5e-06,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"xai/grok-beta": {
"input_cost_per_token": 5e-06,
"litellm_provider": "xai",
pyproject.toml
+18 -14
View File
@@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.84.0"
version = "1.85.0"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.14"
@@ -10,18 +10,22 @@ authors = [
{ name = "BerriAI" },
]
dependencies = [
"fastuuid==0.14.0",
"httpx==0.28.1",
"openai==2.33.0",
"python-dotenv==1.2.2",
"tiktoken==0.12.0",
"importlib-metadata==8.5.0",
"tokenizers==0.23.1",
"click==8.1.8",
"jinja2==3.1.6",
"aiohttp==3.13.4",
"pydantic==2.12.5",
"jsonschema==4.23.0",
# Ranges (not exact pins) so SDK consumers can coexist with their other
# deps. Reproducibility for our Docker/CI comes from `uv.lock`.
# When changing a floor, verify it installs + imports on every supported
# Python with: `uv pip install --resolution=lowest-direct .`
"fastuuid>=0.14.0,<1.0",
"httpx>=0.28.0,<1.0",
"openai>=2.20.0,<3.0.0",
"python-dotenv>=1.0.0,<2.0",
"tiktoken>=0.8.0,<1.0",
"importlib-metadata>=8.0.0,<9.0",
"tokenizers>=0.21.0,<1.0",
"click>=8.0.0,<9.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",
]
[project.urls]
@@ -246,7 +250,7 @@ source-exclude = [
profile = "black"
[tool.commitizen]
version = "1.84.0"
version = "1.85.0"
version_files = [
"pyproject.toml:^version",
]
scripts/mock_bedrock_passthrough_target.py
+276
View File
@@ -0,0 +1,276 @@
#!/usr/bin/env python3
"""
Minimal HTTP target for testing LiteLLM **Bedrock pass-through** (`/bedrock/...` on the proxy).
What it does
- Serves a tiny Converse-shaped JSON (and optional invoke-shaped) response so the proxy can
complete a round trip without calling AWS.
- Does **not** verify SigV4 (Bedrock does); any Authorization header is accepted.
How to run
uv run python scripts/mock_bedrock_passthrough_target.py --host 127.0.0.1 --port 9999
Wire LiteLLM to this host (use **one** of these patterns):
1) model_list (recommended) — set the Bedrock runtime base to the mock:
model_list:
- model_name: mock-bedrock-claude
litellm_params:
model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0
custom_llm_provider: bedrock
aws_region_name: us-west-2
api_base: "http://127.0.0.1:9999"
2) Environment (see litellm BaseAWSLLM.get_runtime_endpoint)::
export AWS_BEDROCK_RUNTIME_ENDPOINT="http://127.0.0.1:9999"
Then call the proxy, e.g. (model_name must match config)::
curl -sS -X POST "http://127.0.0.1:4000/bedrock/model/mock-bedrock-claude/converse" \
-H "Authorization: Bearer $LITELLM_KEY" -H "Content-Type: application/json" \
-d '{"messages":[{"role":"user","content":[{"text":"hi"}]}]}'
The proxy will forward to: {api_base}/model/<resolved model id>/converse (SigV4-signed).
This mock implements POST .../converse and returns a minimal valid Converse response.
Notes
- `invoke-with-response-stream` returns a real **binary** AWS event stream
(`application/vnd.amazon.eventstream`) with Anthropic-style JSON payloads inside each
`PayloadPart`, matching Bedrock's InvokeModelWithResponseStream wire format. See
https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModelWithResponseStream.html
and https://docs.aws.amazon.com/awstreams/latest/devguide/message-formats.html
- `converse-stream` is still JSON-only placeholder (different inner event shapes).
- Use real (or any non-empty) AWS creds in the environment of the **proxy**; signing still runs.
"""
from __future__ import annotations
import argparse
import base64
import json
from binascii import crc32
from struct import pack
from typing import Any, Dict, Iterator, List
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from starlette.responses import StreamingResponse
app = FastAPI(title="Mock Bedrock runtime (pass-through test target)")
# Minimal structure compatible with Converse: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Converse.html
def _converse_response_body() -> Dict[str, Any]:
return {
"output": {
"message": {
"role": "assistant",
"content": [
{"text": "mock: ok from mock_bedrock_passthrough_target.py"}
],
}
},
"stopReason": "end_turn",
"usage": {
"inputTokens": 1,
"outputTokens": 2,
"totalTokens": 3,
},
}
# Minimal invoke (Anthropic messages on bedrock) style — adjust if you test /invoke
def _invoke_response_body() -> Dict[str, Any]:
return {
"id": "msg_mock",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "mock invoke response"}],
"model": "mock",
"stop_reason": "end_turn",
"usage": {"input_tokens": 1, "output_tokens": 2},
}
def _encode_event_stream_message(headers: Dict[str, str], payload: bytes) -> bytes:
"""Single AWS binary event-stream frame (same layout botocore's ``EventStreamBuffer`` parses)."""
header_blob = b""
for name, value in headers.items():
nb = name.encode("utf-8")
vb = value.encode("utf-8")
header_blob += bytes([len(nb)]) + nb + bytes([7]) + pack("!H", len(vb)) + vb
headers_length = len(header_blob)
payload_length = len(payload)
total_length = 12 + headers_length + payload_length + 4
prelude_wo_crc = pack("!II", total_length, headers_length)
prelude_crc_val = crc32(prelude_wo_crc) & 0xFFFFFFFF
prelude = prelude_wo_crc + pack("!I", prelude_crc_val)
wo_msg_crc = prelude + header_blob + payload
msg_crc_val = crc32(wo_msg_crc[8:], prelude_crc_val) & 0xFFFFFFFF
return wo_msg_crc + pack("!I", msg_crc_val)
def _bedrock_payload_part(inner_event: Dict[str, Any]) -> bytes:
"""Outer JSON expected by bedrock-runtime ``ResponseStream`` / ``PayloadPart``."""
inner_bytes = json.dumps(inner_event, separators=(",", ":")).encode("utf-8")
outer = {
"chunk": {
"bytes": base64.b64encode(inner_bytes).decode("ascii"),
}
}
return json.dumps(outer, separators=(",", ":")).encode("utf-8")
def _anthropic_invoke_stream_events(
model_id: str, assistant_text: str
) -> List[Dict[str, Any]]:
"""
Minimal Anthropic Messages stream events as returned inside Bedrock stream chunks.
Mirrors the sequence Amazon emits for Claude on ``invoke-with-response-stream``.
"""
msg_id = "msg_mock_bedrock_stream"
input_tokens = 3
output_tokens = max(1, len(assistant_text) // 4)
events: List[Dict[str, Any]] = [
{
"type": "message_start",
"message": {
"model": model_id,
"id": msg_id,
"type": "message",
"role": "assistant",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {
"input_tokens": input_tokens,
"output_tokens": 1,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"cache_creation": {
"ephemeral_5m_input_tokens": 0,
"ephemeral_1h_input_tokens": 0,
},
},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
]
# Split text into small deltas so downstream streaming behavior is visible.
step = 24
for i in range(0, len(assistant_text), step):
events.append(
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "text_delta",
"text": assistant_text[i : i + step],
},
}
)
events.append({"type": "content_block_stop", "index": 0})
events.append(
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
},
}
)
events.append(
{
"type": "message_stop",
"amazon-bedrock-invocationMetrics": {
"inputTokenCount": input_tokens,
"outputTokenCount": output_tokens,
"invocationLatency": 42,
"firstByteLatency": 10,
},
}
)
return events
def _iter_invoke_with_response_stream(model_id: str) -> Iterator[bytes]:
text = (
"mock streaming: ok from scripts/mock_bedrock_passthrough_target.py "
"(invoke-with-response-stream)."
)
headers = {
":event-type": "chunk",
":content-type": "application/json",
":message-type": "event",
}
for ev in _anthropic_invoke_stream_events(model_id, text):
yield _encode_event_stream_message(headers, _bedrock_payload_part(ev))
@app.get("/health")
def health() -> Dict[str, str]:
return {"status": "ok"}
@app.post("/model/{model_path:path}/converse")
async def converse(model_path: str, request: Request) -> JSONResponse:
# Optional: log body for debugging
_ = await request.body()
return JSONResponse(content=_converse_response_body())
@app.post("/model/{model_path:path}/converse-stream")
async def converse_stream(model_path: str, request: Request) -> JSONResponse:
"""
Not a real AWS event stream — returns JSON for quick smoke tests only.
"""
_ = await request.body()
return JSONResponse(
content={
"note": "This mock does not implement application/vnd.amazon.eventstream; use /converse for basic tests."
}
)
@app.post("/model/{model_path:path}/invoke")
async def invoke(model_path: str, request: Request) -> JSONResponse:
_ = await request.body()
return JSONResponse(content=_invoke_response_body())
@app.post("/model/{model_path:path}/invoke-with-response-stream")
async def invoke_with_response_stream(
model_path: str, request: Request
) -> StreamingResponse:
"""
Binary ``application/vnd.amazon.eventstream`` body compatible with boto3/botocore
``InvokeModelWithResponseStream`` / LiteLLM's Bedrock invoke streaming path.
"""
_ = await request.body()
return StreamingResponse(
_iter_invoke_with_response_stream(model_id=model_path),
media_type="application/vnd.amazon.eventstream",
)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=9999)
args = parser.parse_args()
import uvicorn
uvicorn.run(app, host=args.host, port=args.port, log_level="info")
if __name__ == "__main__":
main()
tests/_vcr_conftest_common.py
+174 -3
View File
@@ -7,7 +7,9 @@ from __future__ import annotations
import atexit
import hashlib
import json
import os
import re
import sys
from typing import Iterable
@@ -74,6 +76,17 @@ FILTERED_RESPONSE_HEADERS = (
"date",
)
# Tiny placeholder used to replace base64 image payloads in cassettes.
# Decodes to b"test" — short, valid base64 so test code that decodes
# the field still succeeds.
VCR_IMAGE_B64_PLACEHOLDER = "dGVzdA=="
# Fixed boundary substituted into multipart request bodies so the
# ``safe_body`` matcher sees the same bytes across record and replay.
# httpx generates a fresh random boundary per request via os.urandom,
# which otherwise turns every multipart cassette into a permanent miss.
VCR_FIXED_MULTIPART_BOUNDARY = "vcr-static-boundary"
def _scrub_response(response):
if not isinstance(response, dict):
@@ -86,8 +99,88 @@ def _scrub_response(response):
return response
def _replace_b64_json_in_place(obj) -> bool:
"""Recursively replace ``b64_json`` string values in a JSON tree.
Returns ``True`` if any value was rewritten. The check on the
existing value's length keeps the function idempotent — once a
value has been swapped to the placeholder, subsequent invocations
are no-ops.
"""
changed = False
if isinstance(obj, dict):
for key, value in obj.items():
if (
key == "b64_json"
and isinstance(value, str)
and len(value) > len(VCR_IMAGE_B64_PLACEHOLDER)
):
obj[key] = VCR_IMAGE_B64_PLACEHOLDER
changed = True
elif _replace_b64_json_in_place(value):
changed = True
elif isinstance(obj, list):
for item in obj:
if _replace_b64_json_in_place(item):
changed = True
return changed
def _strip_image_b64_payloads(response):
"""Replace ``b64_json`` payloads in image-gen responses before save.
Image-edit and image-generation responses carry the full base64
PNG/JPEG (1-10+ MB) in ``data[*].b64_json``. The image_gen tests
only assert response shape — the field decodes, schema validates —
they never inspect pixel content. Swapping to a 4-byte placeholder
preserves all those checks while shrinking cassettes by ~99%.
"""
if not isinstance(response, dict):
return response
body = response.get("body")
if not isinstance(body, dict):
return response
raw = body.get("string")
if raw is None:
return response
if isinstance(raw, (bytes, bytearray)):
try:
text = bytes(raw).decode("utf-8")
except UnicodeDecodeError:
return response
was_bytes = True
elif isinstance(raw, str):
text = raw
was_bytes = False
else:
return response
try:
payload = json.loads(text)
except (ValueError, TypeError):
return response
if not _replace_b64_json_in_place(payload):
return response
new_text = json.dumps(payload, separators=(",", ":"))
body["string"] = new_text.encode("utf-8") if was_bytes else new_text
headers = response.get("headers")
if isinstance(headers, dict):
new_len_value = str(len(new_text.encode("utf-8")))
for key in list(headers):
if str(key).lower() == "content-length":
value = headers[key]
headers[key] = (
[new_len_value] if isinstance(value, list) else new_len_value
)
return response
def _before_record_response(response):
return filter_non_2xx_response(_scrub_response(response))
return filter_non_2xx_response(_scrub_response(_strip_image_b64_payloads(response)))
def _safe_body_matcher(r1, r2) -> None:
@@ -172,8 +265,84 @@ def _strip_headers(headers, names: Iterable[str]) -> None:
pass
def _normalize_multipart_boundary(request) -> None:
"""Rewrite random multipart boundaries to a fixed string in-place.
httpx generates a fresh ``boundary=<random hex>`` for every
multipart request via ``os.urandom``. Without normalization, the
request body bytes differ across runs even when everything else is
identical, the ``safe_body`` matcher misses, and the persister
keeps appending new episodes until ``MAX_EPISODES_PER_CASSETTE``
refuses the save — leaving audio-transcription tests effectively
unmocked. Replacing the boundary in both the Content-Type header
and the body bytes makes the request deterministic.
Idempotent — vcrpy invokes this hook multiple times per request,
so the second invocation sees ``boundary=vcr-static-boundary``
already and short-circuits.
"""
headers = getattr(request, "headers", None)
if headers is None:
return
content_type_key = None
content_type_value = None
try:
for key in list(headers.keys()):
if str(key).lower() == "content-type":
content_type_key = key
value = headers[key]
content_type_value = value if isinstance(value, str) else str(value)
break
except AttributeError:
return
if not content_type_value or "multipart/" not in content_type_value.lower():
return
fixed_param = f"boundary={VCR_FIXED_MULTIPART_BOUNDARY}"
if fixed_param in content_type_value:
return
match = re.search(r"boundary=([^\s;]+)", content_type_value)
if not match:
return
current_boundary = match.group(1).strip('"')
if current_boundary == VCR_FIXED_MULTIPART_BOUNDARY:
return
try:
headers[content_type_key] = content_type_value.replace(
match.group(0), fixed_param
)
except (TypeError, AttributeError):
return
body = getattr(request, "body", None)
if body is None:
return
if isinstance(body, (bytes, bytearray)):
try:
new_body = bytes(body).replace(
current_boundary.encode("utf-8"),
VCR_FIXED_MULTIPART_BOUNDARY.encode("utf-8"),
)
except (TypeError, ValueError):
return
elif isinstance(body, str):
new_body = body.replace(current_boundary, VCR_FIXED_MULTIPART_BOUNDARY)
else:
return
try:
request.body = new_body
except (AttributeError, TypeError):
pass
def _before_record_request(request):
"""Fingerprint API keys, then scrub them.
"""Fingerprint API keys, scrub them, and normalize multipart boundaries.
Order matters in two ways:
@@ -187,7 +356,8 @@ def _before_record_request(request):
auth headers we already stripped, so re-hashing would yield
``"no-key"`` and the stored vs. incoming fingerprints would
diverge. Skip the recompute when the header is already set so
this hook is idempotent.
this hook is idempotent. The boundary normalizer is also
idempotent for the same reason.
"""
headers = getattr(request, "headers", None)
if headers is None:
@@ -199,6 +369,7 @@ def _before_record_request(request):
except (TypeError, AttributeError):
pass
_strip_headers(headers, FILTERED_REQUEST_HEADERS)
_normalize_multipart_boundary(request)
return request
tests/batches_tests/conftest.py
-45
View File
@@ -1,7 +1,6 @@
# conftest.py
import asyncio
import importlib
import os
import sys
@@ -12,16 +11,6 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm # noqa: E402,F401
from tests._vcr_conftest_common import ( # noqa: E402
VerboseReporterState,
apply_vcr_auto_marker_to_items,
record_vcr_outcome,
register_persister_if_enabled,
vcr_config_dict,
)
_verbose_state = VerboseReporterState()
@pytest.fixture(scope="session")
def event_loop():
@@ -31,37 +20,3 @@ def event_loop():
loop = asyncio.new_event_loop()
yield loop
loop.close()
@pytest.fixture(scope="module")
def vcr_config():
return vcr_config_dict()
def pytest_recording_configure(config, vcr):
register_persister_if_enabled(vcr)
@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_makereport(item, call):
outcome = yield
rep = outcome.get_result()
setattr(item, f"rep_{rep.when}", rep)
@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
yield
record_vcr_outcome(request, vcr)
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
def pytest_runtest_logreport(report):
_verbose_state.maybe_emit_verdict(report)
def pytest_collection_modifyitems(config, items):
apply_vcr_auto_marker_to_items(items)

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