scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_registry_audit_2026_09_10b

Browse source
This commit is contained in:
mateo 2026-09-11 13:04:19 +00:00
commit 598e863510
350 changed files with 26608 additions and 1479 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/ISSUE_TEMPLATE/config.yml vendored
View file

@ -1,4 +1,4 @@
blank_issues_enabled: true
blank_issues_enabled: false
contact_links:
- name: Schedule Demo
url: https://enterprise.litellm.ai/demo

28
.github/workflows/test-rust.yml vendored
View file

@ -4,8 +4,22 @@ on:
push:
paths:
- "litellm-rust/**"
- "litellm/rust_bridge/**"
- "tests/test_litellm_rust/**"
- "litellm/integrations/custom_logger.py"
- "litellm/litellm_core_utils/litellm_logging.py"
- "litellm/litellm_core_utils/logging_worker.py"
- "litellm/proxy/guardrails/**"
- "litellm/utils.py"
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/ocr/**"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"
- "uv.lock"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
@ -20,8 +34,22 @@ on:
- "litellm_**"
paths:
- "litellm-rust/**"
- "litellm/rust_bridge/**"
- "tests/test_litellm_rust/**"
- "litellm/integrations/custom_logger.py"
- "litellm/litellm_core_utils/litellm_logging.py"
- "litellm/litellm_core_utils/logging_worker.py"
- "litellm/proxy/guardrails/**"
- "litellm/utils.py"
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/ocr/**"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"
- "uv.lock"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"

40
.github/workflows/test-terraform-modules.yml vendored
View file

@ -25,13 +25,17 @@ concurrency:
cancel-in-progress: true
jobs:
aws-module:
name: fmt, validate, test (aws)
module:
name: fmt, validate, test (${{ matrix.module }})
runs-on: ubuntu-latest
timeout-minutes: 15
strategy:
fail-fast: false
matrix:
module: [aws, gcp]
defaults:
run:
working-directory: terraform/litellm/aws
working-directory: terraform/litellm/${{ matrix.module }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
@ -51,35 +55,7 @@ jobs:
- name: validate
run: terraform validate
# Plan-only, mock_provider-backed: no AWS credentials, no API calls.
# Plan-only, mock_provider-backed: no cloud credentials, no API calls.
- name: test
run: terraform test
gcp-module:
name: fmt, validate, test (gcp)
runs-on: ubuntu-latest
timeout-minutes: 15
defaults:
run:
working-directory: terraform/litellm/gcp
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # v3.1.2
with:
terraform_version: 1.13.3
terraform_wrapper: false
- name: fmt
run: terraform fmt -recursive -check -diff
- name: init
run: terraform init -backend=false -input=false
- name: validate
run: terraform validate
- name: test
run: terraform test

1
.github/workflows/test-unit.yml vendored
View file

@ -200,6 +200,7 @@ jobs:
tests/test_litellm/proxy/types_utils
tests/test_litellm/proxy/logging_endpoints
tests/test_litellm/proxy/test_*.py
tests/test_gateway
workers: 4
reruns: 2
timeout-minutes: 20

2
AGENTS.md
View file

@ -1 +1,3 @@
Read @CLAUDE.md for coding guidelines
Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or lint budgets to satisfy a check

19
Dockerfile
View file

@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.25.2
ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
FROM $UV_IMAGE AS uvbin
FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
ARG PGBOUNCER_VERSION
ARG PGBOUNCER_SHA256
USER root
RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
WORKDIR /build
RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \
tar xzf pgbouncer.tar.gz --strip-components=1 && \
./configure --prefix=/usr/local --with-openssl=/usr && \
make -j"$(nproc)" pgbouncer && \
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
# Admin UI builder. Pinned to the build platform so the architecture-independent
# Next.js static export compiles once natively even in a multi-arch build,
# instead of once per target arch under QEMU.
@ -110,7 +126,8 @@ USER root
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}" \

19
docker/Dockerfile.database
View file

@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.25.2
ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
FROM $UV_IMAGE AS uvbin
FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
ARG PGBOUNCER_VERSION
ARG PGBOUNCER_SHA256
USER root
RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
WORKDIR /build
RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \
tar xzf pgbouncer.tar.gz --strip-components=1 && \
./configure --prefix=/usr/local --with-openssl=/usr && \
make -j"$(nproc)" pgbouncer && \
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
# Admin UI builder. Pinned to the build platform so the architecture-independent
# Next.js static export compiles once natively even in a multi-arch build,
# instead of once per target arch under QEMU.
@ -101,7 +117,8 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}" \

19
docker/Dockerfile.non_root
View file

@ -7,9 +7,25 @@ ARG PROXY_EXTRAS_SOURCE=published
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.25.2
ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
FROM $UV_IMAGE AS uvbin
FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
ARG PGBOUNCER_VERSION
ARG PGBOUNCER_SHA256
USER root
RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
WORKDIR /build
RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \
tar xzf pgbouncer.tar.gz --strip-components=1 && \
./configure --prefix=/usr/local --with-openssl=/usr && \
make -j"$(nproc)" pgbouncer && \
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
# Admin UI builder. Pinned to the build platform so the architecture-independent
# Next.js static export compiles once natively even in a multi-arch build,
# instead of once per target arch under QEMU.
@ -128,8 +144,9 @@ 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 python-3.13 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent && break || sleep 5; \
done
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
# Copy only what runtime needs. The application is installed inside the venv;
# the rest of the builder's /app is source and build metadata that must not

28
gateway/Dockerfile
View file

@ -1,9 +1,25 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.25.2
ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
FROM $UV_IMAGE AS uvbin
FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
ARG PGBOUNCER_VERSION
ARG PGBOUNCER_SHA256
USER root
RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
WORKDIR /build
RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \
tar xzf pgbouncer.tar.gz --strip-components=1 && \
./configure --prefix=/usr/local --with-openssl=/usr && \
make -j"$(nproc)" pgbouncer && \
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
# ---------- Builder ----------
FROM $LITELLM_BUILD_IMAGE AS builder
@ -61,6 +77,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra bedrock-realtime \
--python python3.13
# PYTHONPATH=/app makes the source tree shadow the installed package, so the
# compiled Rust extension must live next to the source or it is never imported.
RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
prisma generate --schema=./schema.prisma
@ -73,7 +93,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN for i in 1 2 3; do \
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -90,15 +110,17 @@ ENV HOME=/home/nonroot \
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder /opt/prisma /opt/prisma
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete && \
chmod -R a+rX /opt/prisma && \
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \
python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()"
USER nonroot
EXPOSE 4000/tcp
ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh python -m gateway.launch --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
CMD ["--host", "0.0.0.0", "--port", "4000"]

77
gateway/launch.py Normal file
View file

@ -0,0 +1,77 @@
"""Gateway supervisor: assemble DATABASE_URL, start the in-container PgBouncer, then run uvicorn.
``gateway/main.py`` assembles ``DATABASE_URL`` inside every uvicorn worker, which
is fine for a plain Postgres URL but not for the pooler: PgBouncer must be
started exactly once per pod, before the workers fork, and the workers must be
handed the loopback URL it listens on. A pre-existing ``DATABASE_URL`` wins in
``DatabaseURLSettings.apply_to_env`` under password auth, and one marked pooled
wins under token auth too, so exporting it here is enough for every worker to
pick the pooled URL up unchanged.
Run with:
python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000
"""
import os
import sys
from collections.abc import Callable, Mapping, Sequence
from typing import Final
from uvicorn.main import main as uvicorn_main
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
from litellm.proxy.db.pgbouncer import (
PgBouncerError,
PgBouncerSettings,
export_pooled_database_url,
start_in_container_pgbouncer,
)
GATEWAY_APP: Final = "gateway.main:app"
KEEPALIVE_FLAG: Final = "--timeout-keep-alive"
def uvicorn_argv(argv: Sequence[str], environ: Mapping[str, str]) -> tuple[str, ...]:
"""Honor ``KEEPALIVE_TIMEOUT`` like ``proxy_cli.py`` does, unless the flag was passed explicitly."""
keepalive: Final = environ.get("KEEPALIVE_TIMEOUT")
if keepalive is None or any(arg == KEEPALIVE_FLAG or arg.startswith(f"{KEEPALIVE_FLAG}=") for arg in argv):
return (GATEWAY_APP, *argv)
return (GATEWAY_APP, *argv, KEEPALIVE_FLAG, keepalive)
def pool_database_url(
settings: DatabaseURLSettings,
pgbouncer: PgBouncerSettings,
environ: Mapping[str, str],
) -> str | PgBouncerError | None:
"""Start the in-container PgBouncer and return its loopback URL, or None when ``pgbouncer.enabled`` is off.
The upstream URL is whatever ``apply_to_env`` assembled from the discrete
``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Under token
auth the pooler mints and renews the upstream token itself.
"""
if not pgbouncer.enabled:
return None
upstream_url: Final = environ.get("DATABASE_URL")
if upstream_url is None:
return PgBouncerError("LITELLM_PGBOUNCER_ENABLED is set but no DATABASE_URL could be assembled")
return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth=settings.token_auth())
def _serve(argv: Sequence[str]) -> None:
uvicorn_main(tuple(argv), prog_name="uvicorn")
def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
settings: Final = DatabaseURLSettings.from_env()
settings.apply_to_env()
pooled_url: Final = pool_database_url(settings, PgBouncerSettings(), os.environ)
if isinstance(pooled_url, PgBouncerError):
sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {pooled_url.reason}")
if pooled_url is not None:
export_pooled_database_url(pooled_url)
serve(uvicorn_argv(argv, os.environ))
if __name__ == "__main__":
main(sys.argv[1:])

160
helm/litellm-helm/templates/_helpers.tpl
View file

@ -161,3 +161,163 @@ taken before the change, which by that point no longer exists.
{{- fail (printf "postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got %q). An unpinned tag can start a different PostgreSQL major against the existing data directory, which makes the database unreadable and is not recoverable in place. Crossing a major version requires a dump and restore." $tag) -}}
{{- end -}}
{{- end -}}
{{/*
Environment shared by the proxy container and the opt-in collector sidecar:
database, pgbouncer, master key, redis, user envVars. Both containers must see
the same DATABASE_URL and REDIS_* so the sidecar reaches the pod's pgbouncer
and the same spend transaction buffer.
*/}}
{{- define "litellm.proxyEnv" -}}
- name: HOST
value: "{{ .Values.listen | default "0.0.0.0" }}"
- name: PORT
value: {{ .Values.service.port | quote}}
{{- if .Values.db.deployStandalone }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: username
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: password
- name: DATABASE_HOST
value: {{ .Release.Name }}-postgresql
- name: DATABASE_NAME
value: litellm
{{- else if .Values.db.useExisting }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.usernameKey }}
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.passwordKey }}
- name: DATABASE_HOST
{{- if .Values.db.secret.endpointKey }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.endpointKey }}
{{- else }}
value: {{ .Values.db.endpoint }}
{{- end }}
- name: DATABASE_NAME
value: {{ .Values.db.database }}
- name: DATABASE_URL
value: {{ .Values.db.url | quote }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
- name: DATABASE_READER_HOST
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaEndpointKey }}
{{- 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 }}
{{- if .Values.db.connectionPool.enabled }}
- name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: {{ .Values.db.connectionPool.maxDbConnections | quote }}
- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: {{ .Values.db.connectionPool.maxClientConn | quote }}
{{- end }}
- name: PROXY_MASTER_KEY
valueFrom:
secretKeyRef:
name: {{ .Values.masterkeySecretName | default (printf "%s-masterkey" (include "litellm.fullname" .)) }}
key: {{ .Values.masterkeySecretKey | default "masterkey" }}
{{- if .Values.redis.enabled }}
- name: REDIS_HOST
value: {{ include "litellm.redis.serviceName" . }}
- name: REDIS_PORT
value: {{ include "litellm.redis.port" . | quote }}
- name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "redis.secretName" .Subcharts.redis }}
key: {{include "redis.secretPasswordKey" .Subcharts.redis }}
{{- end }}
{{- /*
Inject LITELLM_LOG only when envVars does not already define it.
*/}}
{{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }}
- name: LITELLM_LOG
value: {{ .Values.logLevel | quote }}
{{- end }}
{{- if .Values.envVars }}
{{- range $key, $val := .Values.envVars }}
- name: {{ $key }}
value: {{ $val | quote }}
{{- end }}
{{- end }}
{{- with .Values.extraEnvVars }}
{{ toYaml . }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
# one DB on every rollout. Placed last (after envVars and
# extraEnvVars) so this override can't be silently shadowed by a
# user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env
# semantics — same pattern the migrations Job uses.
- name: DISABLE_SCHEMA_UPDATE
value: "true"
{{- end }}
{{- end -}}
{{/*
Proxy-only metering and metrics env. The collector sidecar serves no HTTP
traffic, so it gets neither.
*/}}
{{- define "litellm.proxyMetricsEnv" -}}
{{- if .Values.billingMetrics.enabled }}
{{ include "litellm.billingMetricsEnv" . }}
{{- end }}
{{- if .Values.metricsServer.enabled }}
{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }}
{{- fail "metricsServer.port must differ from service.port" }}
{{- end }}
- name: PROMETHEUS_METRICS_PORT
value: {{ .Values.metricsServer.port | quote }}
{{- end }}
{{- end -}}
{{/*
Directory of the collector's unix socket, shared between the two containers
through an emptyDir. Empty when the sidecar is off or uses 127.0.0.1 TCP.
*/}}
{{- define "litellm.collector.socketDir" -}}
{{- if and .Values.collector.enabled (hasPrefix "unix://" .Values.collector.address) -}}
{{- dir (trimPrefix "unix://" .Values.collector.address) -}}
{{- end -}}
{{- end -}}
{{- define "litellm.collectorEnv" -}}
- name: LITELLM_COLLECTOR_ENABLED
value: "true"
- name: LITELLM_COLLECTOR_ADDRESS
value: {{ .Values.collector.address | quote }}
- name: LITELLM_COLLECTOR_BUFFER_SIZE
value: {{ .Values.collector.bufferSize | quote }}
- name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: {{ .Values.collector.onUnavailable | quote }}
- name: LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS
value: {{ .Values.collector.drainTimeoutSeconds | quote }}
{{- end -}}

172
helm/litellm-helm/templates/deployment.yaml
View file

@ -56,118 +56,10 @@ spec:
image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.image.pullPolicy }}
env:
- name: HOST
value: "{{ .Values.listen | default "0.0.0.0" }}"
- name: PORT
value: {{ .Values.service.port | quote}}
{{- if .Values.db.deployStandalone }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: username
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: password
- name: DATABASE_HOST
value: {{ .Release.Name }}-postgresql
- name: DATABASE_NAME
value: litellm
{{- else if .Values.db.useExisting }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.usernameKey }}
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.passwordKey }}
- name: DATABASE_HOST
{{- if .Values.db.secret.endpointKey }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.endpointKey }}
{{- else }}
value: {{ .Values.db.endpoint }}
{{- end }}
- name: DATABASE_NAME
value: {{ .Values.db.database }}
- name: DATABASE_URL
value: {{ .Values.db.url | quote }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
- name: DATABASE_READER_HOST
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaEndpointKey }}
{{- 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:
name: {{ .Values.masterkeySecretName | default (printf "%s-masterkey" (include "litellm.fullname" .)) }}
key: {{ .Values.masterkeySecretKey | default "masterkey" }}
{{- if .Values.redis.enabled }}
- name: REDIS_HOST
value: {{ include "litellm.redis.serviceName" . }}
- name: REDIS_PORT
value: {{ include "litellm.redis.port" . | quote }}
- name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "redis.secretName" .Subcharts.redis }}
key: {{include "redis.secretPasswordKey" .Subcharts.redis }}
{{- end }}
{{- /*
Inject LITELLM_LOG only when envVars does not already define it.
*/}}
{{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }}
- name: LITELLM_LOG
value: {{ .Values.logLevel | quote }}
{{- end }}
{{- if .Values.envVars }}
{{- range $key, $val := .Values.envVars }}
- name: {{ $key }}
value: {{ $val | quote }}
{{- end }}
{{- end }}
{{- with .Values.extraEnvVars }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.metricsServer.enabled }}
{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }}
{{- fail "metricsServer.port must differ from service.port" }}
{{- end }}
- name: PROMETHEUS_METRICS_PORT
value: {{ .Values.metricsServer.port | quote }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
# one DB on every rollout. Placed last (after envVars and
# extraEnvVars) so this override can't be silently shadowed by a
# user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env
# semantics — same pattern the migrations Job uses.
- name: DISABLE_SCHEMA_UPDATE
value: "true"
{{- include "litellm.proxyEnv" . | nindent 12 }}
{{- include "litellm.proxyMetricsEnv" . | nindent 12 }}
{{- if .Values.collector.enabled }}
{{- include "litellm.collectorEnv" . | nindent 12 }}
{{- end }}
envFrom:
{{- range .Values.environmentSecrets }}
@ -245,6 +137,10 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- if include "litellm.collector.socketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.collector.socketDir" . }}
{{- end }}
{{- with .Values.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
@ -252,6 +148,53 @@ spec:
lifecycle:
{{- toYaml . | nindent 12 }}
{{- end }}
{{- if .Values.collector.enabled }}
- name: {{ include "litellm.name" . }}-collector
securityContext:
{{- toYaml .Values.securityContext | nindent 12 }}
image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.image.pullPolicy }}
command: {{ toYaml .Values.collector.command | nindent 12 }}
env:
{{- include "litellm.proxyEnv" . | nindent 12 }}
{{- include "litellm.collectorEnv" . | nindent 12 }}
- name: LITELLM_JOB_ROLE
value: collector
{{- if not (hasKey (default dict .Values.envVars) "CONFIG_FILE_PATH") }}
- name: CONFIG_FILE_PATH
value: /etc/litellm/config.yaml
{{- end }}
envFrom:
{{- range .Values.environmentSecrets }}
- secretRef:
name: {{ . }}
{{- end }}
{{- range .Values.environmentConfigMaps }}
- configMapRef:
name: {{ . }}
{{- end }}
resources:
{{- toYaml .Values.collector.resources | nindent 12 }}
volumeMounts:
- name: litellm-config
mountPath: /etc/litellm/config.yaml
subPath: config.yaml
{{- if include "litellm.collector.socketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.collector.socketDir" . }}
{{- end }}
{{ if .Values.securityContext.readOnlyRootFilesystem }}
- name: tmp
mountPath: /tmp
- name: cache
mountPath: /.cache
- name: npm
mountPath: /.npm
{{- end }}
{{- with .Values.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- end }}
{{- with .Values.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
@ -280,6 +223,11 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- if include "litellm.collector.socketDir" . }}
- name: collector-socket
emptyDir:
sizeLimit: 1Mi
{{- end }}
{{- with .Values.volumes }}
{{- toYaml . | nindent 8 }}
{{- end }}

28
helm/litellm-helm/templates/hpa.yaml
View file

@ -18,6 +18,15 @@ spec:
{{- end }}
metrics:
{{- if .Values.autoscaling.targetCPUUtilizationPercentage }}
{{- if and .Values.collector.enabled .Values.collector.scaleOnProxyContainerCpu }}
- type: ContainerResource
containerResource:
name: cpu
container: {{ include "litellm.name" . }}
target:
type: Utilization
averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
{{- else }}
- type: Resource
resource:
name: cpu
@ -25,6 +34,7 @@ spec:
type: Utilization
averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
{{- end }}
{{- end }}
{{- if .Values.autoscaling.targetMemoryUtilizationPercentage }}
- type: Resource
resource:
@ -33,4 +43,22 @@ spec:
type: Utilization
averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }}
{{- end }}
{{- with .Values.autoscaling.targetRequestsPerSecond }}
- type: Pods
pods:
metric:
name: litellm_requests_per_second
target:
type: AverageValue
averageValue: {{ toJson . | trimAll "\"" | quote }}
{{- end }}
{{- with .Values.autoscaling.targetTokensPerSecond }}
- type: Pods
pods:
metric:
name: litellm_tokens_per_second
target:
type: AverageValue
averageValue: {{ toJson . | trimAll "\"" | quote }}
{{- end }}
{{- end }}

21
helm/litellm-helm/templates/keda.yaml
View file

@ -23,6 +23,27 @@ spec:
triggers:
{{- with .Values.keda.triggers }}
{{- toYaml . | nindent 2 }}
{{- end }}
{{- $prom := .Values.keda.prometheus }}
{{- if or $prom.requestsPerSecond $prom.tokensPerSecond }}
{{- if not $prom.serverAddress }}
{{- fail "keda.prometheus.serverAddress is required when keda.prometheus.requestsPerSecond or tokensPerSecond is set" }}
{{- end }}
{{- $selector := printf "namespace=%q,job=%q" .Release.Namespace (printf "%s%s" (include "litellm.fullname" .) (ternary "-metrics" "" .Values.metricsServer.enabled)) }}
{{- with $prom.requestsPerSecond }}
- type: prometheus
metadata:
serverAddress: {{ $prom.serverAddress | quote }}
threshold: {{ toJson . | trimAll "\"" | quote }}
query: {{ printf "sum(rate(litellm_proxy_total_requests_metric_total{%s}[1m]))" $selector | quote }}
{{- end }}
{{- with $prom.tokensPerSecond }}
- type: prometheus
metadata:
serverAddress: {{ $prom.serverAddress | quote }}
threshold: {{ toJson . | trimAll "\"" | quote }}
query: {{ printf "sum(rate(litellm_total_tokens_metric_total{%s}[1m]))" $selector | quote }}
{{- end }}
{{- end }}
advanced:
restoreToOriginalReplicaCount: {{ .Values.keda.restoreToOriginalReplicaCount }}

272
helm/litellm-helm/tests/collector_tests.yaml Normal file
View file

@ -0,0 +1,272 @@
suite: test collector sidecar
templates:
- deployment.yaml
- hpa.yaml
- configmap-litellm.yaml
tests:
- it: should run the proxy alone with no collector env by default
template: deployment.yaml
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 1
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ENABLED
value: "true"
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
- it: should add the sidecar on the same image and point both containers at the unix socket
template: deployment.yaml
set:
image.tag: test
db.connectionPool.enabled: true
collector.enabled: true
collector.resources:
requests:
cpu: 500m
memory: 1Gi
limits:
cpu: "1"
memory: 2Gi
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 2
- equal:
path: spec.template.spec.containers[1].name
value: litellm-collector
- equal:
path: spec.template.spec.containers[1].image
value: ghcr.io/berriai/litellm:test
- equal:
path: spec.template.spec.containers[1].command
value: [python, -m, litellm.proxy.collector]
- equal:
path: spec.template.spec.containers[1].resources.requests.cpu
value: 500m
- equal:
path: spec.template.spec.containers[1].resources.limits.memory
value: 2Gi
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_BUFFER_SIZE
value: "1000"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: fallback
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_JOB_ROLE
value: collector
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_JOB_ROLE
value: collector
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
- contains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /etc/litellm/config.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_HOST
value: RELEASE-NAME-postgresql
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: RELEASE-NAME-litellm-dbcredentials
key: password
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: litellm-config
mountPath: /etc/litellm/config.yaml
subPath: config.yaml
- contains:
path: spec.template.spec.volumes
content:
name: collector-socket
emptyDir:
sizeLimit: 1Mi
- it: should skip the socket volume and pass the policy through on tcp transport
template: deployment.yaml
set:
collector.enabled: true
collector.address: tcp://127.0.0.1:4100
collector.onUnavailable: drop
collector.bufferSize: 50
envVars:
CONFIG_FILE_PATH: /custom/config.yaml
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 2
- notContains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /etc/litellm/config.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /custom/config.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: tcp://127.0.0.1:4100
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: drop
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_BUFFER_SIZE
value: "50"
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
- it: should keep metrics and billing env on the proxy container only
template: deployment.yaml
set:
collector.enabled: true
metricsServer.enabled: true
metricsServer.port: 9090
billingMetrics.enabled: true
billingMetrics.endpoint: https://metering.example.com
billingMetrics.secretName: billing-mtls
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
value: "9090"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://metering.example.com
- notContains:
path: spec.template.spec.containers[1].env
content:
name: PROMETHEUS_METRICS_PORT
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
any: true
- notContains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: billing-metrics-mtls
any: true
- it: should give the sidecar the same scratch mounts as the proxy on a read-only root
template: deployment.yaml
set:
collector.enabled: true
securityContext.readOnlyRootFilesystem: true
asserts:
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: npm
mountPath: /.npm
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: cache
mountPath: /.cache
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: tmp
mountPath: /tmp
- it: should keep the pod-wide cpu metric unless asked to scale on the proxy container
template: hpa.yaml
set:
autoscaling.enabled: true
collector.enabled: true
asserts:
- equal: { path: "spec.metrics[0].type", value: Resource }
- equal: { path: "spec.metrics[0].resource.name", value: cpu }
- it: should scale on the proxy container's cpu only when opted in
template: hpa.yaml
set:
autoscaling.enabled: true
collector.enabled: true
collector.scaleOnProxyContainerCpu: true
asserts:
- equal: { path: "spec.metrics[0].type", value: ContainerResource }
- equal: { path: "spec.metrics[0].containerResource.name", value: cpu }
- equal: { path: "spec.metrics[0].containerResource.container", value: litellm }
- equal: { path: "spec.metrics[0].containerResource.target.averageUtilization", value: 60 }
- isNull: { path: "spec.metrics[0].resource" }
- it: should not switch to the container metric while the sidecar is off
template: hpa.yaml
set:
autoscaling.enabled: true
collector.scaleOnProxyContainerCpu: true
asserts:
- equal: { path: "spec.metrics[0].type", value: Resource }

112
helm/litellm-helm/tests/connection_pool_tests.yaml Normal file
View file

@ -0,0 +1,112 @@
suite: test in-container connection pool
templates:
- deployment.yaml
- configmap-litellm.yaml
tests:
- it: should not emit pgbouncer env vars by default
template: deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- it: should enable the pool with the default sizing when connectionPool.enabled is set
template: deployment.yaml
set:
db.connectionPool.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "1000"
- it: should pass custom sizing through as strings next to the worker count
template: deployment.yaml
set:
numWorkers: 4
db.connectionPool.enabled: true
db.connectionPool.maxDbConnections: 8
db.connectionPool.maxClientConn: 400
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "400"
- contains:
path: spec.template.spec.containers[0].args
content: "4"
- it: should give the collector sidecar the same pool env as the proxy container
template: deployment.yaml
set:
collector.enabled: true
db.connectionPool.enabled: true
db.connectionPool.maxDbConnections: 8
db.connectionPool.maxClientConn: 400
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: litellm-collector
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "400"
- it: should give the collector sidecar no pool env when the pool is off
template: deployment.yaml
set:
collector.enabled: true
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: litellm-collector
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
any: true

78
helm/litellm-helm/tests/hpa_tests.yaml
View file

@ -61,6 +61,84 @@ tests:
- equal: { path: "spec.metrics[1].resource.name", value: memory }
- equal: { path: "spec.metrics[1].resource.target.averageUtilization", value: 80 }
- it: "renders no workload metrics by default"
set:
autoscaling.enabled: true
autoscaling.targetMemoryUtilizationPercentage: 80
asserts:
- lengthEqual: { path: spec.metrics, count: 2 }
- notContains: { path: spec.metrics, content: { type: Pods }, any: true }
- it: "adds a requests-per-second Pods metric after the cpu metric"
set:
autoscaling.enabled: true
autoscaling.targetRequestsPerSecond: 90
asserts:
- lengthEqual: { path: spec.metrics, count: 2 }
- equal: { path: "spec.metrics[0].resource.name", value: cpu }
- equal:
path: "spec.metrics[1]"
value:
type: Pods
pods:
metric: { name: litellm_requests_per_second }
target: { type: AverageValue, averageValue: "90" }
- it: "adds a tokens-per-second Pods metric on its own"
set:
autoscaling.enabled: true
autoscaling.targetTokensPerSecond: 6M
asserts:
- lengthEqual: { path: spec.metrics, count: 2 }
- equal:
path: "spec.metrics[1]"
value:
type: Pods
pods:
metric: { name: litellm_tokens_per_second }
target: { type: AverageValue, averageValue: "6M" }
- notContains:
path: spec.metrics
content: { type: Pods, pods: { metric: { name: litellm_requests_per_second } } }
any: true
- it: "renders requests, tokens, cpu and memory metrics together"
set:
autoscaling.enabled: true
autoscaling.targetMemoryUtilizationPercentage: 80
autoscaling.targetRequestsPerSecond: 90
autoscaling.targetTokensPerSecond: 6000000
asserts:
- lengthEqual: { path: spec.metrics, count: 4 }
- equal: { path: "spec.metrics[0].resource.name", value: cpu }
- equal: { path: "spec.metrics[1].resource.name", value: memory }
- equal: { path: "spec.metrics[2].pods.metric.name", value: litellm_requests_per_second }
- equal: { path: "spec.metrics[2].pods.target.averageValue", value: "90" }
- equal: { path: "spec.metrics[3].pods.metric.name", value: litellm_tokens_per_second }
- equal: { path: "spec.metrics[3].pods.target.averageValue", value: "6000000" }
- it: "scales on workload metrics alone when the cpu target is cleared"
set:
autoscaling.enabled: true
autoscaling.targetCPUUtilizationPercentage: null
autoscaling.targetRequestsPerSecond: 90
autoscaling.targetTokensPerSecond: 6000000
asserts:
- lengthEqual: { path: spec.metrics, count: 2 }
- notContains: { path: spec.metrics, content: { type: Resource }, any: true }
- equal: { path: "spec.metrics[0].pods.metric.name", value: litellm_requests_per_second }
- equal: { path: "spec.metrics[1].pods.metric.name", value: litellm_tokens_per_second }
- notMatchRegexRaw: { pattern: per_minute }
- it: "ignores the per-minute keys, which the chart never shipped"
set:
autoscaling.enabled: true
autoscaling.targetRequestsPerMinute: 5400
autoscaling.targetTokensPerMinute: 360000000
asserts:
- lengthEqual: { path: spec.metrics, count: 1 }
- notContains: { path: spec.metrics, content: { type: Pods }, any: true }
- it: "renders no hpa when autoscaling is disabled"
asserts:
- hasDocuments: { count: 0 }

106
helm/litellm-helm/tests/keda_tests.yaml Normal file
View file

@ -0,0 +1,106 @@
suite: "keda"
templates:
- keda.yaml
release:
name: rel
namespace: llm
tests:
- it: "renders no scaled object by default"
asserts:
- hasDocuments: { count: 0 }
- it: "passes user triggers through and adds no prometheus triggers by default"
set:
keda.enabled: true
keda.triggers:
- type: cpu
metricType: Utilization
metadata: { value: "60" }
asserts:
- isKind: { of: ScaledObject }
- equal:
path: spec.triggers
value:
- type: cpu
metricType: Utilization
metadata: { value: "60" }
- it: "scales on release-wide requests per second divided by the per-replica target"
set:
keda.enabled: true
keda.prometheus.serverAddress: http://prometheus-operated.monitoring.svc:9090
keda.prometheus.requestsPerSecond: 90
asserts:
- lengthEqual: { path: spec.triggers, count: 1 }
- equal:
path: "spec.triggers[0]"
value:
type: prometheus
metadata:
serverAddress: http://prometheus-operated.monitoring.svc:9090
threshold: "90"
query: sum(rate(litellm_proxy_total_requests_metric_total{namespace="llm",job="rel-litellm"}[1m]))
- it: "scales on tokens per second on its own"
set:
keda.enabled: true
keda.prometheus.serverAddress: http://prom:9090
keda.prometheus.tokensPerSecond: 6000000
asserts:
- lengthEqual: { path: spec.triggers, count: 1 }
- equal: { path: "spec.triggers[0].type", value: prometheus }
- equal: { path: "spec.triggers[0].metadata.threshold", value: "6000000" }
- equal:
path: "spec.triggers[0].metadata.query"
value: sum(rate(litellm_total_tokens_metric_total{namespace="llm",job="rel-litellm"}[1m]))
- it: "appends requests and tokens triggers after user triggers and selects the metrics service job"
set:
keda.enabled: true
metricsServer.enabled: true
keda.triggers:
- type: cpu
metricType: Utilization
metadata: { value: "60" }
keda.prometheus.serverAddress: http://prom:9090
keda.prometheus.requestsPerSecond: 90
keda.prometheus.tokensPerSecond: 6000000
asserts:
- lengthEqual: { path: spec.triggers, count: 3 }
- equal: { path: "spec.triggers[0].type", value: cpu }
- equal: { path: "spec.triggers[1].metadata.threshold", value: "90" }
- equal:
path: "spec.triggers[1].metadata.query"
value: sum(rate(litellm_proxy_total_requests_metric_total{namespace="llm",job="rel-litellm-metrics"}[1m]))
- equal: { path: "spec.triggers[2].metadata.threshold", value: "6000000" }
- equal:
path: "spec.triggers[2].metadata.query"
value: sum(rate(litellm_total_tokens_metric_total{namespace="llm",job="rel-litellm-metrics"}[1m]))
- notMatchRegexRaw: { pattern: "\\* *60|per_minute|PerMinute" }
- it: "ignores the per-minute keys, which the chart never shipped"
set:
keda.enabled: true
keda.prometheus.serverAddress: http://prom:9090
keda.prometheus.requestsPerMinute: 5400
keda.prometheus.tokensPerMinute: 360000000
asserts:
- isKind: { of: ScaledObject }
- isNullOrEmpty: { path: spec.triggers }
- it: "refuses a workload target without a prometheus server address"
set:
keda.enabled: true
keda.prometheus.requestsPerSecond: 90
asserts:
- failedTemplate:
errorMessage: keda.prometheus.serverAddress is required when keda.prometheus.requestsPerSecond or tokensPerSecond is set
- it: "yields to the hpa when both autoscalers are enabled"
set:
autoscaling.enabled: true
keda.enabled: true
keda.prometheus.serverAddress: http://prom:9090
keda.prometheus.requestsPerSecond: 90
asserts:
- hasDocuments: { count: 0 }

92
helm/litellm-helm/values.yaml
View file

@ -190,6 +190,48 @@ metricsServer:
enabled: false
port: 4001
# Opt-in sidecar that runs the post-response spend pipeline (cost calculation,
# spend logs, spend counters, budget reservation reconciliation) so the proxy's
# uvicorn workers only serialise a compact typed event and go back to serving
# inference. Same image and tag as the proxy, second container in the same pod,
# fed over loopback (a unix socket on a shared emptyDir, or 127.0.0.1 TCP). It
# reuses the pod's in-container pgbouncer (db.connectionPool) and the same Redis
# spend transaction buffer, so the per-pod DB connection budget is unchanged.
# Delivery is at-most-once inside the pod: events already handed to the sidecar
# are lost if it crashes before writing them; events the workers could not hand
# over follow onUnavailable. Both containers drain on SIGTERM within
# terminationGracePeriodSeconds
collector:
enabled: false
# unix:///<dir>/<file>.sock (the <dir> becomes a shared emptyDir) or tcp://127.0.0.1:<port>
address: unix:///var/run/litellm/collector.sock
# Events each uvicorn worker holds in memory while the sidecar is slow or restarting
bufferSize: 1000
# fallback: run the pipeline in the worker when the sidecar is unreachable or the
# buffer is full (spend stays exact, that request costs proxy CPU again)
# drop: count and discard the event instead (spend under-reports)
onUnavailable: fallback
# How long the workers keep pushing buffered events on shutdown, and how long the
# sidecar keeps serving its open connections after SIGTERM
drainTimeoutSeconds: 10
command:
- python
- -m
- litellm.proxy.collector
# Sized independently of the proxy container; the pipeline is CPU bound
resources: {}
# requests:
# cpu: 500m
# memory: 1Gi
# limits:
# cpu: "1"
# memory: 2Gi
# When autoscaling.enabled, swap the pod-wide cpu Resource metric for an
# autoscaling/v2 ContainerResource metric on the proxy container only, so the
# sidecar's CPU never scales inference replicas. Needs Kubernetes 1.30+ (or the
# HPAContainerMetrics feature gate on 1.27 to 1.29)
scaleOnProxyContainerCpu: false
resources:
{}
# Unset by default so the chart installs on small clusters such as Minikube, and so an
@ -222,6 +264,25 @@ autoscaling:
# Memory is a floor to provision under 'resources', not a signal to scale on.
# targetMemoryUtilizationPercentage: 80
# behavior: {}
# Opt-in per-pod workload targets, rendered as autoscaling/v2 `Pods` metrics
# named `litellm_requests_per_second` and `litellm_tokens_per_second` with an
# AverageValue target, alongside whichever resource targets are set (the HPA
# follows the metric asking for the most replicas). A Prometheus Adapter must
# serve those two names on custom.metrics.k8s.io from the proxy's counters,
# grouped by the scrape target's `pod` label (enable serviceMonitor below so
# every pod is scraped on its own):
# litellm_requests_per_second:
# sum(rate(litellm_proxy_total_requests_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
# litellm_tokens_per_second:
# sum(rate(litellm_total_tokens_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
# rate() over [1m] is already per second, so no `* 60`. How fast the HPA
# reacts is set by that window, the scrape interval and the HPA sync period
# (15s by default), not by the unit: keep serviceMonitor.interval at 15s or
# faster so a 1m window holds at least 4 samples. averageValue takes SI
# suffixes, so "6M" is six million tokens per second per pod. Tokens are
# counted when a response completes, so TPS trails long streams.
targetRequestsPerSecond: ""
targetTokensPerSecond: ""
# Autoscaling with keda is mutually exclusive with hpa
keda:
@ -243,6 +304,23 @@ keda:
# metricName: http_requests_total
# threshold: '100'
# query: sum(rate(http_requests_total{deployment="my-deployment"}[2m]))
# First-class Prometheus triggers on the proxy's own request and token
# counters, appended to `triggers`. Each target is the per-second load one
# replica should carry: KEDA divides the release-wide
# `sum(rate(<counter>[1m]))` by it to pick the replica count. Thresholds
# are plain numbers (KEDA parses them as floats, no SI suffixes). The
# queries select samples by the release namespace and the `job` label the
# chart's ServiceMonitor produces (the metrics Service name), so enable
# serviceMonitor below together with metricsServer: the http port serves
# /metrics/ behind virtual-key auth and answers an unauthenticated scrape
# with 401. Reaction time comes from the [1m] window, the scrape interval
# and pollingInterval above, so keep both at 15s or faster. Tokens are
# counted at completion, so TPS trails long streams. serverAddress is
# required once either target is set.
prometheus:
serverAddress: ""
requestsPerSecond: ""
tokensPerSecond: ""
behavior: {}
# scaleDown:
# stabilizationWindowSeconds: 300
@ -319,6 +397,20 @@ db:
# only (e.g. when IAM_TOKEN_DB_AUTH supplies the token at runtime).
readReplicaUrl: ""
# In-container connection pool (PgBouncer, transaction mode) shared by every
# worker in the pod. Without it each --num_workers worker opens its own
# connection_limit connections to Postgres, so a pod's footprint against the
# database's connection ceiling is workers x connection_limit and grows with
# every replica. With it, the pod holds at most maxDbConnections upstream
# connections no matter how many workers run; the workers connect to the pool
# over loopback, with no extra network hop. Migrations still go straight to
# Postgres. Starting profile for numWorkers: 4 is maxDbConnections: 20, so
# a database with a 5000-connection ceiling fits roughly 200 replicas.
connectionPool:
enabled: false
maxDbConnections: 20
maxClientConn: 1000
# Use the Stackgres Helm chart to deploy an instance of a Stackgres cluster.
# The Stackgres Operator must already be installed within the target
# Kubernetes cluster.

45
helm/litellm/templates/_helpers.tpl
View file

@ -360,6 +360,20 @@ harmless no-op for the Job and authoritative for the app pods.
{{- end }}
{{- end -}}
{{/*
In-container PgBouncer env for the gateway container. Under IAM or Entra auth the pooler mints and renews the database token itself.
*/}}
{{- define "litellm.connectionPoolEnv" -}}
{{- with .Values.database.connectionPool -}}
- name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: {{ required "database.connectionPool.maxDbConnections is required when the pool is enabled" .maxDbConnections | quote }}
- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: {{ required "database.connectionPool.maxClientConn is required when the pool is enabled" .maxClientConn | quote }}
{{- end }}
{{- end -}}
{{/*
PodDisruptionBudget shared by gateway, backend, and ui.
@ -443,3 +457,34 @@ ImplementationSpecific
{{- end -}}
{{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}}
{{/*
Directory of the collector's unix socket, shared by the gateway and
collector containers through an emptyDir. Empty when the sidecar is off
or gateway.collector.address is a tcp://127.0.0.1:<port> address.
*/}}
{{- define "litellm.gateway.collectorSocketDir" -}}
{{- if and .Values.gateway.collector.enabled (hasPrefix "unix://" .Values.gateway.collector.address) -}}
{{- dir (trimPrefix "unix://" .Values.gateway.collector.address) -}}
{{- end -}}
{{- end -}}
{{/*
LITELLM_COLLECTOR_* env shared by the producer (gateway container) and the
consumer (collector container), so both agree on the transport and the
shutdown drain window.
*/}}
{{- define "litellm.gateway.collectorEnv" -}}
{{- with .Values.gateway.collector }}
- name: LITELLM_COLLECTOR_ENABLED
value: "true"
- name: LITELLM_COLLECTOR_ADDRESS
value: {{ .address | quote }}
- name: LITELLM_COLLECTOR_BUFFER_SIZE
value: {{ .bufferSize | quote }}
- name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: {{ .onUnavailable | quote }}
- name: LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS
value: {{ .drainTimeoutSeconds | quote }}
{{- end }}
{{- end -}}

62
helm/litellm/templates/gateway/deployment.yaml
View file

@ -61,6 +61,9 @@ spec:
- name: NUM_WORKERS
value: {{ .Values.gateway.numWorkers | quote }}
{{- end }}
{{- if .Values.database.connectionPool.enabled }}
{{- include "litellm.connectionPoolEnv" $ | nindent 12 }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
@ -71,8 +74,11 @@ spec:
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- if .Values.gateway.collector.enabled }}
{{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
{{- end }}
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
@ -83,6 +89,10 @@ spec:
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- if include "litellm.gateway.collectorSocketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
@ -142,10 +152,53 @@ spec:
resources:
{{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }}
{{- end }}
{{- if .Values.gateway.collector.enabled }}
- name: collector
image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
{{- with .Values.gateway.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
command:
- python
- -m
- litellm.proxy.collector
env:
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }}
{{- if .Values.gateway.config.create }}
- name: CONFIG_FILE_PATH
value: /app/config/config.yaml
{{- end }}
{{- if .Values.database.connectionPool.enabled }}
{{- include "litellm.connectionPoolEnv" $ | nindent 12 }}
{{- end }}
{{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
- name: LITELLM_JOB_ROLE
value: collector
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts (include "litellm.gateway.collectorSocketDir" .) }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if include "litellm.gateway.collectorSocketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
{{- end }}
{{- with .Values.gateway.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- end }}
resources:
{{- toYaml .Values.gateway.collector.resources | nindent 12 }}
{{- end }}
{{- with .Values.gateway.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
@ -156,6 +209,11 @@ spec:
- name: prometheus-multiproc
emptyDir: {}
{{- end }}
{{- if include "litellm.gateway.collectorSocketDir" . }}
- name: collector-socket
emptyDir:
sizeLimit: 1Mi
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}

28
helm/litellm/templates/gateway/hpa.yaml
View file

@ -15,6 +15,15 @@ spec:
maxReplicas: {{ .Values.gateway.hpa.maxReplicas }}
metrics:
{{- if .Values.gateway.hpa.targetCPUUtilizationPercentage }}
{{- if and .Values.gateway.collector.enabled .Values.gateway.collector.scaleOnGatewayContainerCpu }}
- type: ContainerResource
containerResource:
name: cpu
container: gateway
target:
type: Utilization
averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }}
{{- else }}
- type: Resource
resource:
name: cpu
@ -22,6 +31,7 @@ spec:
type: Utilization
averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }}
{{- end }}
{{- end }}
{{- if .Values.gateway.hpa.targetMemoryUtilizationPercentage }}
- type: Resource
resource:
@ -30,6 +40,24 @@ spec:
type: Utilization
averageUtilization: {{ .Values.gateway.hpa.targetMemoryUtilizationPercentage }}
{{- end }}
{{- with .Values.gateway.hpa.targetRequestsPerSecond }}
- type: Pods
pods:
metric:
name: litellm_requests_per_second
target:
type: AverageValue
averageValue: {{ toJson . | trimAll "\"" | quote }}
{{- end }}
{{- with .Values.gateway.hpa.targetTokensPerSecond }}
- type: Pods
pods:
metric:
name: litellm_tokens_per_second
target:
type: AverageValue
averageValue: {{ toJson . | trimAll "\"" | quote }}
{{- end }}
{{- with .Values.gateway.hpa.behavior }}
behavior:
{{- toYaml . | nindent 4 }}

28
helm/litellm/templates/gateway/servicemonitor.yaml Normal file
View file

@ -0,0 +1,28 @@
{{- if and .Values.gateway.enabled .Values.gateway.serviceMonitor.enabled }}
{{- if not .Values.gateway.metricsServer.enabled }}
{{- fail "gateway.serviceMonitor.enabled requires gateway.metricsServer.enabled: the http port serves /metrics/ behind virtual-key auth, so an unauthenticated scrape gets 401" }}
{{- end }}
apiVersion: monitoring.coreos.com/v1
kind: ServiceMonitor
metadata:
name: {{ include "litellm.gateway.fullname" . }}
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: gateway
{{- with .Values.gateway.serviceMonitor.labels }}
{{- toYaml . | nindent 4 }}
{{- end }}
spec:
selector:
matchLabels:
{{- include "litellm.gateway.selectorLabels" . | nindent 6 }}
namespaceSelector:
matchNames:
- {{ .Release.Namespace | quote }}
endpoints:
- port: metrics
path: /metrics/
interval: {{ .Values.gateway.serviceMonitor.interval }}
scrapeTimeout: {{ .Values.gateway.serviceMonitor.scrapeTimeout }}
scheme: http
{{- end }}

216
helm/litellm/tests/collector_tests.yaml Normal file
View file

@ -0,0 +1,216 @@
suite: test gateway collector sidecar
templates:
- gateway/configmap.yaml
- gateway/deployment.yaml
- gateway/hpa.yaml
values:
- ./values/required.yaml
tests:
- it: adds no sidecar, env, volume or container metric when the collector is off
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 1
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ENABLED
value: "true"
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
template: gateway/deployment.yaml
- equal:
path: spec.metrics[0].type
value: Resource
template: gateway/hpa.yaml
- it: runs the collector as a sidecar sharing env, config, the pod pool and a unix socket emptyDir, and scales on the gateway container only
set:
gateway.collector.enabled: true
gateway.collector.bufferSize: 250
gateway.collector.onUnavailable: drop
gateway.image.tag: v1.102.0
gateway.numWorkers: 4
database.connectionPool.enabled: true
database.connectionPool.maxDbConnections: 8
database.connectionPool.maxClientConn: 250
gateway.envSecrets:
- litellm-license
gateway.volumes:
- name: redis-ca
secret:
secretName: redis-ca
gateway.volumeMounts:
- name: redis-ca
mountPath: /etc/litellm/redis-ca
readOnly: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_BUFFER_SIZE
value: "250"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: drop
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].name
value: collector
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].image
value: ghcr.io/berriai/litellm-gateway:v1.102.0
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].command
value:
- python
- -m
- litellm.proxy.collector
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_JOB_ROLE
value: collector
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /app/config/config.yaml
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "250"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_HOST
value: postgres.example.com
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[1].env
content:
name: NUM_WORKERS
any: true
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].envFrom
value:
- secretRef:
name: litellm-license
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: redis-ca
mountPath: /etc/litellm/redis-ca
readOnly: true
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].resources.limits.cpu
value: "1"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.volumes
content:
name: collector-socket
emptyDir:
sizeLimit: 1Mi
template: gateway/deployment.yaml
- equal:
path: spec.metrics[0]
value:
type: ContainerResource
containerResource:
name: cpu
container: gateway
target:
type: Utilization
averageUtilization: 70
template: gateway/hpa.yaml
- it: uses loopback tcp without a socket volume and keeps the pod-wide cpu metric when asked
set:
gateway.collector.enabled: true
gateway.collector.address: tcp://127.0.0.1:4010
gateway.collector.scaleOnGatewayContainerCpu: false
asserts:
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: tcp://127.0.0.1:4010
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: collector-socket
any: true
template: gateway/deployment.yaml
- equal:
path: spec.metrics[0].type
value: Resource
template: gateway/hpa.yaml

204
helm/litellm/tests/connection_pool_tests.yaml Normal file
View file

@ -0,0 +1,204 @@
suite: test in-container connection pool env vars
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- backend/configmap.yaml
values:
- ./values/required.yaml
tests:
- it: renders no pool env by default
template: gateway/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
any: true
- it: enabled pool renders the three pgbouncer vars with the configured sizes
template: gateway/deployment.yaml
set:
gateway.numWorkers: 4
database.connectionPool.enabled: true
database.connectionPool.maxDbConnections: 8
database.connectionPool.maxClientConn: 250
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "250"
- contains:
path: spec.template.spec.containers[0].env
content:
name: NUM_WORKERS
value: "4"
- it: enabled pool uses the chart default sizes
template: gateway/deployment.yaml
set:
database.connectionPool.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "1000"
- it: backend never gets the pool env
template: backend/deployment.yaml
set:
database.connectionPool.enabled: true
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- it: collector sidecar gets the same pool env as the gateway container, the metrics sidecar none
template: gateway/deployment.yaml
set:
gateway.collector.enabled: true
gateway.metricsServer.enabled: true
database.connectionPool.enabled: true
database.connectionPool.maxDbConnections: 8
database.connectionPool.maxClientConn: 250
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: metrics
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- equal:
path: spec.template.spec.containers[2].name
value: collector
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "250"
- it: collector sidecar gets no pool env when the pool is off
template: gateway/deployment.yaml
set:
gateway.collector.enabled: true
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: collector
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
any: true
- it: pool with IAM auth renders both the pool and the token auth flag
template: gateway/deployment.yaml
set:
gateway.collector.enabled: true
database.connectionPool.enabled: true
database.writer.useIAMAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- it: pool with Entra auth renders both the pool and the token auth flag
template: gateway/deployment.yaml
set:
database.connectionPool.enabled: true
database.writer.useAzureEntraAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: AZURE_POSTGRESQL_AUTH
value: "true"
- it: IAM auth without the pool still renders
template: gateway/deployment.yaml
set:
database.writer.useIAMAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true

213
helm/litellm/tests/hpa_workload_metrics_tests.yaml Normal file
View file

@ -0,0 +1,213 @@
suite: test gateway HPA per-pod requests-per-second and tokens-per-second targets
templates:
- gateway/hpa.yaml
- gateway/servicemonitor.yaml
values:
- ./values/required.yaml
tests:
- it: scales on CPU and memory only by default
template: gateway/hpa.yaml
asserts:
- equal:
path: spec.metrics
value:
- type: Resource
resource:
name: cpu
target:
type: Utilization
averageUtilization: 70
- type: Resource
resource:
name: memory
target:
type: Utilization
averageUtilization: 80
- it: adds a requests-per-second Pods metric next to the resource metrics
template: gateway/hpa.yaml
set:
gateway.hpa.targetRequestsPerSecond: 90
asserts:
- lengthEqual:
path: spec.metrics
count: 3
- contains:
path: spec.metrics
content:
type: Resource
resource:
name: cpu
target:
type: Utilization
averageUtilization: 70
- equal:
path: spec.metrics[2]
value:
type: Pods
pods:
metric:
name: litellm_requests_per_second
target:
type: AverageValue
averageValue: "90"
- notContains:
path: spec.metrics
content:
type: Pods
pods:
metric:
name: litellm_tokens_per_second
any: true
- it: adds a tokens-per-second Pods metric on its own
template: gateway/hpa.yaml
set:
gateway.hpa.targetTokensPerSecond: 6M
asserts:
- lengthEqual:
path: spec.metrics
count: 3
- equal:
path: spec.metrics[2]
value:
type: Pods
pods:
metric:
name: litellm_tokens_per_second
target:
type: AverageValue
averageValue: "6M"
- notContains:
path: spec.metrics
content:
type: Pods
pods:
metric:
name: litellm_requests_per_second
any: true
- it: renders requests and tokens targets together and keeps CPU and memory
template: gateway/hpa.yaml
set:
gateway.hpa.targetRequestsPerSecond: 90
gateway.hpa.targetTokensPerSecond: 6000000
asserts:
- lengthEqual:
path: spec.metrics
count: 4
- equal:
path: spec.metrics[0].resource.name
value: cpu
- equal:
path: spec.metrics[1].resource.name
value: memory
- equal:
path: spec.metrics[2].pods.metric.name
value: litellm_requests_per_second
- equal:
path: spec.metrics[3].pods.metric.name
value: litellm_tokens_per_second
- equal:
path: spec.metrics[3].pods.target.averageValue
value: "6000000"
- it: scales on workload metrics alone when the resource targets are cleared
template: gateway/hpa.yaml
set:
gateway.hpa.targetCPUUtilizationPercentage: null
gateway.hpa.targetMemoryUtilizationPercentage: null
gateway.hpa.targetRequestsPerSecond: 90
gateway.hpa.targetTokensPerSecond: 6000000
asserts:
- lengthEqual:
path: spec.metrics
count: 2
- notContains:
path: spec.metrics
content:
type: Resource
any: true
- equal:
path: spec.metrics[0].pods.metric.name
value: litellm_requests_per_second
- equal:
path: spec.metrics[1].pods.metric.name
value: litellm_tokens_per_second
- notMatchRegexRaw:
pattern: per_minute
- it: ignores the per-minute keys, which the chart never shipped
template: gateway/hpa.yaml
set:
gateway.hpa.targetRequestsPerMinute: 5400
gateway.hpa.targetTokensPerMinute: 360000000
asserts:
- lengthEqual:
path: spec.metrics
count: 2
- notContains:
path: spec.metrics
content:
type: Pods
any: true
- it: renders no ServiceMonitor by default
template: gateway/servicemonitor.yaml
asserts:
- hasDocuments:
count: 0
- it: refuses a ServiceMonitor without the metrics server, whose http port needs a bearer token
template: gateway/servicemonitor.yaml
set:
gateway.serviceMonitor.enabled: true
asserts:
- failedTemplate:
errorPattern: gateway.serviceMonitor.enabled requires gateway.metricsServer.enabled
- it: scrapes each gateway pod through the metrics port
template: gateway/servicemonitor.yaml
release:
name: rel
namespace: llm
set:
gateway.serviceMonitor.enabled: true
gateway.metricsServer.enabled: true
gateway.serviceMonitor.labels:
release: kube-prometheus-stack
asserts:
- isKind:
of: ServiceMonitor
- equal:
path: metadata.labels.release
value: kube-prometheus-stack
- equal:
path: spec.selector.matchLabels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: rel
app.kubernetes.io/component: gateway
- equal:
path: spec.namespaceSelector.matchNames
value:
- llm
- equal:
path: spec.endpoints
value:
- port: metrics
path: /metrics/
interval: 15s
scrapeTimeout: 10s
scheme: http
- it: honours a custom scrape interval
template: gateway/servicemonitor.yaml
set:
gateway.serviceMonitor.enabled: true
gateway.serviceMonitor.interval: 30s
gateway.metricsServer.enabled: true
asserts:
- equal:
path: spec.endpoints[0].interval
value: 30s

85
helm/litellm/values.yaml
View file

@ -225,6 +225,26 @@ database:
usernameKey: username
passwordKey: password
# In-container connection pool (PgBouncer, transaction mode) shared by every
# gateway worker in the pod. Without it each of the `gateway.numWorkers`
# workers opens its own Prisma pool straight to Postgres, so a pod's
# footprint against the database's connection ceiling is
# numWorkers x connection_limit and grows with every replica. With it, the
# pod holds at most maxDbConnections upstream connections no matter how many
# workers run; the workers connect to the pool over loopback, with no extra
# network hop. The chart emits LITELLM_PGBOUNCER_ENABLED /
# LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / LITELLM_PGBOUNCER_MAX_CLIENT_CONN on
# the gateway container and its collector sidecar only: the backend runs a
# single worker and the migrations Job must keep a direct connection. With
# `database.writer.useIAMAuth` or `useAzureEntraAuth` the pool mints and
# renews the database token itself, so the workers never see it. Starting profile for
# `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a
# 5000-connection ceiling fits roughly 200 gateway replicas.
connectionPool:
enabled: false
maxDbConnections: 20
maxClientConn: 1000
# Optional Redis. Leave host empty to disable.
#
# This is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend
@ -284,6 +304,52 @@ gateway:
memory: 128Mi
limits:
memory: 512Mi
# Prometheus Operator ServiceMonitor for the gateway pods. Scrapes the
# `<gateway>-metrics` Service, so it requires metricsServer above (the http
# port serves /metrics/ behind virtual-key auth). Every pod is its own scrape
# target, so the samples carry the `pod` label the per-pod autoscaling
# queries below group by.
serviceMonitor:
enabled: false
labels: {}
interval: 15s
scrapeTimeout: 10s
# Opt-in `collector` sidecar (same image, `python -m litellm.proxy.collector`)
# that runs the post-response spend pipeline (cost calculation, spend logs,
# spend counters, budget reservation reconciliation) so the uvicorn workers
# only serialise a compact event over loopback and go back to serving
# requests. It shares the pod's env, proxy config, in-container pgbouncer and
# Redis spend buffer, so the per-pod DB connection budget is unchanged.
# Delivery is at-most-once inside the pod: events already handed over are
# lost if the sidecar dies before writing them; events the workers cannot
# hand over follow `onUnavailable`.
collector:
enabled: false
# unix:///<dir>/<file>.sock (the <dir> becomes a shared emptyDir) or
# tcp://127.0.0.1:<port>
address: unix:///var/run/litellm/collector.sock
# Events each uvicorn worker holds in memory while the sidecar is slow or
# restarting.
bufferSize: 1000
# fallback: run the pipeline in the worker when the sidecar is unreachable
# or the buffer is full (spend stays exact, that request costs gateway CPU
# again). drop: count and discard the event instead (spend under-reports).
onUnavailable: fallback
# How long the workers keep pushing buffered events on shutdown, and how
# long the sidecar keeps serving open connections after SIGTERM.
drainTimeoutSeconds: 10
# Sized independently of the gateway container; the pipeline is CPU bound.
resources:
requests:
cpu: 500m
memory: 1Gi
limits:
cpu: "1"
memory: 2Gi
# With hpa.targetCPUUtilizationPercentage set, scale on an autoscaling/v2
# ContainerResource metric of the `gateway` container only, so the
# sidecar's CPU never drives inference replicas. Needs Kubernetes 1.30+.
scaleOnGatewayContainerCpu: true
image:
repository: ghcr.io/berriai/litellm-gateway
tag: "" # defaults to .Chart.AppVersion
@ -340,6 +406,25 @@ gateway:
# policies:
# - { type: Percent, value: 100, periodSeconds: 30 }
behavior: {}
# Opt-in per-pod workload targets, rendered as autoscaling/v2 `Pods` metrics
# named `litellm_requests_per_second` and `litellm_tokens_per_second` with an
# AverageValue target. They coexist with the CPU/memory targets above: the
# HPA scales on whichever metric asks for the most replicas. Kubernetes has
# no idea what a token is, so a Prometheus Adapter must serve those two
# names on custom.metrics.k8s.io from the proxy's counters, grouped by the
# scrape target's `pod` label (enable serviceMonitor above):
# litellm_requests_per_second:
# sum(rate(litellm_proxy_total_requests_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
# litellm_tokens_per_second:
# sum(rate(litellm_total_tokens_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
# rate() over [1m] is already per second, so no `* 60`. How fast the HPA
# reacts is set by that window, the scrape interval and the HPA sync period
# (15s by default), not by the unit: keep serviceMonitor.interval at 15s or
# faster so a 1m window holds at least 4 samples. averageValue takes SI
# suffixes, so "6M" is six million tokens per second per pod. Tokens are
# counted when a response completes, so TPS trails long streams.
targetRequestsPerSecond: ""
targetTokensPerSecond: ""
# PodDisruptionBudget for the gateway pods. Set exactly one of
# `minAvailable` / `maxUnavailable` (minAvailable wins if both are set;
# enabling without either falls back to `maxUnavailable: 1`). Disabled by

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_skills_on_object_permission/migration.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "skills" TEXT[] DEFAULT ARRAY[]::TEXT[];

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -282,6 +282,7 @@ model LiteLLM_ObjectPermissionTable {
mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user
search_tools String[] @default([]) // search_tool_name values this key/team/user may call
mcp_tool_search_enabled Boolean?
skills String[] @default([]) // Claude Code plugin names granted to this key/team beyond the public (enabled) set
teams LiteLLM_TeamTable[]
projects LiteLLM_ProjectTable[]
verification_tokens LiteLLM_VerificationToken[]

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.95"
version = "0.4.96"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.95"
version = "0.4.96"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

5
litellm-rust/AGENTS.md
View file

@ -1,18 +1,19 @@
# AGENTS.md
litellm-rust has five crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
litellm-rust has six crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
## Crates
| Crate | Role |
|-------|------|
| litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. |
| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. |
| litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. |
| litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. |
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`. The interop foundation depends on no LiteLLM domain crate.
Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers, `litellm-token-counter`, and `litellm-python-interop`. The token counter and interop foundations depend on no LiteLLM domain crate.
## Where a route lives

424
litellm-rust/Cargo.lock generated
View file

@ -2,6 +2,20 @@
# It is not intended for manual editing.
version = 4
[[package]]
name = "ahash"
version = "0.8.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75"
dependencies = [
"cfg-if",
"getrandom 0.3.4",
"once_cell",
"serde",
"version_check",
"zerocopy",
]
[[package]]
name = "aho-corasick"
version = "1.1.5"
@ -418,7 +432,7 @@ checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f"
dependencies = [
"async-trait",
"axum-core",
"base64",
"base64 0.22.1",
"bytes",
"futures-util",
"http 1.4.2",
@ -468,6 +482,12 @@ dependencies = [
"tracing",
]
[[package]]
name = "base64"
version = "0.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
[[package]]
name = "base64"
version = "0.22.1"
@ -542,6 +562,15 @@ version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
[[package]]
name = "castaway"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a"
dependencies = [
"rustversion",
]
[[package]]
name = "cc"
version = "1.3.0"
@ -644,6 +673,21 @@ version = "0.5.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a"
[[package]]
name = "compact_str"
version = "0.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab"
dependencies = [
"castaway",
"cfg-if",
"itoa",
"rustversion",
"ryu",
"serde",
"static_assertions",
]
[[package]]
name = "const-oid"
version = "0.10.2"
@ -696,7 +740,7 @@ dependencies = [
"ciborium",
"clap",
"criterion-plot",
"itertools",
"itertools 0.13.0",
"num-traits",
"oorandom",
"page_size",
@ -716,7 +760,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea"
dependencies = [
"cast",
"itertools",
"itertools 0.13.0",
]
[[package]]
@ -778,6 +822,56 @@ dependencies = [
"cmov",
]
[[package]]
name = "daachorse"
version = "3.0.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5614204febbc33cc07a2806aa6440b904ac012b68eecc37f4493ea4a76455a3d"
[[package]]
name = "darling"
version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee"
dependencies = [
"darling_core",
"darling_macro",
]
[[package]]
name = "darling_core"
version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e"
dependencies = [
"fnv",
"ident_case",
"proc-macro2",
"quote",
"strsim",
"syn 2.0.119",
]
[[package]]
name = "darling_macro"
version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
dependencies = [
"darling_core",
"quote",
"syn 2.0.119",
]
[[package]]
name = "dary_heap"
version = "0.3.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe"
dependencies = [
"serde",
]
[[package]]
name = "data-encoding"
version = "2.11.0"
@ -790,6 +884,37 @@ version = "0.5.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c"
[[package]]
name = "derive_builder"
version = "0.20.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947"
dependencies = [
"derive_builder_macro",
]
[[package]]
name = "derive_builder_core"
version = "0.20.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8"
dependencies = [
"darling",
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "derive_builder_macro"
version = "0.20.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c"
dependencies = [
"derive_builder_core",
"syn 2.0.119",
]
[[package]]
name = "digest"
version = "0.10.7"
@ -841,6 +966,12 @@ version = "1.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
[[package]]
name = "esaxx-rs"
version = "0.1.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6"
[[package]]
name = "fastrand"
version = "2.5.0"
@ -964,6 +1095,18 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "getrandom"
version = "0.3.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd"
dependencies = [
"cfg-if",
"libc",
"r-efi 5.3.0",
"wasip2",
]
[[package]]
name = "getrandom"
version = "0.4.3"
@ -973,7 +1116,7 @@ dependencies = [
"cfg-if",
"js-sys",
"libc",
"r-efi",
"r-efi 6.0.0",
"rand_core 0.10.1",
"wasm-bindgen",
]
@ -1220,7 +1363,7 @@ version = "0.1.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0"
dependencies = [
"base64",
"base64 0.22.1",
"bytes",
"futures-channel",
"futures-util",
@ -1319,6 +1462,12 @@ dependencies = [
"zerovec",
]
[[package]]
name = "ident_case"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
[[package]]
name = "idna"
version = "1.1.0"
@ -1348,6 +1497,8 @@ checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9"
dependencies = [
"equivalent",
"hashbrown",
"serde",
"serde_core",
]
[[package]]
@ -1365,6 +1516,15 @@ dependencies = [
"either",
]
[[package]]
name = "itertools"
version = "0.14.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285"
dependencies = [
"either",
]
[[package]]
name = "itoa"
version = "1.0.18"
@ -1409,7 +1569,7 @@ name = "litellm-ai-gateway"
version = "0.1.0"
dependencies = [
"axum",
"base64",
"base64 0.22.1",
"futures-channel",
"futures-util",
"litellm-config",
@ -1447,17 +1607,19 @@ dependencies = [
"aws-sigv4",
"aws-smithy-runtime-api",
"aws-types",
"base64",
"base64 0.22.1",
"rand 0.8.7",
"reqwest",
"rstest",
"serde",
"serde_json",
"serde_path_to_error",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"tracing",
"tracing-subscriber",
"url",
]
[[package]]
@ -1469,6 +1631,7 @@ dependencies = [
"litellm-ai-gateway",
"litellm-core",
"litellm-python-interop",
"litellm-token-counter",
"pyo3",
"pyo3-async-runtimes",
"serde",
@ -1489,6 +1652,22 @@ dependencies = [
"serde_json",
]
[[package]]
name = "litellm-token-counter"
version = "0.1.0"
dependencies = [
"criterion",
"indexmap",
"itoa",
"rand 0.8.7",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokenizers",
"unicode-normalization-alignments",
]
[[package]]
name = "litemap"
version = "0.8.2"
@ -1507,6 +1686,22 @@ version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
[[package]]
name = "macro_rules_attribute"
version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b3ae8f6d608c795738406608304d30a2dfbdc8e58e44f7ba43236da5208ded3c"
dependencies = [
"macro_rules_attribute-proc_macro",
"pastey",
]
[[package]]
name = "macro_rules_attribute-proc_macro"
version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
[[package]]
name = "matchit"
version = "0.7.3"
@ -1535,6 +1730,12 @@ dependencies = [
"unicase",
]
[[package]]
name = "minimal-lexical"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
[[package]]
name = "mio"
version = "1.2.2"
@ -1546,6 +1747,38 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "monostate"
version = "0.1.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67"
dependencies = [
"monostate-impl",
"serde",
"serde_core",
]
[[package]]
name = "monostate-impl"
version = "0.1.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "nom"
version = "7.1.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a"
dependencies = [
"memchr",
"minimal-lexical",
]
[[package]]
name = "num-conv"
version = "0.2.2"
@ -1576,6 +1809,28 @@ version = "1.21.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
[[package]]
name = "onig"
version = "6.5.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2"
dependencies = [
"bitflags",
"libc",
"once_cell",
"onig_sys",
]
[[package]]
name = "onig_sys"
version = "69.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7"
dependencies = [
"cc",
"pkg-config",
]
[[package]]
name = "oorandom"
version = "11.1.5"
@ -1604,6 +1859,18 @@ dependencies = [
"winapi",
]
[[package]]
name = "paste"
version = "1.0.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a"
[[package]]
name = "pastey"
version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4"
[[package]]
name = "percent-encoding"
version = "2.3.2"
@ -1850,6 +2117,12 @@ dependencies = [
"proc-macro2",
]
[[package]]
name = "r-efi"
version = "5.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
[[package]]
name = "r-efi"
version = "6.0.0"
@ -1863,10 +2136,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a"
dependencies = [
"libc",
"rand_chacha",
"rand_chacha 0.3.1",
"rand_core 0.6.4",
]
[[package]]
name = "rand"
version = "0.9.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41"
dependencies = [
"rand_chacha 0.9.0",
"rand_core 0.9.5",
]
[[package]]
name = "rand"
version = "0.10.2"
@ -1888,6 +2171,16 @@ dependencies = [
"rand_core 0.6.4",
]
[[package]]
name = "rand_chacha"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
dependencies = [
"ppv-lite86",
"rand_core 0.9.5",
]
[[package]]
name = "rand_core"
version = "0.6.4"
@ -1897,6 +2190,15 @@ dependencies = [
"getrandom 0.2.17",
]
[[package]]
name = "rand_core"
version = "0.9.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c"
dependencies = [
"getrandom 0.3.4",
]
[[package]]
name = "rand_core"
version = "0.10.1"
@ -1922,6 +2224,17 @@ dependencies = [
"rayon-core",
]
[[package]]
name = "rayon-cond"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f"
dependencies = [
"either",
"itertools 0.14.0",
"rayon",
]
[[package]]
name = "rayon-core"
version = "1.13.0"
@ -1979,7 +2292,7 @@ version = "0.12.28"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
dependencies = [
"base64",
"base64 0.22.1",
"bytes",
"futures-channel",
"futures-core",
@ -2361,12 +2674,36 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "spm_precompiled"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326"
dependencies = [
"base64 0.13.1",
"nom",
"serde",
"unicode-segmentation",
]
[[package]]
name = "stable_deref_trait"
version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596"
[[package]]
name = "static_assertions"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
[[package]]
name = "strsim"
version = "0.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
[[package]]
name = "subtle"
version = "2.6.1"
@ -2535,6 +2872,39 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
[[package]]
name = "tokenizers"
version = "0.23.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7afbf6e88718afcc138bad01d6ccc3051dbbc3b2ce9793d8b8a3aeb610969cfc"
dependencies = [
"ahash",
"compact_str",
"daachorse",
"dary_heap",
"derive_builder",
"esaxx-rs",
"getrandom 0.3.4",
"itertools 0.14.0",
"log",
"macro_rules_attribute",
"monostate",
"onig",
"paste",
"rand 0.9.5",
"rayon",
"rayon-cond",
"regex",
"regex-syntax",
"serde",
"serde_json",
"spm_precompiled",
"thiserror 2.0.19",
"unicode-normalization-alignments",
"unicode-segmentation",
"unicode_categories",
]
[[package]]
name = "tokio"
version = "1.53.0"
@ -2773,6 +3143,27 @@ version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
[[package]]
name = "unicode-normalization-alignments"
version = "0.1.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de"
dependencies = [
"smallvec",
]
[[package]]
name = "unicode-segmentation"
version = "1.13.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8"
[[package]]
name = "unicode_categories"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
[[package]]
name = "untrusted"
version = "0.9.0"
@ -2856,6 +3247,15 @@ version = "0.11.1+wasi-snapshot-preview1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
[[package]]
name = "wasip2"
version = "1.0.4+wasi-0.2.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487"
dependencies = [
"wit-bindgen",
]
[[package]]
name = "wasm-bindgen"
version = "0.2.126"
@ -3081,6 +3481,12 @@ dependencies = [
"memchr",
]
[[package]]
name = "wit-bindgen"
version = "0.57.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e"
[[package]]
name = "writeable"
version = "0.6.3"

4
litellm-rust/Cargo.toml
View file

@ -1,6 +1,7 @@
[workspace]
members = [
"crates/core",
"crates/token-counter",
"crates/config",
"crates/ai-gateway",
"crates/python-interop",
@ -18,6 +19,7 @@ repository = "https://github.com/BerriAI/litellm"
tracing = "0.1"
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
litellm-core = { path = "crates/core" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-config = { path = "crates/config" }
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
litellm-python-interop = { path = "crates/python-interop" }
@ -39,6 +41,8 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
base64 = "0.22"
url = "2.5.8"
criterion = "0.8.2"
[profile.release]
opt-level = 3

5
litellm-rust/crates/ai-gateway/README.md
View file

@ -6,17 +6,18 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame.
## Crates
`litellm-rust` has five crates. A crate is a layer or shared foundation, not a route:
`litellm-rust` has six crates. A crate is a layer or shared foundation, not a route:
| Crate | Role |
|-------|------|
| litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. |
| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. |
| litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. |
| litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. |
Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop.
Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers, token counter, and Python interop.
- **Client endpoint:** `wss://<host>/v1/realtime?model=<model>` (WebSocket)
- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)

2
litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs
View file

@ -269,7 +269,7 @@ fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
fn core_error_kind(error: &Error) -> &'static str {
match error {
Error::Auth(_) => "AuthError",
Error::Auth(_) | Error::MissingApiKey { .. } => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",

2
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
View file

@ -386,7 +386,7 @@ fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
fn core_error_kind(error: &Error) -> &'static str {
match error {
Error::Auth(_) => "AuthError",
Error::Auth(_) | Error::MissingApiKey { .. } => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",

3
litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs
View file

@ -114,7 +114,8 @@ impl IntoResponse for MessagesRouteError {
| Error::Connect(_)
| Error::InvalidResponse(_)
| Error::InvalidType { .. }
| Error::MissingField(_) => (
| Error::MissingField(_)
| Error::MissingApiKey { .. } => (
StatusCode::BAD_GATEWAY,
"messages provider request failed".to_string(),
),

8
litellm-rust/crates/ai-gateway/src/trace_parity.rs
View file

@ -47,10 +47,10 @@ pub async fn messages_request(
.header(CONTENT_TYPE, "application/json")
.body(Body::from(body.to_string()))
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let response = routes::app(state)
.oneshot(request)
.await
.map_err(|error| match error {})?;
let response = match routes::app(state).oneshot(request).await {
Ok(response) => response,
Err(error) => match error {},
};
let status: StatusCode = response.status();
let bytes = to_bytes(response.into_body(), usize::MAX)
.await

8
litellm-rust/crates/core/Cargo.toml
View file

@ -4,6 +4,11 @@ version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
autotests = false
[[test]]
name = "workspace_crate_allowlist"
path = "tests/workspace_crate_allowlist.rs"
[dependencies]
base64.workspace = true
@ -11,10 +16,13 @@ rand.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_path_to_error = "0.1"
tokio.workspace = true
thiserror.workspace = true
tracing.workspace = true
tracing-subscriber = { workspace = true, optional = true }
sha2.workspace = true
url.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }

3
litellm-rust/crates/core/src/constants.rs
View file

@ -43,3 +43,6 @@ pub const EMPTY_TEXT_PLACEHOLDER: &str =
"[System: Empty message content sanitised to satisfy protocol]";
pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";
pub(crate) const OCR_HTTP_TIMEOUT_SECS: u64 = 600;
pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10;
pub(crate) const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1";

107
litellm-rust/crates/core/src/error.rs
View file

@ -17,6 +17,10 @@ pub enum Error {
InvalidRequest(String),
#[error("{0}")]
Auth(String),
#[error(
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error("upstream request failed with status {status}: {body}")]
Http { status: u16, body: String },
#[error("upstream network error: {0}")]
@ -36,6 +40,59 @@ pub enum Error {
Unsupported(&'static str),
}
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum TransportError {
#[error("upstream request failed with status {status}: {body}")]
Http { status: u16, body: String },
#[error("upstream network error: {0}")]
Network(String),
#[error("could not reach the provider: {0}")]
Connect(String),
}
impl TransportError {
pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self {
let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder());
let message = error.without_url().to_string();
if before_dispatch {
Self::Connect(message)
} else {
Self::Network(message)
}
}
}
impl From<reqwest::Error> for TransportError {
fn from(error: reqwest::Error) -> Self {
Self::Network(error.without_url().to_string())
}
}
impl From<crate::ocr::error::OcrRequestError> for Error {
fn from(error: crate::ocr::error::OcrRequestError) -> Self {
match error {
crate::ocr::error::OcrRequestError::MissingField(field) => Self::MissingField(field),
error => Self::InvalidRequest(error.to_string()),
}
}
}
impl From<crate::ocr::error::OcrResponseError> for Error {
fn from(error: crate::ocr::error::OcrResponseError) -> Self {
Self::InvalidResponse(error.to_string())
}
}
impl From<TransportError> for Error {
fn from(error: TransportError) -> Self {
match error {
TransportError::Http { status, body } => Self::Http { status, body },
TransportError::Network(message) => Self::Network(message),
TransportError::Connect(message) => Self::Connect(message),
}
}
}
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
match value {
serde_json::Value::Null => "null",
@ -46,3 +103,53 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str {
serde_json::Value::Object(_) => "object",
}
}
#[cfg(test)]
mod transport_tests {
use super::*;
#[tokio::test]
async fn transport_errors_remove_urls_and_keep_dispatch_context() {
let error = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get("http://localhost:invalid/private?api_key=secret")
.send()
.await
.expect_err("invalid port");
let error = TransportError::from_reqwest_before_dispatch(error);
assert!(matches!(error, TransportError::Connect(_)));
assert!(!error.to_string().contains("secret"));
assert!(!error.to_string().contains("private"));
}
#[tokio::test]
async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() {
use std::time::Duration;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind");
let address = listener.local_addr().expect("address");
let request = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get(format!("http://{address}"))
.timeout(Duration::from_millis(200))
.send();
let (response, accepted) = tokio::join!(
request,
tokio::time::timeout(Duration::from_secs(2), listener.accept())
);
let _connection = accepted
.expect("accept deadline")
.expect("accepted connection");
let error = response.expect_err("server does not respond");
assert!(error.is_timeout());
assert!(matches!(
TransportError::from_reqwest_before_dispatch(error),
TransportError::Network(_)
));
}
}

105
litellm-rust/crates/core/src/http_utils.rs
View file

@ -1,10 +1,43 @@
//! Header and upstream-body helpers shared by every route module.
use serde_json::{Map, Value};
use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
use crate::error::{Error, json_type_name};
#[allow(
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub(crate) enum HeaderPolicy<'a> {
All,
Only(&'a [&'a str]),
Except(&'a [&'a str]),
}
#[allow(
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub(crate) fn with_headers(
builder: reqwest::RequestBuilder,
headers: &[(String, String)],
policy: HeaderPolicy<'_>,
) -> reqwest::RequestBuilder {
headers
.iter()
.filter(|(name, _)| match policy {
HeaderPolicy::All => true,
HeaderPolicy::Only(names) => names
.iter()
.any(|allowed| name.eq_ignore_ascii_case(allowed)),
HeaderPolicy::Except(names) => !names
.iter()
.any(|excluded| name.eq_ignore_ascii_case(excluded)),
})
.fold(builder, |builder, (name, value)| {
builder.header(name, value)
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn http_request(
request: reqwest::RequestBuilder,
@ -12,8 +45,6 @@ pub async fn http_request(
request.send().await
}
/// Bound an upstream error body before it crosses a host boundary, so provider
/// bodies stay data-minimized.
pub fn truncate_error_body(body: &str) -> String {
if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS {
return body.to_string();
@ -61,11 +92,77 @@ pub fn has_bearer_auth(headers: &[(String, String)]) -> bool {
})
}
#[allow(
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub(crate) fn deserialize_optional_param<'de, D, T>(
deserializer: D,
) -> Result<Option<Option<T>>, D::Error>
where
D: serde::Deserializer<'de>,
T: serde::Deserialize<'de>,
{
<Option<T> as serde::Deserialize>::deserialize(deserializer).map(Some)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[rstest::rstest]
#[case(HeaderPolicy::All, true, true)]
#[case(HeaderPolicy::Only(&["authorization"]), true, false)]
#[case(HeaderPolicy::Except(&["authorization"]), false, true)]
fn forwarding_policy_preserves_matching_headers_and_duplicates(
#[case] policy: HeaderPolicy<'_>,
#[case] auth: bool,
#[case] trace: bool,
) {
let request = with_headers(
reqwest::Client::new().get("https://example.com"),
&[
("AuThOrIzAtIoN".into(), "Bearer token".into()),
("X-Trace".into(), "first".into()),
("x-trace".into(), "second".into()),
],
policy,
)
.build()
.unwrap();
assert_eq!(request.headers().contains_key("authorization"), auth);
let traces: Vec<_> = request.headers().get_all("x-trace").iter().collect();
if trace {
assert_eq!(traces, ["first", "second"]);
} else {
assert!(traces.is_empty());
}
}
#[test]
fn multipart_policy_leaves_content_headers_to_reqwest() {
let request = with_headers(
reqwest::Client::new()
.post("https://example.com")
.multipart(reqwest::multipart::Form::new().text("file", "abc")),
&[
("Content-Type".into(), "application/json".into()),
("CONTENT-LENGTH".into(), "0".into()),
],
HeaderPolicy::Except(&["content-type", "content-length"]),
)
.build()
.unwrap();
assert!(
request.headers()["content-type"]
.to_str()
.unwrap()
.starts_with("multipart/form-data; boundary=")
);
assert_ne!(request.headers()["content-length"], "0");
}
#[test]
fn truncate_leaves_short_bodies_untouched() {
assert_eq!(truncate_error_body("short"), "short");

1
litellm-rust/crates/core/src/lib.rs
View file

@ -14,5 +14,6 @@ pub mod realtime;
pub mod responses;
pub mod router;
pub mod routing_utils;
mod url_utils;
pub use error::Error;

147
litellm-rust/crates/core/src/ocr/adapters/mistral.rs Normal file
View file

@ -0,0 +1,147 @@
use super::OcrAdapter;
use crate::Error;
use crate::constants::MISTRAL_OCR_API_BASE;
use crate::ocr::OcrClient;
use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection};
use crate::url_utils::ApiUrl;
const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY";
#[derive(Clone, Debug)]
pub(crate) struct MistralAdapter;
impl OcrAdapter for MistralAdapter {
type ProviderResponse = MistralOcrResponse;
const PROVIDER: OcrProvider = OcrProvider::Mistral;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<MistralOcrParams>(request)?;
let headers = validate_environment(&request.connection, &credential_env)?;
let url = get_complete_url(request.connection.api_base.as_deref())?;
let body =
mistral::transform_ocr_request(&request.model, request.document.clone(), &params)?;
transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
mistral::transform_ocr_response(&request.model, response)
}
}
pub(crate) fn get_complete_url(api_base: Option<&str>) -> Result<String, OcrError> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(MISTRAL_OCR_API_BASE);
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&["v1", "ocr"]))
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
fn validate_environment(
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let api_key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
.ok_or(Error::MissingApiKey {
provider: "Mistral",
})?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {api_key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn complete_url_defaults_and_dedupes_v1() {
assert_eq!(
get_complete_url(None).unwrap(),
"https://api.mistral.ai/v1/ocr"
);
assert_eq!(
get_complete_url(Some("https://example.com/v1?tenant=a")).unwrap(),
"https://example.com/v1/ocr?tenant=a"
);
assert_eq!(
get_complete_url(Some("https://example.com/v1/ocr?tenant=a")).unwrap(),
"https://example.com/v1/ocr?tenant=a"
);
}
#[test]
fn environment_prefers_explicit_key_then_environment() {
let explicit = OcrConnection {
api_key: Some("explicit".into()),
..OcrConnection::default()
};
assert_eq!(
validate_environment(&explicit, &|_| Some("environment".into())).unwrap()[0],
("Authorization".into(), "Bearer explicit".into())
);
assert_eq!(
validate_environment(&OcrConnection::default(), &|_| Some("environment".into()))
.unwrap()[0],
("Authorization".into(), "Bearer environment".into())
);
}
#[test]
fn environment_preserves_forwarded_authorization() {
let connection = OcrConnection {
extra_headers: vec![("authorization".into(), "Bearer forwarded".into())],
..OcrConnection::default()
};
assert_eq!(
validate_environment(&connection, &|_| None).unwrap(),
connection.extra_headers
);
}
#[test]
fn environment_rejects_missing_key() {
assert!(matches!(
validate_environment(&OcrConnection::default(), &|_| None),
Err(OcrError::Public(Error::MissingApiKey {
provider: "Mistral"
}))
));
}
}

69
litellm-rust/crates/core/src/ocr/adapters/mod.rs Normal file
View file

@ -0,0 +1,69 @@
use std::future::Future;
use serde::de::DeserializeOwned;
use super::OcrClient;
use super::error::{OcrError, OcrResponseError};
use super::registry::OcrProvider;
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrResponseFormat};
use super::wire::DecodedOcrResponse;
mod mistral;
pub(crate) use mistral::MistralAdapter;
/// Converts a complete LiteLLM OCR call to provider HTTP and normalizes its response.
pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static {
/// Provider JSON schema; direct and Vertex Mistral share `MistralOcrResponse`.
type ProviderResponse: DeserializeOwned + Send;
const PROVIDER: OcrProvider;
/// Prepares the complete provider HTTP request.
/// `request` contains the model, document, connection, and unmapped caller options.
/// `client` supplies reusable provider and document HTTP clients.
/// Returns the complete HTTP request, whereas Python returns body data.
fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> impl Future<Output = Result<reqwest::Request, OcrError>> + Send;
/// Python: `transform_ocr_response`.
/// `request` supplies caller context, including the fallback model.
/// `response` is the decoded provider payload; the output is the shared LiteLLM schema.
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError>;
/// Decodes provider HTTP; adapters may override this to poll asynchronous operations.
/// Python performs that polling inside `async_transform_ocr_response`.
/// `client` is reused for polling; `response` is the initial HTTP response.
/// `url` and `headers` describe the submitted call; `request` supplies limits and format.
fn read_response(
&self,
_client: &OcrClient,
response: reqwest::Response,
_url: &str,
_headers: &[(String, String)],
request: &LiteLLMOcrRequest,
) -> impl Future<Output = Result<DecodedOcrResponse<Self::ProviderResponse>, OcrError>> + Send
{
let retain_native = request
.response_format()
.map(|format| format == OcrResponseFormat::Native);
async move { super::client::read_json_response(response, retain_native?).await }
}
}
macro_rules! for_each_ocr_adapter {
($callback:ident) => {
$callback! {
Mistral, $crate::ocr::adapters::MistralAdapter, $crate::ocr::adapters::MistralAdapter, Mistral;
}
};
}
pub(crate) use for_each_ocr_adapter;

75
litellm-rust/crates/core/src/ocr/client.rs Normal file
View file

@ -0,0 +1,75 @@
use std::sync::OnceLock;
use std::time::Duration;
use serde::de::DeserializeOwned;
use super::error::OcrError;
use super::handler::perform_ocr_request;
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use super::wire::{DecodedOcrResponse, decode_response};
use crate::Error;
use crate::constants::OCR_CONNECT_TIMEOUT_SECS;
use crate::error::TransportError;
#[derive(Clone)]
pub struct OcrClient {
provider_http: reqwest::Client,
}
impl OcrClient {
pub fn new(provider_http: reqwest::Client) -> Result<Self, TransportError> {
Ok(Self { provider_http })
}
#[tracing::instrument(
name = "ocr",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
pub async fn perform(&self, request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
perform_ocr_request(self, request).await
}
pub(crate) fn provider_http(&self) -> &reqwest::Client {
&self.provider_http
}
#[cfg(test)]
pub(crate) fn for_test(provider_http: reqwest::Client) -> Self {
Self { provider_http }
}
}
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
static CLIENT: OnceLock<Result<OcrClient, TransportError>> = OnceLock::new();
let client = CLIENT
.get_or_init(|| {
reqwest::Client::builder()
.connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS))
.build()
.map_err(TransportError::from)
.and_then(OcrClient::new)
})
.clone()?;
client.perform(request).await
}
pub async fn read_json_response<T: DeserializeOwned>(
response: reqwest::Response,
native: bool,
) -> Result<DecodedOcrResponse<T>, OcrError> {
let status = response.status();
let bytes = response
.bytes()
.await
.map_err(crate::error::TransportError::from)?;
if !status.is_success() {
return Err(crate::error::TransportError::Http {
status: status.as_u16(),
body: crate::http_utils::truncate_error_body(&String::from_utf8_lossy(&bytes)),
}
.into());
}
Ok(decode_response(&bytes, native)?)
}

5
litellm-rust/crates/core/src/ocr/codecs/mistral/mod.rs Normal file
View file

@ -0,0 +1,5 @@
mod transformation;
mod types;
pub(crate) use transformation::{transform_ocr_request, transform_ocr_response};
pub(crate) use types::{MistralOcrParams, MistralOcrRequest, MistralOcrResponse};

102
litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs Normal file
View file

@ -0,0 +1,102 @@
use super::{MistralOcrParams, MistralOcrRequest, MistralOcrResponse};
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) fn transform_ocr_request(
model: &str,
document: OcrDocument,
params: &MistralOcrParams,
) -> Result<MistralOcrRequest, OcrRequestError> {
Ok(MistralOcrRequest {
model: model.to_string(),
document,
params: params.clone(),
})
}
pub(crate) fn transform_ocr_response(
model: &str,
response: MistralOcrResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
Ok(LiteLLMOcrResponse {
pages: response.pages,
model: response.model.unwrap_or_else(|| model.to_string()),
document_annotation: response.document_annotation,
usage_info: response.usage_info,
object: "ocr".to_string(),
extra_fields: response.extra_fields,
provider_native_response: None,
})
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
use serde_json::{Value, json};
#[rstest]
#[case("pages", json!([0, 2]))]
#[case("include_image_base64", json!(true))]
#[case("image_limit", json!(2))]
#[case("image_min_size", json!(100))]
#[case("bbox_annotation_format", json!({"type":"json_schema"}))]
#[case("document_annotation_format", json!({"type":"json_schema"}))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("extract_header", json!(true))]
#[case("extract_footer", json!(false))]
#[case("table_format", json!("html"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("include_blocks", json!(true))]
#[case("id", json!("req-123"))]
fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) {
let params: MistralOcrParams =
serde_json::from_value(json!({name: value.clone()})).unwrap();
let document: OcrDocument = serde_json::from_value(
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
)
.unwrap();
let result =
serde_json::to_value(transform_ocr_request("model", document, &params).unwrap())
.unwrap();
assert_eq!(result["model"], "model");
assert_eq!(result[name], value);
}
#[test]
fn request_mapping_filters_unknown_fields() {
let params: MistralOcrParams = serde_json::from_value(json!({"unknown": true})).unwrap();
let document: OcrDocument = serde_json::from_value(
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
)
.unwrap();
let result =
serde_json::to_value(transform_ocr_request("model", document, &params).unwrap())
.unwrap();
assert!(result.get("unknown").is_none());
}
#[test]
fn response_preserves_provider_fields() {
let response: MistralOcrResponse = serde_json::from_value(json!({
"pages":[{"index":0,"markdown":"hello","header":"head","confidence_scores":{"mean":0.99}}],
"model":"returned-model",
"usage_info":{"pages_processed":1,"future_counter":5},
"future_response_field":"kept"
}))
.unwrap();
let result = transform_ocr_response("model", response)
.unwrap()
.into_json();
assert_eq!(result["pages"][0]["header"], "head");
assert_eq!(result["usage_info"]["future_counter"], 5);
assert_eq!(result["future_response_field"], "kept");
assert_eq!(result["model"], "returned-model");
}
#[test]
fn response_rejects_null_pages() {
assert!(serde_json::from_value::<MistralOcrResponse>(json!({"pages":null})).is_err());
}
}

53
litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs Normal file
View file

@ -0,0 +1,53 @@
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::ocr::types::OcrDocument;
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub(crate) struct MistralOcrParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub pages: Option<Vec<i64>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub include_image_base64: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub image_limit: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub image_min_size: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub bbox_annotation_format: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub document_annotation_format: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub document_annotation_prompt: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub extract_header: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub extract_footer: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub table_format: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub confidence_scores_granularity: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub include_blocks: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct MistralOcrRequest {
pub model: String,
pub document: OcrDocument,
#[serde(flatten)]
pub params: MistralOcrParams,
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct MistralOcrResponse {
#[serde(default)]
pub pages: Vec<Value>,
pub model: Option<String>,
pub document_annotation: Option<Value>,
pub usage_info: Option<Value>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}

1
litellm-rust/crates/core/src/ocr/codecs/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod mistral;

42
litellm-rust/crates/core/src/ocr/error.rs Normal file
View file

@ -0,0 +1,42 @@
use thiserror::Error;
use crate::error::TransportError;
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum OcrRequestError {
#[error("Invalid `req_format`. Expected 'native' or 'litellm'.")]
RequestFormat,
#[error("invalid OCR request field: {path}")]
RequestField { path: String },
#[error("missing required field: {0}")]
MissingField(&'static str),
}
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum OcrResponseError {
#[error("invalid OCR response field: {path}")]
ResponseField { path: String },
}
#[derive(Debug, Error)]
pub enum OcrError {
#[error("{0}")]
Request(#[from] OcrRequestError),
#[error("{0}")]
Response(#[from] OcrResponseError),
#[error("{0}")]
Transport(#[from] TransportError),
#[error("{0}")]
Public(#[from] crate::Error),
}
impl From<OcrError> for crate::Error {
fn from(error: OcrError) -> Self {
match error {
OcrError::Request(error) => error.into(),
OcrError::Response(error) => error.into(),
OcrError::Transport(error) => error.into(),
OcrError::Public(error) => error,
}
}
}

72
litellm-rust/crates/core/src/ocr/handler.rs Normal file
View file

@ -0,0 +1,72 @@
use super::OcrClient;
use super::adapters::OcrAdapter;
use super::hooks::OcrLifecycleHooks;
use super::registry::OcrAdapterKind;
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use crate::Error;
use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext};
pub(crate) async fn perform_ocr_request(
client: &OcrClient,
request: LiteLLMOcrRequest,
) -> Result<LiteLLMOcrResponse, Error> {
let context = CallLifecycleContext::new(
"ocr",
request.model.clone(),
request.adapter.provider().as_str(),
request
.litellm_call_id
.clone()
.unwrap_or_else(|| format!("ocr-{:032x}", rand::random::<u128>())),
);
let hooks = OcrLifecycleHooks {
hooks: request.hooks.clone(),
provider_name: context.custom_llm_provider.clone(),
};
CallLifecycle::default().run(context, request, &hooks, |request| async move {
macro_rules! execute_selected_adapter {
($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => {
match request.adapter {
$( OcrAdapterKind::$variant => execute_ocr_provider_call(client, &$instance, request).await, )+
}
};
}
super::adapters::for_each_ocr_adapter!(execute_selected_adapter)
}).await
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn execute_ocr_provider_call<A: OcrAdapter>(
client: &OcrClient,
adapter: &A,
request: LiteLLMOcrRequest,
) -> Result<LiteLLMOcrResponse, Error> {
let provider_request = adapter.prepare_request(&request, client).await?;
let url = provider_request.url().to_string();
let headers = provider_request
.headers()
.iter()
.map(|(name, value)| {
value
.to_str()
.map(|value| (name.to_string(), value.to_string()))
.map_err(|_| super::error::OcrRequestError::RequestField {
path: "headers".into(),
})
})
.collect::<Result<Vec<_>, _>>()?;
let response = crate::http_utils::http_request(reqwest::RequestBuilder::from_parts(
client.provider_http().clone(),
provider_request,
))
.await
.map_err(crate::error::TransportError::from)?;
let decoded = adapter
.read_response(client, response, &url, &headers, &request)
.await?;
let response = adapter.transform_ocr_response(&request, decoded.data)?;
Ok(LiteLLMOcrResponse {
provider_native_response: decoded.native,
..response
})
}

146
litellm-rust/crates/core/src/ocr/hooks.rs Normal file
View file

@ -0,0 +1,146 @@
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument};
use crate::Error;
use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use serde::Serialize;
use serde_json::Value;
pub type OcrHookFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
pub type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
#[derive(Clone, Debug, Serialize)]
pub struct OcrPreCallRequest {
pub model: String,
pub custom_llm_provider: String,
pub document: OcrDocument,
pub optional_params: Value,
}
#[derive(Clone, Debug, Serialize)]
pub struct OcrDuringCallRequest {
pub model: String,
pub custom_llm_provider: String,
pub url: String,
pub body: Value,
}
pub trait OcrHooks: Send + Sync {
fn has_guardrails(&self) -> bool {
false
}
fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
Box::pin(async move { Ok(request) })
}
fn during_call(
&self,
request: OcrDuringCallRequest,
) -> OcrHookFuture<'_, OcrDuringCallRequest> {
Box::pin(async move { Ok(request) })
}
fn success<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_response: &'a LiteLLMOcrResponse,
_timing: &'a CallLifecycleTiming,
) -> OcrLogFuture<'a> {
Box::pin(async {})
}
fn failure<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_error: &'a Error,
_timing: &'a CallLifecycleTiming,
) -> OcrLogFuture<'a> {
Box::pin(async {})
}
}
pub struct NoopOcrHooks;
impl OcrHooks for NoopOcrHooks {}
pub(crate) struct OcrLifecycleHooks {
pub hooks: Arc<dyn OcrHooks>,
pub provider_name: String,
}
impl CallLifecycleHooks<LiteLLMOcrRequest, LiteLLMOcrRequest, LiteLLMOcrResponse>
for OcrLifecycleHooks
{
type PreCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>;
type DuringCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>;
type SuccessFuture<'a> = OcrLogFuture<'a>;
type FailureFuture<'a> = OcrLogFuture<'a>;
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: LiteLLMOcrRequest,
) -> Self::PreCallFuture<'a> {
Box::pin(async move {
if !self.hooks.has_guardrails() {
return Ok(request);
}
let changed = self
.hooks
.pre_call(OcrPreCallRequest {
model: request.model.clone(),
custom_llm_provider: self.provider_name.clone(),
document: request.document,
optional_params: Value::Object(request.optional_params),
})
.await?;
let Value::Object(optional_params) = changed.optional_params else {
return Err(super::error::OcrRequestError::RequestField {
path: "guardrail.optional_params".into(),
}
.into());
};
Ok(LiteLLMOcrRequest {
document: changed.document,
optional_params,
..request
})
})
}
fn async_during_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: LiteLLMOcrRequest,
) -> Self::DuringCallFuture<'a> {
Box::pin(async move { Ok(request) })
}
#[tracing::instrument(
name = "success_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_success_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
response: &'a LiteLLMOcrResponse,
timing: &'a CallLifecycleTiming,
) -> Self::SuccessFuture<'a> {
self.hooks.success(context, response, timing)
}
#[tracing::instrument(
name = "failure_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a Error,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
self.hooks.failure(context, error, timing)
}
}

19
litellm-rust/crates/core/src/ocr/mod.rs
View file

@ -1,2 +1,21 @@
mod adapters;
pub mod client;
mod codecs;
pub mod error;
mod handler;
pub mod hooks;
mod prepare;
mod registry;
pub mod transformation;
pub mod types;
pub mod wire;
pub use client::{OcrClient, ocr};
pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument};
#[cfg(test)]
#[path = "../../tests/ocr/support.rs"]
pub(crate) mod test_support;
#[cfg(test)]
#[path = "../../tests/ocr.rs"]
pub(crate) mod tests;

142
litellm-rust/crates/core/src/ocr/prepare.rs Normal file
View file

@ -0,0 +1,142 @@
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::{Map, Value};
use super::OcrClient;
use super::error::{OcrError, OcrRequestError};
use super::hooks::OcrDuringCallRequest;
use super::types::LiteLLMOcrRequest;
#[derive(Debug, Deserialize)]
pub(crate) struct ParsedProviderParams<T> {
#[serde(flatten)]
pub known: T,
#[serde(default, flatten)]
pub extra_params: Map<String, Value>,
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) fn _prepare_ocr_request<T: DeserializeOwned>(
request: &LiteLLMOcrRequest,
) -> Result<ParsedProviderParams<T>, OcrRequestError> {
super::wire::decode_request_value(
Value::Object(request.optional_params.clone()),
"optional_params",
)
}
pub(crate) async fn transform_request_body<B>(
client: &OcrClient,
request: &LiteLLMOcrRequest,
url: &str,
headers: &[(String, String)],
body: B,
validate: impl FnOnce(&B) -> Result<(), OcrRequestError>,
) -> Result<reqwest::Request, OcrError>
where
B: Serialize + DeserializeOwned,
{
let body = if request.hooks.has_guardrails() {
let changed = request
.hooks
.during_call(OcrDuringCallRequest {
model: request.model.clone(),
custom_llm_provider: request.adapter.provider().as_str().into(),
url: url.into(),
body: serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField {
path: "body".into(),
})?,
})
.await?;
let body = OcrWireBody::<B>::decode(changed.body)?;
validate(&body.body)?;
body
} else {
OcrWireBody {
body,
extra: Map::new(),
}
};
build_http_request(client, request, url, headers, &body)
}
pub(crate) fn build_http_request<B: Serialize>(
client: &OcrClient,
request: &LiteLLMOcrRequest,
url: &str,
headers: &[(String, String)],
body: &B,
) -> Result<reqwest::Request, OcrError> {
let builder = client
.provider_http()
.post(url)
.json(body)
.timeout(request.connection.timeout);
crate::http_utils::with_headers(builder, headers, crate::http_utils::HeaderPolicy::All)
.build()
.map_err(crate::error::TransportError::from)
.map_err(OcrError::from)
}
#[derive(Serialize)]
struct OcrWireBody<B> {
#[serde(flatten)]
body: B,
#[serde(flatten)]
extra: Map<String, Value>,
}
impl<B: Serialize + DeserializeOwned> OcrWireBody<B> {
fn decode(value: Value) -> Result<Self, OcrRequestError> {
let body: B = super::wire::decode_request_value(value.clone(), "guardrail.body")?;
let Value::Object(fields) = value else {
return Err(OcrRequestError::RequestField {
path: "guardrail.body".into(),
});
};
let known = serde_json::to_value(&body).map_err(|_| OcrRequestError::RequestField {
path: "guardrail.body".into(),
})?;
let extra = fields
.into_iter()
.filter(|(key, _)| known.get(key).is_none())
.collect();
Ok(Self { body, extra })
}
}
pub(crate) fn credential_env(name: &str) -> Option<String> {
std::env::var(name).ok()
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[derive(Debug, Deserialize, PartialEq)]
struct KnownParams {
pages: Option<Vec<i64>>,
}
#[test]
fn parsed_provider_params_separates_known_and_extra_params() {
let parsed: ParsedProviderParams<KnownParams> = super::super::wire::decode_request_value(
json!({
"pages": [0, 2],
"future_ocr_option": true,
"extra_body": {"provider_option": "value"}
}),
"optional_params",
)
.unwrap();
assert_eq!(parsed.known.pages, Some(vec![0, 2]));
assert_eq!(parsed.extra_params["future_ocr_option"], true);
assert_eq!(
parsed.extra_params["extra_body"],
json!({"provider_option": "value"})
);
assert_eq!(parsed.extra_params.len(), 2);
}
}

53
litellm-rust/crates/core/src/ocr/registry.rs Normal file
View file

@ -0,0 +1,53 @@
use super::adapters::OcrAdapter;
use crate::Error;
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
macro_rules! define_adapter_types {
($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => {
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum OcrAdapterKind {
$( $variant, )+
}
impl OcrAdapterKind {
pub(crate) const fn provider(self) -> OcrProvider {
match self {
$( Self::$variant => <$adapter>::PROVIDER, )+
}
}
}
};
}
super::adapters::for_each_ocr_adapter!(define_adapter_types);
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum OcrProvider {
Mistral,
}
impl OcrProvider {
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::Mistral => "mistral",
}
}
}
pub(crate) fn resolve_wire_adapter(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<(String, OcrAdapterKind), Error> {
let provider =
get_custom_llm_provider(model, custom_llm_provider).unwrap_or(CustomLlmProvider {
model,
custom_llm_provider: OcrProvider::Mistral.as_str(),
});
let typed_provider = match provider.custom_llm_provider {
"mistral" => OcrProvider::Mistral,
value => return Err(Error::InvalidProvider(value.to_string())),
};
match typed_provider {
OcrProvider::Mistral => Ok((provider.model.to_string(), OcrAdapterKind::Mistral)),
}
}

6
litellm-rust/crates/core/src/ocr/transformation.rs
View file

@ -1,7 +1,7 @@
use crate::Error;
use serde_json::{Map, Value};
use super::types::{OcrRequestData, OcrResponseData};
use super::types::{LiteLLMOcrResponse, OcrRequestData};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrAuthStrategy {
@ -49,14 +49,14 @@ pub trait OcrProviderConfig: Sync {
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error>;
) -> Result<LiteLLMOcrResponse, Error>;
fn transform_ocr_response_with_params(
&self,
model: &str,
response_json: Value,
_optional_params: &Map<String, Value>,
) -> Result<OcrResponseData, Error> {
) -> Result<LiteLLMOcrResponse, Error> {
self.transform_ocr_response(model, response_json)
}

154
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -1,6 +1,14 @@
use std::sync::Arc;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use super::hooks::{NoopOcrHooks, OcrHooks};
use super::registry::{OcrAdapterKind, resolve_wire_adapter};
use crate::Error;
use crate::constants::OCR_HTTP_TIMEOUT_SECS;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct OcrRequestData {
pub data: Value,
@ -8,31 +16,145 @@ pub struct OcrRequestData {
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct OcrResponseData {
#[serde(tag = "type")]
pub enum OcrDocument {
#[serde(rename = "document_url")]
DocumentUrl {
document_url: String,
#[serde(flatten)]
extra_fields: Map<String, Value>,
},
#[serde(rename = "image_url")]
ImageUrl {
image_url: String,
#[serde(flatten)]
extra_fields: Map<String, Value>,
},
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum OcrResponseFormat {
#[default]
Litellm,
Native,
}
#[derive(Clone)]
pub struct OcrConnection {
pub api_key: Option<String>,
pub api_base: Option<String>,
pub extra_headers: Vec<(String, String)>,
pub timeout: Duration,
}
impl Default for OcrConnection {
fn default() -> Self {
Self {
api_key: None,
api_base: None,
extra_headers: Vec::new(),
timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS),
}
}
}
pub struct LiteLLMOcrRequest {
pub model: String,
pub document: OcrDocument,
pub connection: OcrConnection,
pub hooks: Arc<dyn OcrHooks>,
pub litellm_call_id: Option<String>,
pub optional_params: Map<String, Value>,
pub(crate) adapter: OcrAdapterKind,
}
impl LiteLLMOcrRequest {
pub fn new(
model: String,
document: OcrDocument,
custom_llm_provider: Option<&str>,
optional_params: Map<String, Value>,
) -> Result<Self, Error> {
let (model, adapter_kind) = resolve_wire_adapter(&model, custom_llm_provider)?;
Ok(Self {
model,
document,
connection: OcrConnection::default(),
hooks: Arc::new(NoopOcrHooks),
litellm_call_id: None,
optional_params,
adapter: adapter_kind,
})
}
pub(crate) fn response_format(
&self,
) -> Result<OcrResponseFormat, super::error::OcrRequestError> {
self.optional_params
.get("req_format")
.map(|value| {
serde_json::from_value(value.clone())
.map_err(|_| super::error::OcrRequestError::RequestFormat)
})
.transpose()
.map(|format| format.unwrap_or_default())
}
pub fn with_host_hooks(
self,
hooks: Arc<dyn OcrHooks>,
litellm_call_id: Option<String>,
) -> Self {
Self {
hooks,
litellm_call_id,
..self
}
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct LiteLLMOcrResponse {
pub pages: Vec<Value>,
pub model: String,
pub document_annotation: Option<Value>,
pub usage_info: Option<Value>,
pub object: String,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub provider_native_response: Option<Value>,
}
impl OcrResponseData {
impl LiteLLMOcrResponse {
pub fn into_json(self) -> Value {
let mut response = serde_json::json!({
"pages": self.pages,
"model": self.model,
"document_annotation": self.document_annotation,
"usage_info": self.usage_info,
"object": self.object,
});
if let Value::Object(object) = &mut response {
object.extend(self.extra_fields);
if let Some(native_response) = self.provider_native_response {
object.insert("provider_native_response".to_string(), native_response);
}
}
response
serde_json::to_value(self).expect("OCR response fields are JSON-compatible")
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn response_serialization_flattens_extra_fields_and_omits_absent_native_response() {
let response = LiteLLMOcrResponse {
pages: vec![],
model: "model".into(),
document_annotation: None,
usage_info: None,
object: "ocr".into(),
extra_fields: json!({"provider_field":"kept"})
.as_object()
.unwrap()
.clone(),
provider_native_response: None,
};
let serialized = response.into_json();
assert_eq!(serialized["provider_field"], "kept");
assert!(serialized.get("provider_native_response").is_none());
}
}

154
litellm-rust/crates/core/src/ocr/wire.rs Normal file
View file

@ -0,0 +1,154 @@
use crate::ocr::error::OcrRequestError;
use crate::ocr::error::OcrResponseError;
use std::time::Duration;
use super::hooks::{OcrDuringCallRequest, OcrPreCallRequest};
use super::types::{LiteLLMOcrRequest, OcrConnection, OcrDocument};
use crate::Error;
use serde::{
Deserialize,
de::{DeserializeOwned, IntoDeserializer},
};
use serde_json::{Map, Value};
#[derive(Debug)]
pub struct DecodedOcrResponse<T> {
pub data: T,
pub native: Option<Value>,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
pub struct OcrWireRequest {
pub model: String,
pub document: Value,
pub api_key: Option<String>,
pub api_base: Option<String>,
pub custom_llm_provider: Option<String>,
pub extra_headers: Option<Map<String, Value>>,
#[serde(default)]
pub optional_params: Map<String, Value>,
pub timeout_seconds: Option<f64>,
}
pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> bool {
super::registry::resolve_wire_adapter(model, custom_llm_provider).is_ok()
}
pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error> {
let document = decode_request_value(wire.document, "document")?;
let headers = wire
.extra_headers
.unwrap_or_default()
.into_iter()
.map(|(name, value)| {
let value = value
.as_str()
.ok_or_else(|| OcrRequestError::RequestField {
path: format!("extra_headers.{name}"),
})?;
Ok((name, value.to_string()))
})
.collect::<Result<Vec<_>, OcrRequestError>>()?;
let timeout = wire
.timeout_seconds
.map(|seconds| {
Duration::try_from_secs_f64(seconds).map_err(|_| OcrRequestError::RequestField {
path: "timeout_seconds".into(),
})
})
.transpose()?;
let defaults = OcrConnection::default();
let request = LiteLLMOcrRequest::new(
wire.model,
document,
wire.custom_llm_provider.as_deref(),
wire.optional_params,
)?;
let connection = OcrConnection {
api_key: nonblank(wire.api_key),
api_base: nonblank(wire.api_base),
extra_headers: headers,
timeout: timeout.unwrap_or(defaults.timeout),
};
Ok(LiteLLMOcrRequest {
connection,
..request
})
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
}
pub fn decode_request_value<T: DeserializeOwned>(
value: Value,
prefix: &str,
) -> Result<T, OcrRequestError> {
serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| {
OcrRequestError::RequestField {
path: format!("{prefix}.{}", error.path()),
}
})
}
pub fn decode_response<T: DeserializeOwned>(
bytes: &[u8],
native: bool,
) -> Result<DecodedOcrResponse<T>, OcrResponseError> {
let mut deserializer = serde_json::Deserializer::from_slice(bytes);
let data = serde_path_to_error::deserialize(&mut deserializer).map_err(|error| {
OcrResponseError::ResponseField {
path: error.path().to_string(),
}
})?;
deserializer
.end()
.map_err(|_| OcrResponseError::ResponseField {
path: "response".into(),
})?;
let native = if native {
Some(
serde_json::from_slice(bytes).map_err(|_| OcrResponseError::ResponseField {
path: "response".into(),
})?,
)
} else {
None
};
Ok(DecodedOcrResponse { data, native })
}
pub fn decode_pre_call_result(
original: OcrPreCallRequest,
value: Value,
) -> Result<OcrPreCallRequest, OcrRequestError> {
#[derive(Deserialize)]
struct Changed {
document: OcrDocument,
#[serde(default)]
optional_params: Map<String, Value>,
}
let changed: Changed = decode_request_value(value, "guardrail")?;
Ok(OcrPreCallRequest {
document: changed.document,
optional_params: Value::Object(changed.optional_params),
..original
})
}
pub fn decode_during_call_result(
original: OcrDuringCallRequest,
value: Value,
) -> Result<OcrDuringCallRequest, OcrRequestError> {
#[derive(Deserialize)]
struct Changed {
body: Value,
}
let changed: Changed = decode_request_value(value, "guardrail")?;
Ok(OcrDuringCallRequest {
body: changed.body,
..original
})
}

14
litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs
View file

@ -2,7 +2,7 @@ use std::collections::BTreeSet;
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling};
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData};
use serde_json::{Map, Value, json};
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
@ -440,7 +440,7 @@ fn transform_document_intelligence_response(
model: &str,
response_json: Value,
preserve_native_response: bool,
) -> Result<OcrResponseData, Error> {
) -> Result<LiteLLMOcrResponse, Error> {
let response = response_json
.as_object()
.ok_or_else(|| Error::InvalidType {
@ -488,7 +488,7 @@ fn transform_document_intelligence_response(
})
.collect();
Ok(OcrResponseData {
Ok(LiteLLMOcrResponse {
usage_info: Some(json!({
"pages_processed": pages.len(),
"doc_size_bytes": null,
@ -521,7 +521,7 @@ impl OcrProviderConfig for AzureAiOcrConfig {
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
) -> Result<LiteLLMOcrResponse, Error> {
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
}
@ -599,7 +599,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
) -> Result<LiteLLMOcrResponse, Error> {
transform_document_intelligence_response(model, response_json, false)
}
@ -608,7 +608,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
model: &str,
response_json: Value,
optional_params: &Map<String, Value>,
) -> Result<OcrResponseData, Error> {
) -> Result<LiteLLMOcrResponse, Error> {
transform_document_intelligence_response(
model,
response_json,
@ -718,7 +718,7 @@ mod tests {
})
}
fn assert_native_fields_preserved(response: &OcrResponseData, operation: &Value) {
fn assert_native_fields_preserved(response: &LiteLLMOcrResponse, operation: &Value) {
let analyze_result = &operation["analyzeResult"];
assert_eq!(response.extra_fields["content"], analyze_result["content"]);

11
litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs
View file

@ -1,6 +1,6 @@
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData};
use serde_json::{Map, Value};
const SUPPORTED_OCR_PARAMS: &[&str] = &[
@ -107,7 +107,7 @@ impl OcrProviderConfig for MistralOcrConfig {
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
) -> Result<LiteLLMOcrResponse, Error> {
let response_object = response_json
.as_object()
.ok_or_else(|| Error::InvalidType {
@ -128,7 +128,7 @@ impl OcrProviderConfig for MistralOcrConfig {
let document_annotation = response_object.get("document_annotation").cloned();
let usage_info = response_object.get("usage_info").cloned();
Ok(OcrResponseData {
Ok(LiteLLMOcrResponse {
pages,
model,
document_annotation,
@ -178,7 +178,10 @@ pub fn transform_ocr_request(
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub fn transform_ocr_response(model: &str, response_json: Value) -> Result<OcrResponseData, Error> {
pub fn transform_ocr_response(
model: &str,
response_json: Value,
) -> Result<LiteLLMOcrResponse, Error> {
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
}

10
litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs
View file

@ -6,7 +6,7 @@ use serde_json::{Map, Value, json};
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData};
pub const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
pub const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
@ -283,7 +283,7 @@ fn build_pages(result: &Map<String, Value>) -> Vec<Value> {
pub fn transform_reducto_response(
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
) -> Result<LiteLLMOcrResponse, Error> {
let response = response_json
.as_object()
.ok_or_else(|| Error::InvalidType {
@ -311,7 +311,7 @@ pub fn transform_reducto_response(
"credits": usage.get("credits").cloned().unwrap_or(Value::Null),
}));
Ok(OcrResponseData {
Ok(LiteLLMOcrResponse {
pages: build_pages(result),
model: model.to_string(),
document_annotation: None,
@ -341,7 +341,7 @@ impl OcrProviderConfig for ReductoParseV3Config {
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
) -> Result<LiteLLMOcrResponse, Error> {
transform_reducto_response(model, response_json)
}
@ -383,7 +383,7 @@ impl OcrProviderConfig for ReductoParseLegacyConfig {
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
) -> Result<LiteLLMOcrResponse, Error> {
transform_reducto_response(model, response_json)
}

8
litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs
View file

@ -1,6 +1,6 @@
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData};
use serde_json::{Map, Value, json};
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
@ -226,7 +226,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
) -> Result<LiteLLMOcrResponse, Error> {
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
}
@ -301,7 +301,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
) -> Result<LiteLLMOcrResponse, Error> {
let response = response_json
.as_object()
.ok_or_else(|| Error::InvalidType {
@ -339,7 +339,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
.get("usage_info")
.cloned()
.or_else(|| response.get("usage").cloned());
Ok(OcrResponseData {
Ok(LiteLLMOcrResponse {
pages,
model: object
.get("model")

95
litellm-rust/crates/core/src/url_utils.rs Normal file
View file

@ -0,0 +1,95 @@
use std::marker::PhantomData;
use thiserror::Error;
use url::Url;
#[derive(Debug, Error)]
pub(crate) enum ApiUrlError {
#[error("invalid URL: {0}")]
Parse(#[from] url::ParseError),
#[error("URL cannot be used as a base")]
CannotBeBase,
}
pub(crate) struct Base;
pub(crate) struct Complete;
pub(crate) struct ApiUrl<State> {
url: Url,
state: PhantomData<State>,
}
impl ApiUrl<Base> {
pub(crate) fn parse(value: &str) -> Result<Self, ApiUrlError> {
Ok(Self {
url: Url::parse(value.trim())?,
state: PhantomData,
})
}
pub(crate) fn complete_path(
mut self,
target: &[&str],
) -> Result<ApiUrl<Complete>, ApiUrlError> {
let existing: Vec<String> = self
.url
.path_segments()
.ok_or(ApiUrlError::CannotBeBase)?
.filter(|segment| !segment.is_empty())
.map(str::to_string)
.collect();
let overlap = (0..=existing.len().min(target.len()))
.rev()
.find(|&length| {
existing[existing.len() - length..]
.iter()
.map(String::as_str)
.eq(target[..length].iter().copied())
})
.unwrap_or(0);
self.url
.path_segments_mut()
.map_err(|()| ApiUrlError::CannotBeBase)?
.pop_if_empty()
.extend(target[overlap..].iter().copied());
Ok(ApiUrl {
url: self.url,
state: PhantomData,
})
}
}
impl ApiUrl<Complete> {
pub(crate) fn into_string(self) -> String {
self.url.into()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completion_appends_only_the_missing_path_suffix() {
for (base, expected) in [
("https://example.test", "https://example.test/v1/ocr"),
("https://example.test/v1", "https://example.test/v1/ocr"),
("https://example.test/v1/ocr", "https://example.test/v1/ocr"),
] {
let actual = ApiUrl::parse(base)
.and_then(|url| url.complete_path(&["v1", "ocr"]))
.map(|url| url.into_string())
.expect("url builds");
assert_eq!(actual, expected);
}
}
#[test]
fn completion_places_paths_before_queries() {
let actual = ApiUrl::parse("https://example.test/v1?tenant=a")
.and_then(|url| url.complete_path(&["v1", "ocr"]))
.map(|url| url.into_string())
.expect("url builds");
assert_eq!(actual, "https://example.test/v1/ocr?tenant=a");
}
}

227
litellm-rust/crates/core/tests/ocr.rs Normal file
View file

@ -0,0 +1,227 @@
use std::sync::{Arc, Mutex};
use serde_json::{Value, json};
use super::OcrClient;
use super::hooks::{OcrHookFuture, OcrHooks, OcrLogFuture, OcrPreCallRequest};
use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
use super::wire::{OcrWireRequest, decode_request};
use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming};
#[test]
fn request_boundary_selects_mistral_and_rejects_unknown_providers() {
let request = OcrWireRequest {
model: "mistral/model".into(),
document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}),
api_key: Some("key".into()),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: json!({"extract_header":true,"unknown":42})
.as_object()
.unwrap()
.clone(),
timeout_seconds: None,
};
assert!(decode_request(request).is_ok());
assert!(
decode_request(OcrWireRequest {
model: "model".into(),
document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}),
api_key: Some("key".into()),
api_base: None,
custom_llm_provider: Some("unknown".into()),
extra_headers: None,
optional_params: serde_json::Map::new(),
timeout_seconds: None,
})
.is_err()
);
}
#[tokio::test]
async fn facade_executes_direct_mistral_once() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"pages":[{"index":0,"markdown":"hello","custom":"preserved"}],
"usage_info":{"pages_processed":1}
}))])
.await;
let result = perform_ocr(wire_request(
"mistral/model",
&base,
json!({"extract_header":true,"unknown":"ignored"}),
))
.await
.unwrap();
server.await.unwrap();
assert_eq!(result.pages[0]["markdown"], "hello");
assert_eq!(result.pages[0]["custom"], "preserved");
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with("POST /v1/ocr "));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("authorization: bearer test-key\r\n")
);
let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap();
assert_eq!(
body,
json!({
"model":"model",
"document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
"extract_header":true
})
);
}
#[tokio::test]
async fn facade_retains_native_response_when_requested() {
let provider_response = json!({
"pages":[{"index":0,"markdown":"hello"}],
"usage_info":{"pages_processed":1},
"provider_only":"preserved"
});
let (base, _, server) = mock_server(vec![MockResponse::json(provider_response.clone())]).await;
let response = perform_ocr(wire_request(
"mistral/model",
&base,
json!({"req_format":"native"}),
))
.await
.unwrap();
server.await.unwrap();
assert_eq!(response.provider_native_response, Some(provider_response));
}
#[tokio::test]
async fn facade_uses_the_injected_http_client() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let mut default_headers = reqwest::header::HeaderMap::new();
default_headers.insert(
"x-transport-owner",
reqwest::header::HeaderValue::from_static("host"),
);
let provider_http = reqwest::Client::builder()
.default_headers(default_headers)
.build()
.unwrap();
OcrClient::new(provider_http)
.unwrap()
.perform(wire_request("mistral/model", &base, json!({})))
.await
.unwrap();
server.await.unwrap();
assert!(seen.lock().unwrap()[0].contains("x-transport-owner: host"));
}
struct RecordingHooks {
events: Arc<Mutex<Vec<&'static str>>>,
block: bool,
}
impl OcrHooks for RecordingHooks {
fn has_guardrails(&self) -> bool {
true
}
fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
Box::pin(async move {
self.events.lock().unwrap().push("pre");
if self.block {
return Err(crate::Error::InvalidRequest("blocked".into()));
}
Ok(request)
})
}
fn during_call(
&self,
request: super::hooks::OcrDuringCallRequest,
) -> OcrHookFuture<'_, super::hooks::OcrDuringCallRequest> {
Box::pin(async move {
self.events.lock().unwrap().push("during");
Ok(request)
})
}
fn success<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_response: &'a super::LiteLLMOcrResponse,
_timing: &'a CallLifecycleTiming,
) -> OcrLogFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("success");
})
}
fn failure<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_error: &'a crate::Error,
_timing: &'a CallLifecycleTiming,
) -> OcrLogFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("failure");
})
}
}
#[tokio::test]
async fn lifecycle_orders_hooks_and_emits_one_success() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let events = Arc::new(Mutex::new(Vec::new()));
let request = wire_request("mistral/model", &base, json!({}));
let request = super::LiteLLMOcrRequest {
hooks: Arc::new(RecordingHooks {
events: events.clone(),
block: false,
}),
..request
};
perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(*events.lock().unwrap(), ["pre", "during", "success"]);
assert_eq!(seen.lock().unwrap().len(), 1);
}
#[tokio::test]
async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() {
let events = Arc::new(Mutex::new(Vec::new()));
let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({}));
let request = super::LiteLLMOcrRequest {
hooks: Arc::new(RecordingHooks {
events: events.clone(),
block: true,
}),
..request
};
let error = perform_ocr(request).await.unwrap_err();
assert!(matches!(error, crate::Error::InvalidRequest(_)));
assert_eq!(*events.lock().unwrap(), ["pre", "failure"]);
}
#[tokio::test]
async fn upstream_failure_emits_one_terminal_failure() {
let (base, seen, server) = mock_server(vec![MockResponse {
status: 500,
headers: vec![],
body: json!({"error":"failed"}),
}])
.await;
let events = Arc::new(Mutex::new(Vec::new()));
let request = wire_request("mistral/model", &base, json!({}));
let request = super::LiteLLMOcrRequest {
hooks: Arc::new(RecordingHooks {
events: events.clone(),
block: false,
}),
..request
};
assert!(perform_ocr(request).await.is_err());
server.await.unwrap();
assert_eq!(*events.lock().unwrap(), ["pre", "during", "failure"]);
assert_eq!(seen.lock().unwrap().len(), 1);
}

106
litellm-rust/crates/core/tests/ocr/support.rs Normal file
View file

@ -0,0 +1,106 @@
use std::sync::{Arc, Mutex};
use serde_json::{Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use crate::ocr::wire::{OcrWireRequest, decode_request};
use crate::ocr::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
pub(crate) fn ocr_client() -> OcrClient {
OcrClient::for_test(reqwest::Client::new())
}
pub(crate) async fn perform_ocr(
request: LiteLLMOcrRequest,
) -> Result<LiteLLMOcrResponse, crate::Error> {
ocr_client().perform(request).await
}
pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest {
decode_request(OcrWireRequest {
model: model.into(),
document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}),
api_key: Some("test-key".into()),
api_base: Some(base.into()),
custom_llm_provider: None,
extra_headers: None,
optional_params: options.as_object().unwrap().clone(),
timeout_seconds: Some(2.0),
})
.unwrap()
}
pub(crate) struct MockResponse {
pub status: u16,
pub headers: Vec<(&'static str, String)>,
pub body: Value,
}
impl MockResponse {
pub fn json(body: Value) -> Self {
Self {
status: 200,
headers: vec![],
body,
}
}
}
pub(crate) async fn mock_server(
responses: Vec<MockResponse>,
) -> (String, Arc<Mutex<Vec<String>>>, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let requests = Arc::new(Mutex::new(Vec::new()));
let seen = requests.clone();
let server_base = base.clone();
let task = tokio::spawn(async move {
for response in responses {
let (mut socket, _) = listener.accept().await.unwrap();
let mut bytes = Vec::new();
let mut buffer = [0u8; 4096];
let header_end = loop {
let n = socket.read(&mut buffer).await.unwrap();
assert!(n > 0);
bytes.extend_from_slice(&buffer[..n]);
if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") {
break index + 4;
}
};
let length = String::from_utf8_lossy(&bytes[..header_end])
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().unwrap())
})
.unwrap_or(0);
while bytes.len() < header_end + length {
let n = socket.read(&mut buffer).await.unwrap();
assert!(n > 0);
bytes.extend_from_slice(&buffer[..n]);
}
seen.lock()
.unwrap()
.push(String::from_utf8_lossy(&bytes).into_owned());
let body = serde_json::to_vec(&response.body).unwrap();
let headers = response
.headers
.into_iter()
.map(|(name, value)| {
format!("{name}: {}\r\n", value.replace("{base}", &server_base))
})
.collect::<String>();
let head = format!(
"HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n",
response.status,
body.len(),
headers
);
socket.write_all(head.as_bytes()).await.unwrap();
socket.write_all(&body).await.unwrap();
}
});
(base, requests, task)
}

7
litellm-rust/crates/core/tests/workspace_crate_allowlist.rs
View file

@ -1,6 +1,7 @@
//! Enforcement: the litellm-rust workspace has exactly five crates.
//! Enforcement: the litellm-rust workspace has exactly six crates.
//!
//! `core` (the Rust SDK), `config` (the config-loading boundary),
//! `core` (the Rust SDK), `token-counter` (standalone input token counting),
//! `config` (the config-loading boundary),
//! `ai-gateway` (the HTTP/WebSocket host),
//! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the
//! PyO3 cdylib). Adding or removing a crate must be a
@ -20,6 +21,7 @@ use std::path::{Path, PathBuf};
/// workspace legitimately gains or loses a crate.
const EXPECTED_MEMBERS: &[&str] = &[
"crates/core",
"crates/token-counter",
"crates/config",
"crates/ai-gateway",
"crates/python-interop",
@ -29,6 +31,7 @@ const EXPECTED_MEMBERS: &[&str] = &[
/// The crate subdirectory names that must exist under `crates/`.
const EXPECTED_CRATE_DIRS: &[&str] = &[
"core",
"token-counter",
"config",
"ai-gateway",
"python-interop",

5
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -24,16 +24,17 @@ trace-parity = [
futures-util.workspace = true
tracing = { workspace = true, optional = true }
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-token-counter.workspace = true
litellm-ai-gateway = { workspace = true, default-features = false }
litellm-python-interop.workspace = true
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio.workspace = true
tokio = { workspace = true, features = ["sync"] }
[dev-dependencies]
criterion = "0.8.2"
criterion.workspace = true
tokio-tungstenite.workspace = true
tracing.workspace = true

2
litellm-rust/crates/python-bridge/src/constants.rs Normal file
View file

@ -0,0 +1,2 @@
/// Concurrent token-count encodes allowed when the core count is unavailable.
pub(crate) const TOKEN_COUNT_FALLBACK_PARALLELISM: usize = 1;

1
litellm-rust/crates/python-bridge/src/errors.rs
View file

@ -41,6 +41,7 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr {
| Error::InvalidRequest(_)
| Error::InvalidType { .. }
| Error::MissingField(_)
| Error::MissingApiKey { .. }
| Error::Routing(_)
// Nothing reached the provider, so serving it on Python cannot double
// bill and is the only way the caller gets an answer at all.

39
litellm-rust/crates/python-bridge/src/execution.rs
View file

@ -3,7 +3,6 @@ use std::panic::AssertUnwindSafe;
use std::time::Duration;
use futures_util::FutureExt;
use litellm_core::error::Error;
use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
@ -11,14 +10,15 @@ use serde::Serialize;
use tokio::runtime::{Handle, Runtime};
use tokio::time::{self, MissedTickBehavior};
pub(crate) fn run_sync<T, F>(
pub(crate) fn run_sync<T, E, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
map_error: fn(E) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
E: Send + 'static,
F: Future<Output = Result<T, E>> + Send + 'static,
{
run_sync_on(
py,
@ -28,15 +28,16 @@ where
)
}
fn run_sync_on<T, F>(
fn run_sync_on<T, E, F>(
py: Python<'_>,
runtime: &Runtime,
future: F,
map_error: fn(Error) -> PyErr,
map_error: fn(E) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
E: Send + 'static,
F: Future<Output = Result<T, E>> + Send + 'static,
{
if Handle::try_current().is_ok() {
return Err(PyRuntimeError::new_err(
@ -49,14 +50,15 @@ where
Pythonized(result).into_pyobject(py).map(Bound::unbind)
}
pub(crate) fn run_async<T, F>(
pub(crate) fn run_async<T, E, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
map_error: fn(E) -> PyErr,
) -> PyResult<Bound<'_, PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
E: Send + 'static,
F: Future<Output = Result<T, E>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let result = catch_future_panic(future).await?;
@ -65,7 +67,7 @@ where
})
}
fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
fn map_core_result<T, E>(result: Result<T, E>, map_error: fn(E) -> PyErr) -> PyResult<T> {
match result {
Ok(value) => Ok(value),
Err(error) => Err(
@ -75,9 +77,9 @@ fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -
}
}
async fn catch_future_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
async fn catch_future_panic<T, E, F>(future: F) -> PyResult<Result<T, E>>
where
F: Future<Output = Result<T, Error>>,
F: Future<Output = Result<T, E>>,
{
AssertUnwindSafe(future)
.catch_unwind()
@ -85,9 +87,9 @@ where
.map_err(panic_to_pyerr)
}
async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
async fn wait_for_sync_result<T, E, F>(future: F) -> PyResult<Result<T, E>>
where
F: Future<Output = Result<T, Error>>,
F: Future<Output = Result<T, E>>,
{
let future = catch_future_panic(future);
tokio::pin!(future);
@ -114,6 +116,7 @@ mod tests {
use std::thread;
use std::time::Instant;
use litellm_core::error::Error;
use pyo3::panic::PanicException;
use pyo3::types::{PyDict, PyModule};
use serde::Serializer;
@ -237,7 +240,7 @@ mod tests {
let error = runtime.block_on(async {
Python::attach(|py| {
run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
run_sync::<bool, Error, _>(py, async { Ok(true) }, runtime_error)
.expect_err("sync route should reject a nested Tokio runtime")
})
});
@ -273,7 +276,7 @@ mod tests {
fn sync_runner_maps_a_panicked_future() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
let error = run_sync::<bool, Error, _>(
py,
poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
runtime_error,
@ -289,7 +292,7 @@ mod tests {
fn sync_runner_maps_a_panicked_error_mapper() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
let error = run_sync::<bool, Error, _>(
py,
async { Err(Error::InvalidRequest("invalid".to_string())) },
panicking_error_mapper,

4
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -1,3 +1,4 @@
mod constants;
mod diagnostics;
mod errors;
mod execution;
@ -5,6 +6,7 @@ mod execution;
mod function_trace;
mod marshal;
mod routes;
mod token_counter;
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
use pyo3::prelude::*;
@ -71,6 +73,7 @@ mod _native {
super::errors::register(module)?;
super::routes::register(module)?;
module.add_class::<super::ResponsesWebSocketConnection>()?;
super::token_counter::register(module)?;
super::diagnostics::register(module)
}
}
@ -106,6 +109,7 @@ mod tests {
"chat_completions",
"achat_completions",
"ResponsesWebSocketConnection",
"TokenCounter",
"gil_stats",
];

87
litellm-rust/crates/python-bridge/src/token_counter.rs Normal file
View file

@ -0,0 +1,87 @@
use std::num::NonZero;
use std::sync::Arc;
use std::thread::available_parallelism;
use litellm_python_interop::release_gil;
use litellm_token_counter::{
CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter,
};
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::PyAny;
use tokio::sync::Semaphore;
use crate::constants::TOKEN_COUNT_FALLBACK_PARALLELISM;
use crate::errors::RustBridgeDeclined;
use crate::execution::run_async;
/// Counts the input tokens of a raw request body off the Python event loop with
/// the GIL released. Python owns which requests get here and what to do with
/// the count. At most one encode per core runs at a time; the rest wait in the
/// async task, where a cancelled Python awaiter drops them before any blocking
/// work is scheduled.
#[pyclass(frozen)]
struct TokenCounter {
inner: Arc<CoreTokenCounter>,
encode_slots: Arc<Semaphore>,
}
#[pymethods]
impl TokenCounter {
#[new]
fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
let inner = release_gil(py, || CoreTokenCounter::from_json(tokenizer_json))
.map_err(token_count_error_to_pyerr)?;
Ok(Self {
inner: Arc::new(inner),
encode_slots: Arc::new(Semaphore::new(encode_parallelism())),
})
}
fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
let counter = Arc::clone(&self.inner);
let encode_slots = Arc::clone(&self.encode_slots);
let body = body.to_vec();
run_async(
py,
async move {
let _slot = encode_slots
.acquire_owned()
.await
.map_err(|error| Error::Task(error.to_string()))?;
tokio::task::spawn_blocking(move || count_body(&counter, &body))
.await
.map_err(|error| Error::Task(error.to_string()))?
},
token_count_error_to_pyerr,
)
}
}
fn encode_parallelism() -> usize {
available_parallelism().map_or(TOKEN_COUNT_FALLBACK_PARALLELISM, NonZero::get)
}
fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result<InputTokenCount, Error> {
let request = CountableRequest::parse(body)?;
counter.count_request(&request)
}
fn token_count_error_to_pyerr(error: Error) -> PyErr {
let message = error.to_string();
match error {
Error::Load(_) => PyValueError::new_err(message),
Error::RequestParse(_)
| Error::MissingInput
| Error::FloatText
| Error::ContentBlock
| Error::ArrayItems
| Error::JsonSerialization(_)
| Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message),
Error::Encode(_) | Error::Task(_) => PyRuntimeError::new_err(message),
}
}
pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_class::<TokenCounter>()
}

28
litellm-rust/crates/token-counter/Cargo.toml Normal file
View file

@ -0,0 +1,28 @@
[package]
name = "litellm-token-counter"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
indexmap = { version = "2.14.0", features = ["serde"] }
itoa = "1.0"
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
unicode-normalization-alignments = "0.1.12"
[dev-dependencies]
criterion.workspace = true
rand.workspace = true
rstest.workspace = true
[[bench]]
name = "token_counter"
harness = false
[[bench]]
name = "allocations"
harness = false

128
litellm-rust/crates/token-counter/benches/allocations.rs Normal file
View file

@ -0,0 +1,128 @@
use std::alloc::{GlobalAlloc, Layout, System};
use std::hint::black_box;
use std::sync::atomic::{AtomicUsize, Ordering};
use litellm_token_counter::{CountableRequest, TokenCounter};
struct CountingAllocator;
static ALLOCATIONS: AtomicUsize = AtomicUsize::new(0);
static BYTES: AtomicUsize = AtomicUsize::new(0);
unsafe impl GlobalAlloc for CountingAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
ALLOCATIONS.fetch_add(1, Ordering::Relaxed);
BYTES.fetch_add(layout.size(), Ordering::Relaxed);
// SAFETY: This allocator delegates the unchanged layout to `System`.
unsafe { System.alloc(layout) }
}
unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
// SAFETY: The pointer and layout came from the delegated `System` allocation.
unsafe { System.dealloc(pointer, layout) }
}
unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 {
ALLOCATIONS.fetch_add(1, Ordering::Relaxed);
BYTES.fetch_add(size, Ordering::Relaxed);
// SAFETY: The pointer and layout came from `System`; the new size is unchanged.
unsafe { System.realloc(pointer, layout, size) }
}
}
#[global_allocator]
static ALLOCATOR: CountingAllocator = CountingAllocator;
#[derive(Clone, Copy)]
struct AllocationCount {
allocations: usize,
bytes: usize,
}
impl AllocationCount {
fn assert_max(self, label: &str, maximum: Self) {
eprintln!(
"{label}: {} allocations, {} bytes",
self.allocations, self.bytes
);
assert!(
self.allocations <= maximum.allocations,
"{label} allocation count exceeded {}",
maximum.allocations
);
assert!(
self.bytes <= maximum.bytes,
"{label} allocated bytes exceeded {}",
maximum.bytes
);
}
}
const TOKENIZER_JSON: &str = include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json"
));
const OBJECT_BODY: &[u8] = br#"{"model":"claude-sonnet-4-5","input":{"text":"caf\u00e9","n":3,"ok":true,"list":[1,"a",{"z":[]}]}}"#;
const INTEGER_BODY: &[u8] = br#"{"input":[-9223372036854775808,0,18446744073709551615]}"#;
fn measure(operation: impl FnOnce()) -> AllocationCount {
ALLOCATIONS.store(0, Ordering::Relaxed);
BYTES.store(0, Ordering::Relaxed);
operation();
AllocationCount {
allocations: ALLOCATIONS.load(Ordering::Relaxed),
bytes: BYTES.load(Ordering::Relaxed),
}
}
fn main() {
measure(|| {
black_box(CountableRequest::parse(OBJECT_BODY).expect("object request parses"));
})
.assert_max(
"parse object request",
AllocationCount {
allocations: 16,
bytes: 1_900,
},
);
let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("tokenizer loads");
let object = CountableRequest::parse(OBJECT_BODY).expect("object request parses");
counter
.count_request(&object)
.expect("object warmup succeeds");
measure(|| {
black_box(
counter
.count_request(black_box(&object))
.expect("object counts"),
);
})
.assert_max(
"count object request",
AllocationCount {
allocations: 74,
bytes: 2_200,
},
);
let integers = CountableRequest::parse(INTEGER_BODY).expect("integer request parses");
counter
.count_request(&integers)
.expect("integer warmup succeeds");
measure(|| {
black_box(
counter
.count_request(black_box(&integers))
.expect("integers count"),
);
})
.assert_max(
"count integer list",
AllocationCount {
allocations: 26,
bytes: 1_050,
},
);
}

100
litellm-rust/crates/token-counter/benches/token_counter.rs Normal file
View file

@ -0,0 +1,100 @@
use std::hint::black_box;
use std::time::Duration;
use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use litellm_token_counter::TokenCounter;
use tokenizers::Tokenizer;
const TOKENIZER_JSON: &str = include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json"
));
const FULL_CONTEXT_TOKENS: usize = 1_000_000;
const LARGE_PROMPT_UNIT: &str = "The quick brown fox jumps over the lazy dog. 0123456789\n";
fn full_context_input(tokenizer: &Tokenizer) -> String {
let unit_tokens = tokenizer
.encode_fast(LARGE_PROMPT_UNIT, false)
.expect("reference tokenizer should encode")
.len();
let input = LARGE_PROMPT_UNIT.repeat(FULL_CONTEXT_TOKENS.div_ceil(unit_tokens));
let actual_tokens = tokenizer
.encode_fast(input.as_str(), true)
.expect("reference tokenizer should encode")
.len();
assert!((FULL_CONTEXT_TOKENS..FULL_CONTEXT_TOKENS + unit_tokens).contains(&actual_tokens));
input
}
fn inputs(tokenizer: &Tokenizer) -> Vec<(&'static str, String)> {
vec![
(
"ascii_chat",
"User: Summarize the benefits of statistical benchmarking.\nAssistant:".repeat(8),
),
(
"unicode_nfkc",
"Ａ quick café résumé: مرحبا 世界 🙂 fi Ⅳ.\n".repeat(16),
),
("large_prompt", LARGE_PROMPT_UNIT.repeat(256)),
("full_context_1m_tokens", full_context_input(tokenizer)),
]
}
fn token_counter(c: &mut Criterion) {
let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("token counter should load");
let tokenizer = TOKENIZER_JSON
.parse::<Tokenizer>()
.expect("reference tokenizer should load");
let mut group = c.benchmark_group("anthropic_token_counter");
for (name, input) in inputs(&tokenizer) {
let expected = tokenizer
.encode_fast(input.as_str(), true)
.expect("reference tokenizer should encode")
.len();
let actual = counter
.count_text(input.as_str())
.expect("benchmark path should count");
assert_eq!(
actual, expected,
"benchmark paths should produce the same count"
);
group.throughput(Throughput::Bytes(input.len() as u64));
group.bench_with_input(
BenchmarkId::new("byte_level_fast_path", name),
&input,
|b, input| {
b.iter(|| {
counter
.count_text(black_box(input.as_str()))
.expect("fast path should count")
})
},
);
group.bench_with_input(
BenchmarkId::new("full_encoder", name),
&input,
|b, input| {
b.iter(|| {
tokenizer
.encode_fast(black_box(input.as_str()), true)
.expect("reference tokenizer should encode")
.len()
})
},
);
}
group.finish();
}
criterion_group! {
name = benches;
config = Criterion::default()
.sample_size(20)
.warm_up_time(Duration::from_secs(1))
.measurement_time(Duration::from_secs(4));
targets = token_counter
}
criterion_main!(benches);

650
litellm-rust/crates/token-counter/src/byte_level.rs Normal file
View file

@ -0,0 +1,650 @@
//! Exact token counting for a supported tokenizer configuration: optional
//! NFKC normalization, `ByteLevel` pre-tokenization with the GPT-2 split regex,
//! and no post-processing. A scanner reproduces the regex's piece boundaries
//! and hands each piece to the tokenizer's model. Unsupported configurations
//! and added-token inputs fall back to the full encoder.
use std::borrow::Cow;
use std::iter;
use tokenizers::normalizers::NormalizerWrapper;
use tokenizers::pre_tokenizers::PreTokenizerWrapper;
use tokenizers::{Model, Tokenizer};
use unicode_normalization_alignments::{IsNormalized, UnicodeNormalization, is_nfkc_quick};
use super::unicode_classes::UnicodeClasses;
const CONTRACTIONS: [&str; 7] = ["'s", "'t", "'re", "'ve", "'m", "'ll", "'d"];
pub(super) struct ByteLevelCounter {
nfkc: bool,
normalized_added_tokens: Vec<String>,
unicode_classes: &'static UnicodeClasses,
}
impl ByteLevelCounter {
pub(super) fn detect(tokenizer: &Tokenizer) -> Option<Self> {
let nfkc = match tokenizer.get_normalizer() {
None => false,
Some(NormalizerWrapper::NFKC(_)) => true,
Some(_) => return None,
};
let Some(PreTokenizerWrapper::ByteLevel(byte_level)) = tokenizer.get_pre_tokenizer() else {
return None;
};
let plain = !byte_level.add_prefix_space
&& byte_level.use_regex
&& tokenizer.get_post_processor().is_none()
&& tokenizer.get_truncation().is_none()
&& tokenizer.get_padding().is_none();
if !plain {
return None;
}
let vocabulary = tokenizer.get_added_vocabulary();
let normalized_added_tokens = vocabulary
.get_vocab()
.iter()
.filter_map(|(original, id)| {
vocabulary
.simple_id_to_token(*id)
.filter(|normalized| normalized != original)
})
.collect();
Some(Self {
nfkc,
normalized_added_tokens,
unicode_classes: UnicodeClasses::get()?,
})
}
/// `None` when the text contains an added token or the model rejects a
/// piece; the caller then runs the full encoder.
pub(super) fn count(&self, tokenizer: &Tokenizer, text: &str) -> Option<usize> {
let normalized = self.normalize(text);
let added_tokens = tokenizer.get_added_vocabulary().get_vocab();
if added_tokens
.keys()
.chain(self.normalized_added_tokens.iter())
.any(|token| text.contains(token.as_str()) || normalized.contains(token.as_str()))
{
return None;
}
let model = tokenizer.get_model();
let mapped: String = normalized.bytes().map(byte_char).collect();
pieces(&normalized, self.unicode_classes)
.try_fold((0, 0), |(start, total), piece| {
let end = start + mapped_len(piece);
let tokens = model.tokenize(&mapped[start..end]).ok()?;
Some((end, total + tokens.len()))
})
.map(|(_, total)| total)
}
/// Same crate and Unicode tables as `NormalizedString::nfkc`, so the
/// result is what the full encoder would have tokenized.
fn normalize<'a>(&self, text: &'a str) -> Cow<'a, str> {
if !self.nfkc || text.is_ascii() || is_nfkc_quick(text.chars()) == IsNormalized::Yes {
return Cow::Borrowed(text);
}
Cow::Owned(text.nfkc().map(|(character, _)| character).collect())
}
}
/// GPT-2 `bytes_to_unicode`: printable Latin-1 bytes map to themselves, the
/// rest to U+0100 onwards in byte order.
fn byte_char(byte: u8) -> char {
let code = match byte {
0x21..=0x7E | 0xA1..=0xAC | 0xAE..=0xFF => u32::from(byte),
0x00..=0x20 => 0x100 + u32::from(byte),
0x7F..=0xA0 => 0x121 + u32::from(byte - 0x7F),
0xAD => 0x143,
};
char::from_u32(code).unwrap_or(char::REPLACEMENT_CHARACTER)
}
fn mapped_len(piece: &str) -> usize {
piece.len()
+ piece
.bytes()
.filter(|byte| !byte.is_ascii_graphic())
.count()
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Class {
Letter,
Number,
Space,
Other,
}
fn class(character: char, unicode_classes: &UnicodeClasses) -> Class {
match character {
'A'..='Z' | 'a'..='z' => Class::Letter,
'0'..='9' => Class::Number,
'\t'..='\r' | ' ' => Class::Space,
_ if character.is_ascii() => Class::Other,
_ if unicode_classes.is_letter(character) => Class::Letter,
_ if unicode_classes.is_number(character) => Class::Number,
_ if unicode_classes.is_space(character) => Class::Space,
_ => Class::Other,
}
}
/// The regex matches every character, so the pieces tile the text.
fn pieces<'a>(
text: &'a str,
unicode_classes: &'static UnicodeClasses,
) -> impl Iterator<Item = &'a str> {
iter::successors(split_piece(text, unicode_classes), move |(_, rest)| {
split_piece(rest, unicode_classes)
})
.map(|(piece, _)| piece)
}
fn split_piece<'a>(text: &'a str, unicode_classes: &UnicodeClasses) -> Option<(&'a str, &'a str)> {
let first = text.chars().next()?;
Some(text.split_at(piece_len(text, first, unicode_classes)))
}
fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize {
if let Some(contraction) = CONTRACTIONS.iter().find(|word| text.starts_with(**word)) {
return contraction.len();
}
let first_class = class(first, unicode_classes);
if first_class != Class::Space {
return run_len(text, first_class, unicode_classes);
}
if first != ' ' {
return space_run_len(text, unicode_classes);
}
let after_space = &text[1..];
match after_space
.chars()
.next()
.map(|character| class(character, unicode_classes))
{
None | Some(Class::Space) => space_run_len(text, unicode_classes),
Some(run_class) => 1 + run_len(after_space, run_class, unicode_classes),
}
}
fn run_len(text: &str, run_class: Class, unicode_classes: &UnicodeClasses) -> usize {
text.char_indices()
.find(|(_, character)| class(*character, unicode_classes) != run_class)
.map_or(text.len(), |(index, _)| index)
}
/// `\s+(?!\S)|\s+`: whitespace followed by a non-space leaves its last
/// character to start the next piece (` ?` on the following alternatives).
fn space_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
let run = run_len(text, Class::Space, unicode_classes);
if run == text.len() {
return run;
}
let last = text[..run].chars().next_back().map_or(0, char::len_utf8);
match run - last {
0 => run,
shorter => shorter,
}
}
#[cfg(test)]
mod tests {
use rand::rngs::StdRng;
use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
use rstest::{fixture, rstest};
use tokenizers::normalizers::NFKC;
use tokenizers::pre_tokenizers::byte_level::ByteLevel;
use tokenizers::utils::SysRegex;
use tokenizers::{
NormalizedString, Normalizer, OffsetReferential, OffsetType, PreTokenizedString,
PreTokenizer,
};
use super::*;
#[fixture]
fn anthropic_tokenizer() -> Tokenizer {
let path = concat!(
env!("CARGO_MANIFEST_DIR"),
"/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json"
);
std::fs::read_to_string(path)
.expect("anthropic tokenizer json is in the repo")
.parse()
.expect("anthropic tokenizer loads")
}
fn reference_count(tokenizer: &Tokenizer, text: &str) -> usize {
tokenizer.encode_fast(text, true).expect("encode").len()
}
fn byte_level_counter(nfkc: bool) -> ByteLevelCounter {
ByteLevelCounter {
nfkc,
normalized_added_tokens: Vec::new(),
unicode_classes: UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"),
}
}
const ALPHABET: &[&str] = &[
"a",
"Z",
"e",
"s",
"t",
"d",
"m",
"'",
"'s",
"'re",
"'ll",
"'S",
"0",
"9",
" ",
" ",
"\t",
"\n",
"\r\n",
"\u{b}",
".",
",",
"!",
"-",
"(",
"\"",
"\u{a0}",
"\u{85}",
"\u{2028}",
"\u{3000}",
"\u{200b}",
"\u{200d}",
"é",
"e\u{301}",
"ß",
"漢",
"字",
"ع",
"३",
"½",
"Ⅳ",
"🙂",
"👍🏽",
"Ａ",
"fi",
"㍿",
"㋿",
"ꟲ",
"𐞁",
"a\u{30a}",
"\u{1e0b}\u{323}",
"<",
">",
"EOT",
"<EOT>",
"<META_START>",
];
fn random_text(rng: &mut StdRng) -> String {
let pieces = rng.gen_range(0..40);
(0..pieces)
.map(|_| *ALPHABET.choose(rng).expect("alphabet is not empty"))
.collect()
}
#[rstest]
#[case::plain_text("Hello, how are you today?", true)]
#[case::added_token("stop <EOT> here", false)]
#[case::normalized_added_token("stop ＜ＥＯＴ＞ here", false)]
fn anthropic_tokenizer_takes_the_fast_path(
anthropic_tokenizer: Tokenizer,
#[case] text: &str,
#[case] supported: bool,
) {
let fast =
ByteLevelCounter::detect(&anthropic_tokenizer).expect("anthropic shape is supported");
assert!(fast.nfkc);
let count = fast.count(&anthropic_tokenizer, text);
if supported {
assert_eq!(count, Some(reference_count(&anthropic_tokenizer, text)));
} else {
assert_eq!(count, None);
}
}
#[rstest]
fn counts_match_the_full_encoder(anthropic_tokenizer: Tokenizer) {
let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
let mut rng = StdRng::seed_from_u64(2026);
for _ in 0..4000 {
let text = random_text(&mut rng).replace('<', "(");
let expected = reference_count(&anthropic_tokenizer, &text);
assert_eq!(
fast.count(&anthropic_tokenizer, &text),
Some(expected),
"text {text:?}"
);
}
}
#[rstest]
fn nfkc_matches_the_tokenizer_normalizer_for_every_scalar_value() {
let fast = byte_level_counter(true);
let mut text = String::new();
for character in (0..=0x10FFFFu32).filter_map(char::from_u32) {
text.clear();
text.push(character);
let mut expected = NormalizedString::from(text.as_str());
NFKC.normalize(&mut expected).expect("nfkc");
assert_eq!(
fast.normalize(&text),
expected.get(),
"U+{:04X}",
u32::from(character)
);
}
}
#[rstest]
fn nfkc_matches_the_tokenizer_normalizer_on_random_texts() {
let fast = byte_level_counter(true);
let mut rng = StdRng::seed_from_u64(11);
for _ in 0..4000 {
let text = random_text(&mut rng);
let mut expected = NormalizedString::from(text.as_str());
NFKC.normalize(&mut expected).expect("nfkc");
assert_eq!(fast.normalize(&text), expected.get(), "text {text:?}");
}
}
#[rstest]
fn pieces_match_the_byte_level_pre_tokenizer() {
let byte_level = ByteLevel::new(false, true, true);
let mut rng = StdRng::seed_from_u64(7);
for _ in 0..4000 {
let text = byte_level_counter(true)
.normalize(&random_text(&mut rng))
.into_owned();
let mut pre_tokenized = PreTokenizedString::from(text.as_str());
byte_level
.pre_tokenize(&mut pre_tokenized)
.expect("pre-tokenize");
let expected: Vec<(String, (usize, usize))> = pre_tokenized
.get_splits(OffsetReferential::Original, OffsetType::Byte)
.into_iter()
.map(|(mapped, offsets, _)| (mapped.to_string(), offsets))
.collect();
let actual: Vec<(String, (usize, usize))> = pieces(
&text,
UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"),
)
.map(|piece| {
let start = piece.as_ptr() as usize - text.as_ptr() as usize;
let mapped: String = piece.bytes().map(byte_char).collect();
(mapped, (start, start + piece.len()))
})
.collect();
assert_eq!(actual, expected, "text {text:?}");
}
}
#[rstest]
fn byte_chars_match_the_byte_level_alphabet() {
let byte_level = ByteLevel::new(false, false, false);
let characters: Vec<char> = (0..=0x10FFFFu32).filter_map(char::from_u32).collect();
for chunk in characters.chunks(1024) {
let text: String = chunk.iter().collect();
let mut pre_tokenized = PreTokenizedString::from(text.as_str());
byte_level
.pre_tokenize(&mut pre_tokenized)
.expect("pre-tokenize");
let expected: String = pre_tokenized
.get_splits(OffsetReferential::Original, OffsetType::Byte)
.into_iter()
.map(|(mapped, _, _)| mapped)
.collect();
let actual: String = text.bytes().map(byte_char).collect();
assert_eq!(actual.len(), mapped_len(&text));
assert_eq!(
actual,
expected,
"chunk starting at U+{:04X}",
u32::from(chunk[0])
);
}
}
#[rstest]
fn classes_match_oniguruma() {
let unicode_classes = UnicodeClasses::get().expect("Oniguruma exposes Unicode classes");
let letter = SysRegex::new(r"\p{L}").expect("regex");
let number = SysRegex::new(r"\p{N}").expect("regex");
let space = SysRegex::new(r"\s").expect("regex");
let whole =
|regex: &SysRegex, text: &str| regex.find_iter(text).next() == Some((0, text.len()));
let mut text = String::new();
for character in (0..=0x10FFFFu32).filter_map(char::from_u32) {
text.clear();
text.push(character);
let expected = if whole(&letter, &text) {
Class::Letter
} else if whole(&number, &text) {
Class::Number
} else if whole(&space, &text) {
Class::Space
} else {
Class::Other
};
assert_eq!(
class(character, unicode_classes),
expected,
"U+{:04X}",
u32::from(character)
);
}
}
#[rstest]
#[case("prefix")]
#[case("regex")]
#[case("normalizer")]
#[case("no_pre_tokenizer")]
#[case("other_pre_tokenizer")]
#[case("post_processor")]
#[case("truncation")]
#[case("padding")]
fn other_tokenizer_shapes_are_declined(
mut anthropic_tokenizer: Tokenizer,
#[case] shape: &str,
) {
use tokenizers::{PaddingParams, PaddingStrategy, TruncationParams};
match shape {
"prefix" => {
anthropic_tokenizer.with_pre_tokenizer(Some(ByteLevel::new(true, true, true)));
}
"regex" => {
anthropic_tokenizer.with_pre_tokenizer(Some(ByteLevel::new(false, true, false)));
}
"normalizer" => {
anthropic_tokenizer
.with_normalizer(Some(tokenizers::normalizers::Lowercase))
.expect("normalizer");
}
"no_pre_tokenizer" => {
anthropic_tokenizer.with_pre_tokenizer(None::<PreTokenizerWrapper>);
}
"other_pre_tokenizer" => {
anthropic_tokenizer
.with_pre_tokenizer(Some(tokenizers::pre_tokenizers::whitespace::Whitespace));
}
"post_processor" => {
anthropic_tokenizer.with_post_processor(Some(ByteLevel::default()));
}
"truncation" => {
anthropic_tokenizer
.with_truncation(Some(TruncationParams {
max_length: 2,
..Default::default()
}))
.expect("truncation");
}
"padding" => {
anthropic_tokenizer.with_padding(Some(PaddingParams {
strategy: PaddingStrategy::Fixed(32),
..Default::default()
}));
}
_ => unreachable!(),
}
assert!(ByteLevelCounter::detect(&anthropic_tokenizer).is_none());
let counter = crate::TokenCounter::from_json(
&anthropic_tokenizer.to_string(false).expect("serialize"),
)
.expect("load");
for text in ["", "Hello WORLD! ＡＢ fi Ⅳ", "<EOT> stop"] {
assert_eq!(
counter.count_text(text).expect("count"),
reference_count(&anthropic_tokenizer, text)
);
}
}
#[rstest]
#[case(false)]
#[case(true)]
fn arbitrary_unicode_and_long_inputs_use_fast_path(
mut anthropic_tokenizer: Tokenizer,
#[case] nfkc: bool,
) {
if !nfkc {
anthropic_tokenizer
.with_normalizer(None::<NormalizerWrapper>)
.expect("normalizer");
}
let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
let mut rng = StdRng::seed_from_u64(314159);
for _ in 0..1000 {
let text: String = (0..64)
.filter_map(|_| char::from_u32(rng.gen_range(0..=0x10ffff)))
.collect();
assert_eq!(
fast.count(&anthropic_tokenizer, &text),
Some(reference_count(&anthropic_tokenizer, &text)),
"text {text:?}"
);
}
for text in [
"",
"'s't're've'm'll'd'S'RE",
" a \t\r\n b\u{85}\u{a0}c ",
"\0é漢🙂",
"a\u{30a}\u{301}",
"ＡfiⅣ",
] {
let text = text.repeat(2048);
assert_eq!(
fast.count(&anthropic_tokenizer, &text),
Some(reference_count(&anthropic_tokenizer, &text))
);
}
}
#[rstest]
#[case(false, false, false, false)]
#[case(true, false, false, false)]
#[case(false, true, false, false)]
#[case(false, false, true, false)]
#[case(false, false, false, true)]
fn added_token_options_fall_back(
mut anthropic_tokenizer: Tokenizer,
#[case] special: bool,
#[case] single_word: bool,
#[case] lstrip: bool,
#[case] rstrip: bool,
) {
anthropic_tokenizer
.add_tokens([tokenizers::AddedToken::from("custom token", special)
.single_word(single_word)
.lstrip(lstrip)
.rstrip(rstrip)])
.expect("add token");
let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
let counter = crate::TokenCounter::from_json(
&anthropic_tokenizer.to_string(false).expect("serialize"),
)
.expect("load");
for text in [
"custom token",
"a custom token b",
"acustom tokenb",
" custom token ",
] {
assert_eq!(fast.count(&anthropic_tokenizer, text), None);
assert_eq!(
counter.count_text(text).expect("count"),
reference_count(&anthropic_tokenizer, text)
);
}
}
#[test]
fn model_errors_reach_public_caller() {
let mut tokenizer = Tokenizer::new(tokenizers::models::wordpiece::WordPiece::default());
tokenizer.with_pre_tokenizer(Some(ByteLevel::new(false, true, true)));
let fast = ByteLevelCounter::detect(&tokenizer).expect("supported");
assert_eq!(fast.count(&tokenizer, "hello"), None);
assert!(tokenizer.encode_fast("hello", true).is_err());
let counter =
crate::TokenCounter::from_json(&tokenizer.to_string(false).expect("serialize"))
.expect("load");
assert!(matches!(
counter.count_text("hello"),
Err(crate::Error::Encode(_))
));
}
#[rstest]
fn shared_counter_matches_encoder_across_threads(anthropic_tokenizer: Tokenizer) {
let counter = crate::TokenCounter::from_json(
&anthropic_tokenizer.to_string(false).expect("serialize"),
)
.expect("load");
let inputs = [
"hello world",
"ＡＢ fi\n漢字🙂",
" <EOT> stop",
"\t 're \r\n",
];
let expected = inputs.map(|text| reference_count(&anthropic_tokenizer, text));
std::thread::scope(|scope| {
for _ in 0..8 {
let counter = &counter;
scope.spawn(move || {
for _ in 0..100 {
for (text, count) in inputs.iter().zip(expected) {
assert_eq!(counter.count_text(text).expect("count"), count);
}
}
});
}
});
}
#[rstest]
fn normalized_added_token_spelling_declines_fast_path(mut anthropic_tokenizer: Tokenizer) {
anthropic_tokenizer
.add_tokens([tokenizers::AddedToken::from("ＡＢＣＤ　ＥＦＧＨ", false)])
.expect("add token");
let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
assert_eq!(reference_count(&anthropic_tokenizer, "ABCD EFGH"), 1);
assert_eq!(fast.count(&anthropic_tokenizer, "ABCD EFGH"), None);
let counter = crate::TokenCounter::from_json(
&anthropic_tokenizer.to_string(false).expect("serialize"),
)
.expect("load");
assert_eq!(counter.count_text("ABCD EFGH").expect("count"), 1);
}
}

194
litellm-rust/crates/token-counter/src/counter.rs Normal file
View file

@ -0,0 +1,194 @@
use serde::Serialize;
use crate::Error;
use crate::byte_level::ByteLevelCounter;
use crate::python_json;
use crate::tools::format_function_definitions;
use crate::types::{
ContentBlock, ContentItem, CountableRequest, Message, MessageContent, TextValue, ToolChoice,
ToolDefinition,
};
const TOKENS_PER_MESSAGE: usize = 3;
const TOKENS_PER_NAME: usize = 1;
const REPLY_PRIMING_TOKENS: usize = 3;
const TOOL_DEFINITIONS_TOKENS: usize = 9;
const TOOLS_WITH_SYSTEM_MESSAGE_DISCOUNT: usize = 4;
const TOOL_CHOICE_NONE_TOKENS: usize = 1;
const NAMED_TOOL_CHOICE_TOKENS: usize = 7;
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct InputTokenCount {
pub model: Option<String>,
pub input_tokens: usize,
}
/// A loaded HuggingFace tokenizer plus the message accounting Python applies on
/// top of it. Encoding is CPU-bound and synchronous; hosts run it off their
/// event loop.
pub struct TokenCounter {
tokenizer: tokenizers::Tokenizer,
byte_level: Option<ByteLevelCounter>,
}
impl TokenCounter {
/// Load a HuggingFace `tokenizer.json` document. The host reads the file.
pub fn from_json(tokenizer_json: &str) -> Result<Self, Error> {
let tokenizer = tokenizer_json
.parse::<tokenizers::Tokenizer>()
.map_err(Error::Load)?;
let byte_level = ByteLevelCounter::detect(&tokenizer);
Ok(Self {
tokenizer,
byte_level,
})
}
pub fn count_text(&self, text: &str) -> Result<usize, Error> {
if let Some(count) = self
.byte_level
.as_ref()
.and_then(|counter| counter.count(&self.tokenizer, text))
{
return Ok(count);
}
self.tokenizer
.encode_fast(text, true)
.map(|encoding| encoding.len())
.map_err(Error::Encode)
}
/// Mirrors the host's key precedence: `messages`, then `prompt`, then
/// `input`, then `query` plus `documents`.
pub fn count_request(&self, request: &CountableRequest) -> Result<InputTokenCount, Error> {
let input_tokens = if let Some(messages) = &request.messages {
self.count_messages(request, messages)?
} else if let Some(prompt) = &request.prompt {
self.count_text_value(prompt)?
} else if let Some(input) = &request.input {
self.count_text_value(input)?
} else if request.query.is_some() || request.documents.is_some() {
self.count_optional_text_value(request.query.as_ref())?
+ self.count_optional_text_value(request.documents.as_ref())?
} else {
return Err(Error::MissingInput);
};
Ok(InputTokenCount {
model: request.model.clone(),
input_tokens,
})
}
fn count_messages(
&self,
request: &CountableRequest,
messages: &[Message],
) -> Result<usize, Error> {
let message_tokens = messages
.iter()
.map(|message| self.count_message(message))
.sum::<Result<usize, _>>()?;
let includes_system_message = messages
.iter()
.any(|message| message.role.as_deref() == Some("system"));
let extra_tokens = self.count_extra(
request.tools.as_deref().unwrap_or_default(),
request.tool_choice.as_ref(),
includes_system_message,
)?;
Ok(message_tokens + extra_tokens)
}
fn count_optional_text_value(&self, value: Option<&TextValue>) -> Result<usize, Error> {
value.map_or(Ok(0), |value| self.count_text_value(value))
}
/// `str()` for scalars, `json.dumps()` for objects, lists flattened, nulls
/// skipped. Floats are declined because Python's `repr` and Rust's float
/// formatting disagree on exponents.
fn count_text_value(&self, value: &TextValue) -> Result<usize, Error> {
match value {
TextValue::Null => Ok(0),
TextValue::Bool(true) => self.count_text("True"),
TextValue::Bool(false) => self.count_text("False"),
TextValue::Number(number) => match (number.as_i64(), number.as_u64()) {
(Some(number), _) => self.count_text(itoa::Buffer::new().format(number)),
(_, Some(number)) => self.count_text(itoa::Buffer::new().format(number)),
_ => Err(Error::FloatText),
},
TextValue::Text(text) => self.count_text(text),
TextValue::List(items) => items
.iter()
.map(|item| self.count_text_value(item))
.sum::<Result<usize, _>>(),
TextValue::Object(_) => self.count_text(&python_json::dumps(value)?),
}
}
fn count_message(&self, message: &Message) -> Result<usize, Error> {
let role_tokens = match &message.role {
Some(role) => self.count_text(role)?,
None => 0,
};
let name_tokens = match &message.name {
Some(name) => self.count_text(name)? + TOKENS_PER_NAME,
None => 0,
};
let content_tokens = match &message.content {
Some(MessageContent::Text(text)) => self.count_text(text)?,
Some(MessageContent::Blocks(items)) => items
.iter()
.map(|item| self.count_content_item(item))
.sum::<Result<usize, _>>()?,
None => 0,
};
Ok(TOKENS_PER_MESSAGE + role_tokens + name_tokens + content_tokens)
}
fn count_content_item(&self, item: &ContentItem) -> Result<usize, Error> {
match item {
ContentItem::Text(text) => self.count_text(text),
ContentItem::Block(ContentBlock::Text { text }) => self.count_text(text),
ContentItem::Block(ContentBlock::Thinking { thinking }) => {
if thinking.is_empty() {
return Ok(0);
}
self.count_text(thinking)
}
ContentItem::Block(ContentBlock::ToolReference { tool_name }) => {
match tool_name.as_deref().filter(|name| !name.is_empty()) {
Some(name) => self.count_text(name),
None => Ok(0),
}
}
ContentItem::Block(ContentBlock::Unsupported) => Err(Error::ContentBlock),
}
}
fn count_extra(
&self,
tools: &[ToolDefinition],
tool_choice: Option<&ToolChoice>,
includes_system_message: bool,
) -> Result<usize, Error> {
let tool_tokens = if tools.is_empty() {
0
} else {
let definitions = self.count_text(&format_function_definitions(tools)?)?;
let discount = if includes_system_message {
TOOLS_WITH_SYSTEM_MESSAGE_DISCOUNT
} else {
0
};
definitions + TOOL_DEFINITIONS_TOKENS - discount
};
let choice_tokens = match tool_choice {
Some(ToolChoice::Mode(mode)) if mode == "none" => TOOL_CHOICE_NONE_TOKENS,
Some(ToolChoice::Mode(_)) | None => 0,
Some(ToolChoice::Named(named)) => {
NAMED_TOOL_CHOICE_TOKENS + self.count_text(&named.function.name)?
}
};
Ok(REPLY_PRIMING_TOKENS + tool_tokens + choice_tokens)
}
}

31
litellm-rust/crates/token-counter/src/error.rs Normal file
View file

@ -0,0 +1,31 @@
use std::string::FromUtf8Error;
use thiserror::Error as ThisError;
#[derive(Debug, ThisError)]
pub enum Error {
#[error("failed to load tokenizer: {0}")]
Load(#[source] tokenizers::Error),
#[error("unsupported by the rust token counter: request body could not be parsed: {0}")]
RequestParse(#[source] serde_json::Error),
#[error("unsupported by the rust token counter: request has no countable input")]
MissingInput,
#[error(
"unsupported by the rust token counter: float text values are counted by the python path"
)]
FloatText,
#[error(
"unsupported by the rust token counter: content block type is counted by the python path"
)]
ContentBlock,
#[error("unsupported by the rust token counter: array parameter without items")]
ArrayItems,
#[error("unsupported by the rust token counter: text value could not be serialized: {0}")]
JsonSerialization(#[source] serde_json::Error),
#[error("unsupported by the rust token counter: serialized text value is not UTF-8: {0}")]
JsonUtf8(#[source] FromUtf8Error),
#[error("tokenization failed: {0}")]
Encode(#[source] tokenizers::Error),
#[error("token counting task failed: {0}")]
Task(String),
}

17
litellm-rust/crates/token-counter/src/lib.rs Normal file
View file

@ -0,0 +1,17 @@
//! Input token counting for a request body, mirroring `litellm.token_counter`
//! for the shapes it can count exactly. Everything else is declined so the host
//! keeps its own counter as the reference.
#![forbid(unsafe_code)]
mod byte_level;
mod counter;
mod error;
mod python_json;
mod tools;
mod types;
mod unicode_classes;
pub use counter::{InputTokenCount, TokenCounter};
pub use error::Error;
pub use types::CountableRequest;

155
litellm-rust/crates/token-counter/src/python_json.rs Normal file
View file

@ -0,0 +1,155 @@
//! `json.dumps(value)` with Python's default arguments: `", "` and `": "`
//! separators, `ensure_ascii=True`, and keys in insertion order.
use std::io::{self, Write};
use serde::Serialize;
use serde_json::ser::{Formatter, Serializer};
use super::Error;
use super::types::TextValue;
pub(super) fn dumps(value: &TextValue) -> Result<String, Error> {
let mut output = Vec::with_capacity(serialized_len(value)?);
value
.serialize(&mut Serializer::with_formatter(
&mut output,
PythonFormatter,
))
.map_err(Error::JsonSerialization)?;
debug_assert_eq!(output.len(), output.capacity());
String::from_utf8(output).map_err(Error::JsonUtf8)
}
fn serialized_len(value: &TextValue) -> Result<usize, Error> {
match value {
TextValue::Null => Ok(4),
TextValue::Bool(true) => Ok(4),
TextValue::Bool(false) => Ok(5),
TextValue::Number(number) => match (number.as_i64(), number.as_u64()) {
(Some(number), _) => Ok(unsigned_len(number.unsigned_abs()) + usize::from(number < 0)),
(_, Some(number)) => Ok(unsigned_len(number)),
_ => Err(Error::FloatText),
},
TextValue::Text(text) => Ok(quoted_len(text)),
TextValue::List(items) => items
.iter()
.try_fold(2 + items.len().saturating_sub(1) * 2, |len, item| {
Ok(len + serialized_len(item)?)
}),
TextValue::Object(entries) => entries.iter().try_fold(
2 + entries.len().saturating_sub(1) * 2,
|len, (key, value)| Ok(len + quoted_len(key) + 2 + serialized_len(value)?),
),
}
}
fn unsigned_len(number: u64) -> usize {
if number == 0 {
1
} else {
number.ilog10() as usize + 1
}
}
fn quoted_len(value: &str) -> usize {
value.chars().fold(2, |len, character| {
len + match character {
'"' | '\\' | '\u{0008}' | '\u{000c}' | '\n' | '\r' | '\t' => 2,
'\u{0000}'..='\u{001f}' => 6,
' '..='~' => 1,
_ => character.len_utf16() * 6,
}
})
}
struct PythonFormatter;
impl Formatter for PythonFormatter {
fn begin_array_value<W>(&mut self, writer: &mut W, first: bool) -> io::Result<()>
where
W: ?Sized + Write,
{
if !first {
writer.write_all(b", ")?;
}
Ok(())
}
fn begin_object_key<W>(&mut self, writer: &mut W, first: bool) -> io::Result<()>
where
W: ?Sized + Write,
{
if !first {
writer.write_all(b", ")?;
}
Ok(())
}
fn begin_object_value<W>(&mut self, writer: &mut W) -> io::Result<()>
where
W: ?Sized + Write,
{
writer.write_all(b": ")
}
fn write_string_fragment<W>(&mut self, writer: &mut W, fragment: &str) -> io::Result<()>
where
W: ?Sized + Write,
{
for character in fragment.chars() {
if (' '..='~').contains(&character) {
write!(writer, "{character}")?;
continue;
}
let mut units = [0u16; 2];
for unit in character.encode_utf16(&mut units) {
write!(writer, "\\u{unit:04x}")?;
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::dumps;
#[rstest]
#[case::null("null", "null")]
#[case::boolean("true", "true")]
#[case::signed_integer("-3", "-3")]
#[case::unsigned_integer("18446744073709551615", "18446744073709551615")]
#[case::empty_array("[]", "[]")]
#[case::empty_object("{}", "{}")]
#[case::nested(
r#"{"first":1,"second":{"ok":true,"none":null},"third":[false,2]}"#,
r#"{"first": 1, "second": {"ok": true, "none": null}, "third": [false, 2]}"#
)]
#[case::string_escaping(
r#""caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~ /""#,
r#""caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~ /""#
)]
#[case::short_control_escapes(r#""\b\f\r""#, r#""\b\f\r""#)]
fn matches_python_json_dumps(#[case] input: &str, #[case] expected: &str) {
let value = serde_json::from_str(input).expect("fixture parses");
assert_eq!(dumps(&value).expect("fixture dumps"), expected);
}
#[rstest]
#[case::top_level("1.5")]
#[case::array("[1,2.5]")]
#[case::object(r#"{"nested":{"value":-0.25}}"#)]
fn rejects_floats(#[case] input: &str) {
let value = serde_json::from_str(input).expect("fixture parses");
let error = dumps(&value).expect_err("floats are declined");
assert_eq!(
error.to_string(),
"unsupported by the rust token counter: float text values are counted by the python path"
);
}
}

304
litellm-rust/crates/token-counter/src/tools.rs Normal file
View file

@ -0,0 +1,304 @@
//! Renders tool definitions the way `litellm.token_counter` does before
//! tokenizing them (the TypeScript-like namespace OpenAI appears to use).
use std::fmt::Write;
use super::Error;
use super::types::{EnumValue, Schema, SchemaType, ToolDefinition};
pub(super) fn format_function_definitions(tools: &[ToolDefinition]) -> Result<String, Error> {
ToolFormatter::new(tools.len()).format(tools)
}
struct ToolFormatter {
output: String,
}
impl ToolFormatter {
fn new(tool_count: usize) -> Self {
Self {
output: String::with_capacity(tool_count.saturating_mul(128).saturating_add(48)),
}
}
fn format(mut self, tools: &[ToolDefinition]) -> Result<String, Error> {
self.output.push_str("namespace functions {\n\n");
for tool in tools {
self.write_function(tool)?;
}
self.output.push_str("} // namespace functions");
Ok(self.output)
}
fn write_function(&mut self, tool: &ToolDefinition) -> Result<(), Error> {
let (name, description, parameters) = resolve_function(tool);
let Some(name) = name.filter(|name| !name.is_empty()) else {
return Ok(());
};
if let Some(description) = description.filter(|description| !description.is_empty()) {
self.output.push_str("// ");
self.output.push_str(description);
self.output.push('\n');
}
match parameters.filter(|parameters| {
parameters
.properties
.as_ref()
.is_some_and(|properties| !properties.is_empty())
}) {
Some(parameters) => {
self.output.push_str("type ");
self.output.push_str(name);
self.output.push_str(" = (_: {\n");
self.write_object_parameters(parameters, 0)?;
self.output.push_str("\n}) => any;\n\n");
}
_ => {
self.output.push_str("type ");
self.output.push_str(name);
self.output.push_str(" = () => any;\n\n");
}
}
Ok(())
}
fn write_object_parameters(&mut self, parameters: &Schema, indent: usize) -> Result<(), Error> {
let Some(properties) = parameters
.properties
.as_ref()
.filter(|properties| !properties.is_empty())
else {
return Ok(());
};
let required = parameters.required.as_deref().unwrap_or_default();
for (index, (key, props)) in properties.iter().enumerate() {
if index > 0 {
self.output.push('\n');
}
if let Some(description) = props
.description
.as_deref()
.filter(|description| !description.is_empty())
{
self.write_indent(indent);
self.output.push_str("// ");
self.output.push_str(description);
self.output.push('\n');
}
self.write_indent(indent);
self.output.push_str(key);
if !required.iter().any(|required| required == key) {
self.output.push('?');
}
self.output.push_str(": ");
self.write_type(props, indent)?;
self.output.push(',');
}
Ok(())
}
fn write_type(&mut self, props: &Schema, indent: usize) -> Result<(), Error> {
let Some(SchemaType::Name(schema_type)) = &props.schema_type else {
self.output.push_str("any");
return Ok(());
};
match schema_type.as_str() {
"string" | "integer" | "number" => match &props.enum_values {
Some(values) => self.write_enum(values),
None if schema_type == "string" => self.output.push_str("string"),
None => self.output.push_str("number"),
},
"array" => {
let items = props.items.as_deref().ok_or(Error::ArrayItems)?;
self.write_type(items, indent)?;
self.output.push_str("[]");
}
"object" => {
self.output.push_str("{\n");
self.write_object_parameters(props, indent + 2)?;
self.output.push_str("\n}");
}
"boolean" => self.output.push_str("boolean"),
"null" => self.output.push_str("null"),
_ => self.output.push_str("any"),
}
Ok(())
}
fn write_enum(&mut self, values: &[EnumValue]) {
for (index, value) in values.iter().enumerate() {
if index > 0 {
self.output.push_str(" | ");
}
self.output.push('"');
match value {
EnumValue::Text(text) => self.output.push_str(text),
EnumValue::Integer(number) => {
write!(self.output, "{number}").expect("writing to a String cannot fail");
}
}
self.output.push('"');
}
}
fn write_indent(&mut self, indent: usize) {
for _ in 0..indent {
self.output.push(' ');
}
}
}
fn resolve_function(tool: &ToolDefinition) -> (Option<&str>, Option<&str>, Option<&Schema>) {
match &tool.function {
Some(function) => (
function.name.as_deref(),
function.description.as_deref(),
function.parameters.as_ref(),
),
None => (
tool.name.as_deref(),
tool.description.as_deref(),
tool.input_schema.as_ref().or(tool.parameters.as_ref()),
),
}
}
#[cfg(test)]
mod tests {
use super::*;
fn parse_tools(json: &str) -> Vec<ToolDefinition> {
serde_json::from_str(json).expect("tool fixture parses")
}
#[test]
fn empty_tool_list_renders_an_empty_namespace() {
assert_eq!(
format_function_definitions(&[]).expect("empty tool list renders"),
"namespace functions {\n\n} // namespace functions"
);
}
#[test]
fn unnamed_tools_are_skipped_and_function_shape_takes_precedence() {
let tools = parse_tools(
r#"[
{},
{"name":""},
{"name":"ignored","function":{"description":"missing a function name"}},
{"name":"ping","description":""}
]"#,
);
assert_eq!(
format_function_definitions(&tools).expect("tools render"),
"namespace functions {\n\ntype ping = () => any;\n\n} // namespace functions"
);
}
#[test]
fn empty_or_missing_properties_render_a_no_argument_function() {
let tools = parse_tools(
r#"[
{"name":"missing","input_schema":{"type":"object"}},
{"name":"empty","input_schema":{"type":"object","properties":{}}}
]"#,
);
assert_eq!(
format_function_definitions(&tools).expect("tools render"),
concat!(
"namespace functions {\n\n",
"type missing = () => any;\n\n",
"type empty = () => any;\n\n",
"} // namespace functions"
)
);
}
#[test]
fn anthropic_parameters_render_all_supported_types() {
let tools = parse_tools(
r#"[{
"name":"inspect",
"description":"Inspect a value",
"input_schema":{
"type":"object",
"properties":{
"text":{"type":"string"},
"count":{"type":"integer","description":"Number of attempts"},
"ratio":{"type":"number"},
"enabled":{"type":"boolean"},
"nothing":{"type":"null"},
"unknown":{"type":"custom"},
"union":{"type":["string","null"]},
"labels":{"type":"array","items":{"type":"string"}},
"config":{"type":"object","properties":{"retries":{"type":"integer"}},"required":["retries"]},
"mode":{"type":"string","enum":["fast",2]}
},
"required":["text"]
}
}]"#,
);
assert_eq!(
format_function_definitions(&tools).expect("tool renders"),
concat!(
"namespace functions {\n\n",
"// Inspect a value\n",
"type inspect = (_: {\n",
"text: string,\n",
"// Number of attempts\n",
"count?: number,\n",
"ratio?: number,\n",
"enabled?: boolean,\n",
"nothing?: null,\n",
"unknown?: any,\n",
"union?: any,\n",
"labels?: string[],\n",
"config?: {\n",
" retries: number,\n",
"},\n",
"mode?: \"fast\" | \"2\",\n",
"}) => any;\n\n",
"} // namespace functions"
)
);
}
#[test]
fn input_schema_takes_precedence_over_legacy_parameters() {
let tools = parse_tools(
r#"[{
"name":"choose",
"input_schema":{"type":"object","properties":{"current":{"type":"string"}}},
"parameters":{"type":"object","properties":{"legacy":{"type":"string"}}}
}]"#,
);
let rendered = format_function_definitions(&tools).expect("tool renders");
assert!(rendered.contains("current?: string,"));
assert!(!rendered.contains("legacy"));
}
#[test]
fn array_without_items_returns_an_error() {
let tools = parse_tools(
r#"[{"name":"broken","parameters":{"type":"object","properties":{"values":{"type":"array"}}}}]"#,
);
assert!(matches!(
format_function_definitions(&tools),
Err(Error::ArrayItems)
));
}
}

238
litellm-rust/crates/token-counter/src/types.rs Normal file
View file

@ -0,0 +1,238 @@
use std::fmt;
use indexmap::IndexMap;
use serde::de::{MapAccess, SeqAccess, Visitor};
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::Number;
use super::Error;
/// The parts of a request body the host's budget counter reads. Chat and
/// Anthropic Messages bodies carry `messages`; completions carry `prompt`;
/// Responses and embeddings carry `input`; rerank carries `query` and
/// `documents`. The host checks key presence, not nullness, so an explicit
/// `null` is kept distinct from an absent key. Anything outside this shape is
/// declined so the host can fall back to its own counter instead of silently
/// miscounting.
#[derive(Clone, Debug, Deserialize, PartialEq)]
pub struct CountableRequest {
pub(crate) model: Option<String>,
#[serde(default, deserialize_with = "present_messages")]
pub(crate) messages: Option<Vec<Message>>,
pub(crate) tools: Option<Vec<ToolDefinition>>,
pub(crate) tool_choice: Option<ToolChoice>,
#[serde(default, deserialize_with = "present_text")]
pub(crate) prompt: Option<TextValue>,
#[serde(default, deserialize_with = "present_text")]
pub(crate) input: Option<TextValue>,
#[serde(default, deserialize_with = "present_text")]
pub(crate) query: Option<TextValue>,
#[serde(default, deserialize_with = "present_text")]
pub(crate) documents: Option<TextValue>,
}
impl CountableRequest {
pub fn parse(body: &[u8]) -> Result<Self, Error> {
serde_json::from_slice(body).map_err(Error::RequestParse)
}
}
fn present_messages<'de, D: Deserializer<'de>>(
deserializer: D,
) -> Result<Option<Vec<Message>>, D::Error> {
Option::<Vec<Message>>::deserialize(deserializer)
.map(|messages| Some(messages.unwrap_or_default()))
}
fn present_text<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<TextValue>, D::Error> {
TextValue::deserialize(deserializer).map(Some)
}
/// Free-form JSON the host counts as text: strings and integers via `str()`,
/// objects via `json.dumps()`, lists flattened. Objects keep document order so
/// the dumped text matches Python byte for byte.
#[derive(Clone, Debug, PartialEq, Serialize)]
#[serde(untagged)]
pub(crate) enum TextValue {
Null,
Bool(bool),
Number(Number),
Text(String),
List(Vec<TextValue>),
Object(IndexMap<String, TextValue>),
}
impl<'de> Deserialize<'de> for TextValue {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
deserializer.deserialize_any(TextValueVisitor)
}
}
struct TextValueVisitor;
impl<'de> Visitor<'de> for TextValueVisitor {
type Value = TextValue;
fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("a JSON value")
}
fn visit_unit<E>(self) -> Result<Self::Value, E> {
Ok(TextValue::Null)
}
fn visit_none<E>(self) -> Result<Self::Value, E> {
Ok(TextValue::Null)
}
fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
Ok(TextValue::Bool(value))
}
fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
Ok(TextValue::Number(value.into()))
}
fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
Ok(TextValue::Number(value.into()))
}
fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E>
where
E: serde::de::Error,
{
Number::from_f64(value)
.map(TextValue::Number)
.ok_or_else(|| E::custom("non-finite JSON number"))
}
fn visit_str<E>(self, value: &str) -> Result<Self::Value, E> {
Ok(TextValue::Text(value.to_owned()))
}
fn visit_string<E>(self, value: String) -> Result<Self::Value, E> {
Ok(TextValue::Text(value))
}
fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
where
A: SeqAccess<'de>,
{
let mut items = Vec::with_capacity(sequence.size_hint().unwrap_or(0));
while let Some(item) = sequence.next_element()? {
items.push(item);
}
Ok(TextValue::List(items))
}
fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
where
A: MapAccess<'de>,
{
let mut entries = IndexMap::with_capacity(map.size_hint().unwrap_or(0));
while let Some((key, value)) = map.next_entry()? {
entries.insert(key, value);
}
Ok(TextValue::Object(entries))
}
}
/// Python counts every string-valued key of a message, so any key beyond these
/// makes the shape unsupported rather than silently uncounted.
#[derive(Clone, Debug, Deserialize, PartialEq)]
#[serde(deny_unknown_fields)]
pub(crate) struct Message {
pub(crate) role: Option<String>,
pub(crate) name: Option<String>,
pub(crate) content: Option<MessageContent>,
}
#[derive(Clone, Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub(crate) enum MessageContent {
Text(String),
Blocks(Vec<ContentItem>),
}
#[derive(Clone, Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub(crate) enum ContentItem {
Text(String),
Block(ContentBlock),
}
#[derive(Clone, Debug, Deserialize, PartialEq)]
#[serde(tag = "type")]
pub(crate) enum ContentBlock {
#[serde(rename = "text")]
Text { text: String },
#[serde(rename = "thinking")]
Thinking { thinking: String },
#[serde(rename = "tool_reference")]
ToolReference { tool_name: Option<String> },
/// Images, documents, files and tool use/result blocks price through
/// Python-only helpers, so they stay on the Python counter.
#[serde(other)]
Unsupported,
}
/// Either the OpenAI `{"type": "function", "function": {...}}` shape or the
/// Anthropic `{"name", "description", "input_schema"}` shape.
#[derive(Clone, Debug, Deserialize, PartialEq)]
pub(crate) struct ToolDefinition {
pub(crate) function: Option<FunctionDefinition>,
pub(crate) name: Option<String>,
pub(crate) description: Option<String>,
pub(crate) input_schema: Option<Schema>,
pub(crate) parameters: Option<Schema>,
}
#[derive(Clone, Debug, Deserialize, PartialEq)]
pub(crate) struct FunctionDefinition {
pub(crate) name: Option<String>,
pub(crate) description: Option<String>,
pub(crate) parameters: Option<Schema>,
}
#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
pub(crate) struct Schema {
#[serde(rename = "type")]
pub(crate) schema_type: Option<SchemaType>,
pub(crate) description: Option<String>,
#[serde(rename = "enum")]
pub(crate) enum_values: Option<Vec<EnumValue>>,
pub(crate) items: Option<Box<Schema>>,
pub(crate) properties: Option<IndexMap<String, Schema>>,
pub(crate) required: Option<Vec<String>>,
}
#[derive(Clone, Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub(crate) enum SchemaType {
Name(String),
Union(Vec<String>),
}
#[derive(Clone, Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub(crate) enum EnumValue {
Text(String),
Integer(i64),
}
#[derive(Clone, Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub(crate) enum ToolChoice {
Mode(String),
Named(NamedToolChoice),
}
#[derive(Clone, Debug, Deserialize, PartialEq)]
pub(crate) struct NamedToolChoice {
pub(crate) function: NamedFunction,
}
#[derive(Clone, Debug, Deserialize, PartialEq)]
pub(crate) struct NamedFunction {
pub(crate) name: String,
}

73
litellm-rust/crates/token-counter/src/unicode_classes.rs Normal file
View file

@ -0,0 +1,73 @@
use std::cmp::Ordering;
use std::sync::LazyLock;
use tokenizers::utils::SysRegex;
struct Ranges(Box<[(u32, u32)]>);
pub(super) struct UnicodeClasses {
letters: Ranges,
numbers: Ranges,
spaces: Ranges,
}
static CLASSES: LazyLock<Option<UnicodeClasses>> = LazyLock::new(|| {
let scalars: String = (0..=u32::from(char::MAX))
.filter_map(char::from_u32)
.collect();
Some(UnicodeClasses {
letters: Ranges::load(r"\p{L}+", &scalars)?,
numbers: Ranges::load(r"\p{N}+", &scalars)?,
spaces: Ranges::load(r"\s+", &scalars)?,
})
});
impl Ranges {
fn load(pattern: &str, scalars: &str) -> Option<Self> {
let regex = SysRegex::new(pattern).ok()?;
let ranges = regex
.find_iter(scalars)
.map(|(start, end)| {
let matched = scalars.get(start..end)?;
Some((
u32::from(matched.chars().next()?),
u32::from(matched.chars().next_back()?),
))
})
.collect::<Option<Box<[_]>>>()?;
Some(Self(ranges))
}
fn contains(&self, character: char) -> bool {
let code = u32::from(character);
self.0
.binary_search_by(|(low, high)| {
if *high < code {
Ordering::Less
} else if *low > code {
Ordering::Greater
} else {
Ordering::Equal
}
})
.is_ok()
}
}
impl UnicodeClasses {
pub(super) fn get() -> Option<&'static Self> {
CLASSES.as_ref()
}
pub(super) fn is_letter(&self, character: char) -> bool {
self.letters.contains(character)
}
pub(super) fn is_number(&self, character: char) -> bool {
self.numbers.contains(character)
}
pub(super) fn is_space(&self, character: char) -> bool {
self.spaces.contains(character)
}
}

176
litellm-rust/crates/token-counter/tests/token_counter.rs Normal file
View file

@ -0,0 +1,176 @@
use rstest::rstest;
use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter};
/// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)`
/// so this test also guards Python parity.
fn counter() -> TokenCounter {
let path = concat!(
env!("CARGO_MANIFEST_DIR"),
"/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json"
);
let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo");
TokenCounter::from_json(&json).expect("anthropic tokenizer loads")
}
const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#;
const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[
{"role":"system","content":"You are a terse assistant."},
{"role":"user","name":"alice","content":[
{"type":"text","text":"Summarise this paragraph about ships and harbours."},
"plain string item",
{"type":"thinking","thinking":"pondering"},
{"type":"tool_reference","tool_name":"get_weather"}]},
{"role":"assistant","content":[{"type":"text","text":"Sure.","cache_control":{"type":"ephemeral"}}]}]}"#;
const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"weather?"}],
"tools":[
{"type":"function","function":{"name":"get_weather","description":"Get weather","parameters":{
"type":"object",
"properties":{
"location":{"type":"string","description":"City name"},
"unit":{"type":"string","enum":["celsius","fahrenheit"]},
"days":{"type":"integer"},
"tags":{"type":"array","items":{"type":"string"}},
"opts":{"type":"object","properties":{"verbose":{"type":"boolean"},"level":{"type":"integer","enum":[1,2]}},"required":["verbose"]},
"anything":{}},
"required":["location"]}}},
{"type":"function","function":{"name":"noop"}}],
"tool_choice":{"type":"function","function":{"name":"get_weather"}}}"#;
const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5",
"messages":[{"role":"system","content":"sys"},{"role":"user","content":"weather?"}],
"tools":[{"name":"get_weather","description":"Get weather","input_schema":{
"type":"object","properties":{"location":{"type":["string","null"]}},"required":["location"]}}],
"tool_choice":"none"}"#;
const COMPLETIONS_PROMPT: &str =
r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#;
const COMPLETIONS_PROMPT_LIST: &str =
r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#;
const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[
{"role":"user","content":[{"type":"input_text","text":"Summarise caf\u00e9 menus, na\u00efve \u2014 ok? \"quoted\"\n"}]},
{"role":"assistant","content":"Sure."}],"instructions":"be terse"}"#;
const EMBEDDINGS_TOKEN_IDS: &str =
r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#;
const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour",
"documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#;
/// Expected counts are pinned from
/// `litellm.proxy.spend_tracking.budget_reservation._count_input_tokens(body, "claude-sonnet-4-5")`.
#[rstest]
#[case::text_only(SIMPLE, 14)]
#[case::content_blocks_name_and_system(BLOCKS_AND_SYSTEM, 45)]
#[case::openai_tools_named_choice(TOOLS_OPENAI, 123)]
#[case::anthropic_tools_system_discount_choice_none(TOOLS_ANTHROPIC_SYSTEM, 53)]
#[case::completions_prompt(COMPLETIONS_PROMPT, 7)]
#[case::completions_prompt_list(COMPLETIONS_PROMPT_LIST, 4)]
#[case::responses_input_items(RESPONSES_INPUT, 62)]
#[case::embeddings_token_ids(EMBEDDINGS_TOKEN_IDS, 5)]
#[case::rerank_query_and_documents(RERANK, 41)]
fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expected: usize) {
let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses");
let count = counter().count_request(&request).expect("fixture counts");
assert_eq!(
count,
InputTokenCount {
model: Some("claude-sonnet-4-5".to_string()),
input_tokens: expected,
}
);
}
#[rstest]
#[case::null_messages_win_over_prompt(r#"{"model":"m","messages":null,"prompt":"ignored"}"#, 3)]
#[case::model_from_route(r#"{"prompt":"hi"}"#, 1)]
#[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)]
#[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)]
fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) {
let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses");
let count = counter().count_request(&request).expect("fixture counts");
assert_eq!(count.input_tokens, expected);
}
#[rstest]
#[case::not_json(b"not json" as &[u8])]
#[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)]
#[case::message_with_tool_calls(
br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"#
)]
#[case::dict_content(
br#"{"model":"m","messages":[{"role":"user","content":{"type":"text","text":"x"}}]}"#
)]
#[case::float_enum(
br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"number","enum":[1.5]}}}}]}"#
)]
#[case::anthropic_tool_choice_without_function(
br#"{"model":"m","messages":[],"tool_choice":{"type":"auto"}}"#
)]
fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) {
assert!(matches!(
CountableRequest::parse(body),
Err(Error::RequestParse(_))
));
}
#[rstest]
#[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])]
#[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)]
#[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)]
#[case::image_block(
br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"#
)]
#[case::tool_result_block(
br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"#
)]
#[case::array_without_items(
br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"#
)]
fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) {
let request = CountableRequest::parse(body).expect("shape parses");
assert!(matches!(
counter().count_request(&request),
Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems)
));
}
#[test]
fn tool_choice_and_system_discount_change_the_count() {
let counter = counter();
let count = |body: &str| {
counter
.count_request(&CountableRequest::parse(body.as_bytes()).expect("parses"))
.expect("counts")
.input_tokens
};
let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
assert_eq!(
count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#),
base + 1
);
assert_eq!(
count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"auto"}"#),
base
);
let with_tools = count(
r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}]}"#,
);
let with_tools_and_system = count(
r#"{"model":"m","messages":[{"role":"system","content":"hi"}],"tools":[{"name":"f"}]}"#,
);
assert_eq!(with_tools - with_tools_and_system, 4);
assert_eq!(
count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[]}"#),
base
);
}
#[test]
fn loading_a_bad_tokenizer_is_a_load_error() {
assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_))));
}

65
litellm/_redis.py
View file

@ -24,6 +24,7 @@ from redis.credentials import CredentialProvider
from litellm import get_secret, get_secret_str
from litellm._redis_credential_provider import (
AzureADCredentialProvider,
ElastiCacheIAMCredentialProvider,
GCPIAMCredentialProvider,
_generate_gcp_iam_access_token,
)
@ -38,6 +39,14 @@ from ._logging import verbose_logger
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
_AWS_IAM_KWARG_NAMES: Final = (
"aws_iam_auth",
"aws_iam_user_name",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
)
def _unwrapped_init_args(cls: type) -> frozenset[str]:
"""Every parameter on a single class's own ``__init__``, decorator-unwrapped.
@ -75,6 +84,7 @@ def _get_redis_kwargs():
"azure_client_id",
"azure_tenant_id",
"azure_client_secret",
*_AWS_IAM_KWARG_NAMES,
}
available_args: Final = {x for x in _unwrapped_init_args(redis.Redis) if x not in exclude_args} | include_args
@ -270,6 +280,42 @@ def _redis_kwargs_from_environment():
return return_dict
def _coerces_to_true(value: object | None) -> bool:
return _str_to_bool(value) if isinstance(value, str) else bool(value)
def _uses_tls(redis_kwargs: Mapping[str, object]) -> bool:
if redis_kwargs.get("startup_nodes") is not None:
return _coerces_to_true(redis_kwargs.get("ssl"))
url: Final = redis_kwargs.get("url")
if isinstance(url, str):
return urlsplit(url).scheme.lower() == "rediss"
return _coerces_to_true(redis_kwargs.get("ssl"))
def _build_elasticache_iam_provider(redis_kwargs: Mapping[str, object]) -> ElastiCacheIAMCredentialProvider:
user_name: Final = redis_kwargs.get("aws_iam_user_name")
cache_name: Final = redis_kwargs.get("aws_iam_cache_name")
region: Final = (
redis_kwargs.get("aws_iam_region") or get_secret_str("AWS_REGION") or get_secret_str("AWS_DEFAULT_REGION")
)
required_settings: Final = (
("aws_iam_user_name", user_name),
("aws_iam_cache_name", cache_name),
("aws_iam_region", region),
)
missing_settings: Final = tuple(name for name, value in required_settings if not value)
if missing_settings:
raise ValueError("AWS ElastiCache IAM Redis authentication requires: " + ", ".join(missing_settings))
return ElastiCacheIAMCredentialProvider(
user_name=str(user_name),
cache_name=str(cache_name),
region=str(region),
is_serverless=_coerces_to_true(redis_kwargs.get("aws_iam_serverless")),
)
def create_gcp_iam_redis_connect_func(
service_account: str,
ssl_ca_certs: str | None = None,
@ -540,6 +586,7 @@ def _get_redis_client_logic(**env_overrides):
_azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
_azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
_aws_iam_enabled: Final = _coerces_to_true(redis_kwargs.get("aws_iam_auth"))
if _azure_ad_enabled and _gcp_service_account is not None:
verbose_logger.warning(
@ -567,6 +614,22 @@ def _get_redis_client_logic(**env_overrides):
# credentials via inspection or logging.
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True
if _aws_iam_enabled and _gcp_service_account is not None:
verbose_logger.warning(
"Both GCP IAM (gcp_service_account) and AWS ElastiCache IAM (aws_iam_auth) are configured "
"for Redis. Using GCP IAM. Remove one to avoid misconfiguration."
)
elif _aws_iam_enabled and _azure_ad_enabled:
verbose_logger.warning(
"Both Azure AD (azure_redis_ad_token) and AWS ElastiCache IAM (aws_iam_auth) are configured "
"for Redis. Using Azure AD. Remove one to avoid misconfiguration."
)
elif _aws_iam_enabled:
if not _uses_tls(redis_kwargs):
raise ValueError("AWS ElastiCache IAM Redis authentication requires TLS")
verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.")
redis_kwargs["credential_provider"] = _build_elasticache_iam_provider(redis_kwargs)
redis_kwargs.pop("gcp_service_account", None)
redis_kwargs.pop("gcp_ssl_ca_certs", None)
@ -575,6 +638,8 @@ def _get_redis_client_logic(**env_overrides):
redis_kwargs.pop("azure_client_id", None)
redis_kwargs.pop("azure_tenant_id", None)
redis_kwargs.pop("azure_client_secret", None)
for aws_iam_key in _AWS_IAM_KWARG_NAMES:
redis_kwargs.pop(aws_iam_key, None)
if redis_kwargs.get("credential_provider") is not None:
redis_kwargs.pop("redis_connect_func", None)

85
litellm/_redis_credential_provider.py
View file

@ -1,10 +1,17 @@
from __future__ import annotations
import asyncio
import threading
import time
from typing import Final, Protocol
from collections.abc import Callable
from typing import TYPE_CHECKING, Final, Protocol
from urllib.parse import urlencode
from redis.credentials import CredentialProvider
if TYPE_CHECKING:
from botocore.credentials import Credentials
# Azure AD scope for Redis Cache for Azure.
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
@ -117,6 +124,82 @@ class GCPIAMCredentialProvider(CredentialProvider):
return (token,)
_ELASTICACHE_SERVICE_NAME: Final = "elasticache"
_ELASTICACHE_TOKEN_TTL_SECONDS: Final = 900
_ELASTICACHE_SERVERLESS_RESOURCE_TYPE: Final = "ServerlessCache"
class ElastiCacheIAMCredentialProvider(CredentialProvider):
def __init__(
self,
user_name: str,
cache_name: str,
region: str,
is_serverless: bool = False,
credentials_resolver: Callable[[], Credentials | None] | None = None,
token_lifetime_seconds: int = _ELASTICACHE_TOKEN_TTL_SECONDS,
) -> None:
self._user_name = user_name
self._cache_name = cache_name.lower()
self._region = region
self._is_serverless = is_serverless
self._credentials_resolver = credentials_resolver or self._resolve_credentials
self._credentials: Credentials | None = None
self._token_lifetime_seconds = token_lifetime_seconds
@staticmethod
def _resolve_credentials() -> Credentials | None:
try:
import botocore.session
except ImportError as e:
raise ImportError(
"botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3"
) from e
return botocore.session.get_session().get_credentials()
def _get_credentials(self) -> tuple[str, str]:
credentials: Final = self._credentials if self._credentials is not None else self._credentials_resolver()
if credentials is None:
raise RuntimeError("Unable to resolve AWS credentials for ElastiCache IAM Redis authentication")
self._credentials = credentials
frozen_credentials: Final = credentials.get_frozen_credentials()
try:
from botocore.auth import SigV4QueryAuth
from botocore.awsrequest import AWSRequest
except ImportError as e:
raise ImportError(
"botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3"
) from e
query: Final = urlencode(
(
("Action", "connect"),
("User", self._user_name),
*((("ResourceType", _ELASTICACHE_SERVERLESS_RESOURCE_TYPE),) if self._is_serverless else ()),
)
)
request: Final = AWSRequest(method="GET", url=f"https://{self._cache_name}/?{query}")
SigV4QueryAuth(
frozen_credentials,
_ELASTICACHE_SERVICE_NAME,
self._region,
expires=self._token_lifetime_seconds,
).add_auth(request)
signed_url: Final = request.url
if signed_url is None:
raise RuntimeError("Unable to generate AWS ElastiCache IAM credentials")
return self._user_name, signed_url.removeprefix("https://")
def get_credentials(self) -> tuple[str, str]:
return self._get_credentials()
async def get_credentials_async(self) -> tuple[str, str]:
return await asyncio.to_thread(self._get_credentials)
class AzureADCredentialProvider(CredentialProvider):
"""
redis.credentials.CredentialProvider implementation that supplies Azure AD

9
litellm/caching/caching.py
View file

@ -10,6 +10,7 @@
import ast
import hashlib
import json
import logging
import time
import traceback
from collections.abc import Mapping
@ -32,7 +33,7 @@ from .dual_cache import DualCache # noqa: F401
from .gcs_cache import GCSCache
from .in_memory_cache import InMemoryCache
from .qdrant_semantic_cache import QdrantSemanticCache
from .redis_cache import RedisCache
from .redis_cache import RedisCache, log_redis_failure
from .redis_cluster_cache import RedisClusterCache
from .redis_semantic_cache import RedisSemanticCache
from .s3_cache import S3Cache
@ -678,7 +679,7 @@ class Cache:
cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs)
self.cache.set_cache(cache_key, cached_data, **kwargs)
except Exception as e:
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs):
"""
@ -697,7 +698,7 @@ class Cache:
else:
await self.cache.async_set_cache(cache_key, cached_data, **kwargs)
except Exception as e:
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
def _convert_to_cached_embedding(
self,
@ -876,7 +877,7 @@ class Cache:
else:
await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs)
except Exception as e:
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
def should_use_cache(self, **kwargs):
"""

4
litellm/caching/caching_handler.py
View file

@ -1217,7 +1217,9 @@ class LLMCachingHandler:
}
if litellm.cache is not None:
litellm_params["preset_cache_key"] = litellm.cache._get_preset_cache_key_from_kwargs(**kwargs)
litellm_params["preset_cache_key"] = (
self.preset_cache_key or litellm.cache._get_preset_cache_key_from_kwargs(**kwargs)
)
else:
litellm_params["preset_cache_key"] = None

62
litellm/caching/dual_cache.py
View file

@ -8,8 +8,8 @@ Has 4 primary methods:
- async_get_cache
"""
import logging
import time
import traceback
from collections.abc import Sequence
from threading import Lock
from typing import TYPE_CHECKING, Any, Final
@ -23,7 +23,7 @@ from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
from .base_cache import BaseCache
from .in_memory_cache import InMemoryCache
from .redis_cache import RedisCache
from .redis_cache import RedisCache, RedisCircuitBreakerOpenError, log_redis_failure
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -177,8 +177,10 @@ class DualCache(BaseCache):
print_verbose(f"get cache: cache result: {result}")
return result
except Exception:
verbose_logger.error(traceback.format_exc())
except Exception as e:
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in get_cache", e, with_traceback=True
)
def batch_get_cache(
self,
@ -204,9 +206,12 @@ class DualCache(BaseCache):
redis_result: Final = self.redis_cache.batch_get_cache(
key_list=sublist_keys, parent_otel_span=parent_otel_span
)
except Exception:
except Exception as e:
# Do not throttle subsequent callers if the Redis read fails.
self._rollback_redis_batch_key_reservations(previous_access_times)
if isinstance(e, RedisCircuitBreakerOpenError):
verbose_logger.debug("LiteLLM Cache: batch_get_cache served from memory only: %s", e)
return result
raise
if self.in_memory_cache is not None:
@ -217,8 +222,10 @@ class DualCache(BaseCache):
return list( # mutable-ok: public list contract
redis_result.get(key) if value is None else value for key, value in zip(keys, result)
)
except Exception:
verbose_logger.error(traceback.format_exc())
except Exception as e:
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in batch_get_cache", e, with_traceback=True
)
async def async_get_cache(
self,
@ -250,8 +257,10 @@ class DualCache(BaseCache):
print_verbose(f"get cache: cache result: {result}")
return result
except Exception:
verbose_logger.error(traceback.format_exc())
except Exception as e:
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async_get_cache", e, with_traceback=True
)
def _reserve_redis_batch_keys(
self,
@ -319,9 +328,12 @@ class DualCache(BaseCache):
redis_result: Final = await self.redis_cache.async_batch_get_cache(
sublist_keys, parent_otel_span=parent_otel_span
)
except Exception:
except Exception as e:
# Do not throttle subsequent callers if the Redis read fails.
self._rollback_redis_batch_key_reservations(previous_access_times)
if isinstance(e, RedisCircuitBreakerOpenError):
verbose_logger.debug("LiteLLM Cache: async_batch_get_cache served from memory only: %s", e)
return result
raise
# Short-circuit if redis_result is None or contains only None values
@ -339,8 +351,14 @@ class DualCache(BaseCache):
await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs))
return result
except Exception:
verbose_logger.error(traceback.format_exc())
except Exception as e:
log_redis_failure(
verbose_logger,
logging.ERROR,
"LiteLLM Cache: exception in async_batch_get_cache",
e,
with_traceback=True,
)
async def async_set_cache(self, key, value, local_only: bool = False, **kwargs):
print_verbose(f"async set cache: cache key: {key}; local_only: {local_only}; value: {value}")
@ -353,7 +371,9 @@ class DualCache(BaseCache):
if self.redis_cache is not None and local_only is False:
await self.redis_cache.async_set_cache(key, value, **kwargs)
except Exception as e:
verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e)
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async add_cache", e, with_traceback=True
)
# async_batch_set_cache
async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs):
@ -372,7 +392,9 @@ class DualCache(BaseCache):
cache_list=cache_list, ttl=kwargs.pop("ttl", None), **kwargs
)
except Exception as e:
verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e)
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async add_cache", e, with_traceback=True
)
async def async_increment_cache(
self,
@ -410,8 +432,10 @@ class DualCache(BaseCache):
return result
except Exception as e:
verbose_logger.warning(
"Redis async_increment_cache failed, falling back to in-memory result: %s",
log_redis_failure(
verbose_logger,
logging.WARNING,
"Redis async_increment_cache failed, falling back to in-memory result",
e,
)
return result
@ -439,8 +463,10 @@ class DualCache(BaseCache):
return result
except Exception as e:
verbose_logger.warning(
"Redis async_increment_cache_pipeline failed, falling back to in-memory result: %s",
log_redis_failure(
verbose_logger,
logging.WARNING,
"Redis async_increment_cache_pipeline failed, falling back to in-memory result",
e,
)
return result

73
litellm/caching/redis_cache.py
View file

@ -14,9 +14,11 @@ import functools
import hashlib
import inspect
import json
import logging
import time
from collections.abc import Awaitable, Callable, Sequence
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast
@ -195,8 +197,14 @@ class RedisCircuitBreaker:
self._timeout_streak_started_at: float | None = None
self._opened_at: float | None = None
self._state = self.CLOSED
self._generation = 0
_breaker_metrics().record_state_change(None, self._state)
@property
def generation(self) -> int:
"""Counts state transitions, so a call can tell whether the breaker moved while it ran."""
return self._generation
def is_open(self) -> bool:
"""Returns True if Redis calls should be skipped."""
if not self.enabled:
@ -249,7 +257,7 @@ class RedisCircuitBreaker:
self._set_state(self.OPEN)
def record_success(self) -> None:
if not self.enabled:
if not self.enabled or self._state == self.OPEN:
return
if self._state == self.HALF_OPEN:
verbose_logger.info("Redis circuit breaker CLOSED — Redis recovered")
@ -265,6 +273,7 @@ class RedisCircuitBreaker:
_breaker_metrics().record_transition(state)
_breaker_metrics().record_state_change(self._state, state)
self._state = state
self._generation += 1
_RedisCallResult = TypeVar("_RedisCallResult")
@ -391,21 +400,46 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep
_swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> int:
"""Reject the call if the breaker is open, else return the swallowed-failure count to compare against."""
class RedisCircuitBreakerOpenError(Exception):
pass
def log_redis_failure(
logger: logging.Logger, level: int, message: str, exc: BaseException, with_traceback: bool = False
) -> None:
if isinstance(exc, RedisCircuitBreakerOpenError):
logger.debug("%s: %s", message, exc)
return
logger.log(level, "%s: %s", message, exc, exc_info=exc if with_traceback else None)
@dataclass(frozen=True, slots=True)
class _BreakerAdmission:
swallowed_before: int
generation: int
def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> _BreakerAdmission:
"""Reject the call if the breaker is open, else record what its success may later prove."""
if breaker.is_open():
raise Exception(f"Redis circuit breaker is open — skipping {name}")
return _swallowed_redis_failures.get()
raise RedisCircuitBreakerOpenError(f"Redis circuit breaker is open — skipping {name}")
return _BreakerAdmission(swallowed_before=_swallowed_redis_failures.get(), generation=breaker.generation)
def _exit_circuit_breaker(breaker: RedisCircuitBreaker, swallowed_before: int) -> None:
"""Record success only when nothing failed while the call ran.
def _exit_circuit_breaker(breaker: RedisCircuitBreaker, admission: _BreakerAdmission) -> None:
"""Record success only when nothing failed while the call ran and the breaker has not moved since.
Several Redis methods catch their own connection errors and return a default, so a
method that returned is not on its own proof of a healthy Redis.
method that returned is not on its own proof of a healthy Redis. A success also vouches
only for the breaker state that admitted the call: a call admitted before the breaker
opened, or a probe admitted before a later failure reopened it, finishes knowing nothing
about whether Redis has recovered since, so only the current probe may close the breaker.
"""
if _swallowed_redis_failures.get() == swallowed_before:
breaker.record_success()
if _swallowed_redis_failures.get() != admission.swallowed_before:
return
if breaker.generation != admission.generation:
return
breaker.record_success()
async def _run_under_circuit_breaker(
@ -418,14 +452,14 @@ async def _run_under_circuit_breaker(
Shared by the method decorator and the Lua script executor so both feed the same
health signal.
"""
swallowed_before: Final = _enter_circuit_breaker(breaker, name)
admission: Final = _enter_circuit_breaker(breaker, name)
try:
result: Final = await call()
except Exception as e:
if _is_redis_health_failure(e):
breaker.record_failure(is_timeout=_is_redis_timeout_failure(e))
raise
_exit_circuit_breaker(breaker, swallowed_before)
_exit_circuit_breaker(breaker, admission)
return result
@ -435,14 +469,14 @@ def _run_under_circuit_breaker_sync(
call: Callable[[], _RedisCallResult],
) -> _RedisCallResult:
"""Run one blocking Redis call under a circuit breaker, feeding the same health signal as the async path."""
swallowed_before: Final = _enter_circuit_breaker(breaker, name)
admission: Final = _enter_circuit_breaker(breaker, name)
try:
result: Final = call()
except Exception as e:
if _is_redis_health_failure(e):
breaker.record_failure()
breaker.record_failure(is_timeout=_is_redis_timeout_failure(e))
raise
_exit_circuit_breaker(breaker, swallowed_before)
_exit_circuit_breaker(breaker, admission)
return result
@ -1323,6 +1357,7 @@ class RedisCache(BaseCache):
except Exception:
return ast.literal_eval(decoded)
@_redis_circuit_breaker_guard_sync
def get_cache(self, key, parent_otel_span: Span | None = None, **kwargs):
try:
key = self.check_and_fix_namespace(key=key)
@ -1342,8 +1377,8 @@ class RedisCache(BaseCache):
print_verbose(f"Got Redis Cache: key: {key}, cached_response {cached_response}")
return self._get_cache_logic(cached_response=cached_response)
except Exception as e:
# NON blocking - notify users Redis is throwing an exception
verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: ", e)
verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: %s", e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
"""
@ -1380,12 +1415,12 @@ class RedisCache(BaseCache):
key_value_dict = {}
_key_list: Final = [key for key in key_list if key is not None]
start_time: Final = time.time()
admission: Final = _enter_circuit_breaker(self._circuit_breaker, "batch_get_cache")
try:
swallowed_before: Final = _enter_circuit_breaker(self._circuit_breaker, "batch_get_cache")
_keys: Final = [self.check_and_fix_namespace(key=cache_key or "") for cache_key in _key_list]
results: Final = self._run_redis_mget_operation(keys=_keys)
_exit_circuit_breaker(self._circuit_breaker, swallowed_before)
_exit_circuit_breaker(self._circuit_breaker, admission)
end_time: Final = time.time()
_duration: Final = end_time - start_time
self.service_logger_obj.service_success_hook(

3
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -1201,6 +1201,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
# Cast to Any to match the expected union type for tools list items
tools.append(cast(Any, web_search_tool))
def transform_response_format_to_text_format(self, response_format: object) -> "ResponseText | None":
return self._transform_response_format_to_text_format(response_format)
def _transform_response_format_to_text_format(self, response_format: object) -> "ResponseText | None":
"""
Transform Chat Completion response_format parameter to Responses API text.format parameter.

4
litellm/constants.py
View file

@ -1654,6 +1654,7 @@ CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data"
MAVVRIK_FOCUS_EXPORT_JOB_NAME: Final = "mavvrik_focus_export_usage_data"
CLOUDZERO_MAX_FETCHED_DATA_RECORDS: Final = int(os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000))
SPEND_LOG_CLEANUP_JOB_NAME: Final = "spend_log_cleanup"
BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME: Final = "background_health_check_db_save"
KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job"
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job"
WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job"
@ -1999,6 +2000,9 @@ NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset(
}
)
UNKNOWN_MODEL_SPEND_LOG_MODEL: Final[str] = "unknown-model"
MAX_SPEND_LOG_MODEL_NAME_LENGTH: Final[int] = 256
# PTU reservation rollup writes rows to LiteLLM_DailyTeamSpend with this
# sentinel api_key so PTU flat cost stays distinguishable from real per-request
# spend under the table's composite unique constraint.

41
litellm/integrations/custom_logger.py
View file

@ -3,12 +3,13 @@
import re
import traceback
from collections.abc import AsyncGenerator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
from pydantic import BaseModel
from litellm._logging import verbose_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, EMPTY_MAPPING
from litellm.types.integrations.argilla import ArgillaItem
from litellm.types.integrations.custom_logger import AgenticLoopPlan
from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest
@ -897,10 +898,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
This is useful for logging payloads that contain sensitive information.
"""
from copy import copy
import litellm
from litellm import Choices, Message, ModelResponse
from litellm.litellm_core_utils.classifier_logging import CLASSIFIER_AUDIT_FIELDS, without_classifier_audit
turn_off_message_logging: Final[bool] = getattr(self, "turn_off_message_logging", False)
excluded_fields: Final[list[str] | None] = getattr(litellm, "standard_logging_payload_excluded_fields", None)
@ -909,30 +909,25 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
if turn_off_message_logging is False and not excluded_fields:
return model_call_details
# Only make a shallow copy of the top-level dict to avoid deepcopy issues
# with complex objects like AuthenticationError that may be present
model_call_details_copy: Final = copy(model_call_details)
standard_logging_object: Final = model_call_details.get("standard_logging_object")
if standard_logging_object is None:
return model_call_details_copy
return model_call_details.copy()
# Make a copy of just the standard_logging_object to avoid modifying the original
standard_logging_object_copy: Final = copy(standard_logging_object)
# Handle excluded fields - remove them entirely from the payload
if excluded_fields:
for field in excluded_fields:
if field in standard_logging_object_copy:
del standard_logging_object_copy[field]
standard_logging_object_copy: Final = {
key: value
for key, value in standard_logging_object.items()
if key not in (excluded_fields or ()) and not (turn_off_message_logging and key in CLASSIFIER_AUDIT_FIELDS)
}
# Handle turn_off_message_logging - redact messages and responses (if not already excluded)
if turn_off_message_logging:
redacted_str: Final = "redacted-by-litellm"
if "messages" not in (excluded_fields or []) and standard_logging_object_copy.get("messages") is not None:
if "messages" not in (excluded_fields or ()) and standard_logging_object_copy.get("messages") is not None:
standard_logging_object_copy["messages"] = [Message(content=redacted_str).model_dump()]
if "response" not in (excluded_fields or []) and standard_logging_object_copy.get("response") is not None:
if "response" not in (excluded_fields or ()) and standard_logging_object_copy.get("response") is not None:
response: Final = standard_logging_object_copy["response"]
# Check if this is a ResponsesAPIResponse (has "output" field)
if isinstance(response, dict) and "output" in response:
@ -956,8 +951,18 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
model_response_dict: Final = model_response.model_dump()
standard_logging_object_copy["response"] = model_response_dict
model_call_details_copy["standard_logging_object"] = standard_logging_object_copy
return model_call_details_copy
params: Final = model_call_details.get("litellm_params")
request: Final = params.get("proxy_server_request") if isinstance(params, dict) else None
redacted_params: Final = (
MappingProxyType({"litellm_params": {**params, "proxy_server_request": without_classifier_audit(request)}})
if turn_off_message_logging and isinstance(params, dict) and isinstance(request, dict)
else EMPTY_MAPPING
)
return {
**model_call_details,
**redacted_params,
"standard_logging_object": standard_logging_object_copy,
}
async def get_proxy_server_request_from_cold_storage_with_object_key(
self,

39
litellm/integrations/newrelic/newrelic_metrics.py
View file

@ -5,8 +5,9 @@ NR Reference API: https://docs.newrelic.com/docs/data-apis/ingest-apis/metric-ap
`async_log_success_event` / `async_log_failure_event` queue one record per request;
at flush the queue is aggregated by (team, model group, model, provider, status)
into count/summary metrics. `interval.ms` is the real window between flushes,
computed at flush time.
into count/summary metrics, plus one max/remaining budget gauge pair per team
taken from the team's latest record. `interval.ms` is the real window between
flushes, computed at flush time.
Team-scoped by construction: the ingest key is injected explicitly and there is
deliberately no environment-variable fallback, so a team's metrics are never sent
@ -47,11 +48,14 @@ from litellm.types.integrations.newrelic import (
NEWRELIC_METRIC_PROMPT_TOKENS,
NEWRELIC_METRIC_REQUEST_DURATION_MS,
NEWRELIC_METRIC_REQUESTS,
NEWRELIC_METRIC_TEAM_MAX_BUDGET,
NEWRELIC_METRIC_TEAM_REMAINING_BUDGET,
NEWRELIC_METRIC_TOTAL_TOKENS,
NEWRELIC_METRICS_MAX_BATCH_SIZE,
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
NewRelicCountMetric,
NewRelicGaugeMetric,
NewRelicMetric,
NewRelicMetricCommon,
NewRelicMetricEnvelope,
@ -98,6 +102,8 @@ def _metric_record_from_payload(standard_logging_object: StandardLoggingPayload)
completion_tokens=int(standard_logging_object.get("completion_tokens") or 0),
total_tokens=int(standard_logging_object.get("total_tokens") or 0),
duration_ms=float(standard_logging_object.get("response_time") or 0.0) * 1000.0,
team_max_budget=metadata.get("user_api_key_team_max_budget") if metadata else None,
team_spend=metadata.get("user_api_key_team_spend") if metadata else None,
)
@ -140,6 +146,33 @@ def _bucket_metrics(bucket_records: tuple[NewRelicMetricRecord, ...]) -> tuple[N
return (*count_metrics, summary_metric)
def _team_budget_gauges(record: NewRelicMetricRecord) -> tuple[NewRelicMetric, ...]:
team_max_budget: Final = record.team_max_budget
if team_max_budget is None:
return ()
attributes: Final[Mapping[str, str]] = { # mutable-ok: JSON leaf; safe_dumps stringifies MappingProxyType
key: value[:NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN]
for key, value in (("team_id", record.team_id), ("team_alias", record.team_alias))
if value
}
remaining_budget: Final = team_max_budget - (record.team_spend or 0.0) - record.response_cost
return (
NewRelicGaugeMetric(
name=NEWRELIC_METRIC_TEAM_MAX_BUDGET, type="gauge", value=team_max_budget, attributes=attributes
),
NewRelicGaugeMetric(
name=NEWRELIC_METRIC_TEAM_REMAINING_BUDGET, type="gauge", value=remaining_budget, attributes=attributes
),
)
def _team_budget_metrics(records: tuple[NewRelicMetricRecord, ...]) -> tuple[NewRelicMetric, ...]:
latest_by_team: Final[Mapping[str, NewRelicMetricRecord]] = MappingProxyType(
{record.team_id: record for record in records if record.team_id}
)
return tuple(gauge for record in latest_by_team.values() for gauge in _team_budget_gauges(record))
def build_metric_payload(
records: tuple[NewRelicMetricRecord, ...],
*,
@ -158,7 +191,7 @@ def build_metric_payload(
"timestamp": int(window_start * 1000),
"interval.ms": interval_ms,
}
return (NewRelicMetricEnvelope(common=common, metrics=metrics),)
return (NewRelicMetricEnvelope(common=common, metrics=(*metrics, *_team_budget_metrics(records))),)
class NewRelicMetricsLogger(CustomBatchLogger):

67
litellm/litellm_core_utils/classifier_logging.py Normal file
View file

@ -0,0 +1,67 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from pydantic import JsonValue, TypeAdapter, ValidationError
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN, ClassifierAudit
CLASSIFIER_AUDIT_FIELDS: Final = ("classifier_input", "originating_request_masked")
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def classifier_input_snapshot(value: object, *, openai_sdk: bool = False) -> Mapping[str, JsonValue] | None:
try:
if openai_sdk and isinstance(value, Mapping):
body: Final = MappingProxyType(
{key: item for key, item in value.items() if key not in ("extra_headers", "extra_query", "extra_body")}
)
extra_body: Final = value.get("extra_body")
return _JSON_OBJECT.validate_python(
MappingProxyType({**body, **extra_body}) if isinstance(extra_body, Mapping) else body
)
return (
_JSON_OBJECT.validate_json(value)
if isinstance(value, (str, bytes))
else _JSON_OBJECT.validate_python(value)
)
except ValidationError:
return None
def is_classifier_call(call_type: str, params: Mapping[str, object]) -> bool:
return call_type in ("completion", "acompletion", "responses", "aresponses") and any(
isinstance(metadata := params.get(key), Mapping)
and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == AUTOROUTER_CLASSIFIER_CALL_ORIGIN
for key in ("metadata", "litellm_metadata")
)
def masked_originating_request(request_kwargs: Mapping[str, object] | None) -> Mapping[str, JsonValue] | None:
request: Final = request_kwargs.get("proxy_server_request") if request_kwargs is not None else None
body: Final = request.get("body") if isinstance(request, Mapping) else None
if not isinstance(body, Mapping):
return None
serializable: Final = classifier_input_snapshot(safe_dumps(body))
return classifier_input_snapshot(redact_credentials_in_payload(serializable)) if serializable is not None else None
def classifier_audit_fields(payload: Mapping[str, object]) -> ClassifierAudit:
classifier_input: Final = classifier_input_snapshot(payload.get("classifier_input"))
originating_request: Final = classifier_input_snapshot(payload.get("originating_request_masked"))
if classifier_input is None:
return (
ClassifierAudit(originating_request_masked=originating_request)
if originating_request is not None
else ClassifierAudit()
)
if originating_request is None:
return ClassifierAudit(classifier_input=classifier_input)
return ClassifierAudit(classifier_input=classifier_input, originating_request_masked=originating_request)
def without_classifier_audit(payload: Mapping[str, object]) -> dict[str, object]:
return {key: value for key, value in payload.items() if key not in CLASSIFIER_AUDIT_FIELDS}

8
litellm/litellm_core_utils/core_helpers.py
View file

@ -2,6 +2,7 @@
## Helper utilities
import copy
import logging
import re
from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Any, Final, Literal
@ -21,6 +22,13 @@ else:
Span = Any
_CODEX_CLIENT_PREFIX_RE: Final = re.compile(r"^codex[-_ /]", re.IGNORECASE)
def is_codex_user_agent(user_agent: str) -> bool:
return bool(_CODEX_CLIENT_PREFIX_RE.match(user_agent))
def safe_divide_seconds(seconds: float, denominator: float, default: float | None = None) -> float | None:
"""
Safely divide seconds by denominator, handling zero division.

69
litellm/litellm_core_utils/litellm_logging.py
View file

@ -17,7 +17,7 @@ from types import MappingProxyType, TracebackType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast
from httpx import Response
from pydantic import BaseModel
from pydantic import BaseModel, JsonValue
import litellm
from litellm import (
@ -42,6 +42,7 @@ from litellm.caching.caching_handler import LLMCachingHandler
from litellm.constants import (
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
EMPTY_MAPPING,
PROVIDER_REQUEST_ID_HEADERS,
SENTRY_DENYLIST,
SENTRY_PII_DENYLIST,
@ -64,6 +65,11 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.deepeval.deepeval import DeepEvalLogger
from litellm.integrations.mlflow import MlflowLogger
from litellm.integrations.sqs import SQSLogger
from litellm.litellm_core_utils.classifier_logging import (
classifier_audit_fields,
classifier_input_snapshot,
is_classifier_call,
)
from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.internal_call_metadata import (
@ -89,6 +95,7 @@ from litellm.litellm_core_utils.redact_messages import (
redact_message_input_output_from_custom_logger,
redact_message_input_output_from_logging,
redact_streaming_responses_for_custom_logger,
should_redact_message_logging,
)
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.llms.base_llm.search.transformation import SearchResponse
@ -288,6 +295,9 @@ def _get_provider_request_id(original_exception: Exception) -> str | None:
# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
_MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS
_UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset(
("user_api_key_auth", "user_api_key_budget_reservation")
)
sentry_sdk_instance = None
capture_exception = None
@ -473,6 +483,7 @@ class Logging(LiteLLMLoggingBaseClass):
stream_options = None
litellm_request_debug: bool = False
streamed_anthropic_message_id: str | None = None
classifier_input: Mapping[str, JsonValue] | None = None
def __init__(
self,
@ -1211,6 +1222,14 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["api_key"] = api_key
self.model_call_details["additional_args"] = additional_args
self.model_call_details["log_event_type"] = "pre_api_call"
if is_classifier_call(self.call_type, self.model_call_details.get("litellm_params") or EMPTY_MAPPING):
self.classifier_input = (
None
if should_redact_message_logging(self.model_call_details)
else classifier_input_snapshot(
additional_args.get("complete_input_dict"), openai_sdk=additional_args.get("openai_sdk") is True
)
)
if model: # if model name was changes pre-call, overwrite the initial model call name with the new one
self.model_call_details["model"] = model
self.model_call_details["litellm_params"]["api_base"] = self._get_masked_api_base(
@ -5370,23 +5389,23 @@ class StandardLoggingPayloadSetup:
Returns:
dict: Merged metadata with user API key fields taking precedence
"""
merged_metadata: Final[dict] = {}
# Start with metadata (user API key fields) - but skip non-serializable objects
if litellm_params.get("metadata") and isinstance(litellm_params.get("metadata"), dict):
for key, value in litellm_params["metadata"].items():
# Skip non-serializable objects like UserAPIKeyAuth
if key in {"user_api_key_auth", "user_api_key_budget_reservation"}:
continue
merged_metadata[key] = value
# Then merge litellm_metadata (model-related fields) - this will NOT overwrite existing keys
if litellm_params.get("litellm_metadata") and isinstance(litellm_params.get("litellm_metadata"), dict):
for key, value in litellm_params["litellm_metadata"].items():
if key not in merged_metadata: # Don't overwrite existing keys from metadata
merged_metadata[key] = value
return merged_metadata
metadata: Final = litellm_params.get("metadata")
litellm_metadata: Final = litellm_params.get("litellm_metadata")
user_metadata: Final = MappingProxyType(
{
key: value
for key, value in (metadata.copy().items() if isinstance(metadata, dict) else ())
if key not in _UNSERIALIZABLE_METADATA_KEYS
}
)
model_metadata: Final = MappingProxyType(
{
key: value
for key, value in (litellm_metadata.copy().items() if isinstance(litellm_metadata, dict) else ())
if key not in user_metadata
}
)
return {**user_metadata, **model_metadata} # mutable-ok: function contract returns a plain dict
@staticmethod
def get_standard_logging_metadata(
@ -5644,7 +5663,7 @@ class StandardLoggingPayloadSetup:
additional_logging_headers[key] = additiona_headers[_key]
# Preserve all remaining headers verbatim (e.g. llm_provider-x-request-id)
for k, v in additiona_headers.items():
for k, v in additiona_headers.copy().items():
if k.lower() not in typed_keys:
additional_logging_headers[k] = v
@ -6293,6 +6312,18 @@ def get_standard_logging_object_payload(
)
payload: Final[StandardLoggingPayload] = StandardLoggingPayload(
**(
classifier_audit_fields(
MappingProxyType(
{
"classifier_input": logging_obj.classifier_input,
"originating_request_masked": proxy_server_request.get("originating_request_masked"),
}
)
)
if is_classifier_call(call_type or "", litellm_params) and not should_redact_message_logging(kwargs)
else EMPTY_MAPPING
),
id=str(id),
litellm_call_id=kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"),
trace_id=StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id(

22
litellm/litellm_core_utils/redact_messages.py
View file

@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm.constants import REDACTED_BY_LITELLM
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.classifier_logging import without_classifier_audit
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
)
@ -170,17 +171,11 @@ def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[
payload, but the failure path does not, so a callback that batches both has to redact
the ones it is handed.
"""
redacted: Final = copy.deepcopy(dict(payload)) # mutable-ok: redacted in place below
_redact_standard_logging_object({"standard_logging_object": redacted}) # mutable-ok: the callee's shape
return redacted
return _redact_standard_logging_object(payload)
def _redact_standard_logging_object(model_call_details: dict):
"""Redact messages and response inside standard_logging_object if present."""
standard_logging_object: Final = model_call_details.get("standard_logging_object")
if standard_logging_object is None:
return
def _redact_standard_logging_object(payload: Mapping[str, object]) -> dict[str, object]:
standard_logging_object: Final = copy.deepcopy(without_classifier_audit(payload))
redacted_str: Final = REDACTED_BY_LITELLM
if standard_logging_object.get("messages") is not None:
@ -203,6 +198,7 @@ def _redact_standard_logging_object(model_call_details: dict):
else:
# For other formats (empty dict, None, etc.), use simple text format
standard_logging_object["response"] = {"text": redacted_str}
return standard_logging_object
def _redact_tool_calls_dict(message: Mapping[str, object]) -> None:
@ -254,10 +250,16 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
copy via redact_streaming_responses_for_custom_logger instead.
"""
# Redact model_call_details
params: Final = model_call_details.get("litellm_params")
request: Final = params.get("proxy_server_request") if isinstance(params, dict) else None
if isinstance(params, dict) and isinstance(request, Mapping):
model_call_details["litellm_params"] = {**params, "proxy_server_request": without_classifier_audit(request)}
model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}]
model_call_details["prompt"] = ""
model_call_details["input"] = ""
_redact_standard_logging_object(model_call_details)
standard_logging_object: Final = model_call_details.get("standard_logging_object")
if isinstance(standard_logging_object, Mapping):
model_call_details["standard_logging_object"] = _redact_standard_logging_object(standard_logging_object)
redact_vertex_ai_metadata_from_litellm_params(model_call_details)
# Redact streaming response

1
litellm/litellm_core_utils/sensitive_data_masker.py
View file

@ -15,6 +15,7 @@ _DEFAULT_SENSITIVE_PATTERNS: Final = frozenset(
"token",
"auth",
"authorization",
"cookie",
"credential",
# Plural form: Vertex uses ``vertex_credentials``; segment-exact
# matching otherwise misses it because "credential" != "credentials".

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