diff --git a/.circleci/config.yml b/.circleci/config.yml index 3d1e22eebd3..db8e7d49d71 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -439,7 +439,14 @@ jobs: auth: username: ${DOCKERHUB_USERNAME} password: ${DOCKERHUB_PASSWORD} + - image: cimg/postgres:16.0 + environment: + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + POSTGRES_DB: litellm_test working_directory: ~/project + environment: + DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test" steps: - checkout @@ -463,12 +470,14 @@ jobs: paths: - ./.venv key: v2-dependencies-{{ checksum "uv.lock" }}-{{ checksum ".circleci/config.yml" }} + - wait_for_service: + url: tcp://localhost:5432 + timeout: "60" - run: - name: Run prisma ./docker/entrypoint.sh + name: Seed DB schema via prisma db push command: | set +e - chmod +x docker/entrypoint.sh - ./docker/entrypoint.sh + uv run --no-sync litellm --skip_server_startup --use_prisma_db_push set -e - run: name: Generate Prisma Client @@ -1520,7 +1529,50 @@ jobs: command: | pwd ls - uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py + uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not v2_resolver" + + installing_litellm_on_python_v2_migration_resolver: + docker: + - image: cimg/python:3.11 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + - image: cimg/postgres:16.0 + environment: + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + POSTGRES_DB: litellm_test + working_directory: ~/project + environment: + DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test" + + steps: + - checkout + - setup_google_dns + - run: + name: Install Dependencies + command: | + curl -LsSf -o /tmp/uv-install.sh https://astral.sh/uv/0.10.9/install.sh + echo "7fc46e39cb97290b57169c0c813a17970585ac519139f19006453c99b5f2f45f /tmp/uv-install.sh" | sha256sum -c - + env UV_NO_MODIFY_PATH=1 sh /tmp/uv-install.sh + rm -f /tmp/uv-install.sh + echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$BASH_ENV" + export PATH="$HOME/.local/bin:$PATH" + if [ -f "$HOME/miniconda/etc/profile.d/conda.sh" ]; then + export PATH="$HOME/miniconda/bin:$PATH" + source "$HOME/miniconda/etc/profile.d/conda.sh" + conda activate myenv + fi + uv sync --frozen --all-groups --all-extras --python "$(which python)" + - setup_litellm_enterprise_pip + - wait_for_service: + url: tcp://localhost:5432 + timeout: "60" + - run: + name: Run v2 migration resolver proxy smoke test + command: | + uv run --no-sync python -m pytest -vv \ + tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_v2_resolver installing_litellm_on_python_3_13: docker: @@ -1554,7 +1606,7 @@ jobs: command: | pwd ls - uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py + uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not v2_resolver" helm_chart_testing: machine: image: ubuntu-2204:2023.10.1 # Use machine executor instead of docker @@ -3042,10 +3094,19 @@ jobs: - ui/litellm-dashboard/node_modules - run: name: Build UI from source + # Prior version used `cp -r out/ ../../litellm/proxy/_experimental/out/`. + # GNU cp (used on CircleCI's Ubuntu image) interprets that as "copy the + # source directory as a child of the destination" when the destination + # already exists — silently creating `_experimental/out/out/` instead of + # replacing the served bundle. The proxy continued serving whatever was + # checked into `_experimental/out/*`, so this job was effectively testing + # the pre-build bundle on every run. Replace-and-move guarantees the + # freshly built bundle is what the proxy actually serves. command: | cd ui/litellm-dashboard npm run build - cp -r out/ ../../litellm/proxy/_experimental/out/ + rm -rf ../../litellm/proxy/_experimental/out + mv out ../../litellm/proxy/_experimental/out # Restructure HTML so extensionless routes work (login.html -> login/index.html) find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html" @@ -3526,6 +3587,12 @@ workflows: only: - main - /litellm_.*/ + - installing_litellm_on_python_v2_migration_resolver: + filters: + branches: + only: + - main + - /litellm_.*/ - helm_chart_testing: requires: - build_docker_database_image diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 87e7e17feb7..a631a7c3005 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -31,8 +31,15 @@ jobs: test-path: "tests/proxy_unit_tests/test_auth_checks.py tests/proxy_unit_tests/test_user_api_key_auth.py" workers: 8 timeout: 20 + # test_proxy_utils.py is large (168+ parametrized tests) — run it on its + # own matrix so --dist=loadscope doesn't pin all of it to a single xdist + # worker and push the "remaining" group past the job timeout. + - test-group: proxy-utils + test-path: "tests/proxy_unit_tests/test_proxy_utils.py" + workers: 8 + timeout: 20 - test-group: remaining - test-path: "tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py" + test-path: "tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py --ignore=tests/proxy_unit_tests/test_proxy_utils.py" workers: 8 timeout: 30 uses: ./.github/workflows/_test-unit-services-base.yml diff --git a/ci_cd/run_migration.py b/ci_cd/run_migration.py index 29101bf9505..feec4046ee1 100644 --- a/ci_cd/run_migration.py +++ b/ci_cd/run_migration.py @@ -1,22 +1,231 @@ +import argparse import os -import subprocess -from pathlib import Path -from datetime import datetime -import testing.postgresql +import re import shutil +import subprocess +import sys +from datetime import datetime +from pathlib import Path + +import testing.postgresql -def create_migration(migration_name: str = None): +DESTRUCTIVE_PATTERN = re.compile(r"\bDROP\s+(COLUMN|TABLE|INDEX)\b", re.IGNORECASE) +DEFAULT_BASE_BRANCH = "litellm_internal_staging" + + +def _find_destructive_statements(sql: str) -> list: + """Return SQL lines containing DROP COLUMN, DROP TABLE, or DROP INDEX.""" + return [ + line.strip() for line in sql.splitlines() if DESTRUCTIVE_PATTERN.search(line) + ] + + +def _print_freshness_failure( + base_branch: str, reason: str, stderr_text: str = "" +) -> None: + """Loudly refuse to run when the freshness check can't be completed.""" + banner = "=" * 72 + out = sys.stderr + print(banner, file=out) + print(f" FRESHNESS CHECK FAILED — COULD NOT VERIFY origin/{base_branch}", file=out) + print(banner, file=out) + print("", file=out) + print(f"Reason: {reason}", file=out) + if stderr_text: + print("", file=out) + print("git stderr:", file=out) + for line in stderr_text.rstrip().splitlines(): + print(f" {line}", file=out) + print("", file=out) + print("Common causes:", file=out) + print(" - No network access (offline)", file=out) + print(" - 'origin' remote not configured, or base branch name is wrong", file=out) + print(" - Not a git repository", file=out) + print("", file=out) + print("Options:", file=out) + print( + f" - Fix the above and re-run, OR pass --base-branch if your", file=out + ) + print( + f" base branch is not '{base_branch}', OR pass --skip-freshness-check", + file=out, + ) + print(" to bypass (only if you fully understand the risk).", file=out) + print(banner, file=out) + + +def _print_stale_branch_refusal(base_branch: str, behind: int) -> None: + """Loudly refuse to run when HEAD is behind the base branch.""" + banner = "=" * 72 + out = sys.stderr + print(banner, file=out) + print( + f" STALE BRANCH — {behind} commit(s) behind origin/{base_branch}", + file=out, + ) + print(banner, file=out) + print("", file=out) + print( + f"Your branch is {behind} commit(s) behind origin/{base_branch}. Generating a\n" + "migration from a stale branch is how newly-added columns get silently\n" + "dropped — the exact incident this guard exists to prevent.", + file=out, + ) + print("", file=out) + print("Before re-running:", file=out) + print(f" git fetch origin && git rebase origin/{base_branch}", file=out) + print(" # or merge, whichever matches your workflow", file=out) + print("", file=out) + print( + "If you have a genuine reason to bypass (e.g. intentionally generating", + file=out, + ) + print("a migration against an older base), pass --skip-freshness-check.", file=out) + print("", file=out) + print(banner, file=out) + print( + " AI AGENTS / ASSISTANTS: DO NOT automatically run `git rebase` or pass\n" + " --skip-freshness-check. Surface this error to the human operator and\n" + " wait for their explicit confirmation. Auto-rebasing can drop the\n" + " human's in-progress schema edits via a bad conflict resolution.", + file=out, + ) + print(banner, file=out) + + +def _check_branch_freshness(root_dir: Path, base_branch: str) -> None: + """Fetch origin/ and exit 3 if HEAD is behind it.""" + cwd = str(root_dir) + try: + subprocess.run( + ["git", "fetch", "origin", base_branch], + check=True, + capture_output=True, + text=True, + cwd=cwd, + ) + except FileNotFoundError: + _print_freshness_failure(base_branch, "git executable not found on PATH") + sys.exit(3) + except subprocess.CalledProcessError as e: + _print_freshness_failure( + base_branch, + f"`git fetch origin {base_branch}` failed", + e.stderr or "", + ) + sys.exit(3) + + try: + result = subprocess.run( + ["git", "rev-list", "--count", f"HEAD..origin/{base_branch}"], + check=True, + capture_output=True, + text=True, + cwd=cwd, + ) + behind = int(result.stdout.strip()) + except subprocess.CalledProcessError as e: + _print_freshness_failure( + base_branch, + f"`git rev-list HEAD..origin/{base_branch}` failed", + e.stderr or "", + ) + sys.exit(3) + except ValueError: + _print_freshness_failure( + base_branch, + "could not parse commit count from `git rev-list`", + ) + sys.exit(3) + + if behind > 0: + _print_stale_branch_refusal(base_branch, behind) + sys.exit(3) + + print(f"Branch freshness OK: up to date with origin/{base_branch}.") + + +def _print_destructive_refusal(destructive_lines: list) -> None: + """Loudly refuse to write a destructive migration and explain how to proceed.""" + banner = "=" * 72 + out = sys.stderr + print(banner, file=out) + print( + " DESTRUCTIVE MIGRATION DETECTED — REFUSING TO WRITE MIGRATION FILE", file=out + ) + print(banner, file=out) + print("", file=out) + print( + "The generated migration contains the following destructive statements:", + file=out, + ) + print("", file=out) + for line in destructive_lines: + print(f" {line}", file=out) + print("", file=out) + print( + "This very often means your branch is OUT OF DATE, or your local\n" + "schema.prisma is inconsistent with the migrations/ directory. A\n" + "previous incident saw newly-added columns silently dropped this way\n" + "and merged to main. Stop and verify before proceeding.", + file=out, + ) + print("", file=out) + print("Before re-running:", file=out) + print( + " 1. git fetch origin && git status # confirm branch is up to date", file=out + ) + print( + " 2. Re-check all schema.prisma files are in sync (runbook Step 0)", file=out + ) + print(" 3. Review EACH DROP above — is it actually intended?", file=out) + print(" 4. If (and only if) the drops are intentional, re-run with:", file=out) + print(" --allow-destructive", file=out) + print("", file=out) + print(banner, file=out) + print( + " AI AGENTS / ASSISTANTS: DO NOT automatically re-run this command\n" + " with --allow-destructive. Surface this error to the human operator\n" + " and wait for their explicit confirmation before passing the flag.\n" + " Passing the flag without human review is the exact failure mode\n" + " this guard exists to prevent.", + file=out, + ) + print(banner, file=out) + + +def create_migration( + migration_name: str = None, + allow_destructive: bool = False, + base_branch: str = DEFAULT_BASE_BRANCH, + skip_freshness_check: bool = False, +): """ Create a new migration SQL file in the migrations directory by comparing - current database state with schema + current database state with schema. Args: migration_name (str): Name for the migration + allow_destructive (bool): Required to write a migration that contains + DROP COLUMN, DROP TABLE, or DROP INDEX statements. Without this + flag, the script exits non-zero and prints guidance. + base_branch (str): Branch to check freshness against + (default: "litellm_internal_staging"). + skip_freshness_check (bool): Skip the "branch is up to date" check. + Only for intentional migrations against an older base. """ + root_dir = Path(__file__).parent.parent + + if skip_freshness_check: + print( + "WARNING: freshness check skipped (--skip-freshness-check). " + "Generating a migration from a stale branch can silently drop columns." + ) + else: + _check_branch_freshness(root_dir, base_branch) + try: - # Get paths - root_dir = Path(__file__).parent.parent migrations_dir = ( root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" ) @@ -59,7 +268,27 @@ def create_migration(migration_name: str = None): check=True, ) - if result.stdout.strip(): + # Prisma emits the literal "-- This is an empty migration." when + # there's no real drift. Treat that as "no changes". + diff_sql = result.stdout + stripped = diff_sql.strip() + is_empty_diff = ( + not stripped or stripped == "-- This is an empty migration." + ) + + if not is_empty_diff: + destructive_lines = _find_destructive_statements(diff_sql) + if destructive_lines and not allow_destructive: + _print_destructive_refusal(destructive_lines) + sys.exit(2) + if destructive_lines and allow_destructive: + print( + "WARNING: writing destructive migration " + "(--allow-destructive passed). Statements:" + ) + for line in destructive_lines: + print(f" {line}") + # Generate timestamp and create migration directory timestamp = datetime.now().strftime("%Y%m%d%H%M%S") migration_name = migration_name or "unnamed_migration" @@ -68,7 +297,7 @@ def create_migration(migration_name: str = None): # Write the SQL to migration.sql migration_file = migration_dir / "migration.sql" - migration_file.write_text(result.stdout) + migration_file.write_text(diff_sql) print(f"Created migration in {migration_dir}") return True @@ -90,8 +319,48 @@ def create_migration(migration_name: str = None): if __name__ == "__main__": - # If running directly, can optionally pass migration name as argument - import sys - - migration_name = sys.argv[1] if len(sys.argv) > 1 else None - create_migration(migration_name) + parser = argparse.ArgumentParser( + description=( + "Generate a Prisma migration by diffing the temp DB " + "(existing migrations applied) against schema.prisma." + ) + ) + parser.add_argument( + "migration_name", + nargs="?", + default=None, + help="Name for the migration (used in the generated directory name).", + ) + parser.add_argument( + "--allow-destructive", + action="store_true", + help=( + "Required to write a migration that contains DROP COLUMN, " + "DROP TABLE, or DROP INDEX. Without this flag, destructive " + "diffs are refused." + ), + ) + parser.add_argument( + "--base-branch", + default=DEFAULT_BASE_BRANCH, + help=( + f"Branch to check freshness against (default: {DEFAULT_BASE_BRANCH}). " + "The script fetches origin/ and refuses to run if HEAD " + "is behind it." + ), + ) + parser.add_argument( + "--skip-freshness-check", + action="store_true", + help=( + "Bypass the 'branch is up to date' check. Only for intentional " + "migrations against an older base. Pairs poorly with automation." + ), + ) + args = parser.parse_args() + create_migration( + args.migration_name, + allow_destructive=args.allow_destructive, + base_branch=args.base_branch, + skip_freshness_check=args.skip_freshness_check, + ) diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 5451bff808d..e9161676092 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -15,29 +15,21 @@ COPY --from=uvbin /uv /usr/local/bin/uv COPY --from=uvbin /uvx /usr/local/bin/uvx RUN for i in 1 2 3; do \ - apk add --no-cache \ - python3 \ - python3-dev \ - clang \ - llvm \ - lld \ - gcc \ - linux-headers \ - build-base \ - bash \ - coreutils \ - curl \ - openssl \ - openssl-dev \ - nodejs \ - npm \ - libsndfile && break || sleep 5; \ + apk add --no-cache \ + python3 \ + python3-dev \ + gcc \ + bash \ + coreutils \ + curl \ + openssl \ + libsndfile \ + nodejs && break || sleep 5; \ done ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ UV_LINK_MODE=copy \ - NVM_DIR=/root/.nvm \ - PATH="/root/.nvm/versions/node/v20.20.2/bin:/app/.venv/bin:${PATH}" \ + PATH="/app/.venv/bin:${PATH}" \ LITELLM_NON_ROOT=true \ PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \ PRISMA_CLI_BINARY_TARGETS="debian-openssl-3.0.x" \ @@ -49,7 +41,8 @@ COPY enterprise/pyproject.toml enterprise/ COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/ # Install third-party dependencies (cached unless pyproject.toml/uv.lock change) -RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ +RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ + uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ --extra extra_proxy \ @@ -62,38 +55,12 @@ COPY . . # Set non-root flag for build time consistency ENV LITELLM_NON_ROOT=true -# Build Admin UI once and stage the static output for the runtime image. -# NOTE: .npmrc files (which may set ignore-scripts=true and min-release-age=3d) -# are temporarily renamed during npm install/ci so they don't block lifecycle -# scripts needed by the build. This is safe because npm ci installs from -# package-lock.json with pinned versions + integrity hashes. +# Stage the pre-built Admin UI from the checked-in Next.js static export. +# _experimental/out/ is regenerated as part of the release runbook. +# Restructure extensionless routes (foo.html -> foo/index.html) to match the layout +# proxy_server.py expects, and drop a readiness marker. RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \ - ([ -f /app/.npmrc ] && mv /app/.npmrc /app/.npmrc.bak || true) && \ - NVM_VERSION="v0.40.4" && \ - NVM_CHECKSUM="4b7412c49960c7d31e8df72da90c1fb5b8cccb419ac99537b737028d497aba4f" && \ - NODE_VERSION="v20.20.2" && \ - NVM_SCRIPT="/tmp/install-nvm.sh" && \ - curl -fsSL "https://raw.githubusercontent.com/nvm-sh/nvm/${NVM_VERSION}/install.sh" -o "$NVM_SCRIPT" && \ - echo "${NVM_CHECKSUM} ${NVM_SCRIPT}" | sha256sum -c - && \ - bash "$NVM_SCRIPT" && \ - export NVM_DIR="$HOME/.nvm" && \ - . "$NVM_DIR/nvm.sh" && \ - nvm install "${NODE_VERSION}" && \ - nvm use "${NODE_VERSION}" && \ - npm install -g npm@11.12.1 && \ - npm install -g node-gyp@12.2.0 && \ - ln -sf "$(npm root -g)/node-gyp" "$(npm root -g)/npm/node_modules/node-gyp" && \ - npm cache clean --force && \ - cd /app/ui/litellm-dashboard && \ - if [ -f "/app/enterprise/enterprise_ui/enterprise_colors.json" ]; then \ - cp /app/enterprise/enterprise_ui/enterprise_colors.json ./ui_colors.json; \ - fi && \ - ([ -f .npmrc ] && mv .npmrc .npmrc.bak || true) && \ - npm ci --no-audit --no-fund && \ - ([ -f .npmrc.bak ] && mv .npmrc.bak .npmrc || true) && \ - ([ -f /app/.npmrc.bak ] && mv /app/.npmrc.bak /app/.npmrc || true) && \ - npm run build && \ - cp -r /app/ui/litellm-dashboard/out/* /var/lib/litellm/ui/ && \ + cp -r /app/litellm/proxy/_experimental/out/. /var/lib/litellm/ui/ && \ cp /app/litellm/proxy/logo.jpg /var/lib/litellm/assets/logo.jpg && \ ( cd /var/lib/litellm/ui && \ for html_file in *.html; do \ @@ -103,10 +70,10 @@ RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \ mv "$html_file" "$folder_name/index.html"; \ fi; \ done && \ - touch .litellm_ui_ready ) && \ - cd /app/ui/litellm-dashboard && rm -rf ./out + touch .litellm_ui_ready ) -RUN if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \ +RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ + if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \ uv sync --frozen --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ @@ -123,10 +90,7 @@ RUN if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \ --python python3; \ fi -RUN mkdir -p /app/.cache/npm && \ - prisma generate --schema=./schema.prisma && \ - prisma --version && \ - prisma migrate diff --from-empty --to-schema-datamodel ./schema.prisma --script > /dev/null 2>&1 || true +RUN prisma generate --schema=./schema.prisma RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \ sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh @@ -137,33 +101,11 @@ WORKDIR /app USER root RUN for i in 1 2 3; do \ - apk upgrade --no-cache && break || sleep 5; \ + apk upgrade --no-cache && break || sleep 5; \ done && \ for i in 1 2 3; do \ - apk add --no-cache python3 bash openssl tzdata nodejs npm supervisor libsndfile && break || sleep 5; \ - done && \ - apk upgrade --no-cache nodejs && \ - npm install -g npm@11.12.1 tar@7.5.11 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \ - GLOBAL="$(npm root -g)" && \ - find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ - rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ - done && \ - find "$GLOBAL/npm" -type d -name "glob" -path "*/node_modules/glob" | while read d; do \ - rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \ - done && \ - find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \ - rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \ - done && \ - find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \ - rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \ - done && \ - find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ - rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ - done && \ - find /usr/local/lib /usr/lib -path "*/node_modules/npm/package.json" -exec \ - sed -i 's/"tar": "\^7\.5\.[0-9]*"/"tar": "^7.5.10"/g; s/"minimatch": "\^10\.[0-9.]*"/"minimatch": "^10.2.4"/g' {} + 2>/dev/null && \ - npm cache clean --force && \ - { apk del --no-cache npm 2>/dev/null || true; } + apk add --no-cache python3 bash openssl tzdata supervisor libsndfile nodejs && break || sleep 5; \ + done COPY --from=builder /app /app COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui @@ -179,15 +121,10 @@ ENV PATH="/app/.venv/bin:${PATH}" \ PRISMA_SKIP_POSTINSTALL_GENERATE=1 \ PRISMA_HIDE_UPDATE_MESSAGE=1 \ PRISMA_ENGINES_CHECKSUM_IGNORE_MISSING=1 \ - NPM_CONFIG_CACHE=/app/.cache/npm \ - NPM_CONFIG_PREFER_OFFLINE=true \ PRISMA_OFFLINE_MODE=true -RUN sed -i 's/\r$//' docker/entrypoint.sh && \ - sed -i 's/\r$//' docker/prod_entrypoint.sh && \ - chmod +x docker/entrypoint.sh docker/prod_entrypoint.sh && \ - mkdir -p /nonexistent /.npm /var/lib/litellm/assets /var/lib/litellm/ui /tmp/.npm && \ - chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent /.npm /tmp/.npm && \ +RUN mkdir -p /nonexistent /var/lib/litellm/assets /var/lib/litellm/ui && \ + chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent && \ PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \ chown -R nobody:nogroup "$PRISMA_PATH" && \ LITELLM_PKG_MIGRATIONS_PATH="$(python -c 'import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))' 2>/dev/null || echo '')/migrations" && \ diff --git a/docs/my-website/docs/adaptive_router.md b/docs/my-website/docs/adaptive_router.md new file mode 100644 index 00000000000..1e78ad4647a --- /dev/null +++ b/docs/my-website/docs/adaptive_router.md @@ -0,0 +1,155 @@ +# [BETA] Adaptive Router + +:::info + +Beta feature. Share feedback on [Discord](https://discord.gg/wuPM9dRgDw) or [Slack](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA). + +::: + +**Requirements:** LiteLLM Proxy with a Postgres database. Quality estimates are stored in Postgres and loaded on startup — without a database the router works but forgets everything learned on restart. + +You have a cheap model and an expensive one. You want to use the cheap one when it's good enough, and the expensive one when it actually matters — without hardcoding rules you'll spend months tuning. + +The adaptive router does this automatically. It tracks which model performs best for each type of request (code, writing, analysis, etc.) and routes accordingly, balancing quality against cost based on weights you control. + +## Quick start + +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + model_info: + input_cost_per_token: 0.0000025 + adaptive_router_preferences: + quality_tier: 3 # 1=budget, 2=mid, 3=frontier + strengths: ["code_generation", "analytical_reasoning"] + + - model_name: gpt-4o-mini + litellm_params: + model: openai/gpt-4o-mini + model_info: + input_cost_per_token: 0.00000015 + adaptive_router_preferences: + quality_tier: 2 + strengths: ["factual_lookup"] + + - model_name: my-router + litellm_params: + model: auto_router/adaptive_router + adaptive_router_config: + available_models: ["gpt-4o", "gpt-4o-mini"] + weights: + quality: 0.7 # raise this if quality complaints; lower if bill too high + cost: 0.3 # must sum to 1.0 with quality +``` + +Route to it by setting `model` to your adaptive router's name: + +```bash +curl -X POST {{baseURL}}/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -d '{ + "model": "my-router", + "messages": [ + {"role": "user", "content": "build me a python script that parses CSV"}, + {"role": "assistant", "content": "Here is a script using csv.DictReader..."}, + {"role": "user", "content": "now add error handling for missing files"}, + {"role": "assistant", "content": "Wrap the open() call in a try/except FileNotFoundError..."}, + {"role": "user", "content": "perfect, that worked. thanks!"} + ] + }' +``` + +The response includes a header telling you which model was actually picked: + +``` +x-litellm-adaptive-router-model: gpt-4o +``` + +The "thanks!" turn in the example above fires a satisfaction signal — that's what moves the bandit. + +## Tuning cost vs. quality + +The `weights` are your main lever: + +| Goal | quality | cost | +|---|---|---| +| Minimize cost, quality is secondary | 0.3 | 0.7 | +| Balanced | 0.5 | 0.5 | +| Quality-first (default) | 0.7 | 0.3 | +| Quality non-negotiable | 0.9 | 0.1 | + +The router learns over time. For the first ~10 requests per model, it relies on the tiers you declared. After that, real performance data takes over. + +## Force a minimum quality tier per request + +If a specific request needs a frontier model regardless of cost, pass this header: + +``` +x-litellm-min-quality-tier: 3 +``` + +You can also pass `min_quality_tier` via request metadata instead of a header. + +## What's being learned + +The router classifies each request into one of 7 types and tracks how each model performs on each independently. A model that's great at factual lookup but poor at code will win factual requests and lose code requests — even if it's cheaper overall. + +| Type | Example | +|---|---| +| `code_generation` | "write me a Python sort function" | +| `code_understanding` | "explain what this function does" | +| `technical_design` | "how should I design this API?" | +| `analytical_reasoning` | "calculate the probability that..." | +| `writing` | "draft an email to my team about..." | +| `factual_lookup` | "what is the capital of France?" | +| `general` | anything else | + +[**See classifier code**](https://github.com/BerriAI/litellm/blob/litellm_adaptive_routing/litellm/router_strategy/adaptive_router/classifier.py) + +Learning signals are inspired by [Signals: Trajectory Sampling and Triage for Agentic Interactions](https://arxiv.org/pdf/2604.00356). + +## Inspect the current state + +``` +GET /adaptive_router/{router_name}/state +``` + +Returns current quality estimates per model per request type. Useful for understanding why a model is or isn't being picked. + +```json +{ + "routers": [ + { + "router_name": "smart-cheap-router", + "available_models": ["fast", "smart"], + "weights": { "quality": 0.7, "cost": 0.3 }, + "cells": [ + { + "request_type": "analytical_reasoning", + "model": "fast", + "quality_mean": 0.5, + "samples": 0 + }, + { + "request_type": "analytical_reasoning", + "model": "smart", + "quality_mean": 0.95, + "samples": 0 + } + ] + } + ] +} +``` + +`quality_mean` is the key number — it's the router's current estimate of how well that model handles that request type. `samples` counts how many real observations have moved the prior (starts at 0; the cold-start prior mass is excluded). + +## Known limitations + +- Latency isn't scored — a slow model can still win on quality + cost +- Signals are regex-based and English-biased — no LLM judge +- Hard cap of 200 observations per cell; no decay yet +- Once a model is picked for a session, other models' turns in that session don't contribute to learning diff --git a/docs/my-website/docs/completion/prompt_caching.md b/docs/my-website/docs/completion/prompt_caching.md index 402c7b9f4c7..aaae7e7be76 100644 --- a/docs/my-website/docs/completion/prompt_caching.md +++ b/docs/my-website/docs/completion/prompt_caching.md @@ -10,6 +10,7 @@ Supported Providers: - Vertex AI (`vertex_ai/`, `vertex_ai_beta/`) - Bedrock (`bedrock/`, `bedrock/invoke/`, `bedrock/converse`) ([All models bedrock supports prompt caching on](https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html)) - Deepseek API (`deepseek/`) +- xAI (`xai/`) For the supported providers, LiteLLM follows the OpenAI prompt caching usage object format: diff --git a/docs/my-website/docs/completion/prompt_compression.md b/docs/my-website/docs/completion/prompt_compression.md index 2d999291af6..0d68ea2c101 100644 --- a/docs/my-website/docs/completion/prompt_compression.md +++ b/docs/my-website/docs/completion/prompt_compression.md @@ -8,6 +8,7 @@ The function keeps high-relevance and recent context, replaces low-relevance con ```python import litellm +from litellm.types.utils import CallTypes messages = [ {"role": "system", "content": "You are a coding assistant."}, @@ -19,6 +20,7 @@ messages = [ compressed = litellm.compress( messages=messages, model="gpt-4o", + call_type=CallTypes.completion, compression_trigger=1000, compression_target=500, ) @@ -45,6 +47,7 @@ response = litellm.completion( - `messages` (`List[dict]`, required): input conversation messages - `model` (`str`, required): model name used for token counting +- `call_type` (`CallTypes`, default `CallTypes.completion`): the LiteLLM call type whose message schema these messages follow. Supported values: `CallTypes.completion` / `CallTypes.acompletion` (OpenAI chat-completions shape) and `CallTypes.anthropic_messages` (Anthropic Messages shape) - `compression_trigger` (`int`, default `200000`): compress only if input token count exceeds this - `compression_target` (`Optional[int]`, default `70% of compression_trigger`): desired post-compression token budget - `embedding_model` (`Optional[str]`): if set, combines BM25 + embedding relevance scoring @@ -70,6 +73,28 @@ args = json.loads(tool_call.function.arguments) full_content = compressed["cache"][args["key"]] ``` +## Server-side Callback Loop (`/v1/messages`) + +You can enable callback-based compression interception to make retrieval loops +transparent for Anthropic Messages calls: + +```yaml +litellm_settings: + callbacks: ["compression_interception"] + compression_interception_params: + enabled: true + compression_trigger: 10000 + compression_target: 7000 +``` + +With this enabled, LiteLLM runs the following server-side flow: + +1. Compresses inbound messages before the first provider call. +2. Injects the `litellm_content_retrieve` tool. +3. Detects retrieval `tool_use` blocks in the model response. +4. Resolves retrieval keys from the compression cache. +5. Reruns the model via agentic loop and returns the final answer. + ## Performance Benchmarked on [SWE-bench Lite](https://huggingface.co/datasets/princeton-nlp/SWE-bench_Lite_bm25_27K) (real GitHub issues with ~27k tokens of BM25-retrieved repo context per problem). diff --git a/docs/my-website/docs/providers/scaleway.md b/docs/my-website/docs/providers/scaleway.md index ea57c24db30..8d83a37a3b1 100644 --- a/docs/my-website/docs/providers/scaleway.md +++ b/docs/my-website/docs/providers/scaleway.md @@ -60,3 +60,44 @@ curl http://localhost:4000/chat/completions \ ## Supported features Scaleway provider supports all features in [Generative APIs reference documentation ↗](https://www.scaleway.com/en/developers/api/generative-apis/), such as streaming, structured outputs and tool calling. + +## Audio transcription + +Scaleway's `/audio/transcriptions` endpoint is OpenAI-compatible and works with Whisper models. + +### Python SDK + +```python +import os +from litellm import transcription + +os.environ["SCW_SECRET_KEY"] = "your-scaleway-secret-key" + +with open("speech.mp3", "rb") as audio_file: + response = transcription( + model="scaleway/whisper-large-v3", + file=audio_file, + ) +print(response.text) +``` + +### Proxy config + +```yaml +model_list: + - model_name: scaleway-whisper + litellm_params: + model: scaleway/whisper-large-v3 + api_key: "os.environ/SCW_SECRET_KEY" +``` + +### Proxy request + +```bash +curl http://localhost:4000/v1/audio/transcriptions \ + -H "Authorization: Bearer YOUR_LITELLM_MASTER_KEY" \ + -F model="scaleway-whisper" \ + -F file="@speech.mp3" +``` + +Supported optional params: `language`, `prompt`, `response_format`, `temperature`, `timestamp_granularities`. diff --git a/docs/my-website/docs/proxy/agentic_loop_hook.md b/docs/my-website/docs/proxy/agentic_loop_hook.md new file mode 100644 index 00000000000..054c03228c4 --- /dev/null +++ b/docs/my-website/docs/proxy/agentic_loop_hook.md @@ -0,0 +1,95 @@ +# Agentic Loop Hook + +Build a `CustomLogger` callback that intercepts a model response, fulfills tool calls server-side, and reruns the model — transparently to the caller. + +:::info Supported call types +- `async` only (sync calls do not trigger the hook) +- Non-streaming only (streaming responses cannot be inspected for tool calls) +- Works on both `/v1/messages` and `/v1/chat/completions` +::: + +## Implement the callback + +Override two methods on `CustomLogger`: + +```python +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch + +MY_TOOL = "my_tool" + +class MyToolCallback(CustomLogger): + + async def async_should_run_agentic_loop( + self, response, model, messages, tools, stream, custom_llm_provider, kwargs + ): + # Return (True, context_dict) if there are tool calls to handle + content = getattr(response, "content", None) or [] + calls = [b for b in content if isinstance(b, dict) + and b.get("type") == "tool_use" and b.get("name") == MY_TOOL] + if not calls: + return False, {} + return True, {"tool_calls": calls} + + async def async_build_agentic_loop_plan( + self, tools, model, messages, response, + anthropic_messages_provider_config, + anthropic_messages_optional_request_params, + logging_obj, stream, kwargs, + ): + calls = tools["tool_calls"] + results = [f"result for {c['input']}" for c in calls] # your logic here + + follow_up = messages + [ + {"role": "assistant", "content": [ + {"type": "tool_use", "id": c["id"], "name": c["name"], "input": c["input"]} + for c in calls + ]}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": c["id"], "content": results[i]} + for i, c in enumerate(calls) + ]}, + ] + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch(messages=follow_up), + ) +``` + +For `/v1/chat/completions`, override `async_build_chat_completion_agentic_loop_plan` instead — same idea, `optional_params` replaces `anthropic_messages_optional_request_params`. + +## Register it + +```python +import litellm +litellm.callbacks = [MyToolCallback()] +``` + +Or in `config.yaml`: + +```yaml +litellm_settings: + callbacks: ["my_module.MyToolCallback"] +``` + +## `AgenticLoopPlan` fields + +| Field | Effect | +|---|---| +| `run_agentic_loop=True` + `request_patch` | Reruns the model with the patched request | +| `response_override` | Returns this value directly to the caller (no rerun) | +| `terminate=True` | Stops the loop, returns the current response | +| `run_agentic_loop=False` (default) | Skips; next callback is checked | + +`AgenticLoopRequestPatch` accepts: `model`, `messages`, `tools`, `max_tokens`, `optional_params`, `kwargs`. + +## Loop safety + +- Default max reruns: `3` — override per-request with `kwargs["max_agentic_loops"]` +- Identical tool-call fingerprints abort the loop automatically +- Current depth is in `kwargs["_agentic_loop_depth"]` + +## Examples in this repo + +- `litellm/integrations/compression_interception/handler.py` +- `litellm/integrations/websearch_interception/handler.py` diff --git a/docs/my-website/docs/tutorials/prompt_caching.md b/docs/my-website/docs/tutorials/prompt_caching.md index ab2aa00d773..581d2ba7c36 100644 --- a/docs/my-website/docs/tutorials/prompt_caching.md +++ b/docs/my-website/docs/tutorials/prompt_caching.md @@ -8,6 +8,22 @@ Reduce costs by up to 90% by using LiteLLM to auto-inject prompt caching checkpo +Supported Providers (`cache_control` marker): +- Anthropic API (`anthropic/`) +- AWS Bedrock - Claude (`bedrock/`) +- Vertex AI - Claude and Gemini (`vertex_ai/`) +- Google AI Studio - Gemini (`gemini/`) +- Azure AI - Claude (`azure_ai/`) +- OpenRouter - Claude, Gemini, MiniMax, GLM, z-ai routes (`openrouter/`) +- Databricks - Claude (`databricks/`) +- DashScope / Qwen (`dashscope/`) +- MiniMax (`minimax/`) +- Z.ai / GLM (`zai/`) + +Provider Managed (automatic, no marker needed): +- OpenAI (`openai/`) +- DeepSeek (`deepseek/`) +- xAI (`xai/`) ## How it works diff --git a/docs/my-website/package-lock.json b/docs/my-website/package-lock.json index 77644000aed..3504bb84196 100644 --- a/docs/my-website/package-lock.json +++ b/docs/my-website/package-lock.json @@ -26,6 +26,7 @@ }, "devDependencies": { "@docusaurus/module-type-aliases": "3.8.1", + "ajv": "^8.18.0", "dotenv": "16.6.1" }, "engines": { diff --git a/docs/my-website/package.json b/docs/my-website/package.json index bee7cbca186..6ebe6842e13 100644 --- a/docs/my-website/package.json +++ b/docs/my-website/package.json @@ -32,6 +32,7 @@ }, "devDependencies": { "@docusaurus/module-type-aliases": "3.8.1", + "ajv": "^8.18.0", "dotenv": "16.6.1" }, "browserslist": { diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index c2db54b2237..dbbdb70f6bc 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -536,6 +536,7 @@ const sidebars = { description: "Modify requests, responses, and more", items: [ "proxy/call_hooks", + "proxy/agentic_loop_hook", "proxy/rules", ] }, @@ -1059,6 +1060,7 @@ const sidebars = { }, items: [ "routing", + "adaptive_router", "scheduler", "proxy/auto_routing", "proxy/load_balancing", diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260418000000_add_adaptive_router_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260418000000_add_adaptive_router_tables/migration.sql new file mode 100644 index 00000000000..cdc76a0b915 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260418000000_add_adaptive_router_tables/migration.sql @@ -0,0 +1,39 @@ +-- One row per (router, request_type, model). Hot path on every routing decision. +CREATE TABLE "LiteLLM_AdaptiveRouterState" ( + router_name TEXT NOT NULL, + request_type TEXT NOT NULL, + model_name TEXT NOT NULL, + alpha DOUBLE PRECISION NOT NULL, + beta DOUBLE PRECISION NOT NULL, + total_samples INTEGER NOT NULL DEFAULT 0, + last_updated_at TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (router_name, request_type, model_name) +); + +-- One row per (session, router, model). Updated per turn via the queue. +CREATE TABLE "LiteLLM_AdaptiveRouterSession" ( + session_id TEXT NOT NULL, + router_name TEXT NOT NULL, + model_name TEXT NOT NULL, + classified_type TEXT NOT NULL, + misalignment_count INTEGER NOT NULL DEFAULT 0, + stagnation_count INTEGER NOT NULL DEFAULT 0, + disengagement_count INTEGER NOT NULL DEFAULT 0, + satisfaction_count INTEGER NOT NULL DEFAULT 0, + failure_count INTEGER NOT NULL DEFAULT 0, + loop_count INTEGER NOT NULL DEFAULT 0, + exhaustion_count INTEGER NOT NULL DEFAULT 0, + last_user_content TEXT, + last_assistant_content TEXT, + tool_call_history JSONB NOT NULL DEFAULT '[]', + pending_tool_calls JSONB NOT NULL DEFAULT '{}', + turn_count INTEGER NOT NULL DEFAULT 0, + last_processed_turn INTEGER NOT NULL DEFAULT -1, + clean_credit_awarded BOOLEAN NOT NULL DEFAULT FALSE, + terminal_status INTEGER, + last_activity_at TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (session_id, router_name, model_name) +); + +CREATE INDEX "idx_adaptive_router_session_activity" + ON "LiteLLM_AdaptiveRouterSession" (last_activity_at); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 08aa5645251..7642ad74b20 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1223,3 +1223,46 @@ model LiteLLM_ClaudeCodePluginTable { @@map("LiteLLM_ClaudeCodePluginTable") } + +// Per-(router, request_type, model) Beta posterior for the adaptive router. +model LiteLLM_AdaptiveRouterState { + router_name String + request_type String + model_name String + alpha Float + beta Float + total_samples Int @default(0) + last_updated_at DateTime @default(now()) @updatedAt + + @@id([router_name, request_type, model_name]) +} + +// Per-(session, router, model) signal counters for the adaptive router. +model LiteLLM_AdaptiveRouterSession { + session_id String + router_name String + model_name String + classified_type String + + misalignment_count Int @default(0) + stagnation_count Int @default(0) + disengagement_count Int @default(0) + satisfaction_count Int @default(0) + failure_count Int @default(0) + loop_count Int @default(0) + exhaustion_count Int @default(0) + + last_user_content String? + last_assistant_content String? + tool_call_history Json @default("[]") + pending_tool_calls Json @default("{}") + + turn_count Int @default(0) + last_processed_turn Int @default(-1) + clean_credit_awarded Boolean @default(false) + terminal_status Int? + last_activity_at DateTime @default(now()) @updatedAt + + @@id([session_id, router_name, model_name]) + @@index([last_activity_at], map: "idx_adaptive_router_session_activity") +} diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index c24188cba1d..369b6561931 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -30,6 +30,26 @@ def _get_prisma_env() -> dict: return prisma_env +_MIGRATION_TS_RE = re.compile(r"^(\d{14})_") + + +def _migration_timestamp(name: str) -> int: + """Extract the leading `YYYYMMDDHHMMSS` timestamp from a migration name. + + Returns 0 if the name doesn't match the Prisma pattern — unexpected-format + entries sort as "oldest" and are treated as historical. + """ + m = _MIGRATION_TS_RE.match(name) + return int(m.group(1)) if m else 0 + + +def _max_migration_timestamp(names) -> int: + """Max timestamp in a set/list of migration names (0 if empty).""" + if not names: + return 0 + return max(_migration_timestamp(n) for n in names) + + def _get_prisma_command() -> str: """Get the Prisma command to use, bypassing Python wrapper in offline mode.""" if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): @@ -383,18 +403,301 @@ class ProxyExtrasDBManager: ) @staticmethod - def setup_database(use_migrate: bool = False) -> bool: + def _strip_prisma_query_params(url: str) -> str: + """Remove Prisma-specific query params (connection_limit, pool_timeout, + schema, etc.) from DATABASE_URL so psycopg can parse it.""" + from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode + + parsed = urlparse(url) + if not parsed.query: + return url + libpq_params = { + "sslmode", + "sslcert", + "sslkey", + "sslrootcert", + "sslpassword", + "application_name", + "connect_timeout", + "client_encoding", + "options", + "service", + "gssencmode", + "krbsrvname", + "target_session_attrs", + } + kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params] + return urlunparse(parsed._replace(query=urlencode(kept))) + + @staticmethod + def _warn_if_db_ahead_of_head(migrations_dir: str) -> None: + """ + Log a warning if _prisma_migrations contains applied migrations with + timestamps newer than every migration this build ships. + + This is informational only for the v2 resolver — it tells the operator + the DB was likely migrated by a newer deployment, which is usually a + signal that this (older) version shouldn't run against it. We do NOT + block startup: many users have weird _prisma_migrations state from + prior thrashing bugs, and blocking them would be a breaking change. + + Safe no-op if psycopg isn't installed or DB isn't reachable. + """ + database_url = os.getenv("DATABASE_URL") + if not database_url: + return + + try: + import psycopg + except ImportError: + return + + cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) + known = set(ProxyExtrasDBManager._get_migration_names(migrations_dir)) + + try: + # autocommit=True keeps the SELECT outside a transaction. Without + # it, psycopg3's `with conn` calls COMMIT on clean exit — which + # fails after `UndefinedTable` (fresh DB) leaves the transaction + # in an aborted state. + with psycopg.connect( + cleaned_url, connect_timeout=10, autocommit=True + ) as conn: + try: + rows = conn.execute( + "SELECT migration_name FROM _prisma_migrations " + "WHERE finished_at IS NOT NULL AND rolled_back_at IS NULL" + ).fetchall() + except psycopg.errors.UndefinedTable: + return + except (psycopg.OperationalError, psycopg.DatabaseError): + # Swallow connection failures AND any other DB-layer error + # (e.g. InsufficientPrivilege if the runtime user lacks SELECT + # on _prisma_migrations). This is an informational check — + # never block startup on it. + return + + applied = {r[0] for r in rows} + unknown = applied - known + if not unknown: + return + + head_newest_ts = _max_migration_timestamp(known) + hostile = { + name for name in unknown if _migration_timestamp(name) > head_newest_ts + } + if not hostile: + return + + sorted_hostile = sorted(hostile) + logger.warning( + "Database has %d migration(s) applied that are NEWER than any " + "migration this LiteLLM version ships. This usually means the " + "database was migrated by a newer LiteLLM deployment. Some API " + "endpoints may fail because this proxy's Prisma client does not " + "know about those schema changes. Consider upgrading this " + "deployment. Unknown: %s", + len(hostile), + ", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""), + ) + + @staticmethod + def _setup_database_v2(use_migrate: bool) -> bool: + """ + v2 migration resolver (opt-in via --use_v2_migration_resolver). + + Runs `prisma migrate deploy` and handles standard recovery paths + (P3005 baseline, P3009/P3018 idempotent errors). Critically, it does + NOT call `_resolve_all_migrations` — the diff-and-force recovery that + caused schema thrashing when two LiteLLM versions contended for the + same DB during rolling deploys. + + Ahead-of-HEAD state (DB has migrations newer than this build ships) + is logged as a warning, not a fatal error — users whose DBs got into + weird shapes from the old thrashing should still be able to start. + """ + schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" + migrations_dir = ProxyExtrasDBManager._get_prisma_dir() + + if not use_migrate: + # Preserve `prisma db push` path unchanged. + original_dir = os.getcwd() + os.chdir(migrations_dir) + try: + subprocess.run( + [_get_prisma_command(), "db", "push", "--accept-data-loss"], + timeout=60, + check=True, + env=_get_prisma_env(), + ) + return True + except ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + ) as e: + # Re-raise as RuntimeError so proxy_cli.py's + # `except RuntimeError` catches it and exits cleanly. + raise RuntimeError(f"prisma db push failed.\n\nDetail: {e}") from e + finally: + os.chdir(original_dir) + + # Informational — never blocks. + ProxyExtrasDBManager._warn_if_db_ahead_of_head(migrations_dir) + + original_dir = os.getcwd() + os.chdir(migrations_dir) + try: + for attempt in range(4): + try: + result = subprocess.run( + [_get_prisma_command(), "migrate", "deploy"], + timeout=60, + check=True, + capture_output=True, + text=True, + env=_get_prisma_env(), + ) + logger.info(f"prisma migrate deploy stdout: {result.stdout}") + return True + + except subprocess.TimeoutExpired: + logger.info( + f"prisma migrate deploy attempt {attempt + 1} timed out, retrying" + ) + time.sleep(random.randrange(5, 15)) + continue + + except subprocess.CalledProcessError as e: + stderr = e.stderr or "" + + if "P3005" in stderr and "database schema is not empty" in stderr: + logger.info( + "Schema exists but no migrations ledger — creating baseline" + ) + ProxyExtrasDBManager._create_baseline_migration(schema_path) + continue + + if "P3009" in stderr: + migration_match = re.search(r"`(\d+_\S+?)`", stderr) + if ( + migration_match + and ProxyExtrasDBManager._is_idempotent_error(stderr) + ): + name = migration_match.group(1) + logger.info( + f"Migration {name} failed idempotently — marking applied and retrying" + ) + try: + ProxyExtrasDBManager._roll_back_migration(name) + except ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + ): + pass # may already be rolled-back + try: + ProxyExtrasDBManager._resolve_specific_migration(name) + except ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + ) as resolve_err: + # We're already inside the outer + # `except CalledProcessError` handler — + # re-raising CalledProcessError from here + # would escape as itself, bypassing + # proxy_cli.py's `except RuntimeError`. + raise RuntimeError( + f"Failed to mark migration {name} as applied " + f"after idempotent recovery. Manual " + f"intervention may be required.\n\n" + f"Detail: {resolve_err}" + ) from resolve_err + continue + raise RuntimeError( + "Database migration failed and cannot be auto-recovered. " + f"Manual intervention required.\n\nPrisma error:\n{stderr}" + ) from e + + if "P3018" in stderr: + if ProxyExtrasDBManager._is_permission_error(stderr): + raise RuntimeError( + "Database migration failed due to insufficient " + "permissions. Please grant the required privileges " + f"and retry.\n\nPrisma error:\n{stderr}" + ) from e + + migration_match = re.search( + r"Migration name: (\d+_\S+)", stderr + ) + if ( + migration_match + and ProxyExtrasDBManager._is_idempotent_error(stderr) + ): + name = migration_match.group(1) + logger.info( + f"Migration {name} SQL hit idempotent error — marking applied and retrying" + ) + try: + ProxyExtrasDBManager._roll_back_migration(name) + except ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + ): + pass # may already be rolled-back + try: + ProxyExtrasDBManager._resolve_specific_migration(name) + except ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + ) as resolve_err: + raise RuntimeError( + f"Failed to mark migration {name} as applied " + f"after idempotent recovery. Manual " + f"intervention may be required.\n\n" + f"Detail: {resolve_err}" + ) from resolve_err + continue + + raise RuntimeError( + "Database migration failed and cannot be auto-recovered. " + f"Manual intervention required.\n\nPrisma error:\n{stderr}" + ) from e + + raise RuntimeError( + "Database migration failed and cannot be auto-recovered. " + f"Manual intervention required.\n\nPrisma error:\n{stderr}" + ) from e + + raise RuntimeError( + "Database migration failed after 4 attempts (retry loop " + "exhausted by timeouts or repeated idempotent-recovery " + "continues). Check database connectivity, load, and " + "_prisma_migrations ledger state." + ) + finally: + os.chdir(original_dir) + + @staticmethod + def setup_database( + use_migrate: bool = False, use_v2_resolver: bool = False + ) -> bool: """ Set up the database using either prisma migrate or prisma db push Uses migrations from litellm-proxy-extras package Args: - schema_path (str): Path to the Prisma schema file - use_migrate (bool): Whether to use prisma migrate instead of db push + use_migrate: Whether to use prisma migrate instead of db push + use_v2_resolver: Opt into the v2 migration resolver (safer during + rolling deploys; does not run the diff-and-force recovery + that causes schema thrashing). Defaults to False for + backwards compatibility. Returns: bool: True if setup was successful, False otherwise """ + if use_v2_resolver: + logger.info("Using v2 migration resolver (--use_v2_migration_resolver)") + return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate) + schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" for attempt in range(4): original_dir = os.getcwd() diff --git a/litellm-proxy-extras/migration_runbook.md b/litellm-proxy-extras/migration_runbook.md index 8499bb7ce08..a277441b164 100644 --- a/litellm-proxy-extras/migration_runbook.md +++ b/litellm-proxy-extras/migration_runbook.md @@ -2,6 +2,8 @@ This is a runbook for creating and running database migrations for the LiteLLM proxy. For use for litellm engineers only. +> **AI AGENTS / ASSISTANTS:** If the script refuses with either a "STALE BRANCH" or "DESTRUCTIVE MIGRATION DETECTED" error, **do NOT** bypass it on your own (no `git rebase`, no `--skip-freshness-check`, no `--allow-destructive`). Surface the error to the human operator and wait for their explicit confirmation. See the [Branch freshness](#branch-freshness-check) and [Destructive migrations](#destructive-migrations-drop-column--drop-table) sections below. + ## Step 0: Sync All `schema.prisma` Files Before doing anything else, make sure all `schema.prisma` files in the repo are in sync. There are multiple copies that must match: @@ -46,10 +48,50 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n ## What It Does -1. Creates temp PostgreSQL DB -2. Applies existing migrations -3. Compares with `schema.prisma` -4. Generates new migration if changes found +1. **Verifies the current branch is up to date with `origin/litellm_internal_staging`** (see [Branch freshness](#branch-freshness-check)) +2. Creates temp PostgreSQL DB +3. Applies existing migrations +4. Compares with `schema.prisma` +5. Generates new migration if changes found +6. **Refuses to write if the generated SQL contains `DROP COLUMN`, `DROP TABLE`, or `DROP INDEX`**, unless `--allow-destructive` is passed + +## Branch Freshness Check + +Before generating anything, `run_migration.py` runs `git fetch origin ` and refuses to proceed if `HEAD` is behind `origin/`. Default base is `litellm_internal_staging` (the branch PRs target). A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense. + +Flags: + +- `--base-branch ` — check against a different base (e.g. `main`). Default is `litellm_internal_staging`. +- `--skip-freshness-check` — bypass entirely. Only for intentional migrations against an older base. + +When the guard fires: + +1. Update your branch: + + ```bash + git fetch origin && git rebase origin/litellm_internal_staging + # or git merge origin/litellm_internal_staging — whichever matches your workflow + ``` +2. Re-run `run_migration.py`. + +> **AI AGENTS / ASSISTANTS:** Do **not** auto-rebase or auto-pass `--skip-freshness-check`. A bad conflict resolution during rebase can itself drop the human's in-progress schema edits. Surface the error and wait for explicit confirmation. + +## Destructive Migrations (DROP COLUMN / DROP TABLE / DROP INDEX) + +If the generated diff contains `DROP COLUMN`, `DROP TABLE`, or `DROP INDEX`, `run_migration.py` exits non-zero and refuses to write the migration file. A previous incident saw newly-added columns silently dropped by a stale branch and merged to main — this guard exists to prevent a repeat. + +When the guard fires: + +1. Run `git fetch origin && git status` — confirm your branch is up to date with the base branch. +2. Re-check all `schema.prisma` files are in sync (Step 0). +3. Review EACH `DROP` statement printed in the error — is it actually intended? +4. Only if the drops are genuinely intentional, re-run with the flag: + + ```bash + uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_name" --allow-destructive + ``` + +> **AI AGENTS / ASSISTANTS:** Do **not** automatically re-run the command with `--allow-destructive`. If the guard fires while you are driving the runbook for a human, stop, show them the error, and wait for their explicit confirmation before passing the flag. Auto-passing `--allow-destructive` is the exact failure mode this guard exists to prevent. ## Common Fixes diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py new file mode 100644 index 00000000000..8d66bf872de --- /dev/null +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -0,0 +1,242 @@ +"""Regression tests for ProxyExtrasDBManager v2 migration resolver. + +The v2 resolver is opt-in via `--use_v2_migration_resolver` / the +`use_v2_resolver=True` kwarg. These tests exercise the v2 path; the v1 +(default) behavior is unchanged from pre-fix. +""" + +import subprocess +from unittest.mock import patch + +import pytest + +from litellm_proxy_extras.utils import ( + ProxyExtrasDBManager, + _max_migration_timestamp, + _migration_timestamp, +) + + +def _fake_migrate_deploy_failure(returncode: int, stderr: str): + def _run(*args, **kwargs): + raise subprocess.CalledProcessError( + returncode=returncode, + cmd=args[0], + stderr=stderr, + output="", + ) + + return _run + + +def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): + """v2: a permission failure during migrate deploy raises RuntimeError.""" + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setattr( + ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + + stderr = ( + "Error: P3018\nMigration name: 20250326162113_baseline\n" + "Database error code: 42501\npermission denied for schema public" + ) + with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with pytest.raises(RuntimeError, match="permission"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + +def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path): + """v2: a non-idempotent migration failure raises (no silent recovery).""" + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setattr( + ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + + stderr = ( + "Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n" + 'Reason: syntax error at or near "BRKN" LINE 42' + ) + with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with pytest.raises(RuntimeError, match="cannot be auto-recovered"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + +def test_strip_prisma_query_params_removes_connection_limit(): + """DATABASE_URLs with Prisma-specific params should be parseable by psycopg.""" + url = "postgresql://u:p@h:5432/db?connection_limit=100&pool_timeout=60&sslmode=require" + stripped = ProxyExtrasDBManager._strip_prisma_query_params(url) + assert "connection_limit" not in stripped + assert "pool_timeout" not in stripped + assert "sslmode=require" in stripped + + +def test_strip_prisma_query_params_passthrough_no_query(): + """URLs without query strings are returned unchanged.""" + url = "postgresql://u:p@h:5432/db" + assert ProxyExtrasDBManager._strip_prisma_query_params(url) == url + + +def test_migration_timestamp_extracts_leading_digits(): + assert _migration_timestamp("20260101000000_add_foo") == 20260101000000 + assert _migration_timestamp("20250326162113_baseline") == 20250326162113 + + +def test_migration_timestamp_returns_zero_on_malformed(): + assert _migration_timestamp("0_init") == 0 + assert _migration_timestamp("not_a_migration") == 0 + + +def test_max_migration_timestamp(): + names = {"20250326000000_a", "20260415000000_b", "20251115000000_c"} + assert _max_migration_timestamp(names) == 20260415000000 + + +def test_max_migration_timestamp_empty_set(): + assert _max_migration_timestamp(set()) == 0 + + +def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path): + """v1 (default) continues to call _resolve_all_migrations on the happy path. + + This is the existing buggy behavior — we're not fixing it in v1, only + offering v2 as opt-in. This test pins the default so that a future + inadvertent default flip is caught. + """ + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + + # Stub `prisma migrate deploy` to claim success with pending migrations + # applied, which is the code path that triggers the legacy post-migration + # sanity check (a call to _resolve_all_migrations). + class FakeResult: + stdout = "Applied migration.\n" + stderr = "" + + def fake_run(cmd, *args, **kwargs): + return FakeResult() + + resolve_called = {"n": 0} + + def fake_resolve(*args, **kwargs): + resolve_called["n"] += 1 + + monkeypatch.setattr("subprocess.run", fake_run) + monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set + assert ok is True + assert resolve_called["n"] == 1, "v1 default should still invoke the legacy path" + + +def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path): + """v2: a failing `prisma db push` must raise RuntimeError, not leak + CalledProcessError past proxy_cli.py's `except RuntimeError`.""" + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + + stderr = "db push error" + with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with pytest.raises(RuntimeError, match="prisma db push failed"): + ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) + + +def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path): + """_warn_if_db_ahead_of_head must never raise — it's informational. + + Non-connection DB errors (e.g. InsufficientPrivilege from a user + without SELECT on _prisma_migrations) must be caught, not propagated. + """ + import psycopg + + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + + class _FakeConn: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def execute(self, *a, **kw): + # Simulate an InsufficientPrivilege (subclass of DatabaseError). + raise psycopg.errors.InsufficientPrivilege("permission denied") + + def _fake_connect(*a, **kw): + return _FakeConn() + + monkeypatch.setattr("psycopg.connect", _fake_connect) + + # Must not raise. + ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path)) + + +def test_v2_resolve_specific_migration_failure_raises_runtime_error( + monkeypatch, tmp_path +): + """If marking a migration as applied fails inside P3009 idempotent + recovery, the subprocess error must be re-raised as RuntimeError so + proxy_cli.py catches it cleanly (instead of leaking CalledProcessError).""" + monkeypatch.setattr( + ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setattr( + ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None + ) + + # First call: migrate deploy -> P3009 idempotent error. + # Recovery path tries _resolve_specific_migration; that also raises. + def _failing_resolve(*a, **kw): + raise subprocess.CalledProcessError( + returncode=1, + cmd="prisma migrate resolve --applied", + stderr="resolve failed", + output="", + ) + + monkeypatch.setattr( + ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve + ) + + stderr = ( + "Error: P3009\nMigration `20260101000000_some_migration` failed\n" + "relation already exists" + ) + with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with pytest.raises( + RuntimeError, match="Failed to mark migration .* as applied" + ): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + +def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): + """v2 must never call _resolve_all_migrations — that's the bug it fixes.""" + monkeypatch.setattr( + ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + + class FakeResult: + stdout = "Applied migration.\n" + stderr = "" + + monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult()) + + resolve_called = {"n": 0} + monkeypatch.setattr( + ProxyExtrasDBManager, + "_resolve_all_migrations", + lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1), + ) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True + assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery" diff --git a/litellm/__init__.py b/litellm/__init__.py index 3acbb495356..89cef667c6e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -148,6 +148,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "vantage", "posthog", "levo", + "compression_interception", ] cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None @@ -1501,6 +1502,9 @@ if TYPE_CHECKING: from .llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeMessagesConfig as AmazonAnthropicClaudeMessagesConfig, ) + from .llms.bedrock.messages.mantle_transformation import ( + AmazonMantleMessagesConfig as AmazonMantleMessagesConfig, + ) from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 9164a3c8ae4..119e62a5b38 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -171,6 +171,7 @@ LLM_CONFIG_NAMES = ( "CohereChatConfig", "AnthropicMessagesConfig", "AmazonAnthropicClaudeMessagesConfig", + "AmazonMantleMessagesConfig", "TogetherAIConfig", "NLPCloudConfig", "VertexGeminiConfig", @@ -715,6 +716,10 @@ _LLM_CONFIGS_IMPORT_MAP = { ".llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation", "AmazonAnthropicClaudeMessagesConfig", ), + "AmazonMantleMessagesConfig": ( + ".llms.bedrock.messages.mantle_transformation", + "AmazonMantleMessagesConfig", + ), "TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"), "NLPCloudConfig": (".llms.nlp_cloud.chat.handler", "NLPCloudConfig"), "VertexGeminiConfig": ( diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index b78b04ec43c..45795c9ca15 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -1,9 +1,9 @@ """ -Main compress() function — orchestrates BM25/embedding scoring, message stubbing, -and retrieval tool injection. +Main compress() function — normalizes input messages, orchestrates BM25/embedding +scoring, message stubbing, and retrieval tool injection. """ -from typing import Any, Dict, List, Optional, Set, Union, cast +from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast from litellm.caching.dual_cache import DualCache from litellm.compression.message_stubbing import ( @@ -15,27 +15,196 @@ from litellm.compression.retrieval_tool import build_retrieval_tool from litellm.compression.scoring.bm25 import bm25_score_messages from litellm.litellm_core_utils.token_counter import token_counter from litellm.types.compression import CompressedResult -from litellm.types.utils import AllMessageValues, Message +from litellm.types.utils import CallTypes + +# CallTypes that produce Anthropic-shaped messages (structured content blocks). +# Everything else is treated as OpenAI chat-completions shape. +_ANTHROPIC_CALL_TYPES = frozenset({CallTypes.anthropic_messages.value}) +# CallTypes that are valid targets for compression. Compression operates on +# message-shaped inputs, so we only accept call types whose payload is a list +# of role/content messages. +_SUPPORTED_CALL_TYPES = frozenset( + { + CallTypes.completion.value, + CallTypes.acompletion.value, + CallTypes.anthropic_messages.value, + } +) + + +def _normalize_call_type(call_type: Union[CallTypes, str]) -> str: + """Return the string value for a ``CallTypes`` enum or a raw string.""" + if isinstance(call_type, CallTypes): + return call_type.value + return call_type + + +def _is_anthropic_call_type(call_type: str) -> bool: + return call_type in _ANTHROPIC_CALL_TYPES + + +def _build_retrieval_tools(keys: List[str], call_type: str) -> List[dict]: + """ + Build retrieval tool definitions in the target request schema. + + - Chat-completions call types: keep OpenAI function-tool schema. + - Anthropic messages call type: remap to Anthropic's custom tool schema. + """ + if not keys: + return [] + + openai_tools = [build_retrieval_tool(keys)] + if not _is_anthropic_call_type(call_type): + return openai_tools + + # Lazy import to avoid introducing provider transformation imports during + # module import for non-Anthropic call paths. + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + anthropic_tools, _mcp_servers = AnthropicConfig()._map_tools(openai_tools) + return cast(List[dict], anthropic_tools) + + +def _content_to_text(content: Any) -> str: + """ + Convert OpenAI/Anthropic message content blocks to plain text. + + Text extraction policy: + - Include text-bearing fields only (`text` blocks + string values). + - For `tool_result`, expand into nested `content` items. + - Ignore non-textual blocks (images/documents/tool metadata/thinking metadata). + + Implemented iteratively (stack-based) to avoid unbounded recursion. + """ + parts: List[str] = [] + stack: List[Any] = [content] + while stack: + item = stack.pop() + if isinstance(item, str): + parts.append(item) + elif isinstance(item, list): + # Push list items in reverse order so they are processed left-to-right. + for element in reversed(item): + stack.append(element) + elif isinstance(item, dict): + item_type = item.get("type") + if item_type == "text": + parts.append(str(item.get("text", ""))) + elif item_type == "tool_result": + stack.append(item.get("content", "")) + return " ".join(parts) + + +def _normalize_messages_for_compression( + messages: List[dict], + call_type: str, +) -> Tuple[List[dict], List[dict]]: + """ + Normalize each original message to a text-surrogate content for scoring. + + Returns: + (normalized_messages, original_messages_copy) + """ + if call_type not in _SUPPORTED_CALL_TYPES: + raise ValueError( + f"Unsupported call_type={call_type!r} for compression. " + f"Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}." + ) + + original_messages: List[Dict[str, Any]] = [dict(m) for m in messages] + + normalized_messages: List[dict] = [] + for msg in original_messages: + normalized_messages.append( + { + **msg, + "content": _content_to_text(msg.get("content", "")), + } + ) + return normalized_messages, original_messages def _extract_last_user_message(messages: List[dict]) -> str: """Return the text content of the last user message.""" for msg in reversed(messages): if msg.get("role") == "user": - content = msg.get("content", "") - if isinstance(content, str): - return content - if isinstance(content, list): - parts = [] - for part in content: - if isinstance(part, dict) and part.get("type") == "text": - parts.append(part.get("text", "")) - elif isinstance(part, str): - parts.append(part) - return " ".join(parts) + return _content_to_text(msg.get("content", "")) return "" +def _extract_tool_use_ids(content: Any) -> List[str]: + if not isinstance(content, list): + return [] + tool_use_ids: List[str] = [] + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") != "tool_use": + continue + tool_use_id = part.get("id") + if isinstance(tool_use_id, str) and tool_use_id: + tool_use_ids.append(tool_use_id) + return tool_use_ids + + +def _extract_tool_result_ids(content: Any) -> Set[str]: + if not isinstance(content, list): + return set() + tool_result_ids: Set[str] = set() + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") != "tool_result": + continue + tool_use_id = part.get("tool_use_id") + if isinstance(tool_use_id, str) and tool_use_id: + tool_result_ids.add(tool_use_id) + return tool_result_ids + + +def _extract_anthropic_tool_exchange_spans( + messages: List[dict], +) -> Tuple[List[Set[int]], Optional[str]]: + """ + Return atomic 2-message spans for Anthropic tool exchanges. + + Each assistant message containing `tool_use` must be immediately followed by a + user message containing matching `tool_result` blocks for all tool_use ids. + """ + spans: List[Set[int]] = [] + i = 0 + while i < len(messages): + current = messages[i] + if current.get("role") != "assistant": + i += 1 + continue + + tool_use_ids = _extract_tool_use_ids(current.get("content")) + if not tool_use_ids: + i += 1 + continue + + if i + 1 >= len(messages): + return [], "invalid_anthropic_tool_sequence" + + next_msg = messages[i + 1] + if next_msg.get("role") != "user": + return [], "invalid_anthropic_tool_sequence" + + tool_result_ids = _extract_tool_result_ids(next_msg.get("content")) + if not tool_result_ids: + return [], "invalid_anthropic_tool_sequence" + + for tool_use_id in tool_use_ids: + if tool_use_id not in tool_result_ids: + return [], "invalid_anthropic_tool_sequence" + + spans.append({i, i + 1}) + i += 2 + + return spans, None + + def _get_protected_indices(messages: List[dict]) -> List[int]: """ Return indices of messages that must never be compressed: @@ -87,9 +256,98 @@ def _combine_scores( return [bm25_weight * b + emb_weight * e for b, e in zip(norm_bm25, norm_emb)] +def _select_kept_indices_for_budget( + normalized_messages: List[dict], + original_messages: List[dict], + combined_scores: List[float], + compression_target: int, + model: str, + initial_kept_indices: Set[int], + tool_exchange_spans: List[Set[int]], +) -> Tuple[Set[int], Dict[int, dict]]: + kept_indices = set(initial_kept_indices) + current_tokens = 0 + for i in kept_indices: + current_tokens += token_counter( + model=model, + text=cast(str, normalized_messages[i].get("content", "") or ""), + ) + + # Fill token budget from highest-scoring units. + # A unit is either: + # 1) a single message index, or + # 2) an Anthropic tool-exchange span that must be kept/dropped atomically. + truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict + span_id_by_index: Dict[int, int] = {} + for span_id, span in enumerate(tool_exchange_spans): + for idx in span: + span_id_by_index[idx] = span_id + + # Build single-message candidate units (non-span messages). + candidate_units: List[Tuple[float, Tuple[int, ...], bool]] = [] + for idx in range(len(normalized_messages)): + if idx in span_id_by_index or idx in kept_indices: + continue + candidate_units.append((combined_scores[idx], (idx,), True)) + + # Build span candidate units (atomic keep/drop for tool exchanges). + for span in tool_exchange_spans: + span_indices = tuple(sorted(span)) + if any(idx in kept_indices for idx in span_indices): + continue + span_score = max(combined_scores[idx] for idx in span_indices) + candidate_units.append((span_score, span_indices, False)) + + # Sort by descending relevance score. + candidate_units.sort(key=lambda item: item[0], reverse=True) + + for _score, indices, can_truncate in candidate_units: + if any(idx in kept_indices for idx in indices): + continue + msg_tokens = 0 + for idx in indices: + msg_tokens += token_counter( + model=model, + text=cast(str, normalized_messages[idx].get("content", "") or ""), + ) + remaining = compression_target - current_tokens + + if remaining <= 0: + break # budget exhausted + + if current_tokens + msg_tokens <= compression_target: + # Fits entirely + kept_indices.update(indices) + current_tokens += msg_tokens + elif can_truncate and len(indices) == 1 and remaining >= 100: + # Too large to fit whole single message, but we have budget — truncate it. + idx = indices[0] + truncated = truncate_message(original_messages[idx], remaining) + truncated_tokens = token_counter( + model=model, + text=truncated.get("content", "") or "", + ) + truncated_overrides[idx] = truncated + kept_indices.add(idx) + current_tokens += truncated_tokens + + return kept_indices, truncated_overrides + + +def _get_dropped_tool_span_indices( + kept_indices: Set[int], tool_exchange_spans: List[Set[int]] +) -> Set[int]: + dropped_tool_span_indices: Set[int] = set() + for span in tool_exchange_spans: + if not any(idx in kept_indices for idx in span): + dropped_tool_span_indices.update(span) + return dropped_tool_span_indices + + def compress( messages: List[dict], model: str, + call_type: Union[CallTypes, str] = CallTypes.completion, compression_trigger: int = 200_000, compression_target: Optional[int] = None, embedding_model: Optional[str] = None, @@ -108,6 +366,12 @@ def compress( Parameters: messages: The conversation messages to (potentially) compress. model: The LLM model name — used for token counting. + call_type: The LiteLLM call type whose message schema these messages + follow. Supported values: + - ``CallTypes.completion`` / ``CallTypes.acompletion`` — OpenAI + chat-completions shape (default) + - ``CallTypes.anthropic_messages`` — Anthropic Messages shape + (structured content blocks + atomic tool exchanges) compression_trigger: Only compress if input exceeds this token count. compression_target: Target token count after compression. Defaults to ``compression_trigger // 2``. @@ -122,29 +386,37 @@ def compress( A ``CompressedResult`` dict containing compressed messages, token counts, a cache of original content, and the retrieval tool definition. """ + call_type_str = _normalize_call_type(call_type) + normalized_messages, original_messages = _normalize_messages_for_compression( + messages=messages, + call_type=call_type_str, + ) + if compression_target is None: compression_target = compression_trigger * 7 // 10 original_tokens = token_counter( - model=model, messages=cast(List[Union[AllMessageValues, Message]], messages) + model=model, + messages=cast(List[Any], original_messages), ) # Pass through if below trigger if original_tokens <= compression_trigger: return CompressedResult( - messages=messages, + messages=original_messages, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=0.0, cache={}, tools=[], + compression_skipped_reason="below_trigger", ) # Extract query for relevance scoring - query = _extract_last_user_message(messages) + query = _extract_last_user_message(normalized_messages) # Score each message - bm25_scores = bm25_score_messages(query, messages) + bm25_scores = bm25_score_messages(query, normalized_messages) if embedding_model: from litellm.compression.scoring.embedding_scorer import ( @@ -153,7 +425,7 @@ def compress( emb_scores = embedding_score_messages( query, - messages, + normalized_messages, model=embedding_model, cache=compression_cache, embedding_model_params=embedding_model_params, @@ -162,85 +434,69 @@ def compress( else: combined_scores = bm25_scores - # Sort message indices by score descending - ranked_indices = sorted( - range(len(messages)), - key=lambda i: combined_scores[i], - reverse=True, - ) - # Protected messages are never compressed - protected_indices = _get_protected_indices(messages) + protected_indices = _get_protected_indices(normalized_messages) kept_indices: Set[int] = set(protected_indices) - # Count tokens for protected messages - current_tokens = 0 - for i in kept_indices: - current_tokens += token_counter( - model=model, text=messages[i].get("content", "") or "" + tool_exchange_spans: List[Set[int]] = [] + if _is_anthropic_call_type(call_type_str): + tool_exchange_spans, tool_sequence_error = ( + _extract_anthropic_tool_exchange_spans(original_messages) ) - - # Fill token budget from highest-scoring messages. - # For each candidate (ranked by relevance): - # - If it fits entirely → keep it as-is. - # - If it doesn't fit but there's meaningful remaining budget → truncate it - # to fill as much of the budget as possible. - # - Otherwise → stub it (pointer only, content goes to cache). - # Multiple messages may be truncated so we preserve partial content from - # several high-scoring messages rather than fully stubbing all but one. - truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict - - for idx in ranked_indices: - if idx in kept_indices: - continue - msg_content = messages[idx].get("content", "") or "" - msg_tokens = token_counter(model=model, text=msg_content) - remaining = compression_target - current_tokens - - if remaining <= 0: - break # budget exhausted - - if current_tokens + msg_tokens <= compression_target: - # Fits entirely - kept_indices.add(idx) - current_tokens += msg_tokens - elif remaining >= 100: - # Too large to fit whole, but we have budget — truncate it. - truncated = truncate_message(messages[idx], remaining) - truncated_tokens = token_counter( - model=model, - text=truncated.get("content", "") or "", + if tool_sequence_error is not None: + return CompressedResult( + messages=original_messages, + original_tokens=original_tokens, + compressed_tokens=original_tokens, + compression_ratio=0.0, + cache={}, + tools=[], + compression_skipped_reason=tool_sequence_error, ) - truncated_overrides[idx] = truncated - kept_indices.add(idx) - current_tokens += truncated_tokens + + for span in tool_exchange_spans: + # If any message in the span is protected, keep the whole span. + if any(idx in kept_indices for idx in span): + kept_indices.update(span) + + kept_indices, truncated_overrides = _select_kept_indices_for_budget( + normalized_messages=normalized_messages, + original_messages=original_messages, + combined_scores=combined_scores, + compression_target=compression_target, + model=model, + initial_kept_indices=kept_indices, + tool_exchange_spans=tool_exchange_spans, + ) # Build compressed messages and cache compressed_messages: List[dict] = [] cache: Dict[str, str] = {} used_keys: Set[str] = set() + dropped_tool_span_indices = _get_dropped_tool_span_indices( + kept_indices=kept_indices, tool_exchange_spans=tool_exchange_spans + ) - for i, msg in enumerate(messages): + for i, msg in enumerate(original_messages): + if i in dropped_tool_span_indices: + continue if i in kept_indices: # Use the truncated version if we made one, otherwise the original compressed_messages.append(truncated_overrides.get(i, msg)) else: - key = extract_key(msg, fallback_index=i, used_keys=used_keys) - content = msg.get("content", "") - if isinstance(content, list): - content = " ".join( - p.get("text", "") if isinstance(p, dict) else str(p) - for p in content - ) + key = extract_key( + normalized_messages[i], fallback_index=i, used_keys=used_keys + ) + content = _content_to_text(msg.get("content", "")) cache[key] = content compressed_messages.append(stub_message(msg, key)) - # Build retrieval tool - tools = [build_retrieval_tool(list(cache.keys()))] if cache else [] + # Build retrieval tool in the target request schema + tools = _build_retrieval_tools(list(cache.keys()), call_type=call_type_str) compressed_tokens = token_counter( model=model, - messages=cast(List[Union[AllMessageValues, Message]], compressed_messages), + messages=cast(List[Any], compressed_messages), ) return CompressedResult( diff --git a/litellm/constants.py b/litellm/constants.py index 6c89cf5946d..012599ab6ab 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -164,6 +164,7 @@ MCP_STDIO_ALLOWED_COMMANDS: frozenset = frozenset( LITELLM_UI_ALLOW_HEADERS = [ "x-litellm-semantic-filter", "x-litellm-semantic-filter-tools", + "x-litellm-adaptive-router-model", ] # Gemini model-specific minimal thinking budget constants diff --git a/litellm/integrations/compression_interception/__init__.py b/litellm/integrations/compression_interception/__init__.py new file mode 100644 index 00000000000..14d30af14d8 --- /dev/null +++ b/litellm/integrations/compression_interception/__init__.py @@ -0,0 +1,14 @@ +""" +Compression Interception Module + +Provides server-side prompt compression + retrieval tool fulfillment for +Anthropic Messages agentic loops. +""" + +from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, +) + +__all__ = [ + "CompressionInterceptionLogger", +] diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py new file mode 100644 index 00000000000..c6ae7d9e82b --- /dev/null +++ b/litellm/integrations/compression_interception/handler.py @@ -0,0 +1,399 @@ +""" +Compression Interception Handler + +CustomLogger that compresses inbound Anthropic Messages requests and fulfills +litellm_content_retrieve tool calls server-side via the typed agentic loop plan. +""" + +import time +import uuid +from typing import Any, Dict, List, Optional, Tuple, cast + +from litellm._logging import verbose_logger +from litellm.compression import compress +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.integrations.compression_interception import ( + CompressionInterceptionConfig, +) +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) +from litellm.types.utils import CallTypes + +LITELLM_CONTENT_RETRIEVE_TOOL_NAME = "litellm_content_retrieve" +_CACHE_TTL_SECONDS = 15 * 60 + + +class CompressionInterceptionLogger(CustomLogger): + """ + CustomLogger that implements transparent prompt compression + retrieval loops. + + Flow: + 1. Compress inbound /v1/messages requests in pre-call hook. + 2. Inject litellm_content_retrieve tool and persist compressed cache by call_id. + 3. Detect retrieval tool_use blocks in first model response. + 4. Build typed rerun plan with tool_result blocks from the compressed cache. + """ + + def __init__( + self, + enabled: bool = True, + compression_trigger: int = 200_000, + compression_target: Optional[int] = None, + embedding_model: Optional[str] = None, + embedding_model_params: Optional[Dict[str, Any]] = None, + ): + super().__init__() + self.enabled = enabled + self.compression_trigger = compression_trigger + self.compression_target = compression_target + self.embedding_model = embedding_model + self.embedding_model_params = embedding_model_params + self._compression_cache_by_call_id: Dict[str, Tuple[Dict[str, str], float]] = {} + + @classmethod + def from_config_yaml( + cls, config: CompressionInterceptionConfig + ) -> "CompressionInterceptionLogger": + return cls( + enabled=bool(config.get("enabled", True)), + compression_trigger=int(config.get("compression_trigger", 200_000)), + compression_target=config.get("compression_target"), + embedding_model=config.get("embedding_model"), + embedding_model_params=config.get("embedding_model_params"), + ) + + @staticmethod + def initialize_from_proxy_config( + litellm_settings: Dict[str, Any], + callback_specific_params: Dict[str, Any], + ) -> "CompressionInterceptionLogger": + compression_params: CompressionInterceptionConfig = {} + if "compression_interception_params" in litellm_settings: + compression_params = litellm_settings["compression_interception_params"] + elif "compression_interception" in callback_specific_params: + compression_params = callback_specific_params["compression_interception"] + return CompressionInterceptionLogger.from_config_yaml(compression_params) + + async def async_pre_call_deployment_hook( + self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] + ) -> Optional[dict]: + if not self.enabled: + return None + if call_type is not None and call_type != CallTypes.anthropic_messages: + return None + if int(kwargs.get("_agentic_loop_depth", 0) or 0) > 0: + return None + + messages = kwargs.get("messages") + model = kwargs.get("model") + if not isinstance(messages, list) or not isinstance(model, str): + return None + + if self._has_retrieval_tool(kwargs.get("tools")): + return None + + self._prune_expired_cache() + + compressed = compress( # type: ignore + messages=messages, + model=model, + call_type=CallTypes.anthropic_messages, + compression_trigger=self.compression_trigger, + compression_target=self.compression_target, + embedding_model=self.embedding_model, + embedding_model_params=self.embedding_model_params, + ) + + cache = cast(Dict[str, str], compressed.get("cache", {})) + skip_reason = cast(Optional[str], compressed.get("compression_skipped_reason")) + compressed_tools = cast(List[Dict[str, Any]], compressed.get("tools", [])) + + # Only mutate kwargs when compression actually produced a result. + # If compression was a no-op (below trigger, invalid tool sequence, etc.), + # leave ``messages`` and ``tools`` untouched — injecting an empty + # ``tools: []`` onto a request that originally had no tools breaks + # Anthropic Messages requests. + if cache: + kwargs["messages"] = compressed["messages"] + if compressed_tools: + kwargs["tools"] = self._merge_tools( + existing_tools=cast( + Optional[List[Dict[str, Any]]], kwargs.get("tools") + ), + compressed_tools=compressed_tools, + ) + call_id = cast(Optional[str], kwargs.get("litellm_call_id")) + if not call_id: + call_id = str(uuid.uuid4()) + kwargs["litellm_call_id"] = call_id + self._compression_cache_by_call_id[call_id] = (cache, time.time()) + verbose_logger.debug( + "CompressionInterception: compressed request [call_id=%s original=%d compressed=%d cached_keys=%d]", + call_id, + compressed.get("original_tokens"), + compressed.get("compressed_tokens"), + len(cache), + ) + elif skip_reason is not None: + verbose_logger.debug( + "CompressionInterception: compression skipped [reason=%s original=%d compressed=%d]", + skip_reason, + compressed.get("original_tokens"), + compressed.get("compressed_tokens"), + ) + + return kwargs + + async def async_should_run_agentic_loop( + self, + response: Any, + model: str, + messages: List[Dict], + tools: Optional[List[Dict]], + stream: bool, + custom_llm_provider: str, + kwargs: Dict, + ) -> Tuple[bool, Dict]: + if not self.enabled: + return False, {} + if not self._has_retrieval_tool(tools): + return False, {} + + tool_calls, thinking_blocks = self._extract_retrieval_tool_calls( + response=response + ) + if not tool_calls: + return False, {} + + return True, { + "tool_calls": tool_calls, + "thinking_blocks": thinking_blocks, + "tool_type": "compression_retrieval", + } + + async def async_build_agentic_loop_plan( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict, + logging_obj: Any, + stream: bool, + kwargs: Dict, + ) -> AgenticLoopPlan: + self._prune_expired_cache() + tool_calls = cast(List[Dict[str, Any]], tools.get("tool_calls", [])) + thinking_blocks = cast(List[Dict[str, Any]], tools.get("thinking_blocks", [])) + + call_id = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs) + cache = self._get_cache(call_id=call_id) + retrieval_results = [ + self._resolve_retrieval_content(tc, cache) for tc in tool_calls + ] + + assistant_message = { + "role": "assistant", + "content": thinking_blocks + + [ + { + "type": "tool_use", + "id": tc.get("id"), + "name": tc.get("name", LITELLM_CONTENT_RETRIEVE_TOOL_NAME), + "input": tc.get("input", {}), + } + for tc in tool_calls + ], + } + user_message = { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": tool_calls[i].get("id"), + "content": retrieval_results[i], + } + for i in range(len(tool_calls)) + ], + } + follow_up_messages = messages + [assistant_message, user_message] + + max_tokens = cast( + Optional[int], + anthropic_messages_optional_request_params.get("max_tokens") + or kwargs.get("max_tokens"), + ) + optional_params_without_max_tokens = { + k: v + for k, v in anthropic_messages_optional_request_params.items() + if k != "max_tokens" + } + + full_model_name = model + if logging_obj is not None: + agentic_params = logging_obj.model_call_details.get( + "agentic_loop_params", {} + ) + full_model_name = cast(str, agentic_params.get("model", model)) + + request_patch = AgenticLoopRequestPatch( + model=full_model_name, + messages=follow_up_messages, + max_tokens=max_tokens, + optional_params=optional_params_without_max_tokens, + kwargs=self._prepare_followup_kwargs(kwargs=kwargs), + ) + + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=request_patch, + metadata={"tool_type": "compression_retrieval", "call_id": call_id or ""}, + ) + + def _prune_expired_cache(self) -> None: + now = time.time() + self._compression_cache_by_call_id = { + call_id: (cache, created_at) + for call_id, ( + cache, + created_at, + ) in self._compression_cache_by_call_id.items() + if now - created_at <= _CACHE_TTL_SECONDS + } + + def _get_cache(self, call_id: Optional[str]) -> Dict[str, str]: + if not call_id: + return {} + cache_entry = self._compression_cache_by_call_id.get(call_id) + if cache_entry is None: + return {} + return cache_entry[0] + + def _resolve_call_id( + self, logging_obj: Any, kwargs: Dict[str, Any] + ) -> Optional[str]: + if logging_obj is not None: + logging_call_id = getattr(logging_obj, "litellm_call_id", None) + if isinstance(logging_call_id, str) and logging_call_id: + return logging_call_id + kwargs_call_id = kwargs.get("litellm_call_id") + return cast( + Optional[str], kwargs_call_id if isinstance(kwargs_call_id, str) else None + ) + + def _resolve_retrieval_content( + self, tool_call: Dict[str, Any], cache: Dict[str, str] + ) -> str: + raw_input = tool_call.get("input", {}) + key = "" + if isinstance(raw_input, dict): + key = str(raw_input.get("key", "") or "") + if not key: + return "No retrieval key provided." + if key in cache: + return cache[key] + return f"[compressed content key '{key}' not found]" + + def _extract_retrieval_tool_calls( + self, response: Any + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + if isinstance(response, dict): + content = response.get("content", []) + else: + content = getattr(response, "content", []) or [] + + if not isinstance(content, list): + return [], [] + + tool_calls: List[Dict[str, Any]] = [] + thinking_blocks: List[Dict[str, Any]] = [] + + for block in content: + if isinstance(block, dict): + block_type = block.get("type") + block_name = block.get("name") + if block_type in ("thinking", "redacted_thinking"): + thinking_blocks.append(block) + if ( + block_type == "tool_use" + and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME + ): + tool_calls.append( + { + "id": block.get("id"), + "type": "tool_use", + "name": block_name, + "input": block.get("input", {}), + } + ) + else: + block_type = getattr(block, "type", None) + block_name = getattr(block, "name", None) + if block_type == "thinking": + thinking_blocks.append( + { + "type": "thinking", + "thinking": getattr(block, "thinking", ""), + "signature": getattr(block, "signature", ""), + } + ) + elif block_type == "redacted_thinking": + thinking_blocks.append( + { + "type": "redacted_thinking", + "data": getattr(block, "data", ""), + } + ) + if ( + block_type == "tool_use" + and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME + ): + tool_calls.append( + { + "id": getattr(block, "id", None), + "type": "tool_use", + "name": block_name, + "input": getattr(block, "input", {}) or {}, + } + ) + + return tool_calls, thinking_blocks + + def _prepare_followup_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]: + internal_keys = {"litellm_logging_obj"} + return { + k: v + for k, v in kwargs.items() + if not k.startswith("_compression_interception") and k not in internal_keys + } + + def _has_retrieval_tool(self, tools: Any) -> bool: + if not isinstance(tools, list): + return False + for tool in tools: + if not isinstance(tool, dict): + continue + function = tool.get("function") + if tool.get("type") == "function" and isinstance(function, dict): + if function.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME: + return True + if ( + tool.get("type") == "custom" + and tool.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME + ): + return True + return False + + def _merge_tools( + self, + existing_tools: Optional[List[Dict[str, Any]]], + compressed_tools: List[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + merged = list(existing_tools or []) + if self._has_retrieval_tool(merged): + return merged + merged.extend(compressed_tools) + return merged diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 45c8e2f6262..300c311f36d 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -20,6 +20,7 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER from litellm.types.integrations.argilla import ArgillaItem from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest from litellm.types.prompts.init_prompts import PromptSpec +from litellm.types.integrations.custom_logger import AgenticLoopPlan from litellm.types.utils import ( AdapterCompletionStreamWrapper, CallTypes, @@ -239,7 +240,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac self, model: str, request_kwargs: Dict, - messages: Optional[List[Dict[str, str]]] = None, + messages: Optional[List[Dict[str, Any]]] = None, input: Optional[Union[str, List]] = None, specific_deployment: Optional[bool] = False, ) -> Optional[PreRoutingHookResponse]: @@ -676,6 +677,26 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ pass + async def async_build_agentic_loop_plan( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict, + logging_obj: "LiteLLMLoggingObj", + stream: bool, + kwargs: Dict, + ) -> AgenticLoopPlan: + """ + Build a typed rerun plan for Anthropic Messages agentic loops. + + Override this method to separate callback decision/tool execution from + follow-up request execution (handled by BaseLLMHTTPHandler). + """ + return AgenticLoopPlan(run_agentic_loop=False) + async def async_should_run_chat_completion_agentic_loop( self, response: Any, @@ -707,6 +728,22 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ pass + async def async_build_chat_completion_agentic_loop_plan( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + optional_params: Dict, + logging_obj: "LiteLLMLoggingObj", + stream: bool, + kwargs: Dict, + ) -> AgenticLoopPlan: + """ + Build a typed rerun plan for chat-completions agentic loops. + """ + return AgenticLoopPlan(run_agentic_loop=False) + # Useful helpers for custom logger classes def truncate_standard_logging_payload_content( diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index ecfb42cea7b..b6d91d0b76d 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1615,6 +1615,14 @@ class OpenTelemetry(CustomLogger): value=response_id, ) + litellm_call_id = standard_logging_payload.get("litellm_call_id") + if litellm_call_id: + self.safe_set_attribute( + span=span, + key="litellm.call_id", + value=litellm_call_id, + ) + # The model used to generate the response. if response_obj and response_obj.get("model"): self.safe_set_attribute( @@ -2281,6 +2289,10 @@ class OpenTelemetry(CustomLogger): # Remove trailing slash endpoint = endpoint.rstrip("/") + # Splunk Observability Cloud OTLP/HTTP uses /v2/trace/otlp (not /v1/traces). Do not rewrite. + if signal_type == "traces" and "/v2/trace/otlp" in endpoint: + return endpoint + # Check if endpoint already ends with the correct signal path target_path = f"/v1/{signal_type}" if endpoint.endswith(target_path): diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 1d92a9da073..723b142dfad 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -51,6 +51,7 @@ if TYPE_CHECKING: else: AsyncIOScheduler = Any + class PrometheusLogger(CustomLogger): # Class variables or attributes @@ -991,9 +992,7 @@ class PrometheusLogger(CustomLogger): amount: float = 1.0, ) -> None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name=metric_name - ), + supported_enum_labels=self.get_labels_for_metric(metric_name=metric_name), enum_values=enum_values, label_context=label_context, ) @@ -1118,7 +1117,9 @@ class PrometheusLogger(CustomLogger): user_api_key = hash_token(user_api_key) - label_context = PrometheusLabelFactoryContext(enum_values) #amortized per request. + label_context = PrometheusLabelFactoryContext( + enum_values + ) # amortized per request. # increment total LLM requests and spend metric self._increment_top_level_request_and_spend_metrics( @@ -3490,7 +3491,9 @@ def _prometheus_labels_from_context( } if UserAPIKeyLabelNames.END_USER.value in filtered_labels: - filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ctx.get_resolved_end_user() + filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ( + ctx.get_resolved_end_user() + ) for sk, val in ctx._custom_by_sanitized_key.items(): if sk in supported_enum_labels: diff --git a/litellm/integrations/prometheus_helpers/__init__.py b/litellm/integrations/prometheus_helpers/__init__.py index 34f4855863e..784ab524dd5 100644 --- a/litellm/integrations/prometheus_helpers/__init__.py +++ b/litellm/integrations/prometheus_helpers/__init__.py @@ -51,8 +51,7 @@ class PrometheusLabelFactoryContext: self.enum_values = enum_values enum_dict = enum_values.model_dump() self._sanitized_enum: Dict[str, Optional[str]] = { - k: _sanitize_prometheus_label_value(v) - for k, v in enum_dict.items() + k: _sanitize_prometheus_label_value(v) for k, v in enum_dict.items() } self._custom_by_sanitized_key: Dict[str, Optional[str]] = {} if enum_values.custom_metadata_labels is not None: diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 30fd55a3e9d..41618c72627 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -28,6 +28,10 @@ from litellm.integrations.websearch_interception.transformation import ( from litellm.types.integrations.websearch_interception import ( WebSearchInterceptionConfig, ) +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager @@ -573,6 +577,35 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs=kwargs, ) + async def async_build_agentic_loop_plan( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict, + logging_obj: Any, + stream: bool, + kwargs: Dict, + ) -> AgenticLoopPlan: + tool_calls = tools["tool_calls"] + thinking_blocks = tools.get("thinking_blocks", []) + request_patch = await self._build_anthropic_request_patch( + model=model, + messages=messages, + tool_calls=tool_calls, + thinking_blocks=thinking_blocks, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + kwargs=kwargs, + ) + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=request_patch, + metadata={"tool_type": "websearch", "response_format": "anthropic"}, + ) + async def async_run_chat_completion_agentic_loop( self, tools: Dict, @@ -608,6 +641,33 @@ class WebSearchInterceptionLogger(CustomLogger): response_format=response_format, ) + async def async_build_chat_completion_agentic_loop_plan( + self, + tools: Dict, + model: str, + messages: List[Dict], + response: Any, + optional_params: Dict, + logging_obj: Any, + stream: bool, + kwargs: Dict, + ) -> AgenticLoopPlan: + tool_calls = tools["tool_calls"] + response_format = tools.get("response_format", "openai") + request_patch = await self._build_chat_completion_request_patch( + model=model, + messages=messages, + tool_calls=tool_calls, + optional_params=optional_params, + kwargs=kwargs, + response_format=response_format, + ) + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=request_patch, + metadata={"tool_type": "websearch", "response_format": response_format}, + ) + @staticmethod def _resolve_max_tokens( optional_params: Dict, @@ -672,7 +732,48 @@ class WebSearchInterceptionLogger(CustomLogger): stream: bool, kwargs: Dict, ) -> Any: - """Execute litellm.search() and make follow-up request""" + """Legacy path: execute search + build patch + run follow-up call.""" + request_patch = await self._build_anthropic_request_patch( + model=model, + messages=messages, + tool_calls=tool_calls, + thinking_blocks=thinking_blocks, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + kwargs=kwargs, + ) + if request_patch.messages is None: + raise ValueError("WebSearchInterception: missing follow-up messages") + + optional_params = dict(anthropic_messages_optional_request_params) + optional_params.update(request_patch.optional_params) + max_tokens = request_patch.max_tokens + if max_tokens is None: + max_tokens = cast(Optional[int], optional_params.pop("max_tokens", None)) + else: + optional_params.pop("max_tokens", None) + if max_tokens is None: + max_tokens = cast(int, kwargs.get("max_tokens", 1024)) + + return await anthropic_messages.acreate( + max_tokens=max_tokens, + messages=request_patch.messages, + model=request_patch.model or model, + **optional_params, + **request_patch.kwargs, + ) + + async def _build_anthropic_request_patch( + self, + model: str, + messages: List[Dict], + tool_calls: List[Dict], + thinking_blocks: List[Dict], + anthropic_messages_optional_request_params: Dict, + logging_obj: Any, + kwargs: Dict, + ) -> AgenticLoopRequestPatch: + """Execute litellm.search() and build follow-up request patch.""" # Extract search queries from tool_use blocks search_tasks = [] @@ -721,20 +822,8 @@ class WebSearchInterceptionLogger(CustomLogger): thinking_blocks=thinking_blocks, ) - # Make follow-up request with search results - # Type cast: user_message is a Dict for Anthropic format (default response_format) follow_up_messages = messages + [assistant_message, cast(Dict, user_message)] - verbose_logger.debug( - "WebSearchInterception: Making follow-up request with search results" - ) - verbose_logger.debug( - f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}" - ) - verbose_logger.debug( - f"WebSearchInterception: Last message (tool_result): {user_message}" - ) - # Correlation context for structured logging _call_id = getattr(logging_obj, "litellm_call_id", None) or kwargs.get( "litellm_call_id", "unknown" @@ -742,61 +831,41 @@ class WebSearchInterceptionLogger(CustomLogger): full_model_name = model # safe default before try block - # Use anthropic_messages.acreate for follow-up request - try: - max_tokens = self._resolve_max_tokens( - anthropic_messages_optional_request_params, kwargs - ) + max_tokens = self._resolve_max_tokens( + anthropic_messages_optional_request_params, kwargs + ) - verbose_logger.debug( - f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request" - ) + verbose_logger.debug( + f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request" + ) - # Create a copy of optional params without max_tokens (since we pass it explicitly) - optional_params_without_max_tokens = { - k: v - for k, v in anthropic_messages_optional_request_params.items() - if k != "max_tokens" - } + optional_params_without_max_tokens = { + k: v + for k, v in anthropic_messages_optional_request_params.items() + if k != "max_tokens" + } + kwargs_for_followup = self._prepare_followup_kwargs(kwargs) - kwargs_for_followup = self._prepare_followup_kwargs(kwargs) - - # Get model from logging_obj.model_call_details["agentic_loop_params"] - # This preserves the full model name with provider prefix (e.g., "bedrock/invoke/...") - if logging_obj is not None: - agentic_params = logging_obj.model_call_details.get( - "agentic_loop_params", {} - ) - full_model_name = agentic_params.get("model", model) - verbose_logger.debug( - f"WebSearchInterception: Using model name: {full_model_name}" + if logging_obj is not None: + agentic_params = logging_obj.model_call_details.get( + "agentic_loop_params", {} ) - - final_response = await anthropic_messages.acreate( - max_tokens=max_tokens, - messages=follow_up_messages, - model=full_model_name, - **optional_params_without_max_tokens, - **kwargs_for_followup, - ) - verbose_logger.debug( - f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}" - ) - verbose_logger.debug( - f"WebSearchInterception: Final response: {final_response}" - ) - return final_response - except Exception as e: - verbose_logger.exception( - "WebSearchInterception: Follow-up request failed " - "[call_id=%s model=%s messages=%d searches=%d]: %s", - _call_id, - full_model_name, - len(follow_up_messages), - len(final_search_results), - str(e), - ) - raise + full_model_name = agentic_params.get("model", model) + verbose_logger.debug( + "WebSearchInterception: Built anthropic request patch " + "[call_id=%s model=%s messages=%d searches=%d]", + _call_id, + full_model_name, + len(follow_up_messages), + len(final_search_results), + ) + return AgenticLoopRequestPatch( + model=full_model_name, + messages=follow_up_messages, + max_tokens=max_tokens, + optional_params=optional_params_without_max_tokens, + kwargs=kwargs_for_followup, + ) async def _execute_search(self, query: str) -> str: """Execute a single web search using router's search tools""" @@ -883,7 +952,36 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs: Dict, response_format: str = "openai", ) -> Any: - """Execute litellm.search() and make follow-up chat completion request""" + """Legacy path: execute search + build patch + run follow-up call.""" + request_patch = await self._build_chat_completion_request_patch( + model=model, + messages=messages, + tool_calls=tool_calls, + optional_params=optional_params, + kwargs=kwargs, + response_format=response_format, + ) + if request_patch.messages is None: + raise ValueError("WebSearchInterception: missing follow-up messages") + params = dict(optional_params) + params.update(request_patch.optional_params) + return await litellm.acompletion( + model=request_patch.model or model, + messages=request_patch.messages, + **params, + **request_patch.kwargs, + ) + + async def _build_chat_completion_request_patch( # noqa: PLR0915 + self, + model: str, + messages: List[Dict], + tool_calls: List[Dict], + optional_params: Dict, + kwargs: Dict, + response_format: str = "openai", + ) -> AgenticLoopRequestPatch: + """Execute litellm.search() and build chat-completion rerun patch.""" # Extract search queries from tool_calls search_tasks = [] @@ -963,74 +1061,56 @@ class WebSearchInterceptionLogger(CustomLogger): f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}" ) - # Use litellm.acompletion for follow-up request - try: - # Remove internal parameters that shouldn't be passed to follow-up request - internal_params = { - "_websearch_interception", - "acompletion", - "litellm_logging_obj", - "custom_llm_provider", + # Remove internal parameters that shouldn't be passed to follow-up request + internal_params = { + "_websearch_interception", + "acompletion", + "litellm_logging_obj", + "custom_llm_provider", + "model_alias_map", + "stream_response", + "custom_prompt_dict", + } + kwargs_for_followup = { + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") and k not in internal_params + } + + full_model_name = model + if "custom_llm_provider" in kwargs: + custom_llm_provider = kwargs["custom_llm_provider"] + if not model.startswith(custom_llm_provider) and "/" not in model: + full_model_name = f"{custom_llm_provider}/{model}" + + verbose_logger.debug( + "WebSearchInterception: Built chat completion request patch model=%s messages=%d", + full_model_name, + len(follow_up_messages), + ) + + tools_param = optional_params.get("tools") + optional_params_clean = { + k: v + for k, v in optional_params.items() + if k + not in { + "tools", + "extra_body", "model_alias_map", "stream_response", "custom_prompt_dict", } - kwargs_for_followup = { - k: v - for k, v in kwargs.items() - if not k.startswith("_websearch_interception") - and k not in internal_params - } + } + if tools_param is not None: + optional_params_clean["tools"] = tools_param - # Get full model name from kwargs - full_model_name = model - if "custom_llm_provider" in kwargs: - custom_llm_provider = kwargs["custom_llm_provider"] - # Reconstruct full model name with provider prefix if needed - if not model.startswith(custom_llm_provider): - # Check if model already has a provider prefix - if "/" not in model: - full_model_name = f"{custom_llm_provider}/{model}" - - verbose_logger.debug( - f"WebSearchInterception: Using model name: {full_model_name}" - ) - - # Prepare tools for follow-up request (same as original) - tools_param = optional_params.get("tools") - - # Remove tools and extra_body from optional_params to avoid issues - # extra_body often contains internal LiteLLM params that shouldn't be forwarded - optional_params_clean = { - k: v - for k, v in optional_params.items() - if k - not in { - "tools", - "extra_body", - "model_alias_map", - "stream_response", - "custom_prompt_dict", - } - } - - final_response = await litellm.acompletion( - model=full_model_name, - messages=follow_up_messages, - tools=tools_param, - **optional_params_clean, - **kwargs_for_followup, - ) - - verbose_logger.debug( - f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}" - ) - return final_response - except Exception as e: - verbose_logger.exception( - f"WebSearchInterception: Follow-up request failed: {str(e)}" - ) - raise + return AgenticLoopRequestPatch( + model=full_model_name, + messages=follow_up_messages, + optional_params=optional_params_clean, + kwargs=kwargs_for_followup, + ) async def _create_empty_search_result(self) -> str: """Create an empty search result for tool calls without queries""" diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index b72d7abeae0..9d8bd7523db 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -296,6 +296,15 @@ def get_supported_openai_params( # noqa: PLR0915 return OVHCloudAudioTranscriptionConfig().get_supported_openai_params( model=model ) + elif custom_llm_provider == "scaleway": + if request_type == "transcription": + from litellm.llms.scaleway.audio_transcription.transformation import ( + ScalewayAudioTranscriptionConfig, + ) + + return ScalewayAudioTranscriptionConfig().get_supported_openai_params( + model=model + ) elif custom_llm_provider == "elevenlabs": if request_type == "transcription": from litellm.llms.elevenlabs.audio_transcription.transformation import ( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index fd14f55add3..625cb83724b 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5512,6 +5512,8 @@ def get_standard_logging_object_payload( payload: StandardLoggingPayload = StandardLoggingPayload( 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( logging_obj=logging_obj, litellm_params=litellm_params, diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 3fd913958da..888999504fe 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -684,7 +684,7 @@ def generic_cost_per_token( # noqa: PLR0915 - cache_creation - image_tokens ) - # Clamp to zero: inconsistent streaming usage + # Clamp to zero: inconsistent streaming usage if text_tokens < 0: text_tokens = 0 prompt_tokens_details["text_tokens"] = text_tokens diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index cb430b06940..2bb82f227bb 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -34,6 +34,7 @@ from litellm.types.llms.anthropic import ( ) from litellm.types.llms.openai import ( AllMessageValues, + ChatCompletionRequest, ChatCompletionToolCallChunk, ChatCompletionToolParam, ) @@ -67,6 +68,32 @@ class AnthropicMessagesHandler(BaseTranslation): super().__init__() self.adapter = LiteLLMAnthropicMessagesAdapter() + def _translate_to_openai(self, data: dict) -> ChatCompletionRequest: + """Translate Anthropic request to OpenAI chat completion format.""" + ( + chat_completion_compatible_request, + _tool_name_mapping, + ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + anthropic_message_request=cast(AnthropicMessagesRequest, data.copy()) + ) + return chat_completion_compatible_request + + def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]: + """ + Convert Anthropic messages request data to OpenAI-spec structured messages. + + Uses the Anthropic-to-OpenAI adapter to translate message format. + """ + messages = data.get("messages") + if messages is None: + return None + chat_completion_compatible_request = self._translate_to_openai(data) + result = cast( + List[AllMessageValues], + chat_completion_compatible_request.get("messages", []), + ) + return result if result else None + async def process_input_messages( self, data: dict, @@ -82,13 +109,7 @@ class AnthropicMessagesHandler(BaseTranslation): skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply) - ( - chat_completion_compatible_request, - _tool_name_mapping, - ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( - # Use a shallow copy to avoid mutating request data (pop on litellm_metadata). - anthropic_message_request=cast(AnthropicMessagesRequest, data.copy()) - ) + chat_completion_compatible_request = self._translate_to_openai(data) structured_messages = cast( List[AllMessageValues], @@ -103,8 +124,6 @@ class AnthropicMessagesHandler(BaseTranslation): chat_completion_compatible_request.get("tools", []) ) task_mappings: List[Tuple[int, Optional[int]]] = [] - # Track (message_index, content_index) for each text - # content_index is None for string content, int for list content # Step 1: Extract all text content and images for msg_idx, message in enumerate(messages): diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py new file mode 100644 index 00000000000..d0780c82d06 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -0,0 +1,322 @@ +""" +Agentic Streaming Iterator for Anthropic Messages + +Wraps the raw SSE byte stream from the Anthropic pass-through endpoint, +yields every chunk to the caller (preserving real streaming), collects +all bytes, and on stream exhaustion rebuilds the full Anthropic response +to run through agentic completion hooks. If an agentic hook fires, the +follow-up response is chained as Phase 2 of the same iterator. +""" + +import json +from typing import Any, AsyncIterator, Dict, List, Optional, cast + +from litellm._logging import verbose_logger + + +# --------------------------------------------------------------------------- +# SSE parsing helpers (module-level to keep the class lean) +# --------------------------------------------------------------------------- + + +def _parse_sse_events(raw: bytes) -> List[tuple]: + """Return a list of (event_type, parsed_data_dict) from raw SSE bytes.""" + text = raw.decode("utf-8", errors="replace") + lines = text.split("\n") + events: List[tuple] = [] + current_event_type: Optional[str] = None + + for line in lines: + stripped = line.strip() + if stripped.startswith("event:"): + current_event_type = stripped[len("event:") :].strip() + continue + if not stripped.startswith("data:"): + continue + data_str = stripped[len("data:") :].strip() + try: + data = json.loads(data_str) + except (json.JSONDecodeError, ValueError): + continue + event_type = current_event_type or data.get("type", "") + current_event_type = None + events.append((event_type, data)) + return events + + +def _handle_message_start(data: Dict, response: Dict) -> None: + msg = data.get("message", {}) + response["id"] = msg.get("id", response["id"]) + response["model"] = msg.get("model", response["model"]) + response["role"] = msg.get("role", response["role"]) + usage = msg.get("usage", {}) + if usage: + response["usage"]["input_tokens"] = usage.get("input_tokens", 0) + for key in ("cache_creation_input_tokens", "cache_read_input_tokens"): + if key in usage: + response["usage"][key] = usage[key] + + +def _handle_content_block_start(data: Dict, content_blocks: Dict[int, Dict]) -> None: + idx = data.get("index", len(content_blocks)) + block = data.get("content_block", {}) + block_type = block.get("type", "text") + + _BLOCK_TEMPLATES: Dict[str, Dict] = { + "text": {"type": "text", "text": ""}, + "thinking": {"type": "thinking", "thinking": "", "signature": ""}, + "redacted_thinking": { + "type": "redacted_thinking", + "data": block.get("data", ""), + }, + } + if block_type == "tool_use": + content_blocks[idx] = { + "type": "tool_use", + "id": block.get("id", ""), + "name": block.get("name", ""), + "input": {}, + "_partial_json": "", + } + elif block_type in _BLOCK_TEMPLATES: + content_blocks[idx] = dict(_BLOCK_TEMPLATES[block_type]) + else: + content_blocks[idx] = dict(block) + + +def _handle_content_block_delta(data: Dict, content_blocks: Dict[int, Dict]) -> None: + idx = data.get("index", 0) + delta = data.get("delta", {}) + delta_type = delta.get("type", "") + block = content_blocks.get(idx) + if block is None: + return + + if delta_type == "text_delta": + block["text"] = block.get("text", "") + delta.get("text", "") + elif delta_type == "input_json_delta": + block["_partial_json"] = block.get("_partial_json", "") + delta.get( + "partial_json", "" + ) + elif delta_type == "thinking_delta": + block["thinking"] = block.get("thinking", "") + delta.get("thinking", "") + elif delta_type == "signature_delta": + block["signature"] = delta.get("signature", block.get("signature", "")) + + +def _handle_content_block_stop(data: Dict, content_blocks: Dict[int, Dict]) -> None: + idx = data.get("index", 0) + block = content_blocks.get(idx) + if block and block.get("type") == "tool_use": + partial = block.pop("_partial_json", "") + if partial: + try: + block["input"] = json.loads(partial) + except (json.JSONDecodeError, ValueError): + block["input"] = {"_raw": partial} + + +def _handle_message_delta(data: Dict, response: Dict) -> None: + delta = data.get("delta", {}) + if "stop_reason" in delta: + response["stop_reason"] = delta["stop_reason"] + if "stop_sequence" in delta: + response["stop_sequence"] = delta["stop_sequence"] + usage = data.get("usage", {}) + if usage.get("output_tokens") is not None: + response["usage"]["output_tokens"] = usage["output_tokens"] + for key in ( + "input_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + ): + if key in usage: + response["usage"][key] = usage[key] + + +class AgenticAnthropicStreamingIterator: + """ + Two-phase async iterator that enables agentic hooks on streaming + Anthropic Messages pass-through responses. + + Phase 1: Yield raw SSE bytes from the upstream response while + accumulating them. When the inner iterator is exhausted, + rebuild the full Anthropic response dict and call agentic hooks. + + Phase 2: If an agentic hook fires and returns a follow-up response + (streaming or non-streaming), yield those bytes to the caller. + """ + + def __init__( + self, + completion_stream: AsyncIterator, + http_handler: Any, + model: str, + messages: List[Dict], + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict, + logging_obj: Any, + custom_llm_provider: str, + kwargs: Dict, + ): + self._inner = completion_stream.__aiter__() + self._http_handler = http_handler + self._model = model + self._messages = messages + self._anthropic_messages_provider_config = anthropic_messages_provider_config + self._anthropic_messages_optional_request_params = ( + anthropic_messages_optional_request_params + ) + self._logging_obj = logging_obj + self._custom_llm_provider = custom_llm_provider + self._kwargs = kwargs + + self._collected_bytes: List[bytes] = [] + self._stream_exhausted = False + self._hook_processing_done = False + self._follow_up_iterator: Optional[AsyncIterator] = None + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + # Phase 1: yield from upstream, collect bytes + if not self._stream_exhausted: + try: + chunk = await self._inner.__anext__() + self._collected_bytes.append(chunk) + return chunk + except StopAsyncIteration: + self._stream_exhausted = True + await self._process_agentic_hooks() + # Fall through to Phase 2 + + # Phase 2: yield from follow-up stream if one was created + if self._follow_up_iterator is not None: + chunk = await self._follow_up_iterator.__anext__() + return chunk + + raise StopAsyncIteration + + async def _process_agentic_hooks(self) -> None: + """Rebuild the Anthropic response from collected SSE bytes and call hooks.""" + if self._hook_processing_done: + return + self._hook_processing_done = True + + if not self._collected_bytes: + return + + try: + rebuilt = self._rebuild_anthropic_response_from_sse(self._collected_bytes) + if rebuilt is None: + verbose_logger.debug( + "AgenticStreamingIterator: Could not rebuild response from SSE bytes" + ) + return + + [ + ( + f"{b.get('type')}({b.get('name', '')})" + if b.get("type") == "tool_use" + else b.get("type") + ) + for b in rebuilt.get("content", []) + ] + + result = await self._http_handler._call_agentic_completion_hooks( + response=rebuilt, + model=self._model, + messages=self._messages, + anthropic_messages_provider_config=self._anthropic_messages_provider_config, + anthropic_messages_optional_request_params=self._anthropic_messages_optional_request_params, + logging_obj=self._logging_obj, + stream=True, + custom_llm_provider=self._custom_llm_provider, + kwargs=self._kwargs, + ) + + if result is None: + return + + if hasattr(result, "__aiter__"): + self._follow_up_iterator = result.__aiter__() + elif isinstance(result, dict): + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, + ) + + fake = FakeAnthropicMessagesStreamIterator( + response=cast(AnthropicMessagesResponse, result) + ) + self._follow_up_iterator = fake.__aiter__() + else: + verbose_logger.warning( + "AgenticStreamingIterator: Unexpected result type from hooks: %s", + type(result).__name__, + ) + except Exception as e: + _call_id = getattr(self._logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "AgenticStreamingIterator: Error in agentic hook processing " + "[call_id=%s model=%s]: %s", + _call_id, + self._model, + str(e), + ) + + @staticmethod + def _rebuild_anthropic_response_from_sse( + raw_bytes: List[bytes], + ) -> Optional[Dict[str, Any]]: + """ + Parse collected SSE bytes into an Anthropic Messages response dict. + + Processes SSE events in order: + - message_start -> envelope (id, model, role, usage) + - content_block_start -> new content block + - content_block_delta -> accumulate text/json/thinking deltas + - content_block_stop -> finalize block + - message_delta -> stop_reason, output usage + - message_stop -> end + """ + events = _parse_sse_events(b"".join(raw_bytes)) + + response: Dict[str, Any] = { + "id": "", + "type": "message", + "role": "assistant", + "model": "", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + content_blocks: Dict[int, Dict[str, Any]] = {} + saw_message_start = False + + for event_type, data in events: + if event_type == "message_start": + saw_message_start = True + _handle_message_start(data, response) + elif event_type == "content_block_start": + _handle_content_block_start(data, content_blocks) + elif event_type == "content_block_delta": + _handle_content_block_delta(data, content_blocks) + elif event_type == "content_block_stop": + _handle_content_block_stop(data, content_blocks) + elif event_type == "message_delta": + _handle_message_delta(data, response) + + if not saw_message_start: + return None + + for idx in sorted(content_blocks.keys()): + block = content_blocks[idx] + block.pop("_partial_json", None) + response["content"].append(block) + + return response diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index e1da0dfa29e..1efeb159a3e 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -5,6 +5,7 @@ if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.llms.openai import AllMessageValues class BaseTranslation(ABC): @@ -101,6 +102,16 @@ class BaseTranslation(ABC): """ return responses_so_far + def get_structured_messages(self, data: dict) -> Optional[List["AllMessageValues"]]: + """ + Convert request data to OpenAI-spec structured messages. + + Override in subclasses for format-specific conversion. + + Returns None if no convertible content is found. + """ + return None + def extract_request_tool_names(self, data: dict) -> List[str]: """ Extract tool names from the request body for allowlist/policy checks. diff --git a/litellm/llms/bedrock/chat/mantle/__init__.py b/litellm/llms/bedrock/chat/mantle/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py new file mode 100644 index 00000000000..b9bea77c118 --- /dev/null +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -0,0 +1,91 @@ +""" +Transformation for Bedrock Mantle (Claude Mythos Preview) + +https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-mythos-preview.html + +The bedrock-mantle endpoint uses the Anthropic Messages API format but is served +at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth. +""" + +from typing import TYPE_CHECKING, Any, List, Optional + +from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeConfig, +) +from litellm.types.llms.openai import AllMessageValues + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + +MANTLE_ENDPOINT_TEMPLATE = "https://bedrock-mantle.{region}.api.aws/v1/messages" + + +class AmazonMantleConfig(AmazonAnthropicClaudeConfig): + """ + Config for the bedrock-mantle endpoint (Claude Mythos Preview). + + Uses the Anthropic Messages API format with AWS SigV4 auth, but at a + different endpoint from bedrock-runtime. Model ID goes in the request body. + + Usage: model="bedrock/mantle/anthropic.claude-mythos-preview" + """ + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + region = self._get_aws_region_name(optional_params=optional_params, model=model) + return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + + def transform_request( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + # Strip the "mantle/" routing prefix to get the real model ID + model_id = model.replace("mantle/", "", 1) + + request = self._build_bedrock_anthropic_request_base( + model=model_id, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + # The parent strips "model" from the body (Invoke API puts it in URL). + # The mantle endpoint (Messages API) requires "model" in the body. + request["model"] = model_id + return request + + async def async_transform_request( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + model_id = model.replace("mantle/", "", 1) + + request = self._build_bedrock_anthropic_request_base( + model=model_id, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + await self._async_convert_document_url_sources_to_base64(request) + request["model"] = model_id + return request diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 52697d752be..9a97a134cc4 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -696,6 +696,7 @@ class BedrockModelInfo(BaseLLMModelInfo): "agentcore", "async_invoke", "openai", + "mantle", ]: """ Get the bedrock route for the given model. @@ -710,6 +711,7 @@ class BedrockModelInfo(BaseLLMModelInfo): "agentcore", "async_invoke", "openai", + "mantle", ], ] = { "invoke/": "invoke", @@ -719,6 +721,7 @@ class BedrockModelInfo(BaseLLMModelInfo): "agentcore/": "agentcore", "async_invoke/": "async_invoke", "openai/": "openai", + "mantle/": "mantle", } # Check explicit routes first @@ -770,6 +773,13 @@ class BedrockModelInfo(BaseLLMModelInfo): """ return "agentcore/" in model + @staticmethod + def _explicit_mantle_route(model: str) -> bool: + """ + Check if the model is an explicit mantle route (bedrock-mantle endpoint). + """ + return "mantle/" in model + @staticmethod def _explicit_converse_like_route(model: str) -> bool: """ @@ -809,6 +819,16 @@ class BedrockModelInfo(BaseLLMModelInfo): if BedrockModelInfo._explicit_converse_route(model): return None + ######################################################### + # Mantle route uses the bedrock-mantle endpoint (not bedrock-runtime) + ######################################################### + if BedrockModelInfo._explicit_mantle_route(model): + from litellm.llms.bedrock.messages.mantle_transformation import ( + AmazonMantleMessagesConfig, + ) + + return AmazonMantleMessagesConfig() + ######################################################### # This goes through litellm.AmazonAnthropicClaude3MessagesConfig() # Since bedrock Invoke supports Native Anthropic Messages API @@ -855,6 +875,12 @@ def get_bedrock_chat_config(model: str): ) return AmazonAgentCoreConfig() + elif bedrock_route == "mantle": + from litellm.llms.bedrock.chat.mantle.transformation import ( + AmazonMantleConfig, + ) + + return AmazonMantleConfig() # Handle provider-specific configs if bedrock_invoke_provider == "amazon": diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 31e0e76fd9f..96593b35d0c 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -34,6 +34,7 @@ from litellm.llms.bedrock.common_utils import ( remove_custom_field_from_tools, ) from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER +from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import GenericStreamingChunk @@ -59,6 +60,10 @@ class AmazonAnthropicClaudeMessagesConfig( DEFAULT_BEDROCK_ANTHROPIC_API_VERSION = "bedrock-2023-05-31" + BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset( + BedrockInvokeAnthropicMessagesRequest.__annotations__.keys() + ) + def __init__(self, **kwargs): BaseAnthropicMessagesConfig.__init__(self, **kwargs) AmazonInvokeConfig.__init__(self, **kwargs) @@ -500,10 +505,6 @@ class AmazonAnthropicClaudeMessagesConfig( anthropic_messages_request=anthropic_messages_request, ) - # 5b. Strip `output_config` — Bedrock Invoke doesn't support it - # Fixes: https://github.com/BerriAI/litellm/issues/22797 - anthropic_messages_request.pop("output_config", None) - # 5a. Remove `custom` field from tools (Bedrock doesn't support it) # Claude Code sends `custom: {defer_loading: true}` on tool definitions, # which causes Bedrock to reject the request with "Extra inputs are not permitted" @@ -550,14 +551,43 @@ class AmazonAnthropicClaudeMessagesConfig( if "tool-search-tool-2025-10-19" in beta_set: beta_set.add("tool-examples-2025-10-29") - filtered_auto_betas = filter_and_transform_beta_headers( - beta_headers=list(beta_set - user_beta_set), - provider="bedrock", + filtered_betas = sorted( + filter_and_transform_beta_headers( + beta_headers=list(beta_set), + provider="bedrock", + ) ) - filtered_betas = sorted(user_beta_set.union(set(filtered_auto_betas))) + + dropped_user_betas = sorted( + b + for b in user_beta_set + if not filter_and_transform_beta_headers([b], provider="bedrock") + ) + if dropped_user_betas: + verbose_logger.warning( + "Bedrock Invoke: dropping unsupported anthropic-beta values " + "from client headers: %s. Bedrock has no mapping entry for " + "these; forwarding them would cause a 400.", + dropped_user_betas, + ) + if filtered_betas: anthropic_messages_request["anthropic_beta"] = filtered_betas + # 7. Final safety net: filter top-level fields to the Bedrock Invoke allowlist. + # Catches Anthropic-only extensions (context_management, output_config, speed, + # mcp_servers, ...) and any future additions Claude Code may start sending. + allowed = self.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS + stripped = sorted(k for k in anthropic_messages_request if k not in allowed) + if stripped: + verbose_logger.debug( + "Bedrock Invoke: stripping unsupported top-level request fields: %s", + stripped, + ) + anthropic_messages_request = { + k: v for k, v in anthropic_messages_request.items() if k in allowed + } + return anthropic_messages_request def get_async_streaming_response_iterator( diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py new file mode 100644 index 00000000000..3f04c8a3052 --- /dev/null +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -0,0 +1,69 @@ +""" +Transformation for Bedrock Mantle (Claude Mythos Preview) - /messages endpoint + +Inherits all Messages API request/response transformations from +AmazonAnthropicClaudeMessagesConfig. Overrides only the URL and model-prefix +stripping that are specific to the bedrock-mantle endpoint. +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, +) +from litellm.types.router import GenericLiteLLMParams + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + +MANTLE_ENDPOINT_TEMPLATE = "https://bedrock-mantle.{region}.api.aws/v1/messages" + + +class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): + """ + Config for the bedrock-mantle /messages endpoint (Claude Mythos Preview). + + The mantle endpoint uses the Anthropic Messages API format and requires the + model ID in the request body (unlike Bedrock Invoke which puts it in the URL). + """ + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + region = self._get_aws_region_name(optional_params=optional_params, model=model) + return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + + def transform_anthropic_messages_request( + self, + model: str, + messages: List[Dict], + anthropic_messages_optional_request_params: Dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> Dict: + # Strip "mantle/" routing prefix to get the real model ID + model_id = model.replace("mantle/", "", 1) + + request = super().transform_anthropic_messages_request( + model=model_id, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + + # Parent (AmazonAnthropicClaudeMessagesConfig) removes "model" from the + # body (Bedrock Invoke puts model in the URL). The mantle endpoint + # (Messages API) requires "model" in the request body. + request["model"] = model_id + return request diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ea0c05e7656..8a7043111de 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -78,6 +78,10 @@ from litellm.types.containers.main import ( DeleteContainerResult, ) from litellm.types.files import TwoStepFileUploadConfig +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -2047,7 +2051,23 @@ class BaseLLMHTTPHandler: request_body=request_body, litellm_logging_obj=logging_obj, ) - initial_response = completion_stream + + from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( + AgenticAnthropicStreamingIterator, + ) + + initial_response = AgenticAnthropicStreamingIterator( + completion_stream=completion_stream, + http_handler=self, + model=model, + messages=messages, + anthropic_messages_provider_config=anthropic_messages_provider_config, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ) + return initial_response else: initial_response = anthropic_messages_provider_config.transform_anthropic_messages_response( model=model, @@ -2055,7 +2075,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) - # Call agentic completion hooks + # Call agentic completion hooks (non-streaming path only) final_response = await self._call_agentic_completion_hooks( response=initial_response, model=model, @@ -2063,7 +2083,7 @@ class BaseLLMHTTPHandler: anthropic_messages_provider_config=anthropic_messages_provider_config, anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, logging_obj=logging_obj, - stream=stream or False, + stream=False, custom_llm_provider=custom_llm_provider, kwargs=kwargs, ) @@ -4516,6 +4536,167 @@ class BaseLLMHTTPHandler: return stream, data return stream, data + @staticmethod + def _get_agentic_loop_settings(kwargs: Dict) -> Tuple[int, int, List[str]]: + depth = int(kwargs.get("_agentic_loop_depth", 0) or 0) + max_loops = int(kwargs.get("max_agentic_loops", 3) or 3) + fingerprints = list(kwargs.get("_agentic_loop_fingerprints", []) or []) + return depth, max(max_loops, 1), fingerprints + + @staticmethod + def _check_agentic_loop_safety( + tool_calls: Any, + fingerprints: List[str], + depth: int, + max_loops: int, + model: str, + ) -> str: + """ + Evaluate agentic-loop safety guards (fingerprint cycle / max depth). + + Raises ValueError on abort. Returns the current fingerprint on success. + + These checks must not be swallowed by the per-callback ``except Exception`` + block that wraps callback dispatch — they are bounded-loop / cycle-break + safety rails and must abort the agentic dispatch when they trip. + """ + fingerprint = BaseLLMHTTPHandler._fingerprint_agentic_tools(tool_calls) + if fingerprint in fingerprints: + raise ValueError( + "Agentic loop detected repeated tool-call fingerprint; aborting rerun" + ) + if depth >= max_loops: + raise ValueError( + f"Exceeded max_agentic_loops={max_loops} for model={model}" + ) + return fingerprint + + @staticmethod + def _fingerprint_agentic_tools(tools: Dict) -> str: + try: + return json.dumps(tools, sort_keys=True, default=str) + except Exception: + return str(tools) + + async def _execute_anthropic_agentic_plan( + self, + plan: AgenticLoopPlan, + model: str, + messages: List[Dict], + anthropic_messages_optional_request_params: Dict, + logging_obj: "LiteLLMLoggingObj", + kwargs: Dict, + depth: int, + max_loops: int, + fingerprints: List[str], + fingerprint: str, + stream: bool = False, + ) -> Any: + from litellm.anthropic_interface import messages as anthropic_messages + + patch = plan.request_patch or AgenticLoopRequestPatch() + if patch.messages is None: + raise ValueError("Agentic loop plan missing patched messages") + + full_model_name = model + if logging_obj is not None: + agentic_params = logging_obj.model_call_details.get( + "agentic_loop_params", {} + ) + full_model_name = cast(str, agentic_params.get("model", model)) + + optional_params = dict(anthropic_messages_optional_request_params) + optional_params.update(patch.optional_params) + if patch.tools is not None: + optional_params["tools"] = patch.tools + + max_tokens = patch.max_tokens + if max_tokens is None: + max_tokens = cast(Optional[int], optional_params.pop("max_tokens", None)) + else: + optional_params.pop("max_tokens", None) + if max_tokens is None: + max_tokens = cast(int, kwargs.get("max_tokens", 1024)) + + internal_keys = {"litellm_logging_obj"} + kwargs_for_followup = { + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") + and not k.startswith("_compression_interception") + and k not in internal_keys + and k not in optional_params + } + kwargs_for_followup.update(patch.kwargs) + kwargs_for_followup["_agentic_loop_depth"] = depth + 1 + kwargs_for_followup["max_agentic_loops"] = max_loops + kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] + + return await anthropic_messages.acreate( + **{ + "max_tokens": max_tokens, + "messages": patch.messages, + "model": patch.model or full_model_name, + "stream": stream, + **optional_params, + **kwargs_for_followup, + } + ) + + async def _execute_chat_completion_agentic_plan( + self, + plan: AgenticLoopPlan, + model: str, + messages: List[Dict], + optional_params: Dict, + kwargs: Dict, + custom_llm_provider: str, + depth: int, + max_loops: int, + fingerprints: List[str], + fingerprint: str, + ) -> Any: + patch = plan.request_patch or AgenticLoopRequestPatch() + if patch.messages is None: + raise ValueError("Agentic loop plan missing patched messages") + + full_model_name = patch.model or model + if "/" not in full_model_name: + full_model_name = f"{custom_llm_provider}/{full_model_name}" + + optional_params_for_followup = dict(optional_params) + optional_params_for_followup.update(patch.optional_params) + if patch.tools is not None: + optional_params_for_followup["tools"] = patch.tools + + internal_params = { + "_websearch_interception", + "acompletion", + "litellm_logging_obj", + "custom_llm_provider", + "model_alias_map", + "stream_response", + "custom_prompt_dict", + } + kwargs_for_followup = { + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") + and not k.startswith("_compression_interception") + and k not in internal_params + } + kwargs_for_followup.update(patch.kwargs) + kwargs_for_followup["_agentic_loop_depth"] = depth + 1 + kwargs_for_followup["max_agentic_loops"] = max_loops + kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] + + return await litellm.acompletion( + model=full_model_name, + messages=patch.messages, + **optional_params_for_followup, + **kwargs_for_followup, + ) + async def _call_agentic_completion_hooks( self, response: Any, @@ -4541,45 +4722,111 @@ class BaseLLMHTTPHandler: callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or []) tools = anthropic_messages_optional_request_params.get("tools", []) + depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs) for callback in callbacks: + if not isinstance(callback, CustomLogger): + continue + + should_run: bool = False + tool_calls: Any = None try: - if isinstance(callback, CustomLogger): - # First: Check if agentic loop should run - ( - should_run, - tool_calls, - ) = await callback.async_should_run_agentic_loop( - response=response, + # First: Check if agentic loop should run. Wrap in try/except + # to shield from buggy user callbacks — a callback crash should + # not abort the whole request. + ( + should_run, + tool_calls, + ) = await callback.async_should_run_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_should_run_agentic_loop [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + continue + + if not should_run: + continue + + # Safety guards must run OUTSIDE the callback try/except — they are + # bounded-loop / cycle-break rails that must propagate to the caller. + fingerprint = self._check_agentic_loop_safety( + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + ) + + try: + kwargs_with_provider = kwargs.copy() if kwargs else {} + kwargs_with_provider["custom_llm_provider"] = custom_llm_provider + build_plan_overridden = ( + callback.__class__.async_build_agentic_loop_plan + is not CustomLogger.async_build_agentic_loop_plan + ) + if not build_plan_overridden: + return await callback.async_run_agentic_loop( + tools=tool_calls, model=model, messages=messages, - tools=tools, + response=response, + anthropic_messages_provider_config=anthropic_messages_provider_config, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, stream=stream, - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, + kwargs=kwargs_with_provider, ) - if should_run: - # Second: Execute agentic loop - # Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name - kwargs_with_provider = kwargs.copy() if kwargs else {} - kwargs_with_provider["custom_llm_provider"] = ( - custom_llm_provider - ) - agentic_response = await callback.async_run_agentic_loop( - tools=tool_calls, - model=model, - messages=messages, - response=response, - anthropic_messages_provider_config=anthropic_messages_provider_config, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, - logging_obj=logging_obj, - stream=stream, - kwargs=kwargs_with_provider, - ) - # First hook that runs agentic loop wins - return agentic_response + plan = await callback.async_build_agentic_loop_plan( + tools=tool_calls, + model=model, + messages=messages, + response=response, + anthropic_messages_provider_config=anthropic_messages_provider_config, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs_with_provider, + ) + if plan.response_override is not None: + return plan.response_override + if plan.terminate: + verbose_logger.debug( + "Agentic loop terminated by callback=%s reason=%s", + callback.__class__.__name__, + plan.stop_reason, + ) + return response + if not plan.run_agentic_loop: + continue + + return await self._execute_anthropic_agentic_plan( + plan=plan, + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + kwargs=kwargs_with_provider, + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + stream=stream, + ) except Exception as e: _call_id = getattr(logging_obj, "litellm_call_id", "unknown") verbose_logger.exception( @@ -4653,52 +4900,104 @@ class BaseLLMHTTPHandler: callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or []) tools = optional_params.get("tools", []) + depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs) for callback in callbacks: - try: - if isinstance(callback, CustomLogger): - # Check if callback has the chat completion agentic loop method - if not hasattr( - callback, "async_should_run_chat_completion_agentic_loop" - ): - continue + if not isinstance(callback, CustomLogger): + continue + if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"): + continue - # First: Check if agentic loop should run - ( - should_run, - tool_calls, - ) = await callback.async_should_run_chat_completion_agentic_loop( - response=response, + should_run: bool = False + tool_calls: Any = None + try: + ( + should_run, + tool_calls, + ) = await callback.async_should_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_should_run_chat_completion_agentic_loop: %s", + str(e), + ) + continue + + if not should_run: + continue + + # Safety guards must run OUTSIDE the callback try/except — they are + # bounded-loop / cycle-break rails that must propagate to the caller. + fingerprint = self._check_agentic_loop_safety( + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + ) + + try: + kwargs_with_provider = kwargs.copy() if kwargs else {} + kwargs_with_provider["custom_llm_provider"] = custom_llm_provider + build_plan_overridden = ( + callback.__class__.async_build_chat_completion_agentic_loop_plan + is not CustomLogger.async_build_chat_completion_agentic_loop_plan + ) + if not build_plan_overridden: + return await callback.async_run_chat_completion_agentic_loop( + tools=tool_calls, model=model, messages=messages, - tools=tools, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, stream=stream, - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, + kwargs=kwargs_with_provider, ) - if should_run: - # Second: Execute agentic loop - # Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name - kwargs_with_provider = kwargs.copy() if kwargs else {} - kwargs_with_provider["custom_llm_provider"] = ( - custom_llm_provider - ) - agentic_response = ( - await callback.async_run_chat_completion_agentic_loop( - tools=tool_calls, - model=model, - messages=messages, - response=response, - optional_params=optional_params, - logging_obj=logging_obj, - stream=stream, - kwargs=kwargs_with_provider, - ) - ) - # First hook that runs agentic loop wins - return agentic_response + plan = await callback.async_build_chat_completion_agentic_loop_plan( + tools=tool_calls, + model=model, + messages=messages, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs_with_provider, + ) + if plan.response_override is not None: + return plan.response_override + if plan.terminate: + verbose_logger.debug( + "Agentic chat loop terminated by callback=%s reason=%s", + callback.__class__.__name__, + plan.stop_reason, + ) + return response + if not plan.run_agentic_loop: + continue + + return await self._execute_chat_completion_agentic_plan( + plan=plan, + model=model, + messages=messages, + optional_params=optional_params, + kwargs=kwargs_with_provider, + custom_llm_provider=custom_llm_provider, + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + ) except Exception as e: verbose_logger.exception( f"LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: {str(e)}" diff --git a/litellm/llms/github_copilot/authenticator.py b/litellm/llms/github_copilot/authenticator.py index f4698861edc..9de2987b9f6 100644 --- a/litellm/llms/github_copilot/authenticator.py +++ b/litellm/llms/github_copilot/authenticator.py @@ -294,9 +294,7 @@ class Authenticator: access_token_url = os.getenv( "GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL ) - client_id = os.getenv( - "GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID - ) + client_id = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) for attempt in range(max_attempts): try: diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 2db19dea0b9..86ca6625629 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -48,6 +48,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ + def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]: + """ + Convert chat completions request data to OpenAI-spec structured messages. + + Messages are already in OpenAI format, so this is a simple extraction. + """ + messages = data.get("messages") + if messages is None: + return None + return cast(List[AllMessageValues], messages) + async def process_input_messages( self, data: dict, @@ -68,9 +79,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tool_calls_to_check: List[ChatCompletionToolParam] = [] text_task_mappings: List[Tuple[int, Optional[int]]] = [] tool_call_task_mappings: List[Tuple[int, int]] = [] - # text_task_mappings: Track (message_index, content_index) for each text - # content_index is None for string content, int for list content - # tool_call_task_mappings: Track (message_index, tool_call_index) for each tool call # Step 1: Extract all text content, images, and tool calls for msg_idx, message in enumerate(messages): @@ -92,12 +100,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs["images"] = images_to_check if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check # type: ignore - if messages: - msg_list = cast(List[AllMessageValues], messages) + structured_messages = self.get_structured_messages(data) + if structured_messages: inputs["structured_messages"] = ( - openai_messages_without_system(msg_list) + openai_messages_without_system(structured_messages) if skip_system - else msg_list + else structured_messages ) # Pass tools (function definitions) to the guardrail tools = data.get("tools") diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 76f40eed71f..f7dd68aec55 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -43,6 +43,7 @@ from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) from litellm.types.llms.openai import ( + AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam, ) @@ -70,6 +71,24 @@ class OpenAIResponsesHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ + def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]: + """ + Convert Responses API request data to OpenAI-spec structured messages. + + Transforms `input` (string or ResponseInputParam) and optional + `instructions` into chat completion messages. + """ + input_data = data.get("input") + if input_data is None: + return None + messages = ( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=input_data, + responses_api_request=data, + ) + ) + return cast(List[AllMessageValues], messages) if messages else None + async def process_input_messages( self, data: dict, @@ -86,12 +105,7 @@ class OpenAIResponsesHandler(BaseTranslation): if input_data is None: return data - structured_messages = ( - LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=input_data, - responses_api_request=data, - ) - ) + structured_messages = self.get_structured_messages(data) # Handle simple string input if isinstance(input_data, str): diff --git a/litellm/llms/scaleway/audio_transcription/transformation.py b/litellm/llms/scaleway/audio_transcription/transformation.py new file mode 100644 index 00000000000..b45f287afb4 --- /dev/null +++ b/litellm/llms/scaleway/audio_transcription/transformation.py @@ -0,0 +1,158 @@ +""" +Support for Scaleway's OpenAI-compatible `/v1/audio/transcriptions` endpoint. + +API reference: https://www.scaleway.com/en/developers/api/generative-apis/#path-audio-create-an-audio-transcription +""" + +from typing import List, Optional, Union + +import httpx + +from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.llms.base_llm.audio_transcription.transformation import ( + AudioTranscriptionRequestData, + BaseAudioTranscriptionConfig, +) +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ( + AllMessageValues, + OpenAIAudioTranscriptionOptionalParams, +) +from litellm.types.utils import FileTypes, TranscriptionResponse + + +class ScalewayAudioTranscriptionException(BaseLLMException): + pass + + +class ScalewayAudioTranscriptionConfig(BaseAudioTranscriptionConfig): + def get_supported_openai_params( + self, model: str + ) -> List[OpenAIAudioTranscriptionOptionalParams]: + return [ + "language", + "prompt", + "response_format", + "temperature", + "timestamp_granularities", + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + supported_params = self.get_supported_openai_params(model) + for k, v in non_default_params.items(): + if k in supported_params: + optional_params[k] = v + return optional_params + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + api_base = ( + "https://api.scaleway.ai/v1" if api_base is None else api_base.rstrip("/") + ) + return f"{api_base}/audio/transcriptions" + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] + ) -> BaseLLMException: + return ScalewayAudioTranscriptionException( + message=error_message, + status_code=status_code, + headers=headers, + ) + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + if api_key is None: + api_key = get_secret_str("SCW_SECRET_KEY") + + if not api_key: + raise ScalewayAudioTranscriptionException( + message=( + "Scaleway API key not found. Pass `api_key=...` or set the " + "SCW_SECRET_KEY environment variable." + ), + status_code=401, + headers={}, + ) + + default_headers = { + "Authorization": f"Bearer {api_key}", + "accept": "application/json", + } + default_headers.update(headers or {}) + return default_headers + + def transform_audio_transcription_request( + self, + model: str, + audio_file: FileTypes, + optional_params: dict, + litellm_params: dict, + ) -> AudioTranscriptionRequestData: + processed_audio = process_audio_file(audio_file) + + form_fields: dict = {"model": model} + for key in self.get_supported_openai_params(model): + value = optional_params.get(key) + if value is not None: + form_fields[key] = value + + files = { + "file": ( + processed_audio.filename, + processed_audio.file_content, + processed_audio.content_type, + ) + } + + return AudioTranscriptionRequestData(data=form_fields, files=files) + + def transform_audio_transcription_response( + self, + raw_response: httpx.Response, + ) -> TranscriptionResponse: + content_type = (raw_response.headers.get("content-type") or "").lower() + if "application/json" not in content_type: + return TranscriptionResponse(text=raw_response.text) + + try: + response_json = raw_response.json() + except Exception: + raise ScalewayAudioTranscriptionException( + message=raw_response.text, + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + text = response_json.get("text") or "" + response = TranscriptionResponse(text=text) + + if "segments" in response_json: + response["segments"] = response_json["segments"] + if "language" in response_json: + response["language"] = response_json["language"] + + response._hidden_params = response_json + return response diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 04b68b8f4ec..640607c0748 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22872,6 +22872,22 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k2.6": { + "cache_read_input_token_cost": 1.6e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "moonshot", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://platform.kimi.ai/docs/pricing/chat-k26", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, "moonshot/kimi-latest": { "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 2e-06, diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index 5dde13f0078..d39a0dda152 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -79,7 +79,9 @@ class BasePassthroughUtils: for header_name, header_value in request_headers.items(): if header_name.lower().startswith(PASS_THROUGH_HEADER_PREFIX): # Strip the 'x-pass-' prefix and normalize to lowercase - actual_header_name = header_name[len(PASS_THROUGH_HEADER_PREFIX) :].lower() + actual_header_name = header_name[ + len(PASS_THROUGH_HEADER_PREFIX) : + ].lower() if actual_header_name in _PASS_THROUGH_PROTECTED_HEADERS or any( actual_header_name.startswith(p) for p in _PASS_THROUGH_PROTECTED_HEADER_PREFIXES diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index ed54c707b00..0562b41d2cd 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -1950,7 +1950,7 @@ "responses": true, "embeddings": false, "image_generations": false, - "audio_transcriptions": false, + "audio_transcriptions": true, "audio_speech": false, "moderations": false, "batches": false, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 65e2e3e983d..792a9dace1e 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -323,6 +323,14 @@ async def authorize_with_server( ) parsed = urlparse(redirect_uri) + if parsed.scheme not in ("http", "https"): + raise HTTPException( + status_code=400, + detail={ + "error": "invalid_redirect_uri", + "message": "redirect_uri must use http or https scheme", + }, + ) base_url = urlunparse(parsed._replace(query="")) request_base_url = get_request_base_url(request) encoded_state = encode_state_with_base_url( diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 604e7d5f418..703fe6adc41 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,32 +1,83 @@ +# model_list: +# - model_name: claude-sonnet-4-6 +# litellm_params: {model: anthropic/claude-sonnet-4-6} +# model_info: +# litellm_routing_preferences: +# quality_tier: 1 +# keywords: [tin] +# - model_name: gpt-4o-mini +# litellm_params: {model: openai/gpt-4o-mini} +# model_info: +# litellm_routing_preferences: +# quality_tier: 1 +# keywords: [] +# - model_name: gpt-4o +# litellm_params: {model: openai/gpt-4o} +# model_info: +# litellm_routing_preferences: +# quality_tier: 2 +# keywords: [vision, function_calling] +# - model_name: opus +# litellm_params: {model: anthropic/claude-opus-4-7} +# model_info: +# litellm_routing_preferences: +# quality_tier: 3 +# keywords: ["architecture", "design"] +# - model_name: my-quality-router +# litellm_params: +# model: auto_router/adaptive_router +# adaptive_router_default_model: gpt-4o-mini +# adaptive_router_config: +# available_models: [gpt-4o-mini, gpt-4o, opus, claude-sonnet-4-6] +# Example proxy config for the adaptive router (v0). +# +# Wires one logical router ("smart-cheap-router") that adaptively picks between +# two real deployments ("fast" and "smart") based on per-session feedback signals. +# +# How to use from a client: +# POST /v1/chat/completions { "model": "smart-cheap-router", ... } +# Add { "metadata": { "litellm_session_id": "" } } to enable +# sticky-session routing within a conversation. +# +# Required env vars: OPENAI_API_KEY, DATABASE_URL. + model_list: - - # OpenAI model for /v1/chat/completions test — 200x custom pricing - - model_name: "gpt-4.1-mini" + # ---- The adaptive router "control" deployment ------------------------- + # `model_name` is what clients call. `available_models` lists the underlying + # deployments the router is allowed to pick from (must match other model_name + # entries in this list). + - model_name: smart-cheap-router litellm_params: - model: openai/gpt-4.1-mini - api_key: os.environ/OPENAI_API_KEY - model_info: - id: gpt-4.1-mini-custom-pricing - input_cost_per_token: 0.00004 # 100x standard ($0.40/1M = $0.0000004) - output_cost_per_token: 0.00016 # 100x standard ($1.60/1M = $0.0000016) + model: auto_router/adaptive_router + adaptive_router_config: + available_models: ["fast", "smart"] + weights: + quality: 0.7 + cost: 0.3 - # OpenAI model for /v1/responses test — 100x custom pricing - - model_name: "gpt-5" + # ---- Underlying deployments the router picks from --------------------- + - model_name: fast litellm_params: - model: openai/gpt-5 - api_key: os.environ/OPENAI_API_KEY - model_info: - id: gpt-5-custom-pricing - mode: "chat" - input_cost_per_token: 125 # 100x standard ($1.25/1M = $0.00000125) - output_cost_per_token: 10 # 100x standard ($10.00/1M = $0.00001) - - # Anthropic model for /v1/messages test — 100x custom pricing - - model_name: "claude-sonnet-4-20250514" - litellm_params: - model: anthropic/claude-sonnet-4-20250514 + model: anthropic/claude-sonnet-4-6 api_key: os.environ/ANTHROPIC_API_KEY + input_cost_per_token: 0.00000015 model_info: - id: claude-sonnet-4-custom-pricing - input_cost_per_token: 0.0003 # 100x standard ($0.000003) - output_cost_per_token: 0.0015 # 100x standard ($0.000015) \ No newline at end of file + adaptive_router_preferences: + quality_tier: 2 + strengths: [] + + - model_name: smart + litellm_params: + model: anthropic/claude-opus-4-7 + api_key: os.environ/ANTHROPIC_API_KEY + input_cost_per_token: 0.0000050 + model_info: + adaptive_router_preferences: + quality_tier: 3 + strengths: ["code_generation", "technical_design", "analytical_reasoning"] + +litellm_settings: + drop_params: True + +general_settings: + master_key: sk-1234 # REPLACE in production diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index d245ec53ece..2c8299e77a9 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -626,11 +626,17 @@ async def common_checks( # noqa: PLR0915 and user_object.max_budget is not None ): user_budget = user_object.max_budget - if user_budget < user_object.spend: + from litellm.proxy.proxy_server import get_current_spend + + user_spend = await get_current_spend( + counter_key=f"spend:user:{user_object.user_id}", + fallback_spend=user_object.spend or 0.0, + ) + if user_spend >= user_budget: raise litellm.BudgetExceededError( - current_cost=user_object.spend, + current_cost=user_spend, max_budget=user_budget, - message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}", + message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", ) ## 4.2 check team member budget, if team key @@ -3126,9 +3132,7 @@ async def _virtual_key_max_budget_alert_check( alert_email_config: Optional[Dict[str, List[str]]] = ( _merge_budget_alert_email_configs( global_cfg=litellm.default_key_max_budget_alert_emails, - per_key_cfg=(valid_token.metadata or {}).get( - "max_budget_alert_emails" - ), + per_key_cfg=(valid_token.metadata or {}).get("max_budget_alert_emails"), ) ) @@ -3138,7 +3142,9 @@ async def _virtual_key_max_budget_alert_check( (int(k) for k in alert_email_config if k.isdigit()), default=None, ) - if min_pct is None or valid_token.spend < valid_token.max_budget * (min_pct / 100.0): + if min_pct is None or valid_token.spend < valid_token.max_budget * ( + min_pct / 100.0 + ): return call_info = CallInfo( @@ -3164,8 +3170,7 @@ async def _virtual_key_max_budget_alert_check( else: # Old path: existing single 80% threshold — completely unchanged alert_threshold = ( - valid_token.max_budget - * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE ) if ( @@ -3666,12 +3671,20 @@ async def _organization_max_budget_check( if org_max_budget is None or org_max_budget <= 0: return + # Read spend from cross-pod counter (Redis-first) or cached object (fallback) + from litellm.proxy.proxy_server import get_current_spend + + org_spend = await get_current_spend( + counter_key=f"spend:org:{org_id}", + fallback_spend=org_table.spend or 0.0, + ) + # Check if organization spend exceeds max budget - if org_table.spend >= org_max_budget: + if org_spend >= org_max_budget: # Trigger budget alert call_info = CallInfo( token=valid_token.token, - spend=org_table.spend, + spend=org_spend, max_budget=org_max_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, @@ -3687,9 +3700,9 @@ async def _organization_max_budget_check( ) raise litellm.BudgetExceededError( - current_cost=org_table.spend, + current_cost=org_spend, max_budget=org_max_budget, - message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_table.spend}, Max budget: {org_max_budget}", + message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}", ) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index a206be87a11..e31c76dcac1 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -37,6 +37,20 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 if isinstance(value, list): imported_list: List[Any] = [] for callback in value: # ["presidio", ] + if isinstance(callback, str) and callback == "compression_interception": + from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, + ) + + compression_interception_obj = ( + CompressionInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params, + ) + ) + imported_list.append(compression_interception_obj) + continue + # check if callback is a custom logger compatible callback if isinstance(callback, str): callback = LoggingCallbackManager._add_custom_callback_generic_api_str( diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 114103508ea..73735796eb3 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -403,10 +403,18 @@ class PrismaManager: return dname @staticmethod - def setup_database(use_migrate: bool = False) -> bool: + def setup_database( + use_migrate: bool = False, use_v2_resolver: bool = False + ) -> bool: """ Set up the database using either prisma migrate or prisma db push + Args: + use_migrate: Use `prisma migrate deploy` instead of `db push`. + use_v2_resolver: Opt into the v2 migration resolver that avoids + the diff-and-force recovery behavior (which caused schema + thrashing during rolling deploys). Defaults to False. + Returns: bool: True if setup was successful, False otherwise """ @@ -427,7 +435,10 @@ class PrismaManager: prisma_dir = PrismaManager._get_prisma_dir() - return ProxyExtrasDBManager.setup_database(use_migrate=use_migrate) + return ProxyExtrasDBManager.setup_database( + use_migrate=use_migrate, + use_v2_resolver=use_v2_resolver, + ) else: # Use prisma db push with increased timeout subprocess.run( diff --git a/litellm/proxy/example_config_yaml/adaptive_router_example.yaml b/litellm/proxy/example_config_yaml/adaptive_router_example.yaml new file mode 100644 index 00000000000..58f5398ca57 --- /dev/null +++ b/litellm/proxy/example_config_yaml/adaptive_router_example.yaml @@ -0,0 +1,52 @@ +# Example proxy config for the adaptive router (v0). +# +# Wires one logical router ("smart-cheap-router") that adaptively picks between +# two real deployments ("fast" and "smart") based on per-session feedback signals. +# +# How to use from a client: +# POST /v1/chat/completions { "model": "smart-cheap-router", ... } +# Add { "metadata": { "litellm_session_id": "" } } to enable +# sticky-session routing within a conversation. +# +# Required env vars: OPENAI_API_KEY, DATABASE_URL. + +model_list: + # ---- The adaptive router "control" deployment ------------------------- + # `model_name` is what clients call. `available_models` lists the underlying + # deployments the router is allowed to pick from (must match other model_name + # entries in this list). + - model_name: smart-cheap-router + litellm_params: + model: auto_router/adaptive_router # required prefix -- triggers adaptive-router init + adaptive_router_config: + available_models: ["fast", "smart"] + weights: + quality: 0.7 + cost: 0.3 + + # ---- Underlying deployments the router picks from --------------------- + - model_name: fast + litellm_params: + model: openai/gpt-4o-mini + api_key: os.environ/OPENAI_API_KEY + input_cost_per_token: 0.00000015 + model_info: + adaptive_router_preferences: + quality_tier: 2 + strengths: [] + + - model_name: smart + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + input_cost_per_token: 0.0000050 + model_info: + adaptive_router_preferences: + quality_tier: 3 + strengths: ["code_generation", "technical_design", "analytical_reasoning"] + +litellm_settings: + drop_params: True + +general_settings: + master_key: sk-1234 # REPLACE in production diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 25b46cf3641..77b2f466f2a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -62,6 +62,7 @@ from litellm.types.utils import ( CallTypesLiteral, Choices, GuardrailStatus, + Message, ModelResponse, ModelResponseStream, StreamingChoices, @@ -1563,11 +1564,43 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Bedrock will throw an error if there is no text to process if filtered_messages: - bedrock_response = await self.make_bedrock_api_request( - source="INPUT", - messages=filtered_messages, - request_data=request_data, + # Map the abstract input_type to the Bedrock source parameter. + # "request" -> INPUT (scan user-supplied content) + # "response" -> OUTPUT (scan model-generated content) + # Bedrock guardrail policies are often configured differently + # for Input vs Output (e.g. PII blocking only on Output), so + # the source MUST match where the text originated. + bedrock_source: Literal["INPUT", "OUTPUT"] = ( + "OUTPUT" if input_type == "response" else "INPUT" ) + if bedrock_source == "OUTPUT": + # Build a synthetic ModelResponse whose choices carry the + # text(s) to scan, so _create_bedrock_output_content_request + # can produce the correct Bedrock OUTPUT payload. + synthetic_response = ModelResponse( + choices=[ + Choices( + index=_idx, + message=Message( + role="assistant", + content=str(_msg.get("content") or ""), + ), + finish_reason="stop", + ) + for _idx, _msg in enumerate(filtered_messages) + ] + ) + bedrock_response = await self.make_bedrock_api_request( + source="OUTPUT", + response=synthetic_response, + request_data=request_data, + ) + else: + bedrock_response = await self.make_bedrock_api_request( + source="INPUT", + messages=filtered_messages, + request_data=request_data, + ) # Apply any masking that was applied by the guardrail output_list = bedrock_response.get("output") diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index e0664703d28..7d67750c78f 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -306,7 +306,9 @@ def _health_check_deployment_is_wildcard(litellm_params: dict) -> bool: return "*" in _deployment_model_string_for_health_check(litellm_params) -def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) -> Optional[int]: +def _resolve_health_check_max_tokens( + model_info: dict, litellm_params: dict +) -> Optional[int]: """ Pick max_tokens for the health check request. @@ -341,10 +343,7 @@ def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) -> return int(tokens_reasoning) if not is_reasoning and tokens_non_reasoning is not None: return int(tokens_non_reasoning) - if ( - is_reasoning - and BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING is not None - ): + if is_reasoning and BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING is not None: return int(BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING) if BACKGROUND_HEALTH_CHECK_MAX_TOKENS is not None: diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 8fd19548cbb..b4b5de1746e 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1121,14 +1121,14 @@ async def _db_health_readiness_check(): return db_health_cache except Exception as e: db_health_cache = {"status": "disconnected", "last_updated": datetime.now()} - PrismaDBExceptionHandler.handle_db_exception(e) if PrismaDBExceptionHandler.is_database_transport_error(e): try: verbose_proxy_logger.warning( "_db_health_readiness_check: health_check failed, attempting reconnect" ) - await prisma_client.disconnect() - await prisma_client.connect() + await prisma_client.attempt_db_reconnect( + reason="health_readiness_check" + ) await prisma_client.health_check() verbose_proxy_logger.info( "_db_health_readiness_check: reconnect succeeded" diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 4b59f603d3e..7789fa6a349 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -21,20 +21,30 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): ): try: verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook") - cache_key = f"{user_api_key_dict.user_id}_user_api_key_user_id" - user_row = await cache.async_get_cache( - cache_key, parent_otel_span=user_api_key_dict.parent_otel_span + max_budget = user_api_key_dict.user_max_budget + user_id = user_api_key_dict.user_id + + if max_budget is None or user_id is None: + return + + # Personal budget applies only to non-team requests, matching + # the explicit team-key exemption in common_checks section 4.1. + if user_api_key_dict.team_id is not None: + return + + from litellm.proxy.proxy_server import get_current_spend + + curr_spend = await get_current_spend( + counter_key=f"spend:user:{user_id}", + fallback_spend=user_api_key_dict.user_spend or 0.0, ) - if user_row is None: # value not yet cached - return - max_budget = user_row["max_budget"] - curr_spend = user_row["spend"] - if max_budget is None: - return - - if curr_spend is None: - return + verbose_proxy_logger.debug( + "MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f", + user_id, + curr_spend, + max_budget, + ) # CHECK IF REQUEST ALLOWED if curr_spend >= max_budget: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 5c2b3dfe0ee..f29bbd2d9d5 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1570,9 +1570,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): user_api_key_project_id = standard_logging_metadata.get( "user_api_key_project_id" ) - user_api_key_end_user_id = kwargs.get( - "user" - ) or standard_logging_metadata.get("user_api_key_end_user_id") + user_api_key_end_user_id = kwargs.get("user") or standard_logging_metadata.get( + "user_api_key_end_user_id" + ) model_group = get_model_group_from_litellm_kwargs(kwargs) # Get total tokens from response diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index ea9c92fec6c..c9946f4e26f 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -213,6 +213,7 @@ class _ProxyDBLogger(CustomLogger): team_id=team_id, user_id=user_id, response_cost=response_cost, + org_id=org_id, ) # update cache (fire-and-forget for backward compat: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 7467bbae232..5804e3f8d9f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1,5 +1,6 @@ import asyncio import copy +import re import time from collections import OrderedDict from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union @@ -28,6 +29,14 @@ _SPECIAL_HEADERS_CACHE = frozenset( v.value.lower() for v in SpecialHeaders._member_map_.values() ) +# Matches any header of the form x--session-id (case-insensitive). +# Excludes the two explicit litellm headers which are handled with higher priority. +_GENERIC_SESSION_ID_HEADER_RE = re.compile(r"^x-.+-session-id$", re.IGNORECASE) +_EXPLICIT_SESSION_HEADERS = frozenset({"x-litellm-trace-id", "x-litellm-session-id"}) +# Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores +# (covers UUIDs and most common session-id formats). +_SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") + def _sanitize_for_log(value: Any) -> str: """ @@ -115,13 +124,43 @@ def _get_metadata_variable_name(request: Request) -> str: return "metadata" +def _extract_generic_session_id_from_headers( + normalized: Dict[str, str], +) -> Optional[str]: + """ + Scan a normalised (lower-cased keys) header dict for any header that looks + like ``x--session-id`` and whose value is a plausible session/trace + identifier (alphanumeric + hyphens/underscores, at least 8 chars). + + The two explicit LiteLLM headers (``x-litellm-trace-id`` / + ``x-litellm-session-id``) are excluded here because they are handled with + higher priority by the caller. + + Example: ``x-claude-code-session-id: e96634a3-fa28-4083-b354-55542e2dca01`` + """ + for key, value in normalized.items(): + if ( + key not in _EXPLICIT_SESSION_HEADERS + and _GENERIC_SESSION_ID_HEADER_RE.match(key) + and isinstance(value, str) + and _SESSION_ID_VALUE_RE.match(value) + ): + return value + return None + + def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str]: """ Extract chain id for call chaining from request headers. - x-litellm-trace-id and x-litellm-session-id are interchangeable; when both - are present, x-litellm-trace-id takes precedence. Header keys are matched - case-insensitively so this works with raw header dicts from any transport. + Priority order: + 1. ``x-litellm-trace-id`` (explicit, highest priority) + 2. ``x-litellm-session-id`` (explicit) + 3. Any ``x--session-id`` header whose value looks like a session id + (alphanumeric / UUID, at least 8 chars). E.g. ``x-claude-code-session-id``. + + Header keys are matched case-insensitively so this works with raw header + dicts from any transport. Used by MCP (and other paths that have raw_headers but no Request) to set litellm_trace_id/litellm_session_id for spend logs and logging consistency. @@ -129,8 +168,10 @@ def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str if not headers: return None normalized = {k.lower(): v for k, v in headers.items() if isinstance(k, str)} - return normalized.get("x-litellm-trace-id") or normalized.get( - "x-litellm-session-id" + return ( + normalized.get("x-litellm-trace-id") + or normalized.get("x-litellm-session-id") + or _extract_generic_session_id_from_headers(normalized) ) @@ -649,10 +690,8 @@ class LiteLLMProxyRequestSetup: ######################################################################################### agent_id_from_header = headers.get("x-litellm-agent-id") - # x-litellm-trace-id and x-litellm-session-id are interchangeable for call chaining - chain_id = headers.get("x-litellm-trace-id") or headers.get( - "x-litellm-session-id" - ) + # Explicit litellm headers take precedence; fall back to any x-*-session-id header. + chain_id = get_chain_id_from_headers(dict(headers)) if agent_id_from_header: metadata_from_headers["agent_id"] = agent_id_from_header diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 07286b4fa80..b0ea6b41ac5 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -355,6 +355,7 @@ async def _upsert_budget_and_membership( tpm_limit: Optional[int] = None, rpm_limit: Optional[int] = None, allowed_models: Optional[List[str]] = None, + team_default_budget_id: Optional[str] = None, ): """ Helper function to Create/Update or Delete the budget within the team membership @@ -368,6 +369,11 @@ async def _upsert_budget_and_membership( tpm_limit: Tokens per minute limit for the team member rpm_limit: Requests per minute limit for the team member allowed_models: Per-member model scope. None = don't change. [] = remove restrictions. Non-empty list = enforce. + team_default_budget_id: The team's shared default member budget id (from + team metadata.team_member_budget_id), if any. When the membership's + existing_budget_id matches this, we clone-on-write so editing one + member's budget does not mutate the shared default (and therefore + every other member who still points at it). If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, the user's budget is removed from the team membership. If any of these values exist, a budget is updated or created and linked to the team membership. @@ -385,7 +391,13 @@ async def _upsert_budget_and_membership( ) return - if existing_budget_id is not None: + is_shared_default = ( + existing_budget_id is not None + and team_default_budget_id is not None + and existing_budget_id == team_default_budget_id + ) + + if existing_budget_id is not None and not is_shared_default: # Update the existing budget in-place to preserve fields not being changed. # Only write fields that the caller explicitly provided (non-None). update_data: Dict[str, Any] = { @@ -405,11 +417,40 @@ async def _upsert_budget_and_membership( ) return - # No existing budget — create a new one and link it to the membership. + # Either there is no existing budget, OR the membership is still pointing + # at the team's shared default member budget. In both cases we create a + # NEW private budget for this user and (re)link the membership to it. create_data: Dict[str, Any] = { "created_by": user_api_key_dict.user_id or "", "updated_by": user_api_key_dict.user_id or "", } + + # If we're forking off the shared default, seed the new row with the + # default's values so fields the caller did not change carry over. + if is_shared_default: + default_budget_row = await tx.litellm_budgettable.find_unique( + where={"budget_id": existing_budget_id} + ) + if default_budget_row is not None: + default_budget_dict = default_budget_row.model_dump() + for field in ( + "max_budget", + "soft_budget", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + "model_max_budget", + "budget_duration", + "allowed_models", + ): + value = default_budget_dict.get(field) + if value is None: + continue + if isinstance(value, list) and len(value) == 0: + continue + create_data[field] = value + + # Caller-provided values take precedence over the cloned defaults. if max_budget is not None: create_data["max_budget"] = max_budget if tpm_limit is not None: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 8474f026111..c6d37ace4fe 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -2120,9 +2120,7 @@ async def delete_user( for m in all_target_memberships: if not m.organization_id: continue - target_org_ids_by_user.setdefault(m.user_id, set()).add( - m.organization_id - ) + target_org_ids_by_user.setdefault(m.user_id, set()).add(m.organization_id) # check that all teams passed exist for user_id in data.user_ids: @@ -2141,9 +2139,7 @@ async def delete_user( # Org-admin may only delete users whose entire org membership is # within their admin scope. A target with ANY org outside the # caller's scope (or no org at all) requires PROXY_ADMIN. - if not target_org_ids or not target_org_ids.issubset( - caller_admin_org_ids - ): + if not target_org_ids or not target_org_ids.issubset(caller_admin_org_ids): raise HTTPException( status_code=403, detail={ diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 56bbfe03005..f18e699045f 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1336,7 +1336,9 @@ if MCP_AVAILABLE: return _redact_mcp_credentials(temp_record) - def _get_cached_temporary_mcp_server_or_404(server_id: str) -> MCPServer: + def _get_cached_temporary_mcp_server_or_404( + server_id: str, request: Optional[Request] = None + ) -> MCPServer: server = get_cached_temporary_mcp_server(server_id) if server is None: # Fall back to real DB/config server (e.g. for the user-side OAuth flow @@ -1344,10 +1346,14 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy.auth.ip_address_utils import IPAddressUtils + client_ip = IPAddressUtils.get_mcp_client_ip(request) if request else None server = global_mcp_server_manager.get_mcp_server_by_id( server_id - ) or global_mcp_server_manager.get_mcp_server_by_name(server_id) + ) or global_mcp_server_manager.get_mcp_server_by_name( + server_id, client_ip=client_ip + ) if server is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1358,10 +1364,12 @@ if MCP_AVAILABLE: @router.get( "/server/oauth/{server_id}/authorize", include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], ) async def mcp_authorize( request: Request, server_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), client_id: Optional[str] = None, redirect_uri: str = Query(...), state: str = "", @@ -1370,7 +1378,7 @@ if MCP_AVAILABLE: response_type: Optional[str] = None, scope: Optional[str] = None, ): - mcp_server = _get_cached_temporary_mcp_server_or_404(server_id) + mcp_server = _get_cached_temporary_mcp_server_or_404(server_id, request=request) # Use the server's stored client_id when the caller doesn't supply one resolved_client_id = mcp_server.client_id or client_id or "" if not resolved_client_id: @@ -1399,10 +1407,12 @@ if MCP_AVAILABLE: @router.post( "/server/oauth/{server_id}/token", include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], ) async def mcp_token( request: Request, server_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), grant_type: str = Form(...), code: Optional[str] = Form(None), redirect_uri: Optional[str] = Form(None), @@ -1412,7 +1422,7 @@ if MCP_AVAILABLE: refresh_token: Optional[str] = Form(None), scope: Optional[str] = Form(None), ): - mcp_server = _get_cached_temporary_mcp_server_or_404(server_id) + mcp_server = _get_cached_temporary_mcp_server_or_404(server_id, request=request) resolved_client_id = mcp_server.client_id or client_id or "" if not resolved_client_id: raise HTTPException( @@ -1441,9 +1451,14 @@ if MCP_AVAILABLE: @router.post( "/server/oauth/{server_id}/register", include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], ) - async def mcp_register(request: Request, server_id: str): - mcp_server = _get_cached_temporary_mcp_server_or_404(server_id) + async def mcp_register( + request: Request, + server_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + mcp_server = _get_cached_temporary_mcp_server_or_404(server_id, request=request) request_data = await _read_request_body(request=request) data: dict = {**request_data} diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index a6a1af971e5..442fae2a4fa 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -1078,10 +1078,7 @@ async def organization_member_update( LitellmUserRoles.PROXY_ADMIN.value, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, ): - if ( - user_api_key_dict.user_role - != LitellmUserRoles.PROXY_ADMIN.value - ): + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=403, detail={ diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 8e21b851857..bf912fba4f8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1570,8 +1570,7 @@ async def update_team( # noqa: PLR0915 current_org_id = getattr(existing_team_row, "organization_id", None) if ( data.organization_id != current_org_id - and user_api_key_dict.user_role - != LitellmUserRoles.PROXY_ADMIN.value + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value ): # Is the caller org_admin of the destination org? caller_memberships = ( @@ -2609,6 +2608,15 @@ async def team_member_update( identified_budget_id = tm.budget_id break + # If this membership still points at the team's shared default member + # budget, _upsert_budget_and_membership will clone-on-write so that the + # update only touches this user (not every member sharing the default). + team_default_budget_id: Optional[str] = None + if team_table.metadata is not None: + raw_default_budget_id = team_table.metadata.get("team_member_budget_id") + if isinstance(raw_default_budget_id, str): + team_default_budget_id = raw_default_budget_id + ### upsert new budget async with prisma_client.db.tx() as tx: await _upsert_budget_and_membership( @@ -2621,6 +2629,7 @@ async def team_member_update( tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, allowed_models=data.allowed_models, + team_default_budget_id=team_default_budget_id, ) ### update team member role diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 3e42d392077..5cf53ae06f5 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -140,6 +140,62 @@ async def handle_budget_for_entity( return existing_budget_id +# Fields on LiteLLM_BudgetTable that represent the budget's *configuration* +# (i.e. the values an admin sets). We copy these when cloning a team's +# default member-budget into an individual member-budget so that the new +# row starts with the same limits as the default. +_CLONABLE_BUDGET_FIELDS: Tuple[str, ...] = ( + "max_budget", + "soft_budget", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + "model_max_budget", + "budget_duration", + "allowed_models", +) + + +async def _clone_team_default_budget_for_member( + prisma_client: PrismaClient, + default_team_budget_id: str, + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str, +) -> Optional[str]: + """ + Create a new budget row that copies the values from the team's default + member budget. Returns the new budget_id, or None if the default budget + no longer exists in the DB. + + Used when adding a new team member without an explicit per-member budget, + so the member starts with the team default's values but gets their own + private budget row (which can be edited independently). + """ + default_budget = await prisma_client.db.litellm_budgettable.find_unique( + where={"budget_id": default_team_budget_id} + ) + if default_budget is None: + return None + + default_budget_dict = default_budget.model_dump() + cloned_data: dict = { + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } + for field in _CLONABLE_BUDGET_FIELDS: + value = default_budget_dict.get(field) + if value is None: + continue + # Skip empty list defaults (e.g. allowed_models = []) so the cloned + # row matches the "no value set" shape rather than carrying a default. + if isinstance(value, list) and len(value) == 0: + continue + cloned_data[field] = value + + new_budget = await prisma_client.db.litellm_budgettable.create(data=cloned_data) + return new_budget.budget_id + + async def add_new_member( new_member: Member, max_budget_in_team: Optional[float], @@ -221,8 +277,20 @@ async def add_new_member( response = await prisma_client.db.litellm_budgettable.create(data=budget_data) _budget_id = response.budget_id + elif default_team_budget_id is not None: + # No per-member budget was provided, but the team has a default member + # budget. Clone the default budget into a new row for this user so that + # later edits to one member's budget do not bleed into other members. + # If the default no longer exists in the DB, fall back to no budget. + _budget_id = await _clone_team_default_budget_for_member( + prisma_client=prisma_client, + default_team_budget_id=default_team_budget_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) else: - _budget_id = default_team_budget_id + # No per-member budget and no team default → member gets no budget. + _budget_id = None if _budget_id and returned_user is not None and returned_user.user_id is not None: _returned_team_membership = ( diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index c638e294268..3845203bb9d 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -577,6 +577,22 @@ class ProxyInitializationHelpers: help="Exit with error if database migration fails on startup.", envvar="ENFORCE_PRISMA_MIGRATION_CHECK", ) +@click.option( + "--use_v2_migration_resolver", + is_flag=True, + default=False, + help=( + "Opt into the v2 migration resolver. Avoids the diff-and-force recovery " + "path that can cause schema thrashing during rolling deploys where two " + "LiteLLM versions contend for the same DB. Default is the v1 resolver." + ), +) +@click.option( + "--reload", + is_flag=True, + default=False, + help="Enable uvicorn hot reload (dev only). Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.", +) def run_server( # noqa: PLR0915 host, port, @@ -618,6 +634,8 @@ def run_server( # noqa: PLR0915 keepalive_timeout, max_requests_before_restart, enforce_prisma_migration_check: bool, + use_v2_migration_resolver: bool, + reload: bool, ): if setup: from litellm.setup_wizard import run_setup_wizard @@ -886,9 +904,31 @@ def run_server( # noqa: PLR0915 ): check_prisma_schema_diff(db_url=None) else: - if not PrismaManager.setup_database( - use_migrate=not use_prisma_db_push - ): + if not use_v2_migration_resolver: + print( # noqa + "\033[1;33mLiteLLM Proxy: Using default (v1) migration resolver. " + "If your deployment has seen schema thrashing during rolling " + "deploys, try --use_v2_migration_resolver (safer: avoids the " + "diff-and-force recovery that caused the thrash).\033[0m" + ) + try: + setup_ok = PrismaManager.setup_database( + use_migrate=not use_prisma_db_push, + use_v2_resolver=use_v2_migration_resolver, + ) + except RuntimeError as e: + # v2 resolver raises on unrecoverable migration errors + # (e.g. non-idempotent failures, permission issues). + # v1 never raises here, so this only fires when the + # operator opted into v2. + print( # noqa + "\033[1;31mLiteLLM Proxy: Database migration cannot proceed. " + f"{e}\033[0m", + file=sys.stderr, + flush=True, + ) + sys.exit(2) + if not setup_ok: if enforce_prisma_migration_check: print( # noqa "\033[1;31mLiteLLM Proxy: Database setup failed after multiple retries. " @@ -954,6 +994,9 @@ def run_server( # noqa: PLR0915 if loop_type: uvicorn_args["loop"] = loop_type + if reload: + uvicorn_args["reload"] = True + uvicorn.run( **uvicorn_args, workers=num_workers, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d8354a798b1..aa8122d8fd9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -952,6 +952,17 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 _run_background_health_check() ) # start the background health check coroutine. + # Start adaptive-router queue flusher unconditionally — adaptive routers + # may be added later via `/config/reload`, and the flusher is a no-op when + # `llm_router.adaptive_routers` is empty. Per-router DB state is loaded + # lazily by the flusher on first tick (see `_state_loaded` flag) so + # hot-reloaded routers also get their persisted priors. + if llm_router is not None and getattr(llm_router, "adaptive_routers", None): + for _ar in llm_router.adaptive_routers.values(): + await _ar.load_state_from_db(prisma_client) + _ar._state_loaded = True + asyncio.create_task(_adaptive_router_flusher_loop()) + ## [Optional] Initialize dd tracer ProxyStartupEvent._init_dd_tracer() @@ -1795,6 +1806,7 @@ async def increment_spend_counters( team_id: Optional[str], user_id: Optional[str], response_cost: Optional[float], + org_id: Optional[str] = None, ): """ Atomically increment spend counters for budget enforcement. @@ -1881,6 +1893,20 @@ async def increment_spend_counters( increment=response_cost, ) + if user_id is not None: + await _init_and_increment_spend_counter( + counter_key=f"spend:user:{user_id}", + source_cache_key=user_id, + increment=response_cost, + ) + + if org_id is not None: + await _init_and_increment_spend_counter( + counter_key=f"spend:org:{org_id}", + source_cache_key=f"org_id:{org_id}", + increment=response_cost, + ) + async def _init_and_increment_spend_counter( counter_key: str, @@ -2427,6 +2453,38 @@ def _write_health_state_to_router_cache( ) +_ADAPTIVE_ROUTER_FLUSH_INTERVAL_SECONDS = 10 + + +async def _adaptive_router_flusher_loop(): + """ + Drain every AdaptiveRouter's in-memory state + session aggregators into + Postgres on a fixed cadence. Hot-path writes go to memory; this loop is + the only writer to the adaptive router DB tables. + """ + global llm_router, prisma_client + while True: + try: + await asyncio.sleep(_ADAPTIVE_ROUTER_FLUSH_INTERVAL_SECONDS) + adaptive_routers = getattr(llm_router, "adaptive_routers", None) or {} + if not adaptive_routers or prisma_client is None: + continue + for ar in adaptive_routers.values(): + # Lazy state load: covers adaptive routers registered via + # `/config/reload` after proxy boot. + if not getattr(ar, "_state_loaded", False): + try: + await ar.load_state_from_db(prisma_client) + finally: + ar._state_loaded = True + await ar.queue.flush_state_to_db(prisma_client) + await ar.queue.flush_session_to_db(prisma_client) + except asyncio.CancelledError: + raise + except Exception: + verbose_proxy_logger.exception("adaptive_router flusher iteration failed") + + async def _run_background_health_check(): """ Periodically run health checks in the background on the endpoints. @@ -13953,6 +14011,38 @@ async def home(request: Request): return "LiteLLM: RUNNING" +@router.get( + "/adaptive_router/state", + tags=["adaptive_router"], + dependencies=[Depends(user_api_key_auth)], +) +async def get_adaptive_router_state( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Return live bandit posteriors + queue depth for every configured adaptive router. + + Admin-only. Returns 404 if no adaptive router is configured. + + Response shape: `{"routers": [, ...]}` — one snapshot per + adaptive-router deployment. Each snapshot's `router_name` field identifies + which deployment it came from. + """ + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={"error": CommonProxyErrors.not_allowed_access.value}, + ) + if llm_router is None or not llm_router.adaptive_routers: + raise HTTPException( + status_code=404, + detail={"error": "No adaptive_router is configured on this proxy."}, + ) + snapshots = [ + await ar.get_state_snapshot() for ar in llm_router.adaptive_routers.values() + ] + return {"routers": snapshots} + + @router.get("/routes", dependencies=[Depends(user_api_key_auth)]) async def get_routes(): """ diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 08aa5645251..7642ad74b20 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1223,3 +1223,46 @@ model LiteLLM_ClaudeCodePluginTable { @@map("LiteLLM_ClaudeCodePluginTable") } + +// Per-(router, request_type, model) Beta posterior for the adaptive router. +model LiteLLM_AdaptiveRouterState { + router_name String + request_type String + model_name String + alpha Float + beta Float + total_samples Int @default(0) + last_updated_at DateTime @default(now()) @updatedAt + + @@id([router_name, request_type, model_name]) +} + +// Per-(session, router, model) signal counters for the adaptive router. +model LiteLLM_AdaptiveRouterSession { + session_id String + router_name String + model_name String + classified_type String + + misalignment_count Int @default(0) + stagnation_count Int @default(0) + disengagement_count Int @default(0) + satisfaction_count Int @default(0) + failure_count Int @default(0) + loop_count Int @default(0) + exhaustion_count Int @default(0) + + last_user_content String? + last_assistant_content String? + tool_call_history Json @default("[]") + pending_tool_calls Json @default("{}") + + turn_count Int @default(0) + last_processed_turn Int @default(-1) + clean_credit_awarded Boolean @default(false) + terminal_status Int? + last_activity_at DateTime @default(now()) @updatedAt + + @@id([session_id, router_name, model_name]) + @@index([last_activity_at], map: "idx_adaptive_router_session_activity") +} diff --git a/litellm/router.py b/litellm/router.py index 6572d96f7b9..07053db7d3f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -200,12 +200,20 @@ if TYPE_CHECKING: from litellm.router_strategy.complexity_router.complexity_router import ( ComplexityRouter, ) + from litellm.router_strategy.adaptive_router.adaptive_router import ( + AdaptiveRouter, + ) + from litellm.router_strategy.quality_router.quality_router import ( + QualityRouter, + ) Span = Union[_Span, Any] else: Span = Any AutoRouter = Any ComplexityRouter = Any + AdaptiveRouter = Any + QualityRouter = Any PreRoutingHookResponse = Any @@ -464,6 +472,8 @@ class Router: ) # {"TEAM_ID": PatternMatchRouter} self.auto_routers: Dict[str, "AutoRouter"] = {} self.complexity_routers: Dict[str, "ComplexityRouter"] = {} + self.adaptive_routers: Dict[str, "AdaptiveRouter"] = {} + self.quality_routers: Dict[str, "QualityRouter"] = {} # Initialize model_group_alias early since it's used in set_model_list self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = ( @@ -5884,7 +5894,7 @@ class Router: response = await response ## PROCESS RESPONSE HEADERS response = await self.set_response_headers( - response=response, model_group=model_group + response=response, model_group=model_group, request_kwargs=kwargs ) return response @@ -6810,10 +6820,15 @@ class Router: Check if the deployment is an auto-router deployment (semantic router). Returns True if the litellm_params model starts with "auto_router/" - but NOT "auto_router/complexity_router" (which uses complexity routing). + but NOT "auto_router/complexity_router" or "auto_router/adaptive_router" + (which use the complexity-router and adaptive-router strategies). """ if litellm_params.model.startswith("auto_router/complexity_router"): return False # This is handled by complexity_router + if litellm_params.model.startswith("auto_router/adaptive_router"): + return False # This is handled by adaptive_router + if litellm_params.model.startswith("auto_router/quality_router"): + return False # This is handled by quality_router if litellm_params.model.startswith("auto_router/"): return True return False @@ -6920,6 +6935,196 @@ class Router: ) self.complexity_routers[deployment.model_name] = complexity_router + def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: + """True when this deployment opts in via the `auto_router/adaptive_router` model prefix.""" + return litellm_params.model.startswith("auto_router/adaptive_router") + + def _finalize_adaptive_router_if_configured(self) -> None: + """Locate every adaptive-router deployment in the finalized model_list and + build an AdaptiveRouter for each. Safe no-op when none are configured. + Idempotent: skips any deployment whose model_name is already initialized.""" + # Drop any adaptive-router hooks left over from a previous Router + # instance (e.g. after `/config/reload` replaced `llm_router`). Without + # this, stale AdaptiveRouterPostCallHook callbacks from the old Router + # remain wired up in `litellm.callbacks` and double-fire signal + # recording for every request. + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + + for _cb_list in ( + litellm.callbacks, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ): + litellm.logging_callback_manager.remove_callbacks_by_type( + _cb_list, AdaptiveRouterPostCallHook + ) + + for entry in self.model_list or []: + lp = ( + entry.get("litellm_params") + if isinstance(entry, dict) + else entry.litellm_params + ) + lp_model = ( + (lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None + ) + if not (lp_model and lp_model.startswith("auto_router/adaptive_router")): + continue + model_name = ( + entry.get("model_name") if isinstance(entry, dict) else entry.model_name + ) + if not model_name or not lp: + continue + if model_name in self.adaptive_routers: + continue + deployment = Deployment( + model_name=model_name, + litellm_params=( + lp if not isinstance(lp, dict) else LiteLLM_Params(**lp) + ), + model_info=( + entry.get("model_info") + if isinstance(entry, dict) + else entry.model_info + ), + ) + self.init_adaptive_router_deployment(deployment=deployment) + + def init_adaptive_router_deployment(self, deployment: Deployment) -> None: + """ + Build an AdaptiveRouter instance for this deployment and register its + post-call hook. Multiple adaptive routers can coexist on a single Router, + keyed by `deployment.model_name`. + + `model_to_prefs` and `model_to_cost` are derived from the OTHER models + already registered in `self.model_list` whose `model_name` appears in + `available_models`. Models not yet registered fall back to defaults. + """ + # Local import: AdaptiveRouter -> hooks -> classifier all import litellm + # internals which transitively import this module. (AGENTS.md exception clause.) + from litellm.router_strategy.adaptive_router.adaptive_router import ( + AdaptiveRouter, + ) + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + from litellm.types.router import ( + AdaptiveRouterConfig, + AdaptiveRouterPreferences, + ) + + raw_config = deployment.litellm_params.adaptive_router_config + if raw_config is None: + raise ValueError( + "adaptive_router_config is required for adaptive-router deployments." + ) + + config = AdaptiveRouterConfig(**raw_config) + + model_to_prefs: Dict[str, AdaptiveRouterPreferences] = {} + model_to_cost: Dict[str, float] = {} + # O(k) via the name→indices map: only touch deployments whose name + # is listed in `available_models`, instead of scanning model_list. + for name in config.available_models: + indices = self.model_name_to_deployment_indices.get(name, []) + if not indices: + continue + d = (self.model_list or [])[indices[0]] + mi = d.get("model_info") if isinstance(d, dict) else d.model_info + mi_dict: Dict[str, Any] = ( + mi if isinstance(mi, dict) else (mi.model_dump() if mi else {}) + ) + prefs_raw = mi_dict.get("adaptive_router_preferences") + if prefs_raw is not None: + model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) + + # `input_cost_per_token` is a LiteLLM_Params field per types/router.py. + lp = d.get("litellm_params") if isinstance(d, dict) else d.litellm_params + lp_dict: Dict[str, Any] = ( + lp if isinstance(lp, dict) else (lp.model_dump() if lp else {}) + ) + cost = lp_dict.get("input_cost_per_token") + if cost is not None: + model_to_cost[name] = float(cost) + + if deployment.model_name in self.adaptive_routers: + raise ValueError( + f"Adaptive-router deployment {deployment.model_name} already exists. " + "Please use a different model name." + ) + + adaptive_router = AdaptiveRouter( + router_name=deployment.model_name, + config=config, + model_to_prefs=model_to_prefs, + model_to_cost=model_to_cost, + ) + self.adaptive_routers[deployment.model_name] = adaptive_router + litellm.logging_callback_manager.add_litellm_callback( + AdaptiveRouterPostCallHook(adaptive_router=adaptive_router) + ) + verbose_router_logger.info( + "AdaptiveRouter[%s] initialized with %d models", + deployment.model_name, + len(config.available_models), + ) + + def _is_quality_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: + """ + Check if the deployment is a quality-router deployment. + + Returns True if the litellm_params model starts with "auto_router/quality_router". + """ + if litellm_params.model.startswith("auto_router/quality_router"): + return True + return False + + def init_quality_router_deployment(self, deployment: Deployment): + """ + Initialize the quality-router deployment. + + Resolves the default model from either `quality_router_default_model` or + `quality_router_config["default_model"]`, then instantiates the + QualityRouter and stores it in `self.quality_routers`. + """ + # Import here to mirror the AutoRouter / ComplexityRouter init pattern + # and avoid circular imports. + from litellm.router_strategy.quality_router.quality_router import ( + QualityRouter, + ) + + quality_router_config: Optional[dict] = ( + deployment.litellm_params.quality_router_config + ) + + default_model: Optional[str] = ( + deployment.litellm_params.quality_router_default_model + ) + if default_model is None and quality_router_config: + default_model = quality_router_config.get("default_model") + + if default_model is None: + raise ValueError( + "quality_router_default_model is required for quality-router deployments, " + "or set default_model in quality_router_config. Please configure it in the litellm_params" + ) + + quality_router: QualityRouter = QualityRouter( + model_name=deployment.model_name, + default_model=default_model, + litellm_router_instance=self, + quality_router_config=quality_router_config, + ) + if deployment.model_name in self.quality_routers: + raise ValueError( + f"Quality-router deployment {deployment.model_name} already exists. Please use a different model name." + ) + self.quality_routers[deployment.model_name] = quality_router + def deployment_is_active_for_environment(self, deployment: Deployment) -> bool: """ Function to check if a llm deployment is active for a given environment. Allows using the same config.yaml across multople environments @@ -6966,6 +7171,11 @@ class Router: self.model_id_to_deployment_index_map = {} # Reset the index self.model_name_to_deployment_indices = {} # Reset the model_name index self.team_model_to_deployment_indices = {} # Reset the team_model index + # Reset per-strategy router registries so hot-reload doesn't leave + # stale routers pointing at the old model_list. + self.quality_routers = {} + self.complexity_routers = {} + self.auto_routers = {} self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works @@ -7013,6 +7223,10 @@ class Router: # Note: model_name_to_deployment_indices is already built incrementally # by _create_deployment -> _add_model_to_list_and_index_map + # Deferred: build the AdaptiveRouter strategy now that all underlying + # deployments have been registered. + self._finalize_adaptive_router_if_configured() + def _add_deployment(self, deployment: Deployment) -> Deployment: import os @@ -7140,6 +7354,16 @@ class Router: ): self.init_complexity_router_deployment(deployment=deployment) + # NOTE: adaptive-router deployments are deferred to the end of + # set_model_list() because their init needs visibility into the OTHER + # deployments listed in `available_models` (which may not yet have + # been processed when this one is created). + ######################################################### + # Check if this is a quality-router deployment + ######################################################### + if self._is_quality_router_deployment(litellm_params=deployment.litellm_params): + self.init_quality_router_deployment(deployment=deployment) + return deployment def _initialize_deployment_for_pass_through( @@ -8143,7 +8367,10 @@ class Router: return returned_dict async def set_response_headers( - self, response: Any, model_group: Optional[str] = None + self, + response: Any, + model_group: Optional[str] = None, + request_kwargs: Optional[dict] = None, ) -> Any: """ Add the most accurate rate limit headers for a given model response. @@ -8164,6 +8391,45 @@ class Router: additional_headers = response._hidden_params["additional_headers"] # type: ignore + # Lift QualityRouter routing decision into response headers for + # transparency. The decision is stashed in request_kwargs.metadata + # by QualityRouter.async_pre_routing_hook. + metadata = ( + (request_kwargs.get("metadata") or {}) + if isinstance(request_kwargs, dict) + else {} + ) + decision = ( + metadata.get("quality_router_decision") + if isinstance(metadata, dict) + else None + ) + if isinstance(decision, dict): + # Only emit headers for fields that have a meaningful value. + # `complexity_tier` and `matched_keyword` are mutually exclusive + # (the keyword path short-circuits classification), so each + # request emits one or the other but not both. + if decision.get("routed_model") is not None: + additional_headers["x-litellm-quality-router-model"] = str( + decision["routed_model"] + ) + if decision.get("quality_tier") is not None: + additional_headers["x-litellm-quality-router-tier"] = str( + decision["quality_tier"] + ) + if decision.get("routed_via") is not None: + additional_headers["x-litellm-quality-router-via"] = str( + decision["routed_via"] + ) + if decision.get("matched_keyword") is not None: + additional_headers["x-litellm-quality-router-keyword"] = str( + decision["matched_keyword"] + ) + if decision.get("complexity_tier") is not None: + additional_headers["x-litellm-quality-router-complexity"] = str( + decision["complexity_tier"] + ) + if ( "x-ratelimit-remaining-tokens" not in additional_headers and "x-ratelimit-remaining-requests" not in additional_headers @@ -8708,8 +8974,6 @@ class Router: and self.routing_strategy == "latency-based-routing" ): _settings_to_return[var] = self.lowestlatency_logger.routing_args.json() - elif var == "routing_strategy_args": - _settings_to_return[var] = None return _settings_to_return def update_settings(self, **kwargs): @@ -9620,7 +9884,7 @@ class Router: self, model: str, request_kwargs: Dict, - messages: Optional[List[Dict[str, str]]] = None, + messages: Optional[List[Dict[str, Any]]] = None, input: Optional[Union[str, List]] = None, specific_deployment: Optional[bool] = False, ) -> Optional[PreRoutingHookResponse]: @@ -9653,6 +9917,31 @@ class Router: specific_deployment=specific_deployment, ) + ######################################################### + # Check if an adaptive-router should be used + ######################################################### + adaptive_router = self.adaptive_routers.get(model) + if adaptive_router is not None: + return await adaptive_router.async_pre_routing_hook( + model=model, + request_kwargs=request_kwargs, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) + + ######################################################### + # Check if any quality-router should be used + ######################################################### + if model in self.quality_routers: + return await self.quality_routers[model].async_pre_routing_hook( + model=model, + request_kwargs=request_kwargs, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) + return None def get_available_deployment( diff --git a/litellm/router_strategy/adaptive_router/README.md b/litellm/router_strategy/adaptive_router/README.md new file mode 100644 index 00000000000..7f5d7aa21d0 --- /dev/null +++ b/litellm/router_strategy/adaptive_router/README.md @@ -0,0 +1,95 @@ +# Adaptive Router (v0) + +A request-type-aware routing strategy. For each incoming request, classify the +prompt into one of seven `RequestType` buckets (code generation, writing, +analytical reasoning, …), then Thompson-sample a Beta(α, β) bandit posterior +per `(request_type, model)` cell to pick the best model. Quality estimates are +combined with a normalized cost score via a weighted linear sum. + +A post-call hook reads the response and runs lightweight regex + tool-call +detectors (see `signals.py`) to award per-turn credit/blame to the model that +served the turn. Updates are batched in-memory and flushed to Postgres every +~10s by a background task in `proxy_server.py`. + +## Config example + +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + model_info: + input_cost_per_token: 0.0000025 + adaptive_router_preferences: + quality_tier: 3 + strengths: ["code_generation", "analytical_reasoning"] + + - model_name: gpt-4o-mini + litellm_params: + model: openai/gpt-4o-mini + model_info: + input_cost_per_token: 0.00000015 + adaptive_router_preferences: + quality_tier: 2 + strengths: ["general", "factual_lookup"] + + - model_name: smart-router + litellm_params: + model: auto_router/adaptive_router + adaptive_router_default_model: gpt-4o-mini + adaptive_router_config: + available_models: ["gpt-4o", "gpt-4o-mini"] + weights: + quality: 0.7 + cost: 0.3 +``` + +Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key +`min_quality_tier: 3`) to force selection from tier-3-or-higher models only. + +## Behavior summary + +- **Cold start.** Each `(request_type, model)` cell starts with a + Beta prior whose mean = `BASE_TIER_WEIGHT[tier] (+ STRENGTH_BONUS if declared)` + and total mass = `COLD_START_MASS` (10). About ten real observations move it + meaningfully. +- **Per-request decision.** Sample once per eligible model, score with + `quality_weight·sample + cost_weight·normalized_cost`, pick the argmax. + Routing is stateless per-turn — no sticky lookup. Each call resamples. +- **Owner-cache attribution.** Post-call, the conversation's first picked + model claims an "owner slot" for `OWNER_CACHE_TTL_SECONDS` (24h). Later + turns of the same conversation only fire bandit/state updates if the + same model handled them — mismatches are dropped (no attribution) and + counted in `skipped_updates_total`. Conversation identity is the + client-supplied `litellm_session_id` if present, otherwise a sha256 over + caller identity (api key hash, team, user, end-user) + the first message. +- **Per-turn updates.** `satisfaction → +α`. `misalignment, stagnation, + disengagement, failure → +β` (each). `loop → +0.5β`. `exhaustion → 0` + (uptime, not quality). Skipped if conversation has fewer than + `SIGNAL_GATE_MIN_MESSAGES` messages. +- **Persistence.** Bandit cells: aggregated deltas, eventually consistent. + Session rows: last-write-wins snapshots. + +## Known v0 limitations + +- **Latency is not in the score.** Quality + cost only. A pathologically slow + model can still be picked. +- **Hard sample cap at 200.** Once `α + β > 200`, deltas are silently dropped. + No rescaling — drift is a v1 concern. +- **24h owner-cache TTL.** No explicit eviction below TTL. The in-memory map + can grow if traffic patterns produce many one-shot sessions. +- **Owner-recovery skew.** If model A "owns" a conversation but is then + dethroned in the bandit, later turns served by model B are dropped — so + bandit updates for that conversation flatline until A's TTL expires. + Tracked via `skipped_updates_total`. +- **Signals are regex + tool-call only.** No LLM-judge, no embedding similarity, + no exemplar storage. Signals are best-effort and biased toward English. +- **One AdaptiveRouter per `Router`.** Multiple `adaptive_router/*` deployments + on the same `litellm.Router` raise at init. +- **Bandit-delta mapping is unvalidated.** `_compute_bandit_delta` is a v0 + guess; expect to retune after the first ~1000 sessions of real traffic. +- **`request_type` is classified per turn from the latest user message.** For + non-GENERAL turns, the current-turn type is used for bandit attribution (so + genuine mid-session topic shifts update the correct cell). For GENERAL turns + ("thanks!", "ok", "sounds good"), attribution falls back to the session's + original type to avoid misattributing closing pleasantries. diff --git a/litellm/router_strategy/adaptive_router/__init__.py b/litellm/router_strategy/adaptive_router/__init__.py new file mode 100644 index 00000000000..d7f55ebced9 --- /dev/null +++ b/litellm/router_strategy/adaptive_router/__init__.py @@ -0,0 +1,6 @@ +"""Adaptive router strategy. See README.md for design overview.""" + +from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter +from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook + +__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook"] diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py new file mode 100644 index 00000000000..3bccef36e68 --- /dev/null +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -0,0 +1,454 @@ +""" +Main adaptive router strategy. See README.md for design overview. + +One AdaptiveRouter instance per router_name. Holds in-memory caches: +- _cells: Beta(alpha, beta) bandit posteriors per (request_type, model) +- _owner_cache: session_key -> (owner_model, expires_at) — the first model + picked for a conversation owns its bandit-update slot +- _session_states: (session_key, model) -> SessionState for incremental signal updates + +Owns the AdaptiveRouterUpdateQueue used by the proxy's flusher to persist +state and session snapshots back to Postgres. + +Routing is stateless per-turn (Thompson sample fresh on every call). The +owner cache is consulted only at post-call time to decide whether a turn's +signals should fire a bandit update — turns served by a different model than +the conversation's owner are skipped to avoid cross-model misattribution. +""" + +from __future__ import annotations + +import asyncio +import time +from dataclasses import asdict +from typing import Any, Dict, List, Optional, Tuple, Union, cast + +from litellm._logging import verbose_router_logger +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + get_last_user_message, +) +from litellm.router_strategy.adaptive_router.bandit import ( + BanditCell, + apply_delta, + initial_cell, + pick_best, +) +from litellm.router_strategy.adaptive_router.classifier import classify_prompt +from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + MIN_QUALITY_TIER_HEADER, + MIN_QUALITY_TIER_METADATA_KEY, + OWNER_CACHE_TTL_SECONDS, +) +from litellm.router_strategy.adaptive_router.signals import ( + SessionState, + SignalDelta, + Turn, + apply_turn, +) +from litellm.router_strategy.adaptive_router.update_queue import ( + AdaptiveRouterUpdateQueue, +) + +# Sweep session-state cache when it exceeds this many live entries. Expired +# entries are dropped in bulk; amortizes to O(1) per insert. +_SESSION_STATE_SWEEP_THRESHOLD: int = 1024 +# Same pattern for the owner cache. +_OWNER_CACHE_SWEEP_THRESHOLD: int = 1024 +from litellm.types.llms.openai import AllMessageValues +from litellm.types.router import ( + AdaptiveRouterConfig, + AdaptiveRouterPreferences, + PreRoutingHookResponse, + RequestType, +) + + +def _default_prefs() -> AdaptiveRouterPreferences: + """Tier-2 prior with no declared strengths; used when a model omits prefs.""" + return AdaptiveRouterPreferences(quality_tier=2, strengths=[]) + + +class AdaptiveRouter: + """One instance per router_name. Holds in-memory caches + the update queue.""" + + def __init__( + self, + router_name: str, + config: AdaptiveRouterConfig, + model_to_prefs: Dict[str, AdaptiveRouterPreferences], + model_to_cost: Dict[str, float], + ) -> None: + self.router_name = router_name + self.config = config + self.model_to_prefs = model_to_prefs + self.model_to_cost = model_to_cost + self.queue = AdaptiveRouterUpdateQueue() + + self._cells: Dict[Tuple[RequestType, str], BanditCell] = {} + self._owner_cache: Dict[str, Tuple[str, float]] = {} + self._session_states: Dict[Tuple[str, str], SessionState] = {} + # Parallel expiry map for _session_states, same TTL as _owner_cache. + # Evicted opportunistically in `get_or_create_session_state`. + self._session_states_expiry: Dict[Tuple[str, str], float] = {} + self._skipped_updates_total: int = 0 + # Set to True once the proxy flusher has loaded persisted priors from + # Postgres. Checked to support lazy-load on hot-reloaded routers. + self._state_loaded: bool = False + self._lock = asyncio.Lock() + + self._init_cold_start_cells() + + # ---- Cold-start ------------------------------------------------------ + + def _init_cold_start_cells(self) -> None: + """Populate _cells with cold-start priors for every (rt, model) combination.""" + for rt in RequestType: + for model in self.config.available_models: + prefs = self.model_to_prefs.get(model) or _default_prefs() + self._cells[(rt, model)] = initial_cell(prefs, rt) + + async def load_state_from_db(self, prisma_client: Any) -> None: + """Override cold-start cells with persisted state. Called once at startup.""" + if prisma_client is None: + return + try: + rows = await prisma_client.db.litellm_adaptiverouterstate.find_many( + where={"router_name": self.router_name} + ) + loaded = 0 + for row in rows: + try: + rt = RequestType(row.request_type) + except ValueError: + # Unknown taxonomy entry from an older/newer version. Skip. + continue + if row.model_name not in self.config.available_models: + continue + self._cells[(rt, row.model_name)] = BanditCell( + alpha=row.alpha, beta=row.beta + ) + loaded += 1 + verbose_router_logger.info( + "AdaptiveRouter[%s]: loaded %d cells from DB", + self.router_name, + loaded, + ) + except Exception as e: + verbose_router_logger.exception( + "AdaptiveRouter[%s]: failed to load state from DB: %s", + self.router_name, + e, + ) + + # ---- Pre-routing hook ------------------------------------------------ + + async def async_pre_routing_hook( + self, + model: str, + request_kwargs: Dict[str, Any], + messages: Optional[List[Dict[str, Any]]] = None, + input: Optional[Union[str, List]] = None, + specific_deployment: Optional[bool] = False, + ) -> Optional[PreRoutingHookResponse]: + """ + Plugin entry point invoked by `Router.async_pre_routing_hook` when the + inbound `model` matches this adaptive router's `router_name`. + + Classifies the last user message, picks a logical model via the bandit, + and stashes the chosen model on `request_kwargs["metadata"]` so the + post-call hook can surface it as a response header. + + Routing is stateless per-turn: every call Thompson-samples fresh, + regardless of any prior pick for the same session. Cross-turn + attribution is enforced post-call via the owner cache (see + `claim_or_check_owner`). + """ + user_text = ( + get_last_user_message(cast(List[AllMessageValues], messages or [])) or "" + ) + + request_type = classify_prompt(user_text) + min_quality_tier = self._extract_min_quality_tier(request_kwargs) + chosen_model = await self.pick_model( + request_type=request_type, min_quality_tier=min_quality_tier + ) + verbose_router_logger.debug( + "AdaptiveRouter[%s]: classified=%s -> chose %s", + self.router_name, + request_type.value, + chosen_model, + ) + + # Relay the chosen logical model to the post-call hook, which surfaces + # it as the `x-litellm-adaptive-router-model` response header. We use + # `metadata` (not a top-level kwarg) so the value doesn't leak into + # `litellm.acompletion(**input_kwargs)`. + kwargs_metadata = request_kwargs.setdefault("metadata", {}) + if isinstance(kwargs_metadata, dict): + kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = chosen_model + + return PreRoutingHookResponse(model=chosen_model, messages=messages) + + # ---- Pick model ------------------------------------------------------ + + async def pick_model( + self, + request_type: RequestType, + min_quality_tier: Optional[int] = None, + ) -> str: + """Thompson-sample across eligible models. Stateless per-turn.""" + eligible = self._eligible_models(min_quality_tier) + if not eligible: + raise ValueError( + f"AdaptiveRouter[{self.router_name}]: no models meet " + f"min_quality_tier={min_quality_tier}" + ) + + cells = {m: self._cells[(request_type, m)] for m in eligible} + costs = {m: self.model_to_cost.get(m, 0.0) for m in eligible} + return pick_best( + cells, + costs, + quality_weight=self.config.weights.quality, + cost_weight=self.config.weights.cost, + ) + + def claim_or_check_owner(self, session_key: str, current_model: str) -> bool: + """Resolve attribution for a turn under stateless routing. + + Returns True iff this turn should fire a bandit/state update. The + first call for a `session_key` claims ownership for `current_model` + and returns True. Subsequent calls return True only if the owner is + still live AND matches `current_model`. Mismatches (a different + model handled this turn) and expired owners both increment + `_skipped_updates_total` and return False — no attribution. + """ + now = time.time() + existing = self._owner_cache.get(session_key) + if existing is not None and existing[1] > now: + owner_model, _ = existing + if owner_model == current_model: + return True + self._skipped_updates_total += 1 + return False + + # Opportunistic bulk sweep — sessions that never come back would + # otherwise pile up here forever. Same threshold pattern as the + # session-state cache. + if len(self._owner_cache) >= _OWNER_CACHE_SWEEP_THRESHOLD: + self._evict_expired_owner_cache(now) + + # No live owner -> claim for current_model. + self._owner_cache[session_key] = ( + current_model, + now + OWNER_CACHE_TTL_SECONDS, + ) + return True + + def _evict_expired_owner_cache(self, now: float) -> None: + expired = [k for k, (_, exp) in self._owner_cache.items() if exp <= now] + for k in expired: + self._owner_cache.pop(k, None) + + async def get_state_snapshot(self) -> Dict[str, Any]: + """In-memory snapshot for the introspection endpoint. Cheap; no DB hit.""" + cells = [] + for (rt, model), cell in sorted( + self._cells.items(), key=lambda kv: (kv[0][0].value, kv[0][1]) + ): + total = cell.alpha + cell.beta + cells.append( + { + "request_type": rt.value, + "model": model, + "alpha": cell.alpha, + "beta": cell.beta, + # Net observations that have moved the posterior, excluding + # the cold-start prior mass. `alpha + beta` would show the + # initial COLD_START_MASS (e.g. 10) before any real traffic + # arrives, which confuses operators reading the endpoint. + "samples": cell.total_samples, + "quality_mean": cell.alpha / total if total > 0 else 0.0, + } + ) + queue = await self.queue.queue_size() + now = time.time() + owner_cache_live = sum(1 for _, exp in self._owner_cache.values() if exp > now) + return { + "router_name": self.router_name, + "available_models": list(self.config.available_models), + "weights": { + "quality": self.config.weights.quality, + "cost": self.config.weights.cost, + }, + "model_costs": dict(self.model_to_cost), + "cells": cells, + "owner_cache_live": owner_cache_live, + "skipped_updates_total": self._skipped_updates_total, + "queue": queue, + } + + @staticmethod + def _extract_min_quality_tier( + request_kwargs: Dict[str, Any], + ) -> Optional[int]: + """Pull `min_quality_tier` from request headers or metadata. + + Precedence: headers (`x-litellm-min-quality-tier`) over metadata + (`min_quality_tier`). Headers arrive lowercased from the proxy but we + lookup case-insensitively to be safe. Unparseable values are ignored + (treated as "not set") rather than raising — a bad header shouldn't + fail the request. + """ + headers = request_kwargs.get("headers") or {} + if isinstance(headers, dict): + for k, v in headers.items(): + if isinstance(k, str) and k.lower() == MIN_QUALITY_TIER_HEADER: + try: + return int(v) + except (TypeError, ValueError): + return None + + metadata = request_kwargs.get("metadata") or {} + if isinstance(metadata, dict): + raw = metadata.get(MIN_QUALITY_TIER_METADATA_KEY) + if raw is not None: + try: + return int(raw) + except (TypeError, ValueError): + return None + return None + + def _eligible_models(self, min_quality_tier: Optional[int]) -> List[str]: + if min_quality_tier is None: + return list(self.config.available_models) + return [ + m + for m in self.config.available_models + if (self.model_to_prefs.get(m) or _default_prefs()).quality_tier + >= min_quality_tier + ] + + # ---- Session state --------------------------------------------------- + + def get_or_create_session_state( + self, + session_id: str, + model_name: str, + request_type: RequestType, + ) -> SessionState: + key = (session_id, model_name) + now = time.time() + + # Opportunistic bulk sweep when the cache grows past the threshold. + # Cheap relative to the alternative of a bounded LRU — conversations + # naturally become inactive within OWNER_CACHE_TTL_SECONDS. + if len(self._session_states) >= _SESSION_STATE_SWEEP_THRESHOLD: + self._evict_expired_session_states(now) + + state = self._session_states.get(key) + if state is None: + state = SessionState( + session_id=session_id, + router_name=self.router_name, + model_name=model_name, + classified_type=request_type.value, + ) + self._session_states[key] = state + self._session_states_expiry[key] = now + OWNER_CACHE_TTL_SECONDS + return state + + def _evict_expired_session_states(self, now: float) -> None: + """Drop session states whose TTL has passed. O(n) but amortized O(1) + per insert thanks to `_SESSION_STATE_SWEEP_THRESHOLD`.""" + expired = [k for k, exp in self._session_states_expiry.items() if exp <= now] + for k in expired: + self._session_states.pop(k, None) + self._session_states_expiry.pop(k, None) + + async def record_turn( + self, + session_id: str, + model_name: str, + request_type: RequestType, + turn: Turn, + ) -> SignalDelta: + """Apply one turn, push session snapshot + bandit deltas to the queue.""" + state = self.get_or_create_session_state(session_id, model_name, request_type) + delta = apply_turn(state, turn) + verbose_router_logger.debug( + "AdaptiveRouter[%s]: record_turn delta=%s", self.router_name, delta + ) + + # Strip the raw conversation content before persisting. The + # last_user/assistant_content and tool_call_history fields are only + # needed in-memory for the next turn's incremental signal detection; + # writing user prompts and tool payloads to the DB would store PII + # for every adaptive-router conversation. Counts + bookkeeping is + # all the persisted row needs. + snapshot = asdict(state) + for sensitive in ( + "last_user_content", + "last_assistant_content", + "tool_call_history", + "pending_tool_calls", + ): + snapshot.pop(sensitive, None) + await self.queue.add_session_state( + session_id, self.router_name, model_name, snapshot + ) + + d_alpha, d_beta = self._compute_bandit_delta(delta) + verbose_router_logger.debug( + "AdaptiveRouter[%s]: bandit delta alpha=%.2f beta=%.2f", + self.router_name, + d_alpha, + d_beta, + ) + if d_alpha != 0 or d_beta != 0: + # For non-GENERAL turns, attribute to the current-turn classification + # so genuine mid-session topic shifts (e.g. code → math) update the + # correct cell. For GENERAL turns ("thanks!", "ok", "sounds good"), fall + # back to the session's original type so closing pleasantries don't + # misattribute the reward. + attribution_type = ( + request_type + if request_type != RequestType.GENERAL + else RequestType(state.classified_type) + ) + cell_key = (attribution_type, model_name) + self._cells[cell_key] = apply_delta(self._cells[cell_key], d_alpha, d_beta) + await self.queue.add_state_delta( + self.router_name, + attribution_type.value, + model_name, + d_alpha, + d_beta, + ) + + return delta + + @staticmethod + def _compute_bandit_delta(delta: SignalDelta) -> Tuple[float, float]: + """ + Translate per-turn signal deltas into bandit-cell deltas. + + v0 mapping (UNVALIDATED — D6): + - satisfaction -> +1 alpha + - misalignment, stagnation, + disengagement, failure -> +1 beta each + - loop -> +0.5 beta (weak; could be model OR user) + - exhaustion -> 0 (uptime issue, tracked separately later) + """ + d_alpha = float(delta.satisfaction) + d_beta = ( + float( + delta.misalignment + + delta.stagnation + + delta.disengagement + + delta.failure + ) + + 0.5 * delta.loop + ) + return d_alpha, d_beta diff --git a/litellm/router_strategy/adaptive_router/bandit.py b/litellm/router_strategy/adaptive_router/bandit.py new file mode 100644 index 00000000000..1ab96f0e952 --- /dev/null +++ b/litellm/router_strategy/adaptive_router/bandit.py @@ -0,0 +1,142 @@ +""" +Thompson sampling and prior initialization for the adaptive router bandit. + +Each (router, request_type, model) cell is a Beta(alpha, beta) posterior. +- alpha = pseudo-successes +- beta = pseudo-failures +- mean = alpha / (alpha + beta) +- total samples = alpha + beta - COLD_START_MASS (informative prior, not data) + +Hot path: thompson_sample() — pure function, no I/O. +""" + +import random +from dataclasses import dataclass +from typing import Dict, List, Optional + +from litellm.router_strategy.adaptive_router.config import ( + BASE_TIER_WEIGHT, + COLD_START_MASS, + DEFAULT_COST_WEIGHT, + DEFAULT_QUALITY_WEIGHT, + SAMPLE_CAP, + STRENGTH_BONUS, +) +from litellm.types.router import AdaptiveRouterPreferences, RequestType + + +@dataclass(frozen=True) +class BanditCell: + """Posterior state for a single (router, request_type, model) cell.""" + + alpha: float + beta: float + + @property + def mean(self) -> float: + total = self.alpha + self.beta + return self.alpha / total if total > 0 else 0.5 + + @property + def total_samples(self) -> int: + return max(0, int(self.alpha + self.beta - COLD_START_MASS)) + + +def initial_cell( + prefs: AdaptiveRouterPreferences, request_type: RequestType +) -> BanditCell: + """ + Cold-start prior for a (model, request_type) cell. + + mean = base_tier_weight[tier] + (STRENGTH_BONUS if request_type in strengths else 0) + capped at 0.95 to avoid an over-confident prior. + Total mass = COLD_START_MASS so that ~10 real observations can move it noticeably. + """ + if prefs.quality_tier not in BASE_TIER_WEIGHT: + valid = sorted(BASE_TIER_WEIGHT) + raise ValueError( + f"quality_tier={prefs.quality_tier} is not supported; " + f"valid tiers are {valid}" + ) + base = BASE_TIER_WEIGHT[prefs.quality_tier] + bonus = STRENGTH_BONUS if request_type in prefs.strengths else 0.0 + mean = min(0.95, base + bonus) + alpha = mean * COLD_START_MASS + beta = (1.0 - mean) * COLD_START_MASS + return BanditCell(alpha=alpha, beta=beta) + + +def apply_delta(cell: BanditCell, delta_alpha: float, delta_beta: float) -> BanditCell: + """ + Apply a learning update to a cell, enforcing the sample cap. + + SAMPLE_CAP is a HARD cap on (alpha + beta). When the cap would be exceeded, + we drop the update. (D5: hard cap, no rescaling — keep v0 simple.) + """ + new_alpha = cell.alpha + delta_alpha + new_beta = cell.beta + delta_beta + if new_alpha + new_beta > SAMPLE_CAP: + return cell + return BanditCell(alpha=new_alpha, beta=new_beta) + + +def thompson_sample(cell: BanditCell, rng: Optional[random.Random] = None) -> float: + """Draw a sample from Beta(alpha, beta). Returns a quality estimate in [0, 1].""" + r = rng if rng is not None else random + return r.betavariate(cell.alpha, cell.beta) + + +def normalized_cost(model_cost: float, all_costs: List[float]) -> float: + """ + Map a raw $/1k-token cost into [0, 1] where 0 = most expensive, 1 = cheapest. + Returns 0.5 when there's no spread. + """ + if not all_costs: + return 0.5 + lo, hi = min(all_costs), max(all_costs) + if hi == lo: + return 0.5 + return 1.0 - ((model_cost - lo) / (hi - lo)) + + +def score( + quality_sample: float, + model_cost: float, + all_costs: List[float], + quality_weight: float = DEFAULT_QUALITY_WEIGHT, + cost_weight: float = DEFAULT_COST_WEIGHT, +) -> float: + """ + Multi-objective score. V0 is a weighted linear sum of (quality, normalized_cost). + Higher is better. Both inputs are in [0, 1]. + """ + cost_score = normalized_cost(model_cost, all_costs) + return quality_weight * quality_sample + cost_weight * cost_score + + +def pick_best( + cells: Dict[str, BanditCell], + model_costs: Dict[str, float], + quality_weight: float = DEFAULT_QUALITY_WEIGHT, + cost_weight: float = DEFAULT_COST_WEIGHT, + rng: Optional[random.Random] = None, +) -> str: + """ + Sample once per model, score each, return the model with highest score. + + cells: {model_name: BanditCell} + model_costs: {model_name: $/1k tokens} + """ + if not cells: + raise ValueError("pick_best called with no models") + all_costs = list(model_costs.values()) + best_model: Optional[str] = None + best_score = float("-inf") + for model, cell in cells.items(): + q = thompson_sample(cell, rng=rng) + s = score(q, model_costs[model], all_costs, quality_weight, cost_weight) + if s > best_score: + best_score = s + best_model = model + assert best_model is not None + return best_model diff --git a/litellm/router_strategy/adaptive_router/classifier.py b/litellm/router_strategy/adaptive_router/classifier.py new file mode 100644 index 00000000000..0434dfdb63f --- /dev/null +++ b/litellm/router_strategy/adaptive_router/classifier.py @@ -0,0 +1,140 @@ +""" +Rule-based classifier mapping a user prompt to a RequestType. + +V0 design choice: deterministic regex over the FIRST user message in a session. +Result is cached per session (caller's responsibility, not ours). + +Order matters: we check more specific types first, falling back to GENERAL. +""" + +import re +from typing import List, Pattern, Tuple + +from litellm.types.router import RequestType + +_RULES: List[Tuple[Pattern[str], RequestType]] = [ + ( + re.compile( + r"\b(write|create|generate|implement|build)\s+(?:a |an |the |me )?(?:python|javascript|typescript|java|rust|go|c\+\+|sql|bash|shell)\b", + re.IGNORECASE, + ), + RequestType.CODE_GENERATION, + ), + ( + re.compile( + r"\b(write|create|implement|build)\b(?:\s+\w+){0,4}?\s+(function|class|method|script|program|api|endpoint|microservice)\b", + re.IGNORECASE, + ), + RequestType.CODE_GENERATION, + ), + ( + re.compile( + r"\b(explain|describe|understand|walk me through|what does)\b.*\b(code|function|method|class|algorithm|snippet)\b", + re.IGNORECASE, + ), + RequestType.CODE_UNDERSTANDING, + ), + ( + re.compile( + r"\b(debug|fix|why (?:is|does|isn't)|what.s wrong|trace)\b.*\b(error|bug|exception|stacktrace|stack trace|traceback)\b", + re.IGNORECASE, + ), + RequestType.CODE_UNDERSTANDING, + ), + ( + re.compile( + r"\b(review|critique)\s+(?:this |my |the )?(?:code|pr|pull request|diff|patch)\b", + re.IGNORECASE, + ), + RequestType.CODE_UNDERSTANDING, + ), + ( + re.compile( + r"\b(design|architect|plan|architecture)\b.*\b(system|service|api|database|schema|module|microservice)\b", + re.IGNORECASE, + ), + RequestType.TECHNICAL_DESIGN, + ), + ( + re.compile( + r"\b(should i (?:use|choose|pick)|tradeoffs? between|compare)\b.*\b(library|framework|language|database|protocol|postgres|postgresql|mongodb|dynamodb|mysql|redis|kafka|sql|nosql)\b", + re.IGNORECASE, + ), + RequestType.TECHNICAL_DESIGN, + ), + ( + re.compile( + r"\bhow (?:should|do) i (?:design|structure|organize|model)\b", + re.IGNORECASE, + ), + RequestType.TECHNICAL_DESIGN, + ), + ( + re.compile( + r"\b(solve|compute|calculate|prove|derive)\b.*\b(equation|integral|derivative|theorem|proof|problem)\b", + re.IGNORECASE, + ), + RequestType.ANALYTICAL_REASONING, + ), + ( + re.compile(r"\b(if .+ then|given .+ find|suppose|assume)\b", re.IGNORECASE), + RequestType.ANALYTICAL_REASONING, + ), + ( + re.compile( + r"\b(probability|statistics|combinatorics|optimization problem)\b", + re.IGNORECASE, + ), + RequestType.ANALYTICAL_REASONING, + ), + ( + re.compile( + r"\b(write|draft|compose|rewrite|edit|proofread|polish)\b.*\b(email|essay|blog|post|article|letter|memo|copy|paragraph|sentence)\b", + re.IGNORECASE, + ), + RequestType.WRITING, + ), + ( + re.compile( + r"\b(make (?:this|it)|help me)\s+(?:more |less )?(?:concise|formal|casual|professional|persuasive)\b", + re.IGNORECASE, + ), + RequestType.WRITING, + ), + ( + re.compile( + r"^\s*(who|what|when|where|which)\s+(?:is|was|were|are)\b", re.IGNORECASE + ), + RequestType.FACTUAL_LOOKUP, + ), + ( + re.compile(r"^\s*(define|definition of|meaning of)\b", re.IGNORECASE), + RequestType.FACTUAL_LOOKUP, + ), + ( + re.compile( + r"^\s*how (?:do you spell|to spell|many .* are there|tall is)\b", + re.IGNORECASE, + ), + RequestType.FACTUAL_LOOKUP, + ), +] + + +def classify_prompt(text: str) -> RequestType: + """ + Classify a single user prompt. + + Falls back to GENERAL when no rule matches. Empty/whitespace-only also + returns GENERAL. + """ + if not text or not text.strip(): + return RequestType.GENERAL + + truncated = text[:2000] + + for pattern, request_type in _RULES: + if pattern.search(truncated): + return request_type + + return RequestType.GENERAL diff --git a/litellm/router_strategy/adaptive_router/config.py b/litellm/router_strategy/adaptive_router/config.py new file mode 100644 index 00000000000..e72826cc056 --- /dev/null +++ b/litellm/router_strategy/adaptive_router/config.py @@ -0,0 +1,54 @@ +""" +Configuration constants for the adaptive_router strategy. + +All magic numbers are first-pass guesses (D3-D6 in the handoff plan). +Expect to retune after first 1000 sessions of real traffic. +""" + +from typing import Dict + +from litellm.types.router import RequestType # re-export for convenience # noqa: F401 + +# D3 — Score weights (default; user-overridable via AdaptiveRouterConfig.weights) +DEFAULT_QUALITY_WEIGHT: float = 0.7 # UNVALIDATED — calibrated against [0] sessions +DEFAULT_COST_WEIGHT: float = 0.3 # UNVALIDATED — calibrated against [0] sessions + +# D4 — Cold-start prior: (alpha + beta) total mass = COLD_START_MASS +# Mean of Beta = base_tier_weight + (strength_bonus if declared) +BASE_TIER_WEIGHT: Dict[int, float] = {1: 0.3, 2: 0.5, 3: 0.7} # UNVALIDATED +STRENGTH_BONUS: float = 0.3 # UNVALIDATED +COLD_START_MASS: float = 10.0 + +# D5 — Sample cap. Hard cap, no rescaling (drift handling is v1). +SAMPLE_CAP: int = 200 + +# D6 — Clean-trace credit: minimum turns before α += 1 can fire. +MIN_TURNS_FOR_CLEAN_CREDIT: int = 3 + +# D2 — Owner-cache TTL (seconds). 24h. +# A conversation's first-picked model "owns" the bandit-update slot for +# this long. Subsequent turns of the same conversation only contribute a +# bandit/state update when the same model is re-sampled. +OWNER_CACHE_TTL_SECONDS: int = 24 * 3600 + +# Below this many messages we skip post-call signal recording. Most signals +# (misalignment, stagnation, satisfaction-in-response-to-prior-turn) need at +# least one full prior exchange to be meaningful. +SIGNAL_GATE_MIN_MESSAGES: int = 4 + +# Detector thresholds (from Plano/Chen 2026 paper). +MISALIGNMENT_JACCARD_THRESHOLD: float = 0.45 +STAGNATION_JACCARD_NEAR_DUP: float = 0.50 +LOOP_REPEAT_THRESHOLD: int = 3 +TOOL_CALL_HISTORY_MAX: int = 20 + +# D1 — Caller filter for min quality tier. +MIN_QUALITY_TIER_HEADER: str = "x-litellm-min-quality-tier" +MIN_QUALITY_TIER_METADATA_KEY: str = "min_quality_tier" + +# Pre-routing -> post-call relay: the chosen logical model is stashed on +# request_kwargs["metadata"][ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] by the +# pre-routing hook, then read by the post-call hook to surface as the +# ADAPTIVE_ROUTER_RESPONSE_HEADER response header. +ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY: str = "adaptive_router_chosen_model" +ADAPTIVE_ROUTER_RESPONSE_HEADER: str = "x-litellm-adaptive-router-model" diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py new file mode 100644 index 00000000000..9e346006ac1 --- /dev/null +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -0,0 +1,278 @@ +""" +Post-call hook for the adaptive router. + +On each successful or failed completion, build a Turn from the request/response +and push it through `AdaptiveRouter.record_turn`. The router then updates the +in-memory bandit cell + session state and queues writes for the proxy flusher. + +All work happens after the response has been returned to the caller. Any +exception is swallowed — signal recording must never break a request. +""" + +from __future__ import annotations + +import hashlib +import json +from typing import Any, Dict, List, Optional + +from litellm._logging import verbose_router_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter +from litellm.router_strategy.adaptive_router.classifier import classify_prompt +from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + ADAPTIVE_ROUTER_RESPONSE_HEADER, + SIGNAL_GATE_MIN_MESSAGES, +) +from litellm.router_strategy.adaptive_router.signals import Turn + +# Identity fields hashed into a derived session key so the same conversation +# from the same caller produces a stable key, while different keys/teams/users +# stay segregated even if they happen to send identical first messages. +_IDENTITY_FIELDS = ( + "user_api_key_hash", + "user_api_key_team_id", + "user_api_key_user_id", + "user_api_key_end_user_id", +) + + +def _resolve_session_key(kwargs: Dict[str, Any]) -> Optional[str]: + """Pick a stable per-conversation key for owner-cache attribution. + + Order: + 1. Honor a client-supplied session id (`litellm_session_id` on either + `litellm_params` or `litellm_params.metadata`, or `session_id` on + metadata) — backward compat for callers already wired up. + 2. Otherwise derive a sha256 over (identity fields, first + SIGNAL_GATE_MIN_MESSAGES messages) so the key is stable across turns + and only materialises once there is enough context for the bandit to + act on (matching the gate in the signal-processing path). + + Returns None if the conversation is shorter than SIGNAL_GATE_MIN_MESSAGES. + """ + litellm_params = kwargs.get("litellm_params") or {} + sid = litellm_params.get("litellm_session_id") + if sid: + return str(sid) + metadata = litellm_params.get("metadata") or {} + if isinstance(metadata, dict): + sid = metadata.get("session_id") or metadata.get("litellm_session_id") + if sid: + return str(sid) + + messages = kwargs.get("messages") or [] + if len(messages) < SIGNAL_GATE_MIN_MESSAGES: + # Don't attribute until we have enough turns to match the signal gate — + # ensures the hash is stable (same N messages every time) and avoids + # crediting the bandit for conversations that are too short to signal. + return None + + identity = ":".join( + str(metadata.get(f) or "") if isinstance(metadata, dict) else "" + for f in _IDENTITY_FIELDS + ) + anchor = messages[:SIGNAL_GATE_MIN_MESSAGES] + payload = ( + identity + + "|" + + json.dumps( + [{"role": m.get("role"), "content": m.get("content")} for m in anchor], + sort_keys=True, + default=str, + ) + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _last_user_content(messages: Optional[List[Dict[str, Any]]]) -> Optional[str]: + if not messages: + return None + for msg in reversed(messages): + if msg.get("role") == "user": + content = msg.get("content") + if isinstance(content, str): + return content + if isinstance(content, list): + # OpenAI vision-style content: pick first text part. + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + return part.get("text") + return None + return None + + +def _recent_tool_results( + messages: Optional[List[Dict[str, Any]]] +) -> List[Dict[str, Any]]: + """Extract the current turn's tool result payloads from the request messages. + + Tool results are `role == "tool"` messages that sit at the tail of the + conversation — i.e. after the most recent assistant message with + `tool_calls`, waiting for the model to produce a user-facing reply. Walk + backwards from the end and collect the contiguous run of tool messages; + stop at the first non-tool message. + + Each result is normalized to `{content, is_error}` — the only fields + `signals._detect_failure` / `_detect_exhaustion` actually read. + """ + if not messages: + return [] + results: List[Dict[str, Any]] = [] + for msg in reversed(messages): + if not isinstance(msg, dict): + break + if msg.get("role") != "tool": + break + content = msg.get("content") + # Some providers (Anthropic-style) carry an explicit error flag; OpenAI + # tool results don't, so fall back to an empty/missing content heuristic + # inside `_detect_failure`. + is_error = bool(msg.get("is_error")) + results.append({"content": content, "is_error": is_error}) + results.reverse() + return results + + +def _assistant_content_and_tool_calls(response_obj: Any) -> tuple: + """Return (assistant_text, tool_calls_list) extracted from a ModelResponse-ish object.""" + if response_obj is None: + return None, [] + try: + choices = getattr(response_obj, "choices", None) or response_obj.get("choices") + except Exception: + return None, [] + if not choices: + return None, [] + + msg = choices[0] + msg = getattr(msg, "message", None) or ( + msg.get("message") if isinstance(msg, dict) else None + ) + if msg is None: + return None, [] + + content = getattr(msg, "content", None) + if content is None and isinstance(msg, dict): + content = msg.get("content") + + raw_tool_calls = getattr(msg, "tool_calls", None) + if raw_tool_calls is None and isinstance(msg, dict): + raw_tool_calls = msg.get("tool_calls") + tool_calls: List[Dict[str, Any]] = [] + for tc in raw_tool_calls or []: + if isinstance(tc, dict): + tool_calls.append(tc) + else: + try: + tool_calls.append(tc.model_dump()) + except Exception: + tool_calls.append({"name": getattr(tc, "name", ""), "arguments": ""}) + return content, tool_calls + + +class AdaptiveRouterPostCallHook(CustomLogger): + """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" + + def __init__(self, adaptive_router: AdaptiveRouter) -> None: + self.adaptive_router = adaptive_router + + async def async_post_call_response_headers_hook( + self, + data: Dict[str, Any], + user_api_key_dict: Any, + response: Any, + request_headers: Optional[Dict[str, str]] = None, + litellm_call_info: Optional[Dict[str, Any]] = None, + ) -> Optional[Dict[str, str]]: + """ + Surface the chosen logical model as the `x-litellm-adaptive-router-model` + response header for both streaming and non-streaming responses. + + `async_post_call_success_hook` fires after the stream is fully consumed, + so writing to `_hidden_params["additional_headers"]` there is too late for + streaming — the StreamingResponse headers are already frozen. This hook is + called during header construction (before StreamingResponse is built), so + the header is included for both paths. + """ + metadata = data.get("metadata") or {} + chosen = ( + metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY) + if isinstance(metadata, dict) + else None + ) + if not chosen: + return None + return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen} + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + await self._record(kwargs, response_obj, response_status=200) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + status = kwargs.get("response_status") + if status is None: + exc = kwargs.get("exception") + status = getattr(exc, "status_code", 500) if exc is not None else 500 + await self._record(kwargs, response_obj, response_status=int(status)) + + async def _record( + self, + kwargs: Dict[str, Any], + response_obj: Any, + response_status: int, + ) -> None: + try: + messages = kwargs.get("messages") or [] + if len(messages) < SIGNAL_GATE_MIN_MESSAGES: + # Too few turns for any signal to be meaningful — skip. + return + + session_key = _resolve_session_key(kwargs) + if not session_key: + return + + # The bandit cells are keyed by the *logical* model name from + # `available_models` (e.g. "smart"/"fast"). `kwargs["model"]` at + # post-call time is the physical upstream model + # (e.g. "anthropic/claude-opus-4-7"), so it cannot be used directly. + # The pre-routing hook stashes the logical pick under this key. + litellm_params = kwargs.get("litellm_params") or {} + metadata = litellm_params.get("metadata") or {} + current_model = ( + metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY) + if isinstance(metadata, dict) + else None + ) + if not current_model: + return + + if not self.adaptive_router.claim_or_check_owner( + session_key, current_model + ): + # A different model owns this conversation — skip attribution. + return + + user_text = _last_user_content(messages) + assistant_text, tool_calls = _assistant_content_and_tool_calls(response_obj) + tool_results = _recent_tool_results(messages) + + request_type = classify_prompt(user_text or "") + turn = Turn( + user_content=user_text, + assistant_content=( + assistant_text if isinstance(assistant_text, str) else None + ), + tool_calls=tool_calls, + tool_results=tool_results, + response_status=response_status, + ) + await self.adaptive_router.record_turn( + session_id=session_key, + model_name=current_model, + request_type=request_type, + turn=turn, + ) + except Exception as e: + verbose_router_logger.exception( + "AdaptiveRouterPostCallHook: failed to record turn: %s", e + ) diff --git a/litellm/router_strategy/adaptive_router/signals.py b/litellm/router_strategy/adaptive_router/signals.py new file mode 100644 index 00000000000..a48bdea1eb6 --- /dev/null +++ b/litellm/router_strategy/adaptive_router/signals.py @@ -0,0 +1,287 @@ +""" +Incremental signal detection for the adaptive router. + +Each session maintains a SessionState. On every turn, we call apply_turn(state, turn) +which mutates the state in place and returns a SignalDelta listing which signals +fired on THIS turn. The router then queues the delta to be flushed to DB. + +Design constraint: O(1) work per turn. No re-scanning the full session history. +We keep small bounded windows: last_user_content, last_assistant_content, and a +bounded list of recent tool call signatures. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Set + +from litellm.router_strategy.adaptive_router.config import ( + LOOP_REPEAT_THRESHOLD, + MIN_TURNS_FOR_CLEAN_CREDIT, + MISALIGNMENT_JACCARD_THRESHOLD, + STAGNATION_JACCARD_NEAR_DUP, + TOOL_CALL_HISTORY_MAX, +) + + +# ---- Public types --------------------------------------------------------- + + +@dataclass +class SignalDelta: + """Which signals fired on a single turn. Counts are 0 or 1 (one delta per turn).""" + + misalignment: int = 0 + stagnation: int = 0 + disengagement: int = 0 + satisfaction: int = 0 + failure: int = 0 + loop: int = 0 + exhaustion: int = 0 + + def any_fired(self) -> bool: + return any( + [ + self.misalignment, + self.stagnation, + self.disengagement, + self.satisfaction, + self.failure, + self.loop, + self.exhaustion, + ] + ) + + +@dataclass +class SessionState: + """In-memory rolling state for one session. + + Mirrors the LiteLLM_AdaptiveRouterSession DB row (Wave 0 schema). The flusher + later persists this. We keep this as a plain dataclass — no DB coupling. + """ + + session_id: str + router_name: str + model_name: str + classified_type: str + + misalignment_count: int = 0 + stagnation_count: int = 0 + disengagement_count: int = 0 + satisfaction_count: int = 0 + failure_count: int = 0 + loop_count: int = 0 + exhaustion_count: int = 0 + + last_user_content: Optional[str] = None + last_assistant_content: Optional[str] = None + tool_call_history: List[str] = field(default_factory=list) + pending_tool_calls: Dict[str, str] = field(default_factory=dict) + + turn_count: int = 0 + last_processed_turn: int = -1 + clean_credit_awarded: bool = False + terminal_status: Optional[int] = None + + +@dataclass +class Turn: + """One turn of input. Caller assembles this from the request/response.""" + + user_content: Optional[str] = None + assistant_content: Optional[str] = None + tool_calls: List[Dict[str, Any]] = field(default_factory=list) + tool_results: List[Dict[str, Any]] = field(default_factory=list) + response_status: Optional[int] = None + + +# ---- Detection helpers ---------------------------------------------------- + +_TOKEN_RE = re.compile(r"[A-Za-z0-9]+") + + +def _tokens(text: Optional[str]) -> Set[str]: + if not text: + return set() + return {t.lower() for t in _TOKEN_RE.findall(text)} + + +def _jaccard(a: Set[str], b: Set[str]) -> float: + union = a | b + if not union: + return 0.0 + return len(a & b) / len(union) + + +_DISENGAGEMENT_PATTERNS = [ + re.compile( + r"\b(forget it|never mind|give up|talk to (?:a )?human|cancel)\b", re.IGNORECASE + ), + re.compile(r"\b(this (?:isn'?t|is not) working|stop|abort)\b", re.IGNORECASE), + re.compile(r"\bi'?ll do it (?:myself|manually)\b", re.IGNORECASE), +] + +_SATISFACTION_PATTERNS = [ + re.compile( + r"\b(that worked|that did it|works now|fixed it|solved it|nice)\b", + re.IGNORECASE, + ), + re.compile(r"\b(thanks|thank you|thx|appreciated|appreciate it)\b", re.IGNORECASE), + re.compile(r"\b(perfect|great|excellent|exactly)\b", re.IGNORECASE), +] + + +def _detect_misalignment(prev_user: Optional[str], curr_user: Optional[str]) -> bool: + """Fires when consecutive user messages share *some* topic (jaccard > 0) + but are sufficiently different (jaccard < threshold) — i.e. user is + rephrasing, not changing topic, not repeating.""" + if not prev_user or not curr_user: + return False + j = _jaccard(_tokens(prev_user), _tokens(curr_user)) + return 0.0 < j < MISALIGNMENT_JACCARD_THRESHOLD + + +def _detect_stagnation(prev_asst: Optional[str], curr_asst: Optional[str]) -> bool: + """Fires when consecutive assistant messages are near-duplicates.""" + if not prev_asst or not curr_asst: + return False + j = _jaccard(_tokens(prev_asst), _tokens(curr_asst)) + return j >= STAGNATION_JACCARD_NEAR_DUP + + +def _detect_disengagement(curr_user: Optional[str]) -> bool: + if not curr_user: + return False + return any(p.search(curr_user) for p in _DISENGAGEMENT_PATTERNS) + + +def _detect_satisfaction(curr_user: Optional[str]) -> bool: + if not curr_user: + return False + return any(p.search(curr_user) for p in _SATISFACTION_PATTERNS) + + +def _detect_failure(tool_results: List[Dict[str, Any]]) -> bool: + """Any tool result explicitly flagged as an error. + + We do NOT treat empty content as failure — many tools legitimately return + empty output (zero-result searches, silent bash commands, void writes) and + penalizing the model for those would corrupt the bandit posterior. + """ + for r in tool_results: + if r.get("is_error"): + return True + return False + + +def _signature(call: Dict[str, Any]) -> str: + """Stable signature for loop detection: name + sorted JSON-ish args.""" + name = call.get("name") or call.get("function", {}).get("name", "") + call_args = call.get("arguments") + if call_args is None: + call_args = call.get("function", {}).get("arguments", "") + if isinstance(call_args, dict): + call_args = ",".join(f"{k}={call_args[k]}" for k in sorted(call_args.keys())) + return f"{name}({call_args})" + + +def _detect_loop(history: List[str], new_calls: List[Dict[str, Any]]) -> bool: + """Fires if any new call's signature appears >= LOOP_REPEAT_THRESHOLD-1 times + in recent history (so this call would be the Nth).""" + if not new_calls: + return False + for call in new_calls: + sig = _signature(call) + recent_count = history.count(sig) + if recent_count >= LOOP_REPEAT_THRESHOLD - 1: + return True + return False + + +_EXHAUSTION_STATUSES = {408, 413, 429, 503, 504} + +_EXHAUSTION_KEYWORDS = ( + "context length", + "context window", + "token limit", + "rate limit", + "too many requests", + "timeout", +) + + +def _detect_exhaustion( + status: Optional[int], tool_results: List[Dict[str, Any]] +) -> bool: + if status is not None and status in _EXHAUSTION_STATUSES: + return True + for r in tool_results: + content = str(r.get("content", "")).lower() + if any(kw in content for kw in _EXHAUSTION_KEYWORDS): + return True + return False + + +# ---- Public entrypoint ---------------------------------------------------- + + +def apply_turn(state: SessionState, turn: Turn) -> SignalDelta: + """ + Detect signals on this turn, mutate state, return the delta. + + O(1) per turn (no full-history rescan). Only inspects last_*, recent tool history + (which is bounded at TOOL_CALL_HISTORY_MAX), and the new turn payload. + """ + delta = SignalDelta() + + if _detect_misalignment(state.last_user_content, turn.user_content): + delta.misalignment = 1 + if _detect_stagnation(state.last_assistant_content, turn.assistant_content): + delta.stagnation = 1 + if _detect_disengagement(turn.user_content): + delta.disengagement = 1 + if _detect_satisfaction(turn.user_content): + # Gate: only award satisfaction credit once per session, and only + # after MIN_TURNS_FOR_CLEAN_CREDIT turns of context. Early "thanks" + # on turn 1-2 is noise, not a validated quality signal. + current_turn_index = state.turn_count + 1 + if ( + not state.clean_credit_awarded + and current_turn_index >= MIN_TURNS_FOR_CLEAN_CREDIT + ): + delta.satisfaction = 1 + state.clean_credit_awarded = True + if _detect_failure(turn.tool_results): + delta.failure = 1 + if _detect_loop(state.tool_call_history, turn.tool_calls): + delta.loop = 1 + if _detect_exhaustion(turn.response_status, turn.tool_results): + delta.exhaustion = 1 + + state.misalignment_count += delta.misalignment + state.stagnation_count += delta.stagnation + state.disengagement_count += delta.disengagement + state.satisfaction_count += delta.satisfaction + state.failure_count += delta.failure + state.loop_count += delta.loop + state.exhaustion_count += delta.exhaustion + + if turn.user_content: + state.last_user_content = turn.user_content + if turn.assistant_content: + state.last_assistant_content = turn.assistant_content + + for call in turn.tool_calls: + state.tool_call_history.append(_signature(call)) + if len(state.tool_call_history) > TOOL_CALL_HISTORY_MAX: + state.tool_call_history = state.tool_call_history[-TOOL_CALL_HISTORY_MAX:] + + if turn.response_status is not None: + state.terminal_status = turn.response_status + + state.turn_count += 1 + state.last_processed_turn = state.turn_count + + return delta diff --git a/litellm/router_strategy/adaptive_router/update_queue.py b/litellm/router_strategy/adaptive_router/update_queue.py new file mode 100644 index 00000000000..b667f3a53a7 --- /dev/null +++ b/litellm/router_strategy/adaptive_router/update_queue.py @@ -0,0 +1,213 @@ +""" +In-memory queues for adaptive router state and session updates. + +Pattern follows DailySpendUpdateQueue: hot path is fully in-memory; a background +flusher task drains the aggregator and writes batches to Postgres. + +Two logical queues (one class): + 1. STATE updates: increments to (router, request_type, model) bandit cell. + Aggregator key = (router_name, request_type, model_name) + Aggregated payload = {"delta_alpha": float, "delta_beta": float, "samples_added": int} + 2. SESSION updates: full snapshot of a session row (last-write-wins per session+router+model). + Aggregator key = (session_id, router_name, model_name) + Aggregated payload = the full session state dict. + +Hot-path API is non-blocking and synchronous from the caller's POV (it just appends +to the in-memory aggregator). Flush is async and batched. +""" + +from __future__ import annotations + +import asyncio +from typing import Any, Dict, Tuple + +from litellm._logging import verbose_router_logger + +StateKey = Tuple[str, str, str] # (router_name, request_type, model_name) +SessionKey = Tuple[str, str, str] # (session_id, router_name, model_name) + + +class AdaptiveRouterUpdateQueue: + """ + Single class managing both state-update aggregation and session-snapshot aggregation. + Held by the AdaptiveRouter strategy instance and started by the proxy on boot. + """ + + def __init__(self) -> None: + self._state_agg: Dict[StateKey, Dict[str, float]] = {} + self._session_agg: Dict[SessionKey, Dict[str, Any]] = {} + self._lock = asyncio.Lock() + self._max_state_size_seen = 0 + self._max_session_size_seen = 0 + + # ---- Hot-path: state delta ------------------------------------------- + + async def add_state_delta( + self, + router_name: str, + request_type: str, + model_name: str, + delta_alpha: float, + delta_beta: float, + ) -> None: + """Aggregate a bandit-cell delta. Multiple deltas to the same cell sum.""" + key: StateKey = (router_name, request_type, model_name) + async with self._lock: + current = self._state_agg.get(key) + if current is None: + self._state_agg[key] = { + "delta_alpha": delta_alpha, + "delta_beta": delta_beta, + "samples_added": 1, + } + else: + current["delta_alpha"] += delta_alpha + current["delta_beta"] += delta_beta + current["samples_added"] += 1 + if len(self._state_agg) > self._max_state_size_seen: + self._max_state_size_seen = len(self._state_agg) + + # ---- Hot-path: session snapshot -------------------------------------- + + async def add_session_state( + self, + session_id: str, + router_name: str, + model_name: str, + state_dict: Dict[str, Any], + ) -> None: + """ + Last-write-wins per session row. The state_dict is a snapshot of the + SessionState (signals counts + bookkeeping fields). The flusher will + upsert this into LiteLLM_AdaptiveRouterSession. + """ + key: SessionKey = (session_id, router_name, model_name) + async with self._lock: + self._session_agg[key] = state_dict + if len(self._session_agg) > self._max_session_size_seen: + self._max_session_size_seen = len(self._session_agg) + + # ---- Flushers (called by background task) ---------------------------- + + async def flush_state_to_db(self, prisma_client: Any) -> int: + """ + Drain state aggregator and apply to LiteLLM_AdaptiveRouterState. + Returns number of cells flushed. + """ + async with self._lock: + batch = self._state_agg + self._state_agg = {} + + if not batch: + return 0 + + # Sort keys to give deterministic write order across writers and + # reduce the chance of cross-row deadlocks when other workers race us. + for key in sorted(batch.keys()): + router, rt, model = key + payload = batch[key] + try: + # Atomic increment: push the delta directly into the DB so + # concurrent flushers from multiple pods don't overwrite each + # other. The upsert creates the row with the delta as the + # initial value on first write, then increments on subsequent + # writes — no read-modify-write race. + await prisma_client.db.litellm_adaptiverouterstate.upsert( + where={ + "router_name_request_type_model_name": { + "router_name": router, + "request_type": rt, + "model_name": model, + } + }, + data={ + "create": { + "router_name": router, + "request_type": rt, + "model_name": model, + "alpha": payload["delta_alpha"], + "beta": payload["delta_beta"], + "total_samples": int(payload["samples_added"]), + }, + "update": { + "alpha": {"increment": payload["delta_alpha"]}, + "beta": {"increment": payload["delta_beta"]}, + "total_samples": { + "increment": int(payload["samples_added"]) + }, + }, + }, + ) + except Exception as e: + verbose_router_logger.exception( + "AdaptiveRouterUpdateQueue: failed to flush state for %s: %s", + key, + e, + ) + + return len(batch) + + async def flush_session_to_db(self, prisma_client: Any) -> int: + """ + Drain session aggregator and upsert into LiteLLM_AdaptiveRouterSession. + Returns number of session rows flushed. + """ + async with self._lock: + batch = self._session_agg + self._session_agg = {} + + if not batch: + return 0 + + for key in sorted(batch.keys()): + session_id, router, model = key + payload = batch[key] + try: + # NOTE: Prisma client lower-cases model names, so + # `LiteLLM_AdaptiveRouterSession` -> `litellm_adaptiveroutersession` + # (single 's', not 'litellm_adaptiverouterssession'). + # Strip PK fields from the update payload — Prisma rejects + # writes to fields that are part of the @@id. asdict(state) + # always carries them, so build a separate update dict. + update_payload = { + k: v + for k, v in payload.items() + if k not in ("session_id", "router_name", "model_name") + } + await prisma_client.db.litellm_adaptiveroutersession.upsert( + where={ + "session_id_router_name_model_name": { + "session_id": session_id, + "router_name": router, + "model_name": model, + } + }, + data={ + "create": { + "session_id": session_id, + "router_name": router, + "model_name": model, + **update_payload, + }, + "update": update_payload, + }, + ) + except Exception as e: + verbose_router_logger.exception( + "AdaptiveRouterUpdateQueue: failed to flush session for %s: %s", + key, + e, + ) + + return len(batch) + + # ---- Observability --------------------------------------------------- + + async def queue_size(self) -> Dict[str, int]: + async with self._lock: + return { + "state_pending": len(self._state_agg), + "session_pending": len(self._session_agg), + "max_state_seen": self._max_state_size_seen, + "max_session_seen": self._max_session_size_seen, + } diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index 4ead7225abc..58b2c5a3912 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -82,11 +82,34 @@ class AutoRouter(CustomLogger): ) return auto_router_routes + @staticmethod + def _extract_text_from_messages(messages: List[Dict[str, Any]]) -> str: + """ + Extract text content from the last user message for routing. + + Handles tool-call conversations (where the last message may be an + assistant or tool message with non-string content) and multimodal + messages (where content is a list of content blocks). + """ + for msg in reversed(messages): + if msg.get("role") == "user": + content = msg.get("content") + if content is None: + return "" + if isinstance(content, list): + return " ".join( + block.get("text", "") + for block in content + if isinstance(block, dict) and block.get("type") == "text" + ) + return str(content) + return "" + async def async_pre_routing_hook( self, model: str, request_kwargs: Dict, - messages: Optional[List[Dict[str, str]]] = None, + messages: Optional[List[Dict[str, Any]]] = None, input: Optional[Union[str, List]] = None, specific_deployment: Optional[bool] = False, ) -> Optional["PreRoutingHookResponse"]: @@ -120,8 +143,7 @@ class AutoRouter(CustomLogger): auto_sync=self.auto_sync_value, ) - user_message: Dict[str, str] = messages[-1] - message_content: str = user_message.get("content", "") + message_content = self._extract_text_from_messages(messages) route_choice: Optional[Union[RouteChoice, List[RouteChoice]]] = self.routelayer( text=message_content ) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index e51249b1cb1..aa3bcef6392 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -332,45 +332,68 @@ class ComplexityRouter(CustomLogger): f"No model configured for tier {tier_key} and no default_model set" ) - async def async_pre_routing_hook( + def _resolve_messages( self, - model: str, + messages: Optional[List[Dict[str, Any]]], request_kwargs: Dict, - messages: Optional[List[Dict[str, Any]]] = None, - input: Optional[Union[str, List]] = None, - specific_deployment: Optional[bool] = False, - ) -> Optional["PreRoutingHookResponse"]: + ) -> Optional[List[Dict[str, Any]]]: """ - Pre-routing hook called before the routing decision. + Resolve messages from the request, converting from other formats if needed. - Classifies the request by complexity and returns the appropriate model. - - Args: - model: The original model name requested. - request_kwargs: The request kwargs. - messages: The messages in the request. - input: Optional input for embeddings. - specific_deployment: Whether a specific deployment was requested. - - Returns: - PreRoutingHookResponse with the routed model, or None if no routing needed. + Uses the guardrail translation handler dispatch to convert Responses API + ``input`` (or other non-chat-completions formats) into OpenAI-spec messages. """ - from litellm.types.router import PreRoutingHookResponse + if messages: + return messages - if messages is None or len(messages) == 0: - verbose_router_logger.debug( - "ComplexityRouter: No messages provided, skipping routing" - ) - return None + from litellm.litellm_core_utils.api_route_to_call_types import ( + get_call_types_for_route, + ) + from litellm.llms import load_guardrail_translation_mappings + from litellm.types.utils import CallTypes - # Extract the last user message and the last system prompt + mappings = load_guardrail_translation_mappings() + call_type: Optional[CallTypes] = None + + # 1. Try route-based inference from proxy metadata + route = request_kwargs.get("litellm_metadata", {}).get( + "user_api_key_request_route" + ) + if route: + call_types_list = get_call_types_for_route(route) + if call_types_list: + for ct in call_types_list: + if ct in mappings: + call_type = ct + break + + # 2. Fallback: try each mapped handler until one produces messages + handlers_to_try: List[Any] = [] + if call_type is not None and call_type in mappings: + handlers_to_try.append(mappings[call_type]()) + else: + handlers_to_try.extend(handler_cls() for handler_cls in mappings.values()) + + for handler in handlers_to_try: + structured = handler.get_structured_messages(request_kwargs) + if structured: + return [ + msg if isinstance(msg, dict) else msg.model_dump() # type: ignore + for msg in structured + ] + return None + + @staticmethod + def _extract_user_message_and_system_prompt( + messages: List[Dict[str, Any]], + ) -> Tuple[Optional[str], Optional[str]]: + """Extract the last user message text and last system prompt from messages.""" user_message: Optional[str] = None system_prompt: Optional[str] = None for msg in reversed(messages): role = msg.get("role", "") content = msg.get("content") or "" - # content may be a list of content parts (e.g. [{"type": "text", "text": "..."}]) if isinstance(content, list): text_parts = [ part.get("text", "") @@ -383,6 +406,52 @@ class ComplexityRouter(CustomLogger): user_message = content elif role == "system" and system_prompt is None: system_prompt = content + if user_message is not None and system_prompt is not None: + break + + return user_message, system_prompt + + async def async_pre_routing_hook( + self, + model: str, + request_kwargs: Dict, + messages: Optional[List[Dict[str, Any]]] = None, + input: Optional[Union[str, List]] = None, + specific_deployment: Optional[bool] = False, + ) -> Optional["PreRoutingHookResponse"]: + """ + Pre-routing hook called before the routing decision. + + Classifies the request by complexity and returns the appropriate model. + Supports chat completions (messages), Responses API (input), and other + formats via the guardrail translation handler dispatch. + + Args: + model: The original model name requested. + request_kwargs: The request kwargs. + messages: The messages in the request. + input: Optional input for Responses API or embeddings. + specific_deployment: Whether a specific deployment was requested. + + Returns: + PreRoutingHookResponse with the routed model, or None if no routing needed. + """ + from litellm.types.router import PreRoutingHookResponse + + resolved_messages = self._resolve_messages(messages, request_kwargs) + + if not resolved_messages: + verbose_router_logger.debug( + "ComplexityRouter: No messages could be resolved, skipping routing" + ) + return None + + # Determine whether the original request used messages directly + has_original_messages = messages is not None and len(messages) > 0 + + user_message, system_prompt = self._extract_user_message_and_system_prompt( + resolved_messages + ) if user_message is None: verbose_router_logger.debug( @@ -391,13 +460,10 @@ class ComplexityRouter(CustomLogger): return PreRoutingHookResponse( model=self.config.default_model or self.get_model_for_tier(ComplexityTier.MEDIUM), - messages=messages, + messages=messages if has_original_messages else None, ) - # Classify the request tier, score, signals = self.classify(user_message, system_prompt) - - # Get the model for this tier routed_model = self.get_model_for_tier(tier) verbose_router_logger.info( @@ -407,5 +473,5 @@ class ComplexityRouter(CustomLogger): return PreRoutingHookResponse( model=routed_model, - messages=messages, + messages=messages if has_original_messages else None, ) diff --git a/litellm/router_strategy/quality_router/__init__.py b/litellm/router_strategy/quality_router/__init__.py new file mode 100644 index 00000000000..5728943448a --- /dev/null +++ b/litellm/router_strategy/quality_router/__init__.py @@ -0,0 +1,21 @@ +""" +Quality-tier auto-router. + +Re-uses the ComplexityRouter's classification to decide a request's complexity, +then maps that complexity to an admin-configured quality tier and resolves the +target model from each candidate's `model_info.litellm_routing_preferences`. +""" + +from .config import ( + DEFAULT_COMPLEXITY_TO_QUALITY, + QualityRouterConfig, + RoutingPreferences, +) +from .quality_router import QualityRouter + +__all__ = [ + "QualityRouter", + "QualityRouterConfig", + "RoutingPreferences", + "DEFAULT_COMPLEXITY_TO_QUALITY", +] diff --git a/litellm/router_strategy/quality_router/config.py b/litellm/router_strategy/quality_router/config.py new file mode 100644 index 00000000000..125ecd5bb9b --- /dev/null +++ b/litellm/router_strategy/quality_router/config.py @@ -0,0 +1,74 @@ +""" +Configuration models for the QualityRouter. +""" + +from typing import Dict, List, Optional + +from pydantic import BaseModel, ConfigDict, Field + +# Default mapping from ComplexityTier name (string) to quality tier (int). +# Higher tier = higher capability requirement. +DEFAULT_COMPLEXITY_TO_QUALITY: Dict[str, int] = { + "SIMPLE": 1, + "MEDIUM": 2, + "COMPLEX": 3, + "REASONING": 4, +} + + +class QualityRouterConfig(BaseModel): + """Configuration for the QualityRouter.""" + + available_models: List[str] = Field( + default_factory=list, + description=( + "List of candidate model names this router may route to. Each model " + "must declare its quality_tier in model_info.litellm_routing_preferences." + ), + ) + + default_model: Optional[str] = Field( + default=None, + description="Fallback model when no quality tier resolves.", + ) + + complexity_to_quality: Dict[str, int] = Field( + default_factory=lambda: DEFAULT_COMPLEXITY_TO_QUALITY.copy(), + description="Mapping from ComplexityTier name to quality tier (int).", + ) + + model_config = ConfigDict(extra="allow") + + +class RoutingPreferences(BaseModel): + """Per-deployment routing preferences declared on model_info.""" + + quality_tier: int = Field( + ..., + description="The quality tier this deployment satisfies.", + ) + + keywords: List[str] = Field( + default_factory=list, + description=( + "Substring keywords (case-insensitive) that, when present in the " + "user message, route the request to this deployment. See `order` " + "for explicit collision handling, otherwise ties fall through to " + "(highest quality_tier, then cheapest model_info.input_cost_per_token)." + ), + ) + + order: Optional[int] = Field( + default=None, + description=( + "Explicit priority used to break ties between deployments at the " + "same quality tier. Lower values win. Applies both to keyword " + "collisions and to picking between multiple deployments at the " + "same quality_tier. Tiebreak order is " + "(quality_tier DESC, order ASC, input_cost_per_token ASC, " + "model_name ASC) — quality always wins first, then explicit " + "order, then price." + ), + ) + + model_config = ConfigDict(extra="allow") diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py new file mode 100644 index 00000000000..a79b4384f5e --- /dev/null +++ b/litellm/router_strategy/quality_router/quality_router.py @@ -0,0 +1,446 @@ +""" +Quality-tier Auto Router. + +Routes a request to a model at a target quality tier. The quality tier is +inferred by re-using the existing ComplexityRouter's classification, then +mapped through an admin-configured `complexity_to_quality` table. Each +candidate model declares its own `quality_tier` in +`model_info.litellm_routing_preferences`. + +Optional keyword override: deployments may also declare `keywords` in +`litellm_routing_preferences`. If any declared keyword appears in the user +message (case-insensitive substring match), the router short-circuits the +complexity-classification flow and routes to the matching deployment. When +multiple deployments match, ties are broken by (highest quality_tier first, +then cheapest `model_info.input_cost_per_token`). +""" + +import math +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +from litellm._logging import verbose_router_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_strategy.complexity_router.complexity_router import ( + ComplexityRouter, +) + +from .config import QualityRouterConfig, RoutingPreferences + +if TYPE_CHECKING: + from litellm.router import Router + from litellm.types.router import PreRoutingHookResponse +else: + Router = Any + PreRoutingHookResponse = Any + + +class QualityRouter(CustomLogger): + """ + Routes requests to a model at a target quality tier, with an optional + keyword override. + """ + + def __init__( + self, + model_name: str, + litellm_router_instance: "Router", + default_model: Optional[str] = None, + quality_router_config: Optional[Dict[str, Any]] = None, + ): + self.model_name = model_name + self.litellm_router_instance = litellm_router_instance + + if quality_router_config: + self.config = QualityRouterConfig(**quality_router_config) + else: + self.config = QualityRouterConfig() + + # Explicit default_model arg overrides anything in the config dict. + if default_model: + self.config.default_model = default_model + + # Internal scorer — re-use the existing rule-based classifier. + self._scorer = ComplexityRouter( + model_name=f"{model_name}::scorer", + litellm_router_instance=litellm_router_instance, + ) + + # Per-model indices populated alongside the tier index. `_model_keywords` + # stores keywords lowercased so we can substring-match against the + # lowercased user message in O(total-keyword-count). `_model_quality`, + # `_model_cost`, and `_model_order` drive tiebreaking — `_model_order` + # is the explicit priority (lower wins, unset = +inf). + self._model_keywords: Dict[str, List[str]] = {} + self._model_quality: Dict[str, int] = {} + self._model_cost: Dict[str, Optional[float]] = {} + self._model_order: Dict[str, Optional[int]] = {} + + # Tier → models index. Built lazily on first access so the QualityRouter + # deployment does NOT need to appear after all its referenced models in + # the config — when `_build_tier_index` runs eagerly in `__init__`, the + # router instance's `model_list` is still being assembled incrementally + # by `_create_deployment`, and any `available_models` defined AFTER the + # router entry in config.yaml would silently be reported as missing. + self._tier_to_models_cache: Optional[Dict[int, List[str]]] = None + + verbose_router_logger.debug( + f"QualityRouter initialized for {model_name} with " + f"available_models={self.config.available_models}, " + f"default_model={self.config.default_model}" + ) + + @property + def _tier_to_models(self) -> Dict[int, List[str]]: + """Lazy tier→models index; built on first access.""" + if self._tier_to_models_cache is None: + self._tier_to_models_cache = self._build_tier_index() + return self._tier_to_models_cache + + def _get_routing_preferences(self, deployment: Any) -> Optional[Dict[str, Any]]: + """ + Extract litellm_routing_preferences from a deployment, handling both + dict-shaped and Pydantic-object-shaped deployments. + """ + # Dict-shaped deployment. + if isinstance(deployment, dict): + model_info = deployment.get("model_info") or {} + if isinstance(model_info, dict): + return model_info.get("litellm_routing_preferences") + # Pydantic ModelInfo nested in a dict. + return getattr(model_info, "litellm_routing_preferences", None) + + # Pydantic-object deployment. + model_info = getattr(deployment, "model_info", None) + if model_info is None: + return None + if isinstance(model_info, dict): + return model_info.get("litellm_routing_preferences") + return getattr(model_info, "litellm_routing_preferences", None) + + def _get_deployment_input_cost(self, deployment: Any) -> Optional[float]: + """ + Extract `input_cost_per_token` from a deployment's model_info. + + Returns None when not declared — None is treated as "infinite cost" + for the cheapest-tiebreak ordering, so unpriced models lose ties to + priced ones. (Admins who want a model to win on price must declare it.) + """ + if isinstance(deployment, dict): + model_info = deployment.get("model_info") or {} + else: + model_info = getattr(deployment, "model_info", None) or {} + + if isinstance(model_info, dict): + cost = model_info.get("input_cost_per_token") + else: + cost = getattr(model_info, "input_cost_per_token", None) + + if cost is None: + return None + try: + return float(cost) + except (TypeError, ValueError): + return None + + def _get_deployment_model_name(self, deployment: Any) -> Optional[str]: + """Extract `model_name` from a dict- or object-shaped deployment.""" + if isinstance(deployment, dict): + return deployment.get("model_name") + return getattr(deployment, "model_name", None) + + def _build_tier_index(self) -> Dict[int, List[str]]: + """ + Build {quality_tier: [model_name, ...]} for every model in + `available_models`, plus side indices `_model_keywords`, + `_model_quality`, and `_model_cost`. Raises if any listed model is + missing `litellm_routing_preferences`. + """ + model_list = getattr(self.litellm_router_instance, "model_list", None) or [] + available = set(self.config.available_models) + + # Track which available models we've matched so we can error on missing. + seen: Dict[str, bool] = {name: False for name in available} + tier_to_models: Dict[int, List[str]] = {} + + for deployment in model_list: + name = self._get_deployment_model_name(deployment) + if name is None or name not in available: + continue + + raw_prefs = self._get_routing_preferences(deployment) + if raw_prefs is None: + raise ValueError( + f"QualityRouter: model '{name}' is listed in available_models " + f"but has no model_info.litellm_routing_preferences" + ) + + # Validate via the Pydantic model so we get a clear error for + # missing quality_tier, wrong types, etc. This also means + # `RoutingPreferences` is the single source of truth for the + # accepted shape — readers relied on raw dicts before. + try: + if isinstance(raw_prefs, RoutingPreferences): + prefs = raw_prefs + elif isinstance(raw_prefs, dict): + prefs = RoutingPreferences(**raw_prefs) + else: + # A Pydantic object of some other shape — coerce via its dict. + prefs = RoutingPreferences( + **( + raw_prefs.model_dump() + if hasattr(raw_prefs, "model_dump") + else dict(raw_prefs) + ) + ) + except Exception as e: + raise ValueError( + f"QualityRouter: model '{name}' has invalid " + f"litellm_routing_preferences: {e}" + ) from e + + tier_int = int(prefs.quality_tier) + tier_to_models.setdefault(tier_int, []).append(name) + self._model_keywords[name] = [str(k).lower() for k in prefs.keywords if k] + self._model_quality[name] = tier_int + self._model_cost[name] = self._get_deployment_input_cost(deployment) + self._model_order[name] = prefs.order + seen[name] = True + + missing = [name for name, found in seen.items() if not found] + if missing: + raise ValueError( + f"QualityRouter: the following available_models are not present in " + f"the router's model_list (or are missing routing preferences): {missing}" + ) + + # Sort each tier's model list so `_resolve_model_for_quality_tier` + # (which picks index [0]) honors (order ASC, cost ASC, name ASC). + # Quality is moot within a single tier; keep parity with the keyword + # tiebreak by ordering on (order, cost, name) here. + for models in tier_to_models.values(): + models.sort(key=lambda n: (self._order_key(n), self._cost_key(n), n)) + + return tier_to_models + + def _order_key(self, model_name: str) -> float: + """`order` lookup as a float — unset becomes +inf so explicit wins.""" + order = self._model_order.get(model_name) + return float(order) if order is not None else math.inf + + def _cost_key(self, model_name: str) -> float: + """`input_cost_per_token` as a float — unset becomes +inf.""" + cost = self._model_cost.get(model_name) + return float(cost) if cost is not None else math.inf + + def _keyword_override(self, user_message: str) -> Optional[Tuple[str, str]]: + """ + Find a deployment whose declared keywords appear in `user_message`. + + Returns (model_name, matched_keyword) or None when no keyword matches. + When multiple deployments match, sorts by: + 1. quality_tier DESC (best quality always wins first) + 2. `order` ASC (explicit priority — unset = +inf so explicit wins + within the same tier) + 3. input_cost_per_token ASC (unpriced = +inf so priced wins) + 4. model_name ASC (deterministic stability) + """ + # Touch the lazy index so `_model_keywords` / `_model_quality` / + # `_model_cost` / `_model_order` are populated. + _ = self._tier_to_models + + text = user_message.lower() + + matches: List[Tuple[str, str]] = [] # (model_name, matched_keyword) + for model_name, keywords in self._model_keywords.items(): + for kw in keywords: + if kw and kw in text: + matches.append((model_name, kw)) + break # one match per model is enough + + if not matches: + return None + + def sort_key(match: Tuple[str, str]) -> Tuple[int, float, float, str]: + name = match[0] + quality = self._model_quality.get(name, 0) + order_val = self._order_key(name) + cost = self._model_cost.get(name) + cost_val = cost if cost is not None else math.inf + # Negate quality so higher tier sorts first under ASC sort. + return (-quality, order_val, cost_val, name) + + matches.sort(key=sort_key) + return matches[0] + + def _resolve_model_for_quality_tier(self, tier: int) -> str: + """ + Resolve a quality tier to a concrete model name. + + Strategy: + 1. Exact tier match → first model registered at that tier. + 2. Round UP to the next higher tier that has a model (closer to a + request we might lack capacity for). + 3. Round DOWN to the closest lower tier that has a model (degrade + gracefully instead of jumping straight to `default_model`, + which may be off-tier). + 4. Fall back to `config.default_model`. + 5. Otherwise raise. + """ + tier_index = self._tier_to_models + if tier in tier_index and tier_index[tier]: + return tier_index[tier][0] + + # Round up. + higher_tiers = sorted(t for t in tier_index if t > tier) + for t in higher_tiers: + if tier_index[t]: + return tier_index[t][0] + + # Round down — closest lower tier first. + lower_tiers = sorted((t for t in tier_index if t < tier), reverse=True) + for t in lower_tiers: + if tier_index[t]: + return tier_index[t][0] + + if self.config.default_model: + return self.config.default_model + + raise ValueError( + f"QualityRouter: no model available for quality tier {tier} and " + f"no default_model configured" + ) + + def _stash_decision( + self, + request_kwargs: Optional[Dict[str, Any]], + decision: Dict[str, Any], + ) -> None: + """ + Stash the routing decision in request_kwargs.metadata so the Router can + lift it into response headers (`x-litellm-quality-router-*`). The same + dict object flows from here through to `make_call.set_response_headers`. + """ + if request_kwargs is None: + return + metadata = request_kwargs.setdefault("metadata", {}) + if isinstance(metadata, dict): + metadata["quality_router_decision"] = decision + + async def async_pre_routing_hook( + self, + model: str, + request_kwargs: Dict, + messages: Optional[List[Dict[str, Any]]] = None, + input: Optional[Union[str, List]] = None, + specific_deployment: Optional[bool] = False, + ) -> Optional["PreRoutingHookResponse"]: + """Try keyword override first; fall back to complexity-tier routing.""" + from litellm.types.router import PreRoutingHookResponse + + if messages is None or len(messages) == 0: + verbose_router_logger.debug( + "QualityRouter: No messages provided, skipping routing" + ) + return None + + # Extract last user message and last system prompt — same rules as + # ComplexityRouter.async_pre_routing_hook. + user_message: Optional[str] = None + system_prompt: Optional[str] = None + + for msg in reversed(messages): + role = msg.get("role", "") + content = msg.get("content") or "" + if isinstance(content, list): + text_parts = [ + part.get("text", "") + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ] + content = " ".join(text_parts).strip() + if isinstance(content, str) and content: + if role == "user" and user_message is None: + user_message = content + elif role == "system" and system_prompt is None: + system_prompt = content + + if user_message is None: + verbose_router_logger.debug( + "QualityRouter: No user message found, routing to default model" + ) + if not self.config.default_model: + raise ValueError( + "QualityRouter: no user message and no default_model configured" + ) + return PreRoutingHookResponse( + model=self.config.default_model, + messages=messages, + ) + + # Try keyword override first — it short-circuits complexity classification. + keyword_match = self._keyword_override(user_message) + if keyword_match is not None: + routed_model, matched_keyword = keyword_match + verbose_router_logger.info( + f"QualityRouter: keyword override matched='{matched_keyword}' " + f"routed_model={routed_model} " + f"(quality_tier={self._model_quality.get(routed_model)}, " + f"input_cost_per_token={self._model_cost.get(routed_model)})" + ) + self._stash_decision( + request_kwargs, + { + "router_model_name": self.model_name, + "routed_model": routed_model, + "routed_via": "keyword", + "matched_keyword": matched_keyword, + "quality_tier": self._model_quality.get(routed_model), + "complexity_tier": None, + }, + ) + return PreRoutingHookResponse( + model=routed_model, + messages=messages, + ) + + # No keyword match → complexity classification flow. + complexity_tier, score, signals = self._scorer.classify( + user_message, system_prompt + ) + complexity_name = ( + complexity_tier.value + if hasattr(complexity_tier, "value") + else str(complexity_tier) + ) + + quality_tier = self.config.complexity_to_quality.get(complexity_name) + if quality_tier is None: + raise ValueError( + f"QualityRouter: complexity tier '{complexity_name}' not present " + f"in complexity_to_quality mapping {self.config.complexity_to_quality}" + ) + + routed_model = self._resolve_model_for_quality_tier(int(quality_tier)) + + verbose_router_logger.info( + f"QualityRouter: complexity={complexity_name}, score={score:.3f}, " + f"signals={signals}, quality_tier={quality_tier}, " + f"routed_model={routed_model}" + ) + + self._stash_decision( + request_kwargs, + { + "router_model_name": self.model_name, + "routed_model": routed_model, + "routed_via": "quality_tier", + "matched_keyword": None, + "quality_tier": int(quality_tier), + "complexity_tier": complexity_name, + }, + ) + + return PreRoutingHookResponse( + model=routed_model, + messages=messages, + ) diff --git a/litellm/types/compression.py b/litellm/types/compression.py index 01d5a6dd4d6..5dae0c397f0 100644 --- a/litellm/types/compression.py +++ b/litellm/types/compression.py @@ -2,7 +2,14 @@ Type definitions for litellm.compress(). """ -from typing import Dict, List, TypedDict +import sys + +if sys.version_info >= (3, 11): + from typing import Dict, List, NotRequired, TypedDict +else: + from typing import Dict, List, TypedDict + + from typing_extensions import NotRequired class CompressedResult(TypedDict): @@ -12,3 +19,4 @@ class CompressedResult(TypedDict): compression_ratio: float # fraction reduced, e.g. 0.6 means 60% reduction cache: Dict[str, str] # key -> original content (for retrieval tool responses) tools: List[dict] # [litellm_content_retrieve tool definition] + compression_skipped_reason: NotRequired[str] diff --git a/litellm/types/integrations/compression_interception.py b/litellm/types/integrations/compression_interception.py new file mode 100644 index 00000000000..fe52d2ad0d5 --- /dev/null +++ b/litellm/types/integrations/compression_interception.py @@ -0,0 +1,27 @@ +""" +Type definitions for Compression Interception integration. +""" + +from typing import Any, Dict, Optional, TypedDict + + +class CompressionInterceptionConfig(TypedDict, total=False): + """ + Configuration parameters for CompressionInterceptionLogger. + + Used in proxy_config.yaml under litellm_settings: + litellm_settings: + compression_interception_params: + enabled: true + compression_trigger: 100000 + compression_target: 70000 + embedding_model: "text-embedding-3-small" + embedding_model_params: + dimensions: 512 + """ + + enabled: bool + compression_trigger: int + compression_target: Optional[int] + embedding_model: Optional[str] + embedding_model_params: Optional[Dict[str, Any]] diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 06989409229..b5726a11ca0 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -1,6 +1,6 @@ -from typing import Optional +from typing import Any, Dict, List, Optional -from pydantic import BaseModel +from pydantic import BaseModel, Field class StandardCustomLoggerInitParams(BaseModel): @@ -9,3 +9,29 @@ class StandardCustomLoggerInitParams(BaseModel): """ turn_off_message_logging: Optional[bool] = False + + +class AgenticLoopRequestPatch(BaseModel): + """ + Patch returned by callbacks to request a follow-up LLM call. + """ + + model: Optional[str] = None + messages: Optional[List[Dict[str, Any]]] = None + tools: Optional[List[Dict[str, Any]]] = None + max_tokens: Optional[int] = None + optional_params: Dict[str, Any] = Field(default_factory=dict) + kwargs: Dict[str, Any] = Field(default_factory=dict) + + +class AgenticLoopPlan(BaseModel): + """ + Typed callback response for agentic-loop reruns. + """ + + run_agentic_loop: bool = False + request_patch: Optional[AgenticLoopRequestPatch] = None + response_override: Optional[Any] = None + terminate: bool = False + stop_reason: Optional[str] = None + metadata: Dict[str, Any] = Field(default_factory=dict) diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 338c5a79ce6..43a287f29bc 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -784,7 +784,7 @@ class UserAPIKeyLabelValues: org_id: Optional[str] = None org_alias: Optional[str] = None - #Added for test compatibility. + # Added for test compatibility. def __init__(self, **kwargs: Any) -> None: """ Match former Pydantic behavior: unknown keys are ignored; ``api_key_hash`` maps to diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 6830d95d36f..9ffb52ef88d 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -997,3 +997,47 @@ class BedrockToolBlock(TypedDict, total=False): toolSpec: Optional[ToolSpecBlock] systemTool: Optional[SystemToolBlock] # For Nova grounding cachePoint: Optional[CachePointBlock] + + +class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False): + """ + Top-level request body accepted by AWS Bedrock `InvokeModel` / + `InvokeModelWithResponseStream` when calling an Anthropic Claude model with + the Messages API format. The LiteLLM /v1/messages → Bedrock Invoke + transformation filters outgoing requests to the keys of this TypedDict; any + other field (Anthropic-only extension, internal metadata, future addition) + is dropped before signing so Bedrock doesn't 400 with + "Extra inputs are not permitted". + + Reference: + https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html + https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + + Editing this type is the single source of truth — the runtime allowlist in + `AmazonAnthropicClaudeMessagesConfig.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS` + is derived from `__annotations__`, and a test asserts the resolved set + exactly, so any edit forces a conscious review. + + Value types are intentionally loose (`list`, `dict`) — this type exists to + pin the allowed field names, not to validate nested structure. + """ + + # Required by Bedrock + anthropic_version: str + max_tokens: int + messages: list + + # Documented optional fields + anthropic_beta: List[str] + system: object # str or list[TextBlock] + stop_sequences: List[str] + temperature: float + top_p: float + top_k: int + tools: list + tool_choice: dict + + # `thinking` is required for Opus 4.5 / Sonnet 4 extended thinking, + # `metadata` is part of the common Anthropic Messages API shape. + thinking: dict + metadata: dict diff --git a/litellm/types/router.py b/litellm/types/router.py index 6bd64915d79..58b64b667dd 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from typing import Any, Dict, List, Literal, Optional, Tuple, Union, get_type_hints import httpx -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing_extensions import Required, TypedDict from litellm._uuid import uuid @@ -221,6 +221,13 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): complexity_router_config: Optional[Dict] = None complexity_router_default_model: Optional[str] = None + # adaptive-router params + adaptive_router_default_model: Optional[str] = None + adaptive_router_config: Optional[Dict] = None + # quality-router params + quality_router_config: Optional[Dict] = None + quality_router_default_model: Optional[str] = None + # Batch/File API Params s3_bucket_name: Optional[str] = None s3_encryption_key_id: Optional[str] = None @@ -790,3 +797,44 @@ class PreRoutingHookResponse(BaseModel): model: str messages: Optional[List[Dict[str, Any]]] + + +class RequestType(str, enum.Enum): + """Fixed v0 taxonomy. User-extensible types come in v1.""" + + CODE_GENERATION = "code_generation" + CODE_UNDERSTANDING = "code_understanding" + TECHNICAL_DESIGN = "technical_design" + ANALYTICAL_REASONING = "analytical_reasoning" + WRITING = "writing" + FACTUAL_LOOKUP = "factual_lookup" + GENERAL = "general" + + +class AdaptiveRouterWeights(BaseModel): + quality: float = Field(default=0.7, ge=0.0, le=1.0) + cost: float = Field(default=0.3, ge=0.0, le=1.0) + + @field_validator("cost") + @classmethod + def _weights_sum_to_one(cls, v, info): + q = info.data.get("quality", 0.7) + if abs(q + v - 1.0) > 0.001: + raise ValueError( + f"weights must sum to 1.0, got quality={q} + cost={v} = {q + v}" + ) + return v + + +class AdaptiveRouterConfig(BaseModel): + available_models: List[str] + weights: AdaptiveRouterWeights = Field(default_factory=AdaptiveRouterWeights) + + +class AdaptiveRouterPreferences(BaseModel): + """model_info.adaptive_router_preferences — declared by each model.""" + + model_config = ConfigDict(use_enum_values=False) + + quality_tier: int = Field(ge=1, le=3) + strengths: List[RequestType] = Field(default_factory=list) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d0bc9b78941..e3058d106a6 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2851,6 +2851,7 @@ class StandardAuditLogPayload(TypedDict): class StandardLoggingPayload(TypedDict): id: str trace_id: str # Trace multiple LLM calls belonging to same overall request (e.g. fallbacks/retries) + litellm_call_id: Optional[str] # UUID returned in x-litellm-call-id response header call_type: str stream: Optional[bool] response_cost: float @@ -3290,6 +3291,7 @@ class LlmProviders(str, Enum): MANUS = "manus" WANDB = "wandb" OVHCLOUD = "ovhcloud" + SCALEWAY = "scaleway" LEMONADE = "lemonade" AMAZON_NOVA = "amazon_nova" A2A_AGENT = "a2a_agent" diff --git a/litellm/utils.py b/litellm/utils.py index 2125875ee1d..c4aee792972 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8472,6 +8472,12 @@ class ProviderConfigManager: ) return OVHCloudAudioTranscriptionConfig() + elif litellm.LlmProviders.SCALEWAY == provider: + from litellm.llms.scaleway.audio_transcription.transformation import ( + ScalewayAudioTranscriptionConfig, + ) + + return ScalewayAudioTranscriptionConfig() elif litellm.LlmProviders.MISTRAL == provider: from litellm.llms.mistral.audio_transcription.transformation import ( MistralAudioTranscriptionConfig, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 04b68b8f4ec..4e629bbd947 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1148,6 +1148,20 @@ "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, + "anthropic.claude-mythos-preview": { + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_tool_choice": true + }, "global.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, "cache_read_input_token_cost": 5e-07, @@ -22872,6 +22886,22 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k2.6": { + "cache_read_input_token_cost": 1.6e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "moonshot", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://platform.kimi.ai/docs/pricing/chat-k26", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, "moonshot/kimi-latest": { "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 2e-06, @@ -25135,6 +25165,28 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346 }, + "openrouter/anthropic/claude-opus-4.7": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "tool_use_system_prompt_tokens": 346 + }, "openrouter/bytedance/ui-tars-1.5-7b": { "input_cost_per_token": 1e-07, "litellm_provider": "openrouter", @@ -33302,6 +33354,7 @@ "output_cost_per_token": 1.5e-05, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true @@ -33317,6 +33370,7 @@ "output_cost_per_token": 1.5e-05, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true @@ -33332,6 +33386,7 @@ "output_cost_per_token": 2.5e-05, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true @@ -33347,6 +33402,7 @@ "output_cost_per_token": 2.5e-05, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true @@ -33362,6 +33418,7 @@ "output_cost_per_token": 1.5e-05, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, "supports_web_search": true @@ -33378,6 +33435,7 @@ "output_cost_per_token": 5e-07, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33395,6 +33453,7 @@ "output_cost_per_token": 5e-07, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33411,6 +33470,7 @@ "output_cost_per_token": 4e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33427,6 +33487,7 @@ "output_cost_per_token": 4e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33443,6 +33504,7 @@ "output_cost_per_token": 4e-06, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33459,6 +33521,7 @@ "output_cost_per_token": 5e-07, "source": "https://x.ai/api#pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": false, "supports_tool_choice": true, @@ -33474,38 +33537,41 @@ "output_cost_per_token": 1.5e-05, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true }, "xai/grok-4-fast-reasoning": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_128k_tokens": 4e-07, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, "max_tokens": 2000000.0, "mode": "chat", - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, "output_cost_per_token": 5e-07, "output_cost_per_token_above_128k_tokens": 1e-06, - "cache_read_input_token_cost": 5e-08, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true }, "xai/grok-4-fast-non-reasoning": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_128k_tokens": 4e-07, "litellm_provider": "xai", "max_input_tokens": 2000000.0, "max_output_tokens": 2000000.0, - "cache_read_input_token_cost": 5e-08, "max_tokens": 2000000.0, "mode": "chat", - "input_cost_per_token": 2e-07, - "input_cost_per_token_above_128k_tokens": 4e-07, "output_cost_per_token": 5e-07, "output_cost_per_token_above_128k_tokens": 1e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true }, @@ -33521,6 +33587,7 @@ "output_cost_per_token_above_128k_tokens": 3e-05, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true }, @@ -33536,6 +33603,7 @@ "output_cost_per_token_above_128k_tokens": 3e-05, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_web_search": true }, @@ -33553,6 +33621,7 @@ "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -33573,6 +33642,7 @@ "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -33593,6 +33663,7 @@ "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -33613,6 +33684,7 @@ "source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning", "supports_audio_input": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -33632,6 +33704,7 @@ "source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning", "supports_audio_input": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -33648,6 +33721,7 @@ "output_cost_per_token": 6e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, @@ -33664,6 +33738,7 @@ "output_cost_per_token": 6e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, @@ -33696,6 +33771,7 @@ "output_cost_per_token": 6e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true @@ -33724,6 +33800,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, @@ -33738,6 +33815,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, @@ -33752,6 +33830,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 2f3302bb574..6f23c87f911 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1968,7 +1968,7 @@ "responses": true, "embeddings": false, "image_generations": false, - "audio_transcriptions": false, + "audio_transcriptions": true, "audio_speech": false, "moderations": false, "batches": false, diff --git a/pyproject.toml b/pyproject.toml index d5d238473b1..75aec08c99b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.83.10" +version = "1.83.11" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -208,7 +208,7 @@ build-backend = "uv_build" [tool.uv] default-groups = ["dev"] -required-version = "==0.10.9" +required-version = ">=0.10.9" exclude-newer = "3 days" [tool.uv.sources] @@ -236,7 +236,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.83.10" +version = "1.83.11" version_files = [ "pyproject.toml:^version", ] diff --git a/schema.prisma b/schema.prisma index 08aa5645251..7642ad74b20 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1223,3 +1223,46 @@ model LiteLLM_ClaudeCodePluginTable { @@map("LiteLLM_ClaudeCodePluginTable") } + +// Per-(router, request_type, model) Beta posterior for the adaptive router. +model LiteLLM_AdaptiveRouterState { + router_name String + request_type String + model_name String + alpha Float + beta Float + total_samples Int @default(0) + last_updated_at DateTime @default(now()) @updatedAt + + @@id([router_name, request_type, model_name]) +} + +// Per-(session, router, model) signal counters for the adaptive router. +model LiteLLM_AdaptiveRouterSession { + session_id String + router_name String + model_name String + classified_type String + + misalignment_count Int @default(0) + stagnation_count Int @default(0) + disengagement_count Int @default(0) + satisfaction_count Int @default(0) + failure_count Int @default(0) + loop_count Int @default(0) + exhaustion_count Int @default(0) + + last_user_content String? + last_assistant_content String? + tool_call_history Json @default("[]") + pending_tool_calls Json @default("{}") + + turn_count Int @default(0) + last_processed_turn Int @default(-1) + clean_credit_awarded Boolean @default(false) + terminal_status Int? + last_activity_at DateTime @default(now()) @updatedAt + + @@id([session_id, router_name, model_name]) + @@index([last_activity_at], map: "idx_adaptive_router_session_activity") +} diff --git a/scripts/adaptive_router_demo/README.md b/scripts/adaptive_router_demo/README.md new file mode 100644 index 00000000000..1965dbbf168 --- /dev/null +++ b/scripts/adaptive_router_demo/README.md @@ -0,0 +1,157 @@ +# Adaptive Router — Live Demo + +A 5-minute demo of LiteLLM's adaptive router learning, in real time, that +the smart model wins for code while the fast model is fine for facts. + +``` +┌─ traffic.py ──┐ ┌─ litellm proxy ──────────┐ ┌─ dashboard.html ─┐ +│ synthetic │──▶│ adaptive_router strategy │──▶│ bandit bars + │ +│ chat sessions │ │ /adaptive_router/state │ │ cost meter + │ +└───────────────┘ └──────────┬───────────────┘ │ activity log │ + │ └───────────────────┘ + ┌─────────▼───────────┐ + │ chat.html │ + │ interactive chat │ + │ with preset │ + │ scenarios │ + └─────────────────────┘ +``` + +## Files + +| File | What it does | +|---|---| +| `dashboard.html` | Live bandit dashboard — polls `/adaptive_router/state` every 500ms | +| `chat.html` | Interactive chat with preset scenarios — sends real requests through the router | +| `traffic.py` | Synthetic traffic generator — drives labeled sessions for automated demo | + +## What you're watching + +- **Bandit posteriors** — one Beta(α, β) bar per `(request_type, model)` + cell. Bars fill up as α grows from positive feedback signals. +- **Pick share** — softmax estimate of how often the router would currently + pick each model for that request type. +- **Cost meter** — total spend so far compared to "always use the most + expensive model". The savings line is the headline number. +- **Activity log** — every signal that moves the bandit, in real time. + +## 1. Start the proxy + +The repo ships with a working example config: + +```bash +export OPENAI_API_KEY=sk-... # underlying models hit OpenAI +uv run litellm \ + --config litellm/proxy/example_config_yaml/adaptive_router_example.yaml \ + --port 4000 +``` + +`DATABASE_URL` is optional — the proxy falls back to a bundled Neon dev DB. +Wait ~15s until you see `Application startup complete`. + +## 2. Chat interactively with the router + +Open `chat.html` in a browser (same `file://` or `python3 -m http.server` approach as the dashboard): + +- Click **Connect** after filling in the proxy URL and API key. +- Pick a preset scenario: + - **🐛 Debug my code** — paste broken code and get a fix + - **💡 Brainstorm a feature** — ideate on a product capability + - **📚 Explain a concept** — get a clear technical explanation + - **✍️ Write something** — draft emails, docs, or any prose +- A starter message is pre-filled — edit it or send as-is. +- Each response shows which model the router picked and the inferred request type (from the `x-litellm-adaptive-router-model` and `x-litellm-request-type` response headers). +- A sidebar gate indicator tells you when the session has accumulated enough messages for the bandit to start updating (4+ turns). + +> **Note on headers:** The model/type headers are only readable in the browser if the proxy sets `Access-Control-Expose-Headers`. LiteLLM defaults to exposing them. If the info panel shows `check dashboard`, the router still works — you can verify picks in `dashboard.html`. + +## 4. Open the dashboard + +The dashboard is a single static HTML file. Either: + +- **Easy:** double-click `dashboard.html`. Most browsers will load it from + `file://` and the LiteLLM proxy's CORS defaults (`*`) will accept it. +- **If your browser blocks `file://` fetches:** + + ```bash + cd scripts/adaptive_router_demo + python3 -m http.server 8080 + ``` + + Then open . + +In the connect bar, fill in: + +- **Proxy URL:** `http://localhost:4000` +- **Master Key:** the `master_key` from your config (`sk-1234` in the example). + +Click **Connect**. The dashboard polls `GET /adaptive_router/state` every +500ms (admin-only endpoint, returns one snapshot per configured router). + +## 5. Drive synthetic traffic + +In a second terminal: + +```bash +uv run python scripts/adaptive_router_demo/traffic.py \ + --proxy-url http://localhost:4000 \ + --api-key sk-1234 \ + --router smart-cheap-router \ + --rounds 100 \ + --rate 0.5 +``` + +What it does: + +- Picks a random `(request_type, prompt)` per round from a small labeled corpus. +- Sends a 5-message conversation (passes the `SIGNAL_GATE_MIN_MESSAGES=4` gate + in one round-trip) so the post-call hook runs and updates the bandit. +- Reads the `x-litellm-adaptive-router-model` response header to see what + the router picked. +- Rolls Bernoulli against a hard-coded oracle: + ``` + code_generation : smart=0.92 fast=0.35 + factual_lookup : smart=0.90 fast=0.85 + writing : smart=0.85 fast=0.55 + ``` +- On success → sends a follow-up engineered to match the satisfaction + regex (and re-classify into the same type). Bandit cell gets +α. +- On failure → sends a neutral follow-up. No signal fires. + +After 50–80 rounds you'll see `code_generation` decisively favor `smart` +while `factual_lookup` stays near a coin flip — the router learned the +asymmetry from the oracle. + +## Tuning knobs + +| Knob | Where | What changes | +|---|---|---| +| Quality vs. cost weight | `adaptive_router_config.weights` in proxy yaml | Bias toward quality or savings | +| Per-cell cold-start mass | `litellm/router_strategy/adaptive_router/config.py` `COLD_START_MASS` | How long until the prior is overwritten | +| Avg tokens per request | dashboard input box | How the cost meter estimates spend | +| Oracle | `traffic.py` `ORACLE` dict | Which model "should" win for which type | +| Sessions to drive | `--rounds` | Total learning budget | +| Throttle | `--rate` | Seconds between sessions | + +## Multi-router + +If your proxy has more than one `auto_router/adaptive_router` deployment, +the dashboard shows a router dropdown above the bars. Each router is +independent; the cost meter is per-router (and resets when you switch). + +## Troubleshooting + +- **"Disconnected" / HTTP 401 in the dashboard** — wrong master key. +- **HTTP 403** — your key isn't `proxy_admin`. The state endpoint is + admin-only. Use the master key. +- **HTTP 404 from `/adaptive_router/state`** — proxy started, but no + `auto_router/adaptive_router` deployment is in the model list. +- **Bars don't move** — check the proxy logs for `record_turn` activity. + Common cause: requests are not including 4+ messages, so the signal + gate skips them. `traffic.py` already builds 5-message conversations, + so this only happens if you've changed the script. +- **Cost meter stays at $0** — your model deployments don't have + `input_cost_per_token` set in `litellm_params`. Add it. +- **CORS error in the dashboard console** — set `LITELLM_CORS_ORIGINS=*` + on the proxy (the default), or serve `dashboard.html` from + `python3 -m http.server` instead of `file://`. diff --git a/scripts/adaptive_router_demo/chat.html b/scripts/adaptive_router_demo/chat.html new file mode 100644 index 00000000000..9e7237847c0 --- /dev/null +++ b/scripts/adaptive_router_demo/chat.html @@ -0,0 +1,838 @@ + + + + + Adaptive Router — Chat + + + + +
+

⚡ Adaptive Router — Chat

+ Disconnected + → Open live dashboard +
+ +
+ + + + +
+ +
+ + + + + +
+ +
+
+
+
+

Pick a scenario to start

+

Choose one of the presets above or connect to the proxy and type your own message. The adaptive router will pick the best model for each turn.

+
+
+
+
+ + +
+
Connect first to start chatting.
+
+
+ + +
+ + + + + diff --git a/scripts/adaptive_router_demo/dashboard.html b/scripts/adaptive_router_demo/dashboard.html new file mode 100644 index 00000000000..6652aa19805 --- /dev/null +++ b/scripts/adaptive_router_demo/dashboard.html @@ -0,0 +1,635 @@ + + + + + Adaptive Router — Live + + + + +
+

⚡ Adaptive Router — Live

+ Disconnected + +
+ +
+ + + + + + +
+ +
+
+

How well each model performs, by request type

+
+ Each bar shows the fraction of recent feedback that was positive + for that model on that kind of request. Wider = better. The number + next to it ("N signals") is how much real feedback the bar is + based on — more signals means the router is more confident. + It picks higher-quality bars first, with cost as a tiebreaker. +
+
Connect to see live bandit state.
+
+ + +
+ + + + + diff --git a/scripts/adaptive_router_demo/eval.py b/scripts/adaptive_router_demo/eval.py new file mode 100644 index 00000000000..b02e4a37d31 --- /dev/null +++ b/scripts/adaptive_router_demo/eval.py @@ -0,0 +1,271 @@ +# ruff: noqa: T201 +""" +Adaptive router evaluator — LLM-as-judge harness. + +For each test case: + 1. Sends the prompt to the adaptive router. + 2. Reads which model was picked (x-litellm-adaptive-router-model header). + 3. Asks the judge model whether the response meets the ideal criteria. + 4. Prints PASS or FAIL with one line of reasoning. + +Run: + uv run python scripts/adaptive_router_demo/eval.py \ + --proxy-url http://localhost:4000 \ + --api-key sk-1234 \ + --router smart-cheap-router \ + --judge-model smart +""" + +from __future__ import annotations + +import argparse +import asyncio +import sys +import uuid +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import httpx + + +# --------------------------------------------------------------------------- +# Test cases +# --------------------------------------------------------------------------- +@dataclass +class EvalCase: + category: str + prompt: str + ideal: str # criteria the judge checks the response against + + +EVAL_CASES: List[EvalCase] = [ + # code_generation + EvalCase( + category="code_generation", + prompt="Write a Python function that flattens a nested list of arbitrary depth.", + ideal=( + "A Python function (def flatten(...)) that accepts a list which may " + "contain nested lists to arbitrary depth and returns a single flat list " + "with all elements in order. Must handle at least two levels of nesting." + ), + ), + EvalCase( + category="code_generation", + prompt="Write a Python decorator that retries a function up to 3 times on exception.", + ideal=( + "A Python decorator that wraps a callable, catches exceptions, and " + "retries the call up to 3 times before re-raising. Should use functools.wraps " + "or equivalent to preserve the wrapped function's metadata." + ), + ), + EvalCase( + category="code_generation", + prompt="Write a SQL query that returns the top 5 customers by total order value.", + ideal=( + "A valid SQL SELECT query that JOINs an orders or order_items table with a " + "customers table, groups by customer, sums order value, orders descending, " + "and limits to 5 rows." + ), + ), + # factual_lookup + EvalCase( + category="factual_lookup", + prompt="What is the capital of New Zealand?", + ideal="The answer must state Wellington as the capital of New Zealand.", + ), + EvalCase( + category="factual_lookup", + prompt="In what year did World War II end?", + ideal="The answer must state 1945 as the year World War II ended.", + ), + EvalCase( + category="factual_lookup", + prompt="What is the chemical symbol for gold?", + ideal="The answer must include 'Au' as the chemical symbol for gold.", + ), + # writing + EvalCase( + category="writing", + prompt=( + "Write a short, polite email declining a meeting request because of " + "a scheduling conflict." + ), + ideal=( + "A professional email that: (1) thanks the sender for the invitation, " + "(2) clearly declines, (3) mentions a scheduling conflict as the reason, " + "and (4) offers to reschedule or an alternative. Tone must be polite." + ), + ), + EvalCase( + category="writing", + prompt="Write a one-paragraph product description for noise-cancelling headphones.", + ideal=( + "A marketing paragraph for noise-cancelling headphones that mentions " + "noise cancellation as a feature, highlights at least one other benefit " + "(comfort, audio quality, battery life, or similar), and ends with a " + "persuasive call to action or closing statement." + ), + ), +] + +# Matches the satisfaction regex in signals.py (_SATISFACTION_PATTERNS). +SATISFY_FOLLOWUP = "great, thanks!" +NEUTRAL_FOLLOWUP = "ok, noted" +FAB_ASSISTANT = "Got it. Working on that now." + +JUDGE_SYSTEM = ( + "You are a strict but fair evaluator. Your job is to decide whether a model " + "response meets the stated requirements. Reply with exactly two lines:\n" + "Line 1: PASS or FAIL\n" + "Line 2: One sentence of reasoning (≤ 25 words)." +) + + +def _judge_user(prompt: str, ideal: str, actual: str) -> str: + return ( + f"Question sent to model:\n{prompt}\n\n" + f"Requirements the response must meet:\n{ideal}\n\n" + f"Actual model response:\n{actual}\n\n" + "Does the response meet the requirements? Reply PASS or FAIL." + ) + + +# --------------------------------------------------------------------------- +# HTTP helpers +# --------------------------------------------------------------------------- +async def _chat( + client: httpx.AsyncClient, + proxy_url: str, + api_key: str, + model: str, + messages: List[Dict[str, str]], + session_id: Optional[str] = None, +) -> Tuple[str, str]: + """ + Returns (response_text, chosen_model_header). + chosen_model_header is empty for non-router calls. + """ + body: Dict = {"model": model, "messages": messages} + if session_id: + body["metadata"] = {"litellm_session_id": session_id} + + resp = await client.post( + f"{proxy_url}/v1/chat/completions", + json=body, + headers={"Authorization": f"Bearer {api_key}"}, + timeout=60.0, + ) + resp.raise_for_status() + data = resp.json() + text = data["choices"][0]["message"]["content"] + chosen = resp.headers.get("x-litellm-adaptive-router-model", "") + return text, chosen + + +# --------------------------------------------------------------------------- +# Evaluation loop +# --------------------------------------------------------------------------- +async def evaluate( + proxy_url: str, + api_key: str, + router: str, + judge_model: str, +) -> None: + passed = 0 + failed = 0 + + async with httpx.AsyncClient() as client: + for i, case in enumerate(EVAL_CASES, 1): + print(f"\n[{i}/{len(EVAL_CASES)}] category={case.category}") + print(f" prompt : {case.prompt[:80]}{'…' if len(case.prompt) > 80 else ''}") + + session_id = f"eval-{uuid.uuid4()}" + + # Round 1: single-turn real request — get the actual LLM response to judge. + try: + response, chosen = await _chat( + client, proxy_url, api_key, router, + [{"role": "user", "content": case.prompt}], + session_id=session_id, + ) + except Exception as exc: # noqa: BLE001 + print(f" ERROR calling router: {exc}", file=sys.stderr) + failed += 1 + continue + + print(f" model : {chosen or router}") + print(f" response : {response[:120].replace(chr(10), ' ')}{'…' if len(response) > 120 else ''}") + + # Judge the real response. + judge_msgs = [ + {"role": "system", "content": JUDGE_SYSTEM}, + {"role": "user", "content": _judge_user(case.prompt, case.ideal, response)}, + ] + try: + verdict, _ = await _chat( + client, proxy_url, api_key, judge_model, judge_msgs, + ) + except Exception as exc: # noqa: BLE001 + print(f" ERROR calling judge: {exc}", file=sys.stderr) + failed += 1 + continue + + # Parse verdict — first non-empty line should be PASS or FAIL. + lines = [ln.strip() for ln in verdict.splitlines() if ln.strip()] + first = lines[0].upper() if lines else "" + reason = lines[1] if len(lines) > 1 else "" + is_pass = "PASS" in first + + if is_pass: + passed += 1 + print(f" verdict : \033[32mPASS\033[0m {reason}") + else: + failed += 1 + print(f" verdict : \033[31mFAIL\033[0m {reason}") + + # Round 2: 5-message conversation on the same session_id so the bandit fires. + # On PASS → satisfaction follow-up (+alpha). On FAIL → neutral (no signal). + follow_up = SATISFY_FOLLOWUP if is_pass else NEUTRAL_FOLLOWUP + bandit_msgs = [ + {"role": "user", "content": case.prompt}, + {"role": "assistant", "content": response}, + {"role": "user", "content": "ok continue"}, + {"role": "assistant", "content": FAB_ASSISTANT}, + {"role": "user", "content": follow_up}, + ] + try: + await _chat( + client, proxy_url, api_key, router, bandit_msgs, + session_id=session_id, + ) + except Exception as exc: # noqa: BLE001 + print(f" WARNING: bandit update failed: {exc}", file=sys.stderr) + + total = passed + failed + print(f"\n{'='*60}") + print(f"Results: {passed}/{total} passed ({failed} failed)") + if passed == total: + print("All test cases passed — the adaptive router is working well!") + elif passed >= total * 0.8: + print("Most test cases passed — minor issues to investigate.") + else: + print("Significant failures — check router config and model availability.") + print("=" * 60) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- +def main() -> None: + ap = argparse.ArgumentParser(description="Evaluate the adaptive router with LLM-as-judge.") + ap.add_argument("--proxy-url", default="http://localhost:4000") + ap.add_argument("--api-key", required=True, help="proxy API key") + ap.add_argument("--router", default="smart-cheap-router", help="adaptive router model name") + ap.add_argument("--judge-model", default="smart", help="model name for the judge (via proxy)") + args = ap.parse_args() + + asyncio.run(evaluate(args.proxy_url, args.api_key, args.router, args.judge_model)) + + +if __name__ == "__main__": + main() diff --git a/scripts/adaptive_router_demo/traffic.py b/scripts/adaptive_router_demo/traffic.py new file mode 100644 index 00000000000..eae5506eaee --- /dev/null +++ b/scripts/adaptive_router_demo/traffic.py @@ -0,0 +1,227 @@ +""" +Synthetic traffic generator for the adaptive_router demo dashboard. + +What it does: + - Sends labeled multi-turn chat requests to the proxy's adaptive router. + - For each turn, peeks at the `x-litellm-adaptive-router-model` response + header to learn which underlying model was picked. + - Draws a Bernoulli outcome from a hard-coded ORACLE table that says + "model M succeeds at request type T with probability p". + - Sends a final follow-up turn whose user message is engineered to + BOTH classify into the same RequestType AND match the + satisfaction regex on success (so the bandit's `(type, model)` cell + gets +alpha). On failure we send a neutral follow-up so no signal + fires — over time, models the oracle favors accumulate alpha faster. + +Why this shape: + - The post-call hook gates signal recording on len(messages) >= 4. + A single 5-message request passes the gate in one round-trip, which + keeps the demo cheap. + - Mock responses (`mock_response=...`) skip the real LLM call but still + flow through routing + post-call hooks, so no API keys / no spend. + +Run: + uv run python scripts/adaptive_router_demo/traffic.py \\ + --proxy-url http://localhost:4000 \\ + --api-key sk-1234 \\ + --router smart-cheap-router \\ + --rounds 100 \\ + --rate 0.5 + +Open `dashboard.html` in a browser alongside this and watch the bars move. +""" + +from __future__ import annotations + +import argparse +import asyncio +import random +import sys +import uuid +from typing import Dict, List, Tuple + +import httpx + +# ---- prompts (paired with the RequestType the classifier will assign) ---- +# Each prompt is engineered to (a) classify into the listed type and (b) make +# sense as a user request. Keep prompts short to limit token cost. +PROMPTS: Dict[str, List[str]] = { + "code_generation": [ + "Write a Python function that flattens a nested list", + "Create a TypeScript function that debounces another function", + "Build a Rust function that parses a CSV string", + "Generate a SQL function that returns running totals", + ], + "factual_lookup": [ + "What is the capital of New Zealand?", + "When was the Treaty of Westphalia signed?", + "Who is the current Secretary General of the UN?", + "Where is Mount Kilimanjaro located?", + ], + "writing": [ + "Write an email declining a meeting politely", + "Draft a paragraph introducing a product launch", + "Compose a short blog post about morning routines", + "Rewrite this sentence to be more concise: ...", + ], +} + +# Engineered satisfaction follow-ups — each one is designed to: +# (1) match the satisfaction regex (thanks/great/works/perfect/etc.), AND +# (2) re-classify into the SAME RequestType as the first prompt +# so that signals attribute to the right (type, model) bandit cell. +SATISFY: Dict[str, str] = { + "code_generation": "thanks, that works! now write me a python function that does the inverse", + "factual_lookup": "perfect, thanks! who is the current prime minister?", + "writing": "great, thanks! now write a follow-up email confirming attendance", +} + +# Neutral follow-up — does not match any signal regex, does not move the bandit. +NEUTRAL_FOLLOWUP = "ok, noted" + +# Oracle: P(success | request_type, model). Tunable. +# Defaults: smart dominates code/writing; both are fine for factual_lookup. +ORACLE: Dict[str, Dict[str, float]] = { + "code_generation": {"smart": 0.92, "fast": 0.35}, + "factual_lookup": {"smart": 0.90, "fast": 0.85}, + "writing": {"smart": 0.85, "fast": 0.55}, +} + +# Fabricated assistant turn — content doesn't matter for the hook, only the role. +FAB_ASSISTANT = "Got it. Working on that now." + + +def _build_messages(prompt: str, last_user: str) -> List[Dict[str, str]]: + """5-message conversation that passes the SIGNAL_GATE_MIN_MESSAGES=4 gate.""" + return [ + {"role": "user", "content": prompt}, + {"role": "assistant", "content": FAB_ASSISTANT}, + {"role": "user", "content": "ok continue"}, + {"role": "assistant", "content": FAB_ASSISTANT}, + {"role": "user", "content": last_user}, + ] + + +async def _send( + client: httpx.AsyncClient, + proxy_url: str, + api_key: str, + router: str, + session_id: str, + messages: List[Dict[str, str]], + mock_response: str, +) -> Tuple[bool, str]: + """Returns (ok, chosen_model).""" + body = { + "model": router, + "messages": messages, + "metadata": {"litellm_session_id": session_id}, + "mock_response": mock_response, + } + try: + r = await client.post( + f"{proxy_url}/v1/chat/completions", + json=body, + headers={"Authorization": f"Bearer {api_key}"}, + timeout=15.0, + ) + r.raise_for_status() + except Exception as e: # noqa: BLE001 + print(f" request failed: {e}", file=sys.stderr) + return False, "" + chosen = r.headers.get("x-litellm-adaptive-router-model", "") + return True, chosen + + +async def _drive_one_session( + client: httpx.AsyncClient, + proxy_url: str, + api_key: str, + router: str, + request_type: str, + prompt: str, +) -> str: + """Run one labeled session. Returns the chosen model (for logging).""" + session_id = f"demo-{uuid.uuid4()}" + + # Send the engineered 5-message conversation. The follow-up is chosen + # AFTER we observe what model the router would pick — but since the + # router is sticky-per-session, the model on this single round-trip + # IS the model we're crediting. + # + # Pre-decide success based on the oracle for whichever model gets picked. + # We can't know the pick before sending, so: send a neutral follow-up + # first to learn the pick, then send a second round with credit attached. + # + # Round 1: neutral follow-up → no signal fires, but we learn the pick. + ok, chosen = await _send( + client, proxy_url, api_key, router, session_id, + _build_messages(prompt, NEUTRAL_FOLLOWUP), + mock_response=FAB_ASSISTANT, + ) + if not ok or not chosen: + return "" + + # Decide outcome from oracle. + p = ORACLE.get(request_type, {}).get(chosen, 0.5) + success = random.random() < p + follow_up = SATISFY[request_type] if success else NEUTRAL_FOLLOWUP + + # Round 2: include the round-1 turns + a new follow-up. On success the + # follow-up matches satisfaction → +alpha for (request_type, chosen). + history = _build_messages(prompt, NEUTRAL_FOLLOWUP) + [ + {"role": "assistant", "content": FAB_ASSISTANT}, + {"role": "user", "content": follow_up}, + ] + await _send( + client, proxy_url, api_key, router, session_id, history, + mock_response=FAB_ASSISTANT, + ) + return chosen + + +async def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--proxy-url", default="http://localhost:4000") + ap.add_argument("--api-key", required=True, help="proxy key with /v1/chat/completions perms") + ap.add_argument("--router", default="smart-cheap-router") + ap.add_argument("--rounds", type=int, default=100) + ap.add_argument("--rate", type=float, default=0.5, + help="seconds between sessions; lower = faster") + ap.add_argument("--types", default="code_generation,factual_lookup,writing", + help="comma-separated subset of request types to drive") + args = ap.parse_args() + + types = [t.strip() for t in args.types.split(",") if t.strip() in PROMPTS] + if not types: + print(f"ERROR: no valid types. Choose from: {list(PROMPTS)}", file=sys.stderr) + sys.exit(2) + + print(f"driving {args.rounds} sessions across types: {types}") + print(f"oracle: {ORACLE}") + print(f"proxy: {args.proxy_url} router: {args.router}\n") + + counts: Dict[Tuple[str, str], int] = {} + async with httpx.AsyncClient() as client: + for i in range(args.rounds): + rt = random.choice(types) + prompt = random.choice(PROMPTS[rt]) + chosen = await _drive_one_session( + client, args.proxy_url, args.api_key, args.router, rt, prompt, + ) + if chosen: + counts[(rt, chosen)] = counts.get((rt, chosen), 0) + 1 + if (i + 1) % 10 == 0: + summary = ", ".join( + f"{rt}/{m}={n}" for (rt, m), n in sorted(counts.items()) + ) + print(f" round {i + 1}/{args.rounds} picks: {summary}") + await asyncio.sleep(args.rate) + + print("\nfinal pick distribution:") + for (rt, m), n in sorted(counts.items()): + print(f" {rt:22s} → {m:8s} {n}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/scripts/eval_compression.py b/scripts/eval_compression.py index d7d90dacc2e..a169cc02d74 100644 --- a/scripts/eval_compression.py +++ b/scripts/eval_compression.py @@ -33,6 +33,7 @@ from dataclasses import asdict, dataclass, field from typing import Optional import litellm +from litellm.types.utils import CallTypes # --------------------------------------------------------------------------- # Problem definitions (HumanEval-style) @@ -880,6 +881,7 @@ def eval_problem( result = litellm.compress( messages=messages, model=model, + call_type=CallTypes.completion, compression_trigger=compression_trigger, embedding_model=embedding_model, ) diff --git a/scripts/verify_adaptive_router.py b/scripts/verify_adaptive_router.py new file mode 100644 index 00000000000..fde9dc51a15 --- /dev/null +++ b/scripts/verify_adaptive_router.py @@ -0,0 +1,216 @@ +""" +End-to-end verification script for the adaptive router. + +Requires: + - LiteLLM proxy running on http://localhost:4000 with adaptive_router configured + (see litellm/proxy/example_config_yaml/adaptive_router_example.yaml). + - Postgres reachable via DATABASE_URL (same one the proxy uses). + - LITELLM_PROXY_KEY env var set (a valid key with permission to send requests). + - Two model deployments configured under one adaptive_router: + * "fast" (cheap, lower quality) + * "smart" (expensive, higher quality) + +Run: + uv run python scripts/verify_adaptive_router.py + +Optional env: + LITELLM_PROXY_URL (default: http://localhost:4000) + ADAPTIVE_ROUTER_NAME (default: smart-cheap-router) + EXPECTED_WINNER (default: smart) -- model expected to dominate after training + TRAIN_SESSIONS (default: 20) -- training sessions in phase 1 + CONVERGE_SESSIONS (default: 10) -- cold sessions in phase 2 + WIN_THRESHOLD (default: 0.7) -- min share for EXPECTED_WINNER in phase 2 +""" + +from __future__ import annotations + +import asyncio +import os +import sys +import time +import uuid +from typing import List, Optional + +import httpx + +PROXY_URL: str = os.environ.get("LITELLM_PROXY_URL", "http://localhost:4000") +try: + PROXY_KEY: str = os.environ["LITELLM_PROXY_KEY"] +except KeyError: + print( + "ERROR: LITELLM_PROXY_KEY env var must be set (a proxy key with /chat/completions perms).", + file=sys.stderr, + ) + sys.exit(2) + +ROUTER_NAME: str = os.environ.get("ADAPTIVE_ROUTER_NAME", "smart-cheap-router") +EXPECTED_WINNER: str = os.environ.get("EXPECTED_WINNER", "smart") +TRAIN_SESSIONS: int = int(os.environ.get("TRAIN_SESSIONS", "20")) +CONVERGE_SESSIONS: int = int(os.environ.get("CONVERGE_SESSIONS", "10")) +WIN_THRESHOLD: float = float(os.environ.get("WIN_THRESHOLD", "0.7")) + +REQUEST_TIMEOUT_SECONDS: float = 30.0 +RETRY_ATTEMPTS: int = 3 +RETRY_BACKOFF_SECONDS: float = 1.0 +FLUSHER_DRAIN_WAIT_SECONDS: float = 30.0 # proxy flusher loop is 10s; pad with margin + +PROMPTS: List[str] = [ + "Write a Python function that reverses a binary tree", + "Explain the time complexity of quicksort", + "Design an API for a chat application", +] +SATISFACTION_PROMPT: str = "thanks, that worked!" + + +async def _post_chat( + client: httpx.AsyncClient, session_id: str, prompt: str +) -> Optional[dict]: + """POST a chat completion with retry + timeout. Returns response JSON or None.""" + body = { + "model": ROUTER_NAME, + "messages": [{"role": "user", "content": prompt}], + "metadata": {"litellm_session_id": session_id}, + } + last_exc: Optional[Exception] = None + for attempt in range(1, RETRY_ATTEMPTS + 1): + try: + r = await client.post( + f"{PROXY_URL}/v1/chat/completions", + json=body, + headers={"Authorization": f"Bearer {PROXY_KEY}"}, + timeout=REQUEST_TIMEOUT_SECONDS, + ) + r.raise_for_status() + return r.json() + except Exception as e: # noqa: BLE001 + last_exc = e + if attempt < RETRY_ATTEMPTS: + await asyncio.sleep(RETRY_BACKOFF_SECONDS * attempt) + print( + f" request failed after {RETRY_ATTEMPTS} attempts (session={session_id}): {last_exc}", + file=sys.stderr, + ) + return None + + +async def send_session( + client: httpx.AsyncClient, + session_id: str, + prompts: List[str], + satisfy: bool = True, +) -> Optional[str]: + """Send a session of N turns. Returns the model that handled the last turn.""" + last_model: Optional[str] = None + for prompt in prompts: + resp = await _post_chat(client, session_id, prompt) + if resp is None: + return None + last_model = resp.get("model") or last_model + if satisfy: + await _post_chat(client, session_id, SATISFACTION_PROMPT) + return last_model + + +async def _proxy_health_check(client: httpx.AsyncClient) -> bool: + """Confirm the proxy is reachable before doing anything else.""" + try: + r = await client.get(f"{PROXY_URL}/health/liveliness", timeout=5.0) + return r.status_code == 200 + except Exception as e: # noqa: BLE001 + print(f"proxy unreachable at {PROXY_URL}: {e}", file=sys.stderr) + return False + + +async def main() -> None: + print("=== verify_adaptive_router.py ===") + print(f"proxy: {PROXY_URL}") + print(f"router: {ROUTER_NAME}") + print(f"expected winner: {EXPECTED_WINNER}") + print(f"train sessions: {TRAIN_SESSIONS}") + print(f"converge runs: {CONVERGE_SESSIONS}\n") + + async with httpx.AsyncClient() as client: + if not await _proxy_health_check(client): + print("FAIL: proxy health check did not return 200.", file=sys.stderr) + sys.exit(1) + + # ---- Phase 1: training ------------------------------------------- + print( + f"Phase 1: training ({TRAIN_SESSIONS} sessions of 3 turns + satisfaction)..." + ) + for i in range(TRAIN_SESSIONS): + sid = f"verify-train-{uuid.uuid4()}" + await send_session(client, sid, PROMPTS, satisfy=True) + if (i + 1) % 5 == 0: + print(f" trained {i + 1}/{TRAIN_SESSIONS} sessions") + + print( + f"\nWaiting {FLUSHER_DRAIN_WAIT_SECONDS:.0f}s for flusher to drain queue..." + ) + await asyncio.sleep(FLUSHER_DRAIN_WAIT_SECONDS) + + # ---- Phase 2: convergence ---------------------------------------- + print(f"\nPhase 2: convergence test ({CONVERGE_SESSIONS} cold sessions)...") + picks: List[str] = [] + for i in range(CONVERGE_SESSIONS): + sid = f"verify-test-{uuid.uuid4()}" + m = await send_session(client, sid, [PROMPTS[0]], satisfy=False) + if m: + picks.append(m) + print(f" session {i + 1}: picked {m}") + + if not picks: + print("\nFAIL: no successful picks in convergence phase.", file=sys.stderr) + sys.exit(1) + winner_share = picks.count(EXPECTED_WINNER) / len(picks) + print( + f"\n{EXPECTED_WINNER} share: {winner_share:.0%} " + f"({picks.count(EXPECTED_WINNER)}/{len(picks)})" + ) + + # ---- Phase 3: sticky session ------------------------------------- + print("\nPhase 3: sticky session test...") + sid = f"verify-sticky-{uuid.uuid4()}" + models: List[str] = [] + for _ in range(3): + m = await send_session(client, sid, [PROMPTS[0]], satisfy=False) + if m: + models.append(m) + if len(models) == 3 and len(set(models)) == 1: + print(f" PASS: same model {models[0]} across 3 turns of session {sid}") + else: + print( + f" FAIL: models differed within session: {models}", + file=sys.stderr, + ) + sys.exit(1) + + # ---- Phase 4: latency benchmark ---------------------------------- + print("\nPhase 4: routing latency (5 picks, p50)...") + latencies: List[float] = [] + for _ in range(5): + t0 = time.perf_counter() + await send_session( + client, f"verify-lat-{uuid.uuid4()}", [PROMPTS[0]], satisfy=False + ) + latencies.append(time.perf_counter() - t0) + latencies.sort() + p50 = latencies[len(latencies) // 2] + print(f" p50 e2e roundtrip: {p50 * 1000:.0f}ms") + + # ---- Verdict ----------------------------------------------------- + if winner_share >= WIN_THRESHOLD: + print( + f"\nPASS: convergence ({winner_share:.0%} >= {WIN_THRESHOLD:.0%}) + " + f"sticky + latency checks all green." + ) + sys.exit(0) + print( + f"\nFAIL: convergence too weak ({winner_share:.0%} < {WIN_THRESHOLD:.0%}).", + file=sys.stderr, + ) + sys.exit(1) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/eval_swe_bench.py b/tests/eval_swe_bench.py index 9c986283abd..6ae99f83ca1 100644 --- a/tests/eval_swe_bench.py +++ b/tests/eval_swe_bench.py @@ -40,6 +40,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) import litellm # noqa: E402 from litellm.compression import compress as litellm_compress # noqa: E402 +from litellm.types.utils import CallTypes # noqa: E402 # --------------------------------------------------------------------------- # Prompts @@ -445,7 +446,7 @@ def eval_instance( compress_kwargs: dict = { "messages": messages, "model": model, - "input_type": "openai_chat_completions", + "call_type": CallTypes.completion, "compression_trigger": compression_trigger, "embedding_model": embedding_model, } diff --git a/tests/llm_translation/test_bedrock_mantle.py b/tests/llm_translation/test_bedrock_mantle.py new file mode 100644 index 00000000000..d545f78bc43 --- /dev/null +++ b/tests/llm_translation/test_bedrock_mantle.py @@ -0,0 +1,149 @@ +""" +E2E tests for Bedrock Mantle (Claude Mythos Preview) integration. + +Tests use a fake/mocked HTTP layer to verify the full request pipeline: +- correct endpoint URL +- model ID in the request body +- AWS SigV4 Authorization header present +- response parsing +""" + +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +MODEL = "bedrock/mantle/anthropic.claude-mythos-preview" +REGION = "us-east-1" +EXPECTED_URL = f"https://bedrock-mantle.{REGION}.api.aws/v1/messages" + +FAKE_ANTHROPIC_RESPONSE = { + "id": "msg_fake123", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-mythos-preview", + "content": [{"type": "text", "text": "Hello from Mythos!"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, +} + + +def _make_fake_response(body: dict) -> MagicMock: + mock_resp = MagicMock(spec=httpx.Response) + mock_resp.status_code = 200 + mock_resp.headers = httpx.Headers({"content-type": "application/json"}) + mock_resp.text = json.dumps(body) + mock_resp.json.return_value = body + mock_resp.is_error = False + mock_resp.raise_for_status = MagicMock() + return mock_resp + + +def test_mantle_request_url_and_body(): + """Verify the correct URL is called and model appears in the request body.""" + client = HTTPHandler() + + with patch.object( + client, "post", return_value=_make_fake_response(FAKE_ANTHROPIC_RESPONSE) + ) as mock_post: + try: + litellm.completion( + model=MODEL, + messages=[{"role": "user", "content": "Hello"}], + max_tokens=50, + aws_region_name=REGION, + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + client=client, + ) + except Exception: + pass # response parsing may fail on mock; we only care about the outgoing call + + mock_post.assert_called_once() + call_kwargs = mock_post.call_args.kwargs + + # Correct endpoint + assert ( + call_kwargs["url"] == EXPECTED_URL + ), f"Expected {EXPECTED_URL}, got {call_kwargs['url']}" + + # Request body has model ID (without "mantle/" prefix) + raw_data = call_kwargs.get("data") or call_kwargs.get("json") + body = json.loads(raw_data) if isinstance(raw_data, (str, bytes)) else raw_data + assert ( + body["model"] == "anthropic.claude-mythos-preview" + ), f"body['model'] = {body.get('model')}" + assert "messages" in body + assert body["max_tokens"] == 50 + + # AWS SigV4 Authorization header must be present + headers = call_kwargs.get("headers", {}) + assert "Authorization" in headers, f"No Authorization header in {headers}" + assert headers["Authorization"].startswith( + "AWS4-HMAC-SHA256" + ), f"Expected SigV4 auth, got: {headers['Authorization'][:50]}" + + +def test_mantle_request_does_not_include_mantle_prefix_in_body(): + """Ensure 'mantle/' never leaks into the request body.""" + client = HTTPHandler() + + with patch.object( + client, "post", return_value=_make_fake_response(FAKE_ANTHROPIC_RESPONSE) + ) as mock_post: + try: + litellm.completion( + model=MODEL, + messages=[{"role": "user", "content": "Hi"}], + max_tokens=10, + aws_region_name=REGION, + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + client=client, + ) + except Exception: + pass + + call_kwargs = mock_post.call_args.kwargs + raw_data = call_kwargs.get("data") or call_kwargs.get("json") + body = json.loads(raw_data) if isinstance(raw_data, (str, bytes)) else raw_data + + body_str = json.dumps(body) + assert "mantle/" not in body_str, f"'mantle/' leaked into body: {body_str}" + + +def test_mantle_region_reflected_in_url(): + """The region from aws_region_name must appear in the endpoint URL.""" + client = HTTPHandler() + + for region in ["us-east-1", "us-west-2", "eu-west-1"]: + with patch.object( + client, "post", return_value=_make_fake_response(FAKE_ANTHROPIC_RESPONSE) + ) as mock_post: + try: + litellm.completion( + model=MODEL, + messages=[{"role": "user", "content": "Hi"}], + max_tokens=10, + aws_region_name=region, + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + client=client, + ) + except Exception: + pass + + call_kwargs = mock_post.call_args.kwargs + expected = f"https://bedrock-mantle.{region}.api.aws/v1/messages" + assert ( + call_kwargs["url"] == expected + ), f"region={region}: expected URL {expected}, got {call_kwargs['url']}" diff --git a/tests/local_testing/test_basic_python_version.py b/tests/local_testing/test_basic_python_version.py index 37e23d64677..8308e0d6033 100644 --- a/tests/local_testing/test_basic_python_version.py +++ b/tests/local_testing/test_basic_python_version.py @@ -100,8 +100,12 @@ import pytest import requests -def test_litellm_proxy_server_config_no_general_settings(): - # Sync the local litellm packages into the project environment +def _run_proxy_server_smoke_test(extra_proxy_args=None): + """Sync deps, generate Prisma client, start proxy with optional extra args, + send a health check + chat/completions request, and tear down.""" + if extra_proxy_args is None: + extra_proxy_args = [] + server_process = None try: _run_uv( @@ -144,6 +148,7 @@ def test_litellm_proxy_server_config_no_general_settings(): "litellm.proxy.proxy_cli", "--config", config_fp, + *extra_proxy_args, ], cwd=PROJECT_ROOT, ) @@ -182,3 +187,17 @@ def test_litellm_proxy_server_config_no_general_settings(): # Additional assertions can be added here assert True + + +def test_litellm_proxy_server_config_no_general_settings(): + """Exercises the default (v1) migration resolver.""" + _run_proxy_server_smoke_test() + + +def test_litellm_proxy_server_config_no_general_settings_v2_resolver(): + """Exercises the opt-in v2 migration resolver. + + Runs in a separate CI job against a local Postgres to avoid collisions + with the v1 variant when they share a database. + """ + _run_proxy_server_smoke_test(extra_proxy_args=["--use_v2_migration_resolver"]) diff --git a/tests/local_testing/test_batch_completions.py b/tests/local_testing/test_batch_completions.py index 2125a998f84..95bfe5e6e2b 100644 --- a/tests/local_testing/test_batch_completions.py +++ b/tests/local_testing/test_batch_completions.py @@ -72,7 +72,7 @@ def test_batch_completions_models(): def test_batch_completion_models_all_responses(): try: responses = batch_completion_models_all_responses( - models=["gemini/gemini-2.5-flash-lite", "claude-3-haiku-20240307"], + models=["gemini/gemini-2.5-flash-lite", "claude-haiku-4-5-20251001"], messages=[{"role": "user", "content": "write a poem"}], max_tokens=10, ) diff --git a/tests/local_testing/test_function_call_parsing.py b/tests/local_testing/test_function_call_parsing.py index 0351ce70572..f9582fcc574 100644 --- a/tests/local_testing/test_function_call_parsing.py +++ b/tests/local_testing/test_function_call_parsing.py @@ -142,7 +142,7 @@ def trade(model_name: str) -> List[Trade]: # type: ignore @pytest.mark.parametrize( - "model", ["claude-3-haiku-20240307", "anthropic.claude-3-haiku-20240307-v1:0"] + "model", ["claude-haiku-4-5-20251001", "anthropic.claude-3-haiku-20240307-v1:0"] ) @pytest.mark.flaky(retries=6, delay=10) def test_function_call_parsing(model): diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index 13adb163d5f..b52805c0664 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -47,7 +47,7 @@ def get_current_weather(location, unit="fahrenheit"): [ "gpt-3.5-turbo-1106", "mistral/mistral-large-latest", - "claude-3-haiku-20240307", + "claude-haiku-4-5-20251001", "gemini/gemini-2.5-flash-lite", "anthropic.claude-3-sonnet-20240229-v1:0", ], @@ -275,7 +275,7 @@ from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message "anthropic.claude-3-sonnet-20240229-v1:0", "bedrock", ), - ("claude-3-haiku-20240307", "anthropic"), + ("claude-haiku-4-5-20251001", "anthropic"), ], ) @pytest.mark.parametrize( diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 383ad104577..a14e53adbc4 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -1509,7 +1509,7 @@ def test_router_fallbacks_with_wildcard_model_name(): { "model_name": "claude-3-haiku", "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), "mock_response": "Hi this is claude!", }, @@ -1555,7 +1555,7 @@ def test_fallbacks_with_different_messages(): { "model_name": "claude-3-haiku", "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, }, diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index 3aed0699603..ecac2cfe40e 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -1727,7 +1727,7 @@ def test_openai_chat_completion_complete_response_call(): "model", [ "gpt-3.5-turbo", - "claude-3-haiku-20240307", + "claude-haiku-4-5-20251001", "o1", ], ) @@ -2247,7 +2247,7 @@ def streaming_and_function_calling_format_tests(idx, chunk): [ # "gpt-3.5-turbo", # "anthropic.claude-3-sonnet-20240229-v1:0", - "claude-3-haiku-20240307", + "claude-haiku-4-5-20251001", ], ) def test_streaming_and_function_calling(model): diff --git a/tests/logging_callback_tests/test_otel_logging.py b/tests/logging_callback_tests/test_otel_logging.py index fdb333899cc..ea1c884c324 100644 --- a/tests/logging_callback_tests/test_otel_logging.py +++ b/tests/logging_callback_tests/test_otel_logging.py @@ -253,6 +253,7 @@ def validate_redacted_message_span_attributes(span): or attr.startswith("gen_ai.cost.") or attr.startswith("gen_ai.operation.") or attr.startswith("gen_ai.request.") + or attr.startswith("litellm.") ), f"Non-metadata attribute found: {attr}" pass diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py index ae2d66eb955..84b14f9508b 100644 --- a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py +++ b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py @@ -77,7 +77,7 @@ class TestAnthropicDirectAPI(BaseAnthropicMessagesTest): @property def model_config(self) -> Dict[str, Any]: return { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), } @@ -86,7 +86,7 @@ class TestAnthropicDirectAPI(BaseAnthropicMessagesTest): """ This is the model name that is expected to be in the logging payload """ - return "claude-3-haiku-20240307" + return "claude-haiku-4-5-20251001" class TestAnthropicBedrockAPI(BaseAnthropicMessagesTest): @@ -140,7 +140,7 @@ async def test_anthropic_messages_streaming_with_bad_request(): response = await litellm.anthropic.messages.acreate( messages=[{"role": "user", "content": "hi"}], api_key=os.getenv("ANTHROPIC_API_KEY"), - model="claude-3-haiku-20240307", + model="claude-haiku-4-5-20251001", max_tokens=100, stream=True, ) @@ -168,7 +168,7 @@ async def test_anthropic_messages_router_streaming_with_bad_request(): { "model_name": "claude-special-alias", "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, } @@ -205,7 +205,7 @@ async def test_anthropic_messages_litellm_router_non_streaming(): { "model_name": "claude-special-alias", "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, } @@ -243,7 +243,7 @@ async def test_anthropic_messages_litellm_router_routing_strategy(): { "model_name": "claude-special-alias", "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, } @@ -341,7 +341,7 @@ async def test_anthropic_messages_litellm_router_latency_metadata_tracking(): "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Here's a joke for you!"}], - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "stop_reason": "end_turn", "usage": {"input_tokens": 10, "output_tokens": 20}, } @@ -355,7 +355,7 @@ async def test_anthropic_messages_litellm_router_latency_metadata_tracking(): { "model_name": MODEL_GROUP, "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, } @@ -419,7 +419,7 @@ async def test_anthropic_messages_litellm_router_latency_metadata_tracking(): assert "model_info" in litellm_metadata # Verify other call parameters - assert call_kwargs["model"] == "claude-3-haiku-20240307" + assert call_kwargs["model"] == "claude-haiku-4-5-20251001" assert call_kwargs["messages"] == messages assert call_kwargs["max_tokens"] == 100 assert call_kwargs["metadata"] == {"user_id": "hello"} @@ -459,7 +459,7 @@ async def test_anthropic_messages_litellm_router_non_streaming_with_logging(): { "model_name": MODEL_GROUP, "litellm_params": { - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "api_key": os.getenv("ANTHROPIC_API_KEY"), }, } @@ -496,7 +496,7 @@ async def test_anthropic_messages_litellm_router_non_streaming_with_logging(): assert test_custom_logger.logged_standard_logging_payload["response"] is not None assert ( test_custom_logger.logged_standard_logging_payload["model"] - == "claude-3-haiku-20240307" + == "claude-haiku-4-5-20251001" ) # check logged usage + spend @@ -543,7 +543,7 @@ async def test_anthropic_messages_with_extra_headers(): "text": "Why did the chicken cross the road? To get to the other side!", } ], - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "stop_reason": "end_turn", "usage": {"input_tokens": 10, "output_tokens": 20}, } @@ -556,7 +556,7 @@ async def test_anthropic_messages_with_extra_headers(): response = await litellm.anthropic.messages.acreate( messages=messages, api_key=api_key, - model="claude-3-haiku-20240307", + model="claude-haiku-4-5-20251001", max_tokens=100, client=mock_client, provider_specific_header={ @@ -689,7 +689,7 @@ async def test_anthropic_messages_with_thinking(): "text": "Why did the chicken cross the road? To get to the other side!", } ], - "model": "claude-3-haiku-20240307", + "model": "claude-haiku-4-5-20251001", "stop_reason": "end_turn", "usage": {"input_tokens": 10, "output_tokens": 20}, } @@ -702,7 +702,7 @@ async def test_anthropic_messages_with_thinking(): response = await litellm.anthropic.messages.acreate( messages=messages, api_key=api_key, - model="claude-3-haiku-20240307", + model="claude-haiku-4-5-20251001", max_tokens=100, client=mock_client, thinking={"budget_tokens": 100}, @@ -717,7 +717,7 @@ async def test_anthropic_messages_with_thinking(): request_body = json.loads(call_kwargs.get("data", {})) print("REQUEST BODY", request_body) assert request_body["max_tokens"] == 100 - assert request_body["model"] == "claude-3-haiku-20240307" + assert request_body["model"] == "claude-haiku-4-5-20251001" assert request_body["messages"] == messages assert request_body["thinking"] == {"budget_tokens": 100} diff --git a/tests/proxy_unit_tests/conftest.py b/tests/proxy_unit_tests/conftest.py index 1421700c9a8..a0326f64ed7 100644 --- a/tests/proxy_unit_tests/conftest.py +++ b/tests/proxy_unit_tests/conftest.py @@ -1,48 +1,143 @@ # conftest.py -import importlib +import asyncio +import copy +import inspect import os import sys +import warnings import pytest sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path + import litellm +import litellm.proxy.proxy_server + + +# Top-level assignments of these types are the ones importlib.reload(litellm) +# would have effectively reset. We snapshot them at conftest import time and +# deep-copy the snapshot back before every test. +_SNAPSHOT_TYPES = (list, dict, set, tuple, str, int, float, bool, bytes) + + +def _snapshot_mutable_state(module): + """Capture a per-module snapshot of primitive and collection attributes.""" + snapshot = {} + for attr in list(vars(module)): + if attr.startswith("_"): + continue + try: + value = getattr(module, attr) + except Exception as exc: + warnings.warn( + f"conftest: could not read {module.__name__}.{attr} during snapshot: {exc}", + stacklevel=2, + ) + continue + if value is None or isinstance(value, _SNAPSHOT_TYPES): + try: + snapshot[attr] = copy.deepcopy(value) + except Exception as exc: + warnings.warn( + f"conftest: could not snapshot {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + return snapshot + + +def _restore_mutable_state(module, snapshot): + for attr, default in snapshot.items(): + try: + setattr(module, attr, copy.deepcopy(default)) + except Exception as exc: + warnings.warn( + f"conftest: could not restore {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + + +def _collect_flushable_caches(): + """Return (module, attr) pairs whose values expose flush_cache().""" + targets = [] + for module in (litellm, litellm.proxy.proxy_server): + for attr in list(vars(module)): + if attr.startswith("_"): + continue + try: + value = getattr(module, attr) + except Exception: + continue + # Only instances — a class reference has an unbound flush_cache + # that can't be called without a self argument. + if inspect.isclass(value) or inspect.ismodule(value): + continue + if callable(getattr(value, "flush_cache", None)): + targets.append((module, attr)) + return targets + + +def _flush_caches(targets): + for module, attr in targets: + try: + value = getattr(module, attr) + except Exception: + continue + flush = getattr(value, "flush_cache", None) + if callable(flush): + try: + flush() + except Exception as exc: + warnings.warn( + f"conftest: flush_cache failed on {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + + +# Snapshot once at conftest import — these are the "clean" module states. +_LITELLM_STATE = _snapshot_mutable_state(litellm) +_PROXY_SERVER_STATE = _snapshot_mutable_state(litellm.proxy.proxy_server) +_FLUSHABLE_CACHES = _collect_flushable_caches() @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(): + """Reset mutable module state on litellm and proxy_server before each test. + + Replaces a previous importlib.reload(litellm) approach that cost ~17s + per test (re-executing the full litellm __init__ import chain). + + What IS reset: + - Top-level module attributes of type list / dict / set / tuple + / str / int / float / bool / bytes, and None-valued attributes. + These cover callback lists, general_settings, master_key, + premium_user, prisma_client, etc. — anything the old reload() reset + by re-executing the module body. + - Any module-level object instance that exposes flush_cache() (the + DualCache and LLMClientCache family), which handles cache state + that can't round-trip through deepcopy because of internal locks. + + What is NOT reset: + - Class instances without flush_cache() (e.g. ProxyLogging, + JWTHandler, FastAPI routers, loggers). If a test mutates such an + instance in-place (setattr on the instance, appending to one of + its internal lists, etc.), the mutation will leak into later tests. + Use pytest's monkeypatch.setattr() or a local fixture for those + cases — don't rely on this autouse fixture to undo them. """ - This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. - """ - curr_dir = os.getcwd() # Get the current working directory - sys.path.insert( - 0, os.path.abspath("../..") - ) # Adds the project directory to the system path - - import litellm - from litellm import Router - - importlib.reload(litellm) - try: - if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): - importlib.reload(litellm.proxy.proxy_server) - except Exception as e: - print(f"Error reloading litellm.proxy.proxy_server: {e}") - - import asyncio + _restore_mutable_state(litellm, _LITELLM_STATE) + _restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE) + _flush_caches(_FLUSHABLE_CACHES) loop = asyncio.get_event_loop_policy().new_event_loop() asyncio.set_event_loop(loop) - print(litellm) - # from litellm import Router, completion, aembedding, acompletion, embedding - yield - - # Teardown code (executes after the yield point) - loop.close() # Close the loop created earlier - asyncio.set_event_loop(None) # Remove the reference to the loop + try: + yield + finally: + loop.close() + asyncio.set_event_loop(None) def pytest_collection_modifyitems(config, items): diff --git a/tests/router_unit_tests/test_router_index_management.py b/tests/router_unit_tests/test_router_index_management.py index f67edac494c..43718590808 100644 --- a/tests/router_unit_tests/test_router_index_management.py +++ b/tests/router_unit_tests/test_router_index_management.py @@ -244,6 +244,7 @@ class TestRouterIndexManagement: # Methods that are allowed to iterate through self.model_list ALLOWED_METHODS = [ "_get_deployment_by_litellm_model", # Edge case: lookup by litellm_params.model (not indexed) + "_finalize_adaptive_router_if_configured", # Init-time prefix scan for "auto_router/adaptive_router" (no index for prefix match) ] # Get path to router.py diff --git a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py new file mode 100644 index 00000000000..56e5a94cd49 --- /dev/null +++ b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py @@ -0,0 +1,364 @@ +""" +Unit tests for Compression Interception Handler. +""" + +from unittest.mock import MagicMock + +import pytest + +from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, +) +from litellm.types.utils import CallTypes + + +def test_initialize_from_proxy_config(): + """Test initialization from proxy config with litellm_settings.""" + litellm_settings = { + "compression_interception_params": { + "enabled": True, + "compression_trigger": 1234, + "compression_target": 789, + } + } + + logger = CompressionInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params={}, + ) + + assert logger.enabled is True + assert logger.compression_trigger == 1234 + assert logger.compression_target == 789 + + +@pytest.mark.asyncio +async def test_pre_call_hook_compresses_messages_and_injects_tool(monkeypatch): + """Test pre-call hook compresses and stores per-call cache.""" + logger = CompressionInterceptionLogger() + compressed_result = { + "messages": [{"role": "user", "content": "stubbed"}], + "original_tokens": 12000, + "compressed_tokens": 5000, + "compression_ratio": 0.58, + "cache": {"auth.py": "full file content"}, + "tools": [ + { + "type": "function", + "function": { + "name": "litellm_content_retrieve", + "parameters": { + "type": "object", + "properties": {"key": {"type": "string"}}, + }, + }, + } + ], + } + + def _fake_compress(**kwargs): + return compressed_result + + # The handler does ``from litellm.compression import compress`` at module + # scope, so we must patch the binding on the handler module — patching + # ``litellm.compress`` has no effect on the already-bound reference. + monkeypatch.setattr( + "litellm.integrations.compression_interception.handler.compress", + _fake_compress, + ) + + kwargs = { + "model": "bedrock/us.anthropic.claude-sonnet-4-5", + "messages": [{"role": "user", "content": "very large context"}], + "tools": [ + { + "type": "function", + "function": {"name": "existing_tool", "parameters": {"type": "object"}}, + } + ], + } + + result = await logger.async_pre_call_deployment_hook( + kwargs=kwargs, call_type=CallTypes.anthropic_messages + ) + + assert result is not None + assert result["messages"] == compressed_result["messages"] + tool_names = [t.get("function", {}).get("name") for t in result["tools"]] + assert "existing_tool" in tool_names + assert "litellm_content_retrieve" in tool_names + assert result["litellm_call_id"] in logger._compression_cache_by_call_id + + +@pytest.mark.asyncio +async def test_pre_call_hook_below_trigger_does_not_inject_empty_tools(monkeypatch): + """ + When compression is a no-op (below trigger / invalid tool sequence), the + hook must NOT replace ``messages`` or inject an empty ``tools: []`` onto + a request that originally had no tools — Anthropic Messages rejects + ``tools: []``. + """ + logger = CompressionInterceptionLogger() + original_messages = [{"role": "user", "content": "short prompt"}] + + def _fake_compress_noop(**kwargs): + return { + "messages": original_messages, + "original_tokens": 42, + "compressed_tokens": 42, + "compression_ratio": 0.0, + "cache": {}, + "tools": [], + "compression_skipped_reason": "below_trigger", + } + + monkeypatch.setattr( + "litellm.integrations.compression_interception.handler.compress", + _fake_compress_noop, + ) + + kwargs = { + "model": "bedrock/us.anthropic.claude-sonnet-4-5", + "messages": original_messages, + } + + result = await logger.async_pre_call_deployment_hook( + kwargs=kwargs, call_type=CallTypes.anthropic_messages + ) + + assert result is not None + # Original request had no ``tools`` — skipped compression must leave it that way. + assert "tools" not in result + # Cache must not be populated for a no-op. + assert result.get("litellm_call_id") not in logger._compression_cache_by_call_id + + +@pytest.mark.asyncio +async def test_should_run_agentic_loop_detects_retrieval_tool_use(): + """Test should-run hook returns tool calls for retrieval tool_use blocks.""" + logger = CompressionInterceptionLogger() + response = { + "content": [ + { + "type": "tool_use", + "id": "toolu_123", + "name": "litellm_content_retrieve", + "input": {"key": "auth.py"}, + } + ] + } + + should_run, tools_dict = await logger.async_should_run_agentic_loop( + response=response, + model="bedrock/claude", + messages=[], + tools=[ + { + "type": "function", + "function": { + "name": "litellm_content_retrieve", + "parameters": {"type": "object"}, + }, + } + ], + stream=False, + custom_llm_provider="bedrock", + kwargs={}, + ) + + assert should_run is True + assert len(tools_dict["tool_calls"]) == 1 + assert tools_dict["tool_calls"][0]["input"]["key"] == "auth.py" + + +@pytest.mark.asyncio +async def test_build_agentic_loop_plan_returns_request_patch(): + """Callback should return typed patch with tool_result content.""" + logger = CompressionInterceptionLogger() + call_id = "call_123" + logger._compression_cache_by_call_id[call_id] = ( + {"auth.py": "full auth file"}, + 9999999999.0, + ) + + logging_obj = MagicMock() + logging_obj.litellm_call_id = call_id + logging_obj.model_call_details = { + "agentic_loop_params": {"model": "bedrock/invoke/claude-3-5-sonnet"} + } + + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "toolu_abc", + "type": "tool_use", + "name": "litellm_content_retrieve", + "input": {"key": "auth.py"}, + } + ] + }, + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "read auth.py"}], + response=None, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={ + "max_tokens": 1024, + "tools": [{"name": "litellm_content_retrieve"}], + }, + logging_obj=logging_obj, + stream=False, + kwargs={ + "temperature": 0.1, + "_compression_interception_internal": True, + "litellm_logging_obj": object(), + }, + ) + + assert plan.run_agentic_loop is True + assert plan.request_patch is not None + assert plan.request_patch.model == "bedrock/invoke/claude-3-5-sonnet" + assert plan.request_patch.max_tokens == 1024 + assert plan.request_patch.messages is not None + assert len(plan.request_patch.messages) == 3 + tool_result_content = plan.request_patch.messages[-1]["content"][0]["content"] + assert tool_result_content == "full auth file" + assert "_compression_interception_internal" not in plan.request_patch.kwargs + assert "litellm_logging_obj" not in plan.request_patch.kwargs + assert plan.request_patch.kwargs["temperature"] == 0.1 + assert "max_tokens" not in plan.request_patch.optional_params + + +@pytest.mark.asyncio +async def test_should_run_agentic_loop_with_custom_type_tools(): + """Test that async_should_run_agentic_loop returns True when tools contain + litellm_content_retrieve as a custom-typed tool (e.g. Claude Code tool list) + and the model response includes a matching tool_use block.""" + logger = CompressionInterceptionLogger() + + # Exact tools payload produced by Claude Code – litellm_content_retrieve is + # the final entry and uses type="custom" (not type="function"). + tools = [ + { + "name": "Agent", + "description": "Launch a new agent to handle complex, multi-step tasks.", + "input_schema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "description": {"type": "string"}, + "prompt": {"type": "string"}, + }, + "required": ["description", "prompt"], + "additionalProperties": False, + }, + }, + { + "name": "AskUserQuestion", + "description": "Use this tool when you need to ask the user questions.", + "input_schema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "questions": {"type": "array", "items": {"type": "object"}}, + }, + "required": ["questions"], + "additionalProperties": False, + }, + }, + { + "name": "Bash", + "description": "Executes a given bash command and returns its output.", + "input_schema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + "additionalProperties": False, + }, + }, + { + "name": "litellm_content_retrieve", + "description": "Retrieve the full content of a file or message that was compressed to save tokens.", + "input_schema": { + "type": "object", + "properties": { + "key": { + "type": "string", + "description": "The identifier of the content to retrieve", + "enum": [ + "message_0", + "HA_UPTIME_ROUTER_SPEC.md", + "message_159", + "message_160", + ], + } + }, + "required": ["key"], + }, + "type": "custom", + }, + ] + + response = { + "content": [ + { + "type": "tool_use", + "id": "toolu_abc", + "name": "litellm_content_retrieve", + "input": {"key": "message_0"}, + } + ] + } + + should_run, tools_dict = await logger.async_should_run_agentic_loop( + response=response, + model="claude-3-5-sonnet", + messages=[], + tools=tools, + stream=False, + custom_llm_provider="anthropic", + kwargs={}, + ) + + assert should_run is True + assert tools_dict["tool_type"] == "compression_retrieval" + assert len(tools_dict["tool_calls"]) == 1 + assert tools_dict["tool_calls"][0]["input"]["key"] == "message_0" + + +@pytest.mark.asyncio +async def test_build_agentic_loop_plan_missing_key_fallback(): + """Missing cache keys should produce deterministic fallback content.""" + logger = CompressionInterceptionLogger() + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "missing_call" + logging_obj.model_call_details = {"agentic_loop_params": {}} + + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "toolu_missing", + "type": "tool_use", + "name": "litellm_content_retrieve", + "input": {"key": "not_found.py"}, + } + ] + }, + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "read file"}], + response=None, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=logging_obj, + stream=False, + kwargs={}, + ) + + assert plan.request_patch is not None + assert ( + plan.request_patch.messages[-1]["content"][0]["content"] + == "[compressed content key 'not_found.py' not found]" + ) diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 725836e1340..f7106471894 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -1047,6 +1047,36 @@ class TestOpenTelemetryEndpointNormalization(unittest.TestCase): result = otel._normalize_otel_endpoint("http://collector:4318/", "traces") self.assertEqual(result, "http://collector:4318/v1/traces") + @parameterized.expand( + [ + ( + "https://ingest.eu1.observability.splunkcloud.com/v2/trace/otlp", + "https://ingest.eu1.observability.splunkcloud.com/v2/trace/otlp", + ), + ( + "https://ingest.us0.observability.splunkcloud.com/v2/trace/otlp/", + "https://ingest.us0.observability.splunkcloud.com/v2/trace/otlp", + ), + ( + "https://ingest.eu0.signalfx.com/v2/trace/otlp", + "https://ingest.eu0.signalfx.com/v2/trace/otlp", + ), + ( + "https://example.com/prefix/v2/trace/otlp", + "https://example.com/prefix/v2/trace/otlp", + ), + ] + ) + def test_normalize_traces_nonstandard_otlp_ingest_urls_unchanged( + self, input_url: str, expected: str + ) -> None: + """Splunk-style /v2/trace/otlp endpoints must not get /v1/traces appended.""" + otel = OpenTelemetry() + self.assertEqual( + otel._normalize_otel_endpoint(input_url, "traces"), + expected, + ) + def test_normalize_endpoint_none(self): """Test that None endpoint returns None""" otel = OpenTelemetry() @@ -1315,7 +1345,7 @@ class TestOpenTelemetryProtocolSelection(unittest.TestCase): @patch.dict( os.environ, { - "OTEL_EXPORTER": "otlp_http", + "OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf", "OTEL_EXPORTER_OTLP_ENDPOINT": "http://collector:4318", }, clear=False, @@ -1339,7 +1369,7 @@ class TestOpenTelemetryProtocolSelection(unittest.TestCase): @patch.dict( os.environ, { - "OTEL_EXPORTER": "otlp_grpc", + "OTEL_EXPORTER_OTLP_PROTOCOL": "grpc", "OTEL_EXPORTER_OTLP_ENDPOINT": "http://collector:4317", }, clear=False, @@ -1360,6 +1390,60 @@ class TestOpenTelemetryProtocolSelection(unittest.TestCase): self.assertIsInstance(processor, BatchSpanProcessor) self.assertIsInstance(processor.span_exporter, OTLPSpanExporterGRPC) + @patch.dict( + os.environ, + { + "OTEL_EXPORTER": "otlp_http", + "OTEL_EXPORTER_OTLP_ENDPOINT": "http://collector:4318", + }, + clear=False, + ) + def test_protocol_selection_from_otel_exporter_fallback_http(self): + """OTEL_EXPORTER drives protocol when OTEL_EXPORTER_OTLP_PROTOCOL is unset.""" + from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter as OTLPSpanExporterHTTP, + ) + from opentelemetry.sdk.trace.export import BatchSpanProcessor + + popped_protocol = os.environ.pop("OTEL_EXPORTER_OTLP_PROTOCOL", None) + try: + config = OpenTelemetryConfig.from_env() + self.assertEqual(config.exporter, "otlp_http") + otel = OpenTelemetry(config=config) + processor = otel._get_span_processor() + self.assertIsInstance(processor, BatchSpanProcessor) + self.assertIsInstance(processor.span_exporter, OTLPSpanExporterHTTP) + finally: + if popped_protocol is not None: + os.environ["OTEL_EXPORTER_OTLP_PROTOCOL"] = popped_protocol + + @patch.dict( + os.environ, + { + "OTEL_EXPORTER": "otlp_grpc", + "OTEL_EXPORTER_OTLP_ENDPOINT": "http://collector:4317", + }, + clear=False, + ) + def test_protocol_selection_from_otel_exporter_fallback_grpc(self): + """OTEL_EXPORTER drives protocol when OTEL_EXPORTER_OTLP_PROTOCOL is unset.""" + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( + OTLPSpanExporter as OTLPSpanExporterGRPC, + ) + from opentelemetry.sdk.trace.export import BatchSpanProcessor + + popped_protocol = os.environ.pop("OTEL_EXPORTER_OTLP_PROTOCOL", None) + try: + config = OpenTelemetryConfig.from_env() + self.assertEqual(config.exporter, "otlp_grpc") + otel = OpenTelemetry(config=config) + processor = otel._get_span_processor() + self.assertIsInstance(processor, BatchSpanProcessor) + self.assertIsInstance(processor.span_exporter, OTLPSpanExporterGRPC) + finally: + if popped_protocol is not None: + os.environ["OTEL_EXPORTER_OTLP_PROTOCOL"] = popped_protocol + def test_http_exporter_endpoint_normalization_for_traces(self): """Test that HTTP trace exporter gets properly normalized endpoint""" config = OpenTelemetryConfig( @@ -2752,3 +2836,26 @@ class TestResponseIdFallback(unittest.TestCase): mock_span.set_attribute.assert_any_call( "gen_ai.response.id", "litellm-img-call-101" ) + + def test_litellm_call_id_emitted_as_span_attribute(self): + """litellm.call_id must be set on the span from standard_logging_payload.""" + otel = OpenTelemetry() + mock_span = MagicMock() + + call_id = "my-litellm-call-uuid-456" + kwargs = { + "model": "gpt-4o", + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai"}, + "standard_logging_object": { + "id": "chatcmpl-provider-id", + "litellm_call_id": call_id, + "call_type": "completion", + "metadata": {}, + }, + } + response_obj = {"id": "chatcmpl-provider-id", "model": "gpt-4o"} + + otel.set_attributes(mock_span, kwargs, response_obj) + + mock_span.set_attribute.assert_any_call("litellm.call_id", call_id) diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index c8617a3c1b1..10951265115 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -4,7 +4,7 @@ Unit tests for WebSearch Interception Handler Tests the WebSearchInterceptionLogger class and helper functions. """ -from unittest.mock import MagicMock, Mock +from unittest.mock import AsyncMock, MagicMock, Mock import pytest @@ -69,6 +69,61 @@ async def test_async_should_run_agentic_loop(): assert tools_dict == {} +@pytest.mark.asyncio +async def test_async_build_agentic_loop_plan_returns_request_patch(): + """Callback should return a typed patch for base handler reruns.""" + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + logger._execute_search = AsyncMock( # type: ignore + return_value="Title: LiteLLM\nURL: docs\nSnippet: test" + ) + + tools_dict = { + "tool_calls": [ + { + "id": "toolu_123", + "type": "tool_use", + "name": "litellm_web_search", + "input": {"query": "what is litellm"}, + } + ], + "response_format": "anthropic", + } + logging_obj = MagicMock() + logging_obj.model_call_details = { + "agentic_loop_params": {"model": "bedrock/invoke/claude-3-5-sonnet"} + } + kwargs = { + "temperature": 0.2, + "_websearch_interception_converted_stream": True, + "litellm_logging_obj": object(), + } + + plan = await logger.async_build_agentic_loop_plan( + tools=tools_dict, + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "search LiteLLM"}], + response=None, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={ + "max_tokens": 1024, + "tools": [{"name": "litellm_web_search"}], + }, + logging_obj=logging_obj, + stream=False, + kwargs=kwargs, + ) + + assert plan.run_agentic_loop is True + assert plan.request_patch is not None + assert plan.request_patch.model == "bedrock/invoke/claude-3-5-sonnet" + assert plan.request_patch.max_tokens == 1024 + assert plan.request_patch.messages is not None + assert len(plan.request_patch.messages) == 3 + assert "_websearch_interception_converted_stream" not in plan.request_patch.kwargs + assert "litellm_logging_obj" not in plan.request_patch.kwargs + assert plan.request_patch.kwargs["temperature"] == 0.2 + + @pytest.mark.asyncio async def test_internal_flags_filtered_from_followup_kwargs(): """Test that internal _websearch_interception flags are filtered from follow-up request kwargs. diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index c3849e5869a..cf7be6bf1c7 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2410,3 +2410,29 @@ def test_get_additional_headers_reset_fields_preserved(): assert result is not None assert result["x_ratelimit_reset_requests"] == "1s" # type: ignore assert result["x_ratelimit_reset_tokens"] == "100ms" # type: ignore + + +# ── litellm_call_id propagation ─────────────────────────────────────────────── + + +def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_obj): + """litellm_call_id from kwargs must appear in the returned StandardLoggingPayload.""" + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + call_id = "test-call-id-abc-123" + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={"litellm_call_id": call_id, "model": "gpt-4o", "messages": []}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["litellm_call_id"] == call_id diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py new file mode 100644 index 00000000000..b9bda07336f --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py @@ -0,0 +1,792 @@ +""" +Tests for AgenticAnthropicStreamingIterator and SSE rebuild helpers. +""" + +import json +import os +import sys +from typing import Any, Dict, List, Optional, Tuple +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( + AgenticAnthropicStreamingIterator, + _handle_content_block_delta, + _handle_content_block_start, + _handle_content_block_stop, + _handle_message_delta, + _handle_message_start, + _parse_sse_events, +) + + +# --------------------------------------------------------------------------- +# Helpers to build SSE byte payloads +# --------------------------------------------------------------------------- + + +def _sse_event(event_type: str, data: dict) -> bytes: + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() + + +def _build_simple_text_stream() -> List[bytes]: + """Produce SSE bytes for a simple text response (no tool calls).""" + chunks = [] + chunks.append( + _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-20250514", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 0}, + }, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello, world!"}, + }, + ) + ) + chunks.append( + _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}) + ) + chunks.append( + _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ) + ) + chunks.append(_sse_event("message_stop", {"type": "message_stop"})) + return chunks + + +def _build_tool_use_stream() -> List[bytes]: + """Produce SSE bytes for a response with a tool_use block.""" + chunks = [] + chunks.append( + _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_tool_456", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-20250514", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 50, "output_tokens": 0}, + }, + }, + ) + ) + # thinking block + chunks.append( + _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "thinking", + "thinking": "", + "signature": "", + }, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "thinking_delta", + "thinking": "I need to retrieve...", + }, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "signature_delta", "signature": "sig_abc"}, + }, + ) + ) + chunks.append( + _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}) + ) + # tool_use block + chunks.append( + _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "tool_use", + "id": "toolu_001", + "name": "litellm_content_retrieve", + "input": {}, + }, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 1, + "delta": { + "type": "input_json_delta", + "partial_json": '{"key": "section_', + }, + }, + ) + ) + chunks.append( + _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": '1"}'}, + }, + ) + ) + chunks.append( + _sse_event("content_block_stop", {"type": "content_block_stop", "index": 1}) + ) + chunks.append( + _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use"}, + "usage": {"output_tokens": 20}, + }, + ) + ) + chunks.append(_sse_event("message_stop", {"type": "message_stop"})) + return chunks + + +# --------------------------------------------------------------------------- +# Mock async stream +# --------------------------------------------------------------------------- + + +class MockAsyncStream: + """Async iterator that yields a list of byte chunks.""" + + def __init__(self, chunks: List[bytes]): + self._chunks = list(chunks) + self._idx = 0 + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + if self._idx >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._idx] + self._idx += 1 + return chunk + + +# --------------------------------------------------------------------------- +# Tests for _parse_sse_events +# --------------------------------------------------------------------------- + + +class TestParseSSEEvents: + def test_should_parse_single_event(self): + raw = _sse_event( + "message_start", {"type": "message_start", "message": {"id": "1"}} + ) + events = _parse_sse_events(raw) + assert len(events) == 1 + assert events[0][0] == "message_start" + assert events[0][1]["message"]["id"] == "1" + + def test_should_parse_multiple_events(self): + raw = b"".join(_build_simple_text_stream()) + events = _parse_sse_events(raw) + event_types = [e[0] for e in events] + assert "message_start" in event_types + assert "content_block_start" in event_types + assert "content_block_delta" in event_types + assert "content_block_stop" in event_types + assert "message_delta" in event_types + assert "message_stop" in event_types + + def test_should_skip_malformed_json(self): + raw = b"event: message_start\ndata: {invalid json}\n\n" + events = _parse_sse_events(raw) + assert len(events) == 0 + + def test_should_handle_empty_bytes(self): + events = _parse_sse_events(b"") + assert events == [] + + +# --------------------------------------------------------------------------- +# Tests for _handle_* helpers +# --------------------------------------------------------------------------- + + +class TestHandleMessageStart: + def test_should_populate_envelope(self): + response: Dict[str, Any] = { + "id": "", + "model": "", + "role": "assistant", + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + data = { + "message": { + "id": "msg_abc", + "model": "claude-sonnet-4-20250514", + "role": "assistant", + "usage": { + "input_tokens": 42, + "cache_creation_input_tokens": 100, + }, + } + } + _handle_message_start(data, response) + assert response["id"] == "msg_abc" + assert response["model"] == "claude-sonnet-4-20250514" + assert response["usage"]["input_tokens"] == 42 + assert response["usage"]["cache_creation_input_tokens"] == 100 + + +class TestHandleContentBlockStart: + def test_should_create_text_block(self): + blocks: Dict[int, Dict] = {} + data = {"index": 0, "content_block": {"type": "text", "text": ""}} + _handle_content_block_start(data, blocks) + assert blocks[0] == {"type": "text", "text": ""} + + def test_should_create_tool_use_block(self): + blocks: Dict[int, Dict] = {} + data = { + "index": 1, + "content_block": { + "type": "tool_use", + "id": "toolu_x", + "name": "my_tool", + "input": {}, + }, + } + _handle_content_block_start(data, blocks) + assert blocks[1]["type"] == "tool_use" + assert blocks[1]["name"] == "my_tool" + assert blocks[1]["_partial_json"] == "" + + def test_should_create_thinking_block(self): + blocks: Dict[int, Dict] = {} + data = { + "index": 0, + "content_block": {"type": "thinking", "thinking": "", "signature": ""}, + } + _handle_content_block_start(data, blocks) + assert blocks[0]["type"] == "thinking" + + +class TestHandleContentBlockDelta: + def test_should_accumulate_text(self): + blocks = {0: {"type": "text", "text": "Hello"}} + _handle_content_block_delta( + {"index": 0, "delta": {"type": "text_delta", "text": " World"}}, + blocks, + ) + assert blocks[0]["text"] == "Hello World" + + def test_should_accumulate_json(self): + blocks = {0: {"type": "tool_use", "_partial_json": '{"key":'}} + _handle_content_block_delta( + { + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '"val"}'}, + }, + blocks, + ) + assert blocks[0]["_partial_json"] == '{"key":"val"}' + + def test_should_ignore_missing_block(self): + blocks: Dict[int, Dict] = {} + _handle_content_block_delta( + {"index": 99, "delta": {"type": "text_delta", "text": "x"}}, + blocks, + ) + assert 99 not in blocks + + +class TestHandleContentBlockStop: + def test_should_parse_tool_input_json(self): + blocks = { + 0: { + "type": "tool_use", + "input": {}, + "_partial_json": '{"key": "section_1"}', + } + } + _handle_content_block_stop({"index": 0}, blocks) + assert blocks[0]["input"] == {"key": "section_1"} + assert "_partial_json" not in blocks[0] + + def test_should_handle_invalid_json_gracefully(self): + blocks = { + 0: { + "type": "tool_use", + "input": {}, + "_partial_json": "not valid json", + } + } + _handle_content_block_stop({"index": 0}, blocks) + assert blocks[0]["input"] == {"_raw": "not valid json"} + + +class TestHandleMessageDelta: + def test_should_set_stop_reason_and_usage(self): + response: Dict[str, Any] = { + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + _handle_message_delta( + { + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 15}, + }, + response, + ) + assert response["stop_reason"] == "end_turn" + assert response["usage"]["output_tokens"] == 15 + + +# --------------------------------------------------------------------------- +# Tests for _rebuild_anthropic_response_from_sse +# --------------------------------------------------------------------------- + + +class TestRebuildAnthropicResponse: + def test_should_rebuild_simple_text_response(self): + raw_bytes = _build_simple_text_stream() + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + raw_bytes + ) + assert result is not None + assert result["id"] == "msg_123" + assert result["model"] == "claude-sonnet-4-20250514" + assert result["stop_reason"] == "end_turn" + assert len(result["content"]) == 1 + assert result["content"][0]["type"] == "text" + assert result["content"][0]["text"] == "Hello, world!" + assert result["usage"]["input_tokens"] == 10 + assert result["usage"]["output_tokens"] == 5 + + def test_should_rebuild_tool_use_response(self): + raw_bytes = _build_tool_use_stream() + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + raw_bytes + ) + assert result is not None + assert result["id"] == "msg_tool_456" + assert result["stop_reason"] == "tool_use" + assert len(result["content"]) == 2 + + thinking = result["content"][0] + assert thinking["type"] == "thinking" + assert thinking["thinking"] == "I need to retrieve..." + assert thinking["signature"] == "sig_abc" + + tool = result["content"][1] + assert tool["type"] == "tool_use" + assert tool["id"] == "toolu_001" + assert tool["name"] == "litellm_content_retrieve" + assert tool["input"] == {"key": "section_1"} + + def test_should_return_none_without_message_start(self): + raw_bytes = [ + _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text"}, + }, + ) + ] + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + raw_bytes + ) + assert result is None + + def test_should_handle_empty_bytes(self): + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + [] + ) + assert result is None + + def test_should_handle_multi_event_chunks(self): + """When multiple SSE events arrive in a single bytes chunk.""" + combined = b"".join(_build_simple_text_stream()) + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + [combined] + ) + assert result is not None + assert result["content"][0]["text"] == "Hello, world!" + + def test_should_preserve_cache_usage_fields(self): + raw_bytes = [ + _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_cache", + "model": "claude-sonnet-4-20250514", + "role": "assistant", + "usage": { + "input_tokens": 100, + "cache_creation_input_tokens": 50, + "cache_read_input_tokens": 30, + }, + }, + }, + ), + _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 10}, + }, + ), + _sse_event("message_stop", {"type": "message_stop"}), + ] + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + raw_bytes + ) + assert result is not None + assert result["usage"]["cache_creation_input_tokens"] == 50 + assert result["usage"]["cache_read_input_tokens"] == 30 + + def test_should_handle_redacted_thinking_block(self): + raw_bytes = [ + _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_redact", + "model": "claude-sonnet-4-20250514", + "role": "assistant", + "usage": {"input_tokens": 5}, + }, + }, + ), + _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "redacted_thinking", "data": "abc123"}, + }, + ), + _sse_event( + "content_block_stop", + {"type": "content_block_stop", "index": 0}, + ), + _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 1}, + }, + ), + _sse_event("message_stop", {"type": "message_stop"}), + ] + result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse( + raw_bytes + ) + assert result is not None + assert result["content"][0]["type"] == "redacted_thinking" + + +# --------------------------------------------------------------------------- +# Tests for AgenticAnthropicStreamingIterator (Phase 1 / Phase 2) +# --------------------------------------------------------------------------- + + +class TestAgenticStreamingIteratorPhase1: + @pytest.mark.asyncio + async def test_should_yield_all_chunks_when_no_hook_fires(self): + """When hooks return None, the wrapper should yield all original chunks.""" + chunks = _build_simple_text_stream() + mock_stream = MockAsyncStream(chunks) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert len(collected) == len(chunks) + for orig, got in zip(chunks, collected): + assert orig == got + + mock_handler._call_agentic_completion_hooks.assert_awaited_once() + + @pytest.mark.asyncio + async def test_should_pass_rebuilt_response_to_hooks(self): + """The rebuilt dict passed to hooks should match the original stream content.""" + chunks = _build_tool_use_stream() + mock_stream = MockAsyncStream(chunks) + + captured_response = {} + + async def mock_hooks(**kwargs): + captured_response.update(kwargs["response"]) + return None + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = mock_hooks + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + async for _ in iterator: + pass + + assert captured_response["id"] == "msg_tool_456" + assert captured_response["stop_reason"] == "tool_use" + assert captured_response["content"][1]["name"] == "litellm_content_retrieve" + + +class TestAgenticStreamingIteratorPhase2: + @pytest.mark.asyncio + async def test_should_chain_follow_up_async_iterator(self): + """When hooks return an async iterator, Phase 2 should yield from it.""" + phase1_chunks = _build_simple_text_stream() + phase2_chunks = [b"follow-up-chunk-1", b"follow-up-chunk-2"] + + mock_stream = MockAsyncStream(phase1_chunks) + follow_up = MockAsyncStream(phase2_chunks) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=follow_up) + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert len(collected) == len(phase1_chunks) + len(phase2_chunks) + assert collected[-2:] == phase2_chunks + + @pytest.mark.asyncio + async def test_should_convert_dict_response_to_fake_stream(self): + """When hooks return a dict, it should be wrapped in FakeAnthropicMessagesStreamIterator.""" + phase1_chunks = _build_simple_text_stream() + mock_stream = MockAsyncStream(phase1_chunks) + + fake_response = { + "id": "msg_followup", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-20250514", + "content": [{"type": "text", "text": "follow-up answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 100, "output_tokens": 20}, + } + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + return_value=fake_response + ) + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + # Phase 1 chunks + Phase 2 fake-stream chunks + assert len(collected) > len(phase1_chunks) + # The follow-up chunks should contain the text from the dict response + phase2_bytes = b"".join(collected[len(phase1_chunks) :]) + assert b"follow-up answer" in phase2_bytes + + +class TestAgenticStreamingIteratorErrorHandling: + @pytest.mark.asyncio + async def test_should_swallow_hook_errors(self): + """Errors in hook processing should be swallowed; Phase 1 chunks are still yielded.""" + chunks = _build_simple_text_stream() + mock_stream = MockAsyncStream(chunks) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + side_effect=RuntimeError("hook exploded") + ) + + mock_logging = MagicMock() + mock_logging.litellm_call_id = "test_call_123" + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=mock_logging, + custom_llm_provider="anthropic", + kwargs={}, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + # All Phase 1 chunks should still have been yielded + assert len(collected) == len(chunks) + + @pytest.mark.asyncio + async def test_should_handle_empty_stream(self): + """An empty upstream stream should not crash.""" + mock_stream = MockAsyncStream([]) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + collected = [] + async for chunk in iterator: + collected.append(chunk) + + assert collected == [] + # hooks should not be called since no bytes were collected + mock_handler._call_agentic_completion_hooks.assert_not_awaited() + + @pytest.mark.asyncio + async def test_should_pass_stream_true_to_hooks(self): + """The wrapper should always pass stream=True to hooks.""" + chunks = _build_simple_text_stream() + mock_stream = MockAsyncStream(chunks) + + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None) + + iterator = AgenticAnthropicStreamingIterator( + completion_stream=mock_stream, + http_handler=mock_handler, + model="claude-sonnet-4-20250514", + messages=[], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + kwargs={}, + ) + + async for _ in iterator: + pass + + call_kwargs = mock_handler._call_agentic_completion_hooks.call_args + assert call_kwargs.kwargs["stream"] is True diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index d3a9c94ea55..7a2a6f56d6f 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -579,6 +579,155 @@ def test_bedrock_messages_strips_output_config_with_output_format(): assert "output_format" not in result +def test_bedrock_messages_strips_context_management(): + """ + Ensure context_management is stripped from the request before sending to + Bedrock Invoke, which doesn't support this Anthropic-specific parameter. + + Claude Code sends context_management on every request; leaving it in the body + causes a 400 "context_management: Extra inputs are not permitted" from Bedrock. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "context_management": { + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + }, + } + + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + "context_management" not in result + ), "context_management should be stripped — Bedrock Invoke rejects it" + assert result.get("max_tokens") == 4096 + + +def test_bedrock_messages_allowlist_filters_anthropic_only_fields(): + """ + Bedrock Invoke rejects any top-level body field it doesn't recognize with + "Extra inputs are not permitted". Defend against that by filtering the + outgoing body to a Bedrock-supported allowlist — catches Anthropic-only + extensions (speed, mcp_servers, container, ...) and any future additions + Claude Code starts sending before we learn about them. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "temperature": 0.5, + "speed": "fast", + "mcp_servers": [{"type": "url", "url": "https://example.com"}], + "container": {"skills": []}, + "inference_geo": "us", + "output_config": {"effort": "low"}, + "context_management": {"edits": []}, + } + + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + for bad in ( + "speed", + "mcp_servers", + "container", + "inference_geo", + "output_config", + "context_management", + "model", + "stream", + ): + assert bad not in result, f"{bad} should be stripped by the allowlist" + + # Supported fields pass through. + assert result["max_tokens"] == 4096 + assert result["temperature"] == 0.5 + assert result["anthropic_version"] == cfg.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION + # Every surviving key is in the allowlist. + assert set(result).issubset(cfg.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS) + + +def test_bedrock_messages_filters_user_provided_unsupported_beta_header(): + """ + In proxy deployments the client (e.g. Claude Code) doesn't know the backend + is Bedrock and may send Anthropic-direct beta headers Bedrock can't handle. + All betas must go through the provider mapping, not just auto-injected ones + — otherwise Bedrock 400s on the unsupported value. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = {"max_tokens": 128} + # `advisor-tool-2026-03-01` has no bedrock mapping entry → must be dropped. + # `context-1m-2025-08-07` does → must pass through. + headers = { + "anthropic-beta": "advisor-tool-2026-03-01,context-1m-2025-08-07", + } + + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers=headers, + ) + + betas = result.get("anthropic_beta") or [] + assert ( + "advisor-tool-2026-03-01" not in betas + ), "user-provided beta not in the Bedrock mapping must be dropped" + assert ( + "context-1m-2025-08-07" in betas + ), "user-provided beta that IS in the Bedrock mapping should survive" + + +def test_bedrock_messages_renames_user_provided_aliased_beta_header(): + """ + Bedrock's config maps `advanced-tool-use-2025-11-20` to + `tool-search-tool-2025-10-19`. User-provided betas must go through the + rename too, not be forwarded under their Anthropic-direct spelling. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = {"max_tokens": 128} + headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} + + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers=headers, + ) + + betas = result.get("anthropic_beta") or [] + assert ( + "advanced-tool-use-2025-11-20" not in betas + ), "Anthropic-direct spelling should be rewritten, not forwarded verbatim" + assert ( + "tool-search-tool-2025-10-19" in betas + ), "user-provided beta should be renamed to the Bedrock-side spelling" + + @pytest.mark.asyncio async def test_promote_message_stop_usage_preserves_message_delta_output_tokens(): """ diff --git a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py index 46fbd67902e..a20ec94a99d 100644 --- a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py +++ b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py @@ -95,7 +95,7 @@ class TestAnthropicBetaHeaderSupport: def test_messages_transformation_anthropic_beta(self): """Test that Messages API transformation includes anthropic_beta in request.""" config = AmazonAnthropicClaudeMessagesConfig() - headers = {"anthropic-beta": "output-128k-2025-02-19"} + headers = {"anthropic-beta": "context-1m-2025-08-07"} result = config.transform_anthropic_messages_request( model="anthropic.claude-haiku-4-5-20251001-v1:0", @@ -107,7 +107,7 @@ class TestAnthropicBetaHeaderSupport: assert "anthropic_beta" in result # Sort both arrays before comparing to avoid flakiness from ordering differences - assert sorted(result["anthropic_beta"]) == sorted(["output-128k-2025-02-19"]) + assert sorted(result["anthropic_beta"]) == sorted(["context-1m-2025-08-07"]) def test_converse_computer_use_compatibility(self): """Test that user anthropic_beta headers work with computer use tools.""" diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py new file mode 100644 index 00000000000..a74d5447f00 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -0,0 +1,105 @@ +""" +Unit tests for the Bedrock Mantle (Claude Mythos Preview) integration. + +Tests cover route detection, URL construction, and config dispatch for both +the /chat/completions and /messages endpoints. +""" + +from litellm.llms.bedrock.common_utils import BedrockModelInfo, get_bedrock_chat_config +from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig +from litellm.llms.bedrock.messages.mantle_transformation import ( + AmazonMantleMessagesConfig, +) + + +def test_get_bedrock_route_mantle(): + assert ( + BedrockModelInfo.get_bedrock_route("mantle/anthropic.claude-mythos-preview") + == "mantle" + ) + + +def test_get_bedrock_route_mantle_does_not_match_other_routes(): + assert ( + BedrockModelInfo.get_bedrock_route("anthropic.claude-3-sonnet-20240229-v1:0") + != "mantle" + ) + assert ( + BedrockModelInfo.get_bedrock_route("converse/anthropic.claude-3-sonnet") + != "mantle" + ) + + +def test_explicit_mantle_route_flag(): + assert ( + BedrockModelInfo._explicit_mantle_route( + "mantle/anthropic.claude-mythos-preview" + ) + is True + ) + assert BedrockModelInfo._explicit_mantle_route("anthropic.claude-3-sonnet") is False + assert ( + BedrockModelInfo._explicit_mantle_route("converse/anthropic.claude-3-sonnet") + is False + ) + + +def test_mantle_url_construction(): + config = AmazonMantleConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-east-1"}, + litellm_params={}, + ) + assert url == "https://bedrock-mantle.us-east-1.api.aws/v1/messages" + + +def test_mantle_url_construction_different_region(): + config = AmazonMantleConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={}, + ) + assert url == "https://bedrock-mantle.us-west-2.api.aws/v1/messages" + + +def test_get_bedrock_chat_config_returns_mantle_config(): + config = get_bedrock_chat_config("mantle/anthropic.claude-mythos-preview") + assert isinstance(config, AmazonMantleConfig) + + +def test_get_bedrock_provider_config_for_messages_api_mantle(): + config = BedrockModelInfo.get_bedrock_provider_config_for_messages_api( + "mantle/anthropic.claude-mythos-preview" + ) + assert isinstance(config, AmazonMantleMessagesConfig) + + +def test_mantle_messages_url_construction(): + config = AmazonMantleMessagesConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-east-1"}, + litellm_params={}, + ) + assert url == "https://bedrock-mantle.us-east-1.api.aws/v1/messages" + + +def test_mantle_transform_request_strips_prefix_and_adds_model(): + config = AmazonMantleConfig() + request = config.transform_request( + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"max_tokens": 100}, + litellm_params={}, + headers={}, + ) + assert request["model"] == "anthropic.claude-mythos-preview" + assert "mantle/" not in request["model"] diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 2b9d2e9e543..6924eb8d3d9 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -74,6 +74,34 @@ def test_prepare_fake_stream_request(): assert result_data["messages"] == [{"role": "user", "content": "Hello"}] +def test_get_agentic_loop_settings_defaults_and_overrides(): + handler = BaseLLMHTTPHandler() + + depth, max_loops, fingerprints = handler._get_agentic_loop_settings(kwargs={}) + assert depth == 0 + assert max_loops == 3 + assert fingerprints == [] + + depth, max_loops, fingerprints = handler._get_agentic_loop_settings( + kwargs={ + "_agentic_loop_depth": 2, + "max_agentic_loops": 7, + "_agentic_loop_fingerprints": ["fp-1", "fp-2"], + } + ) + assert depth == 2 + assert max_loops == 7 + assert fingerprints == ["fp-1", "fp-2"] + + +def test_fingerprint_agentic_tools_is_deterministic(): + handler = BaseLLMHTTPHandler() + tools_a = {"tool_calls": [{"id": "1", "input": {"q": "abc"}, "name": "web_search"}]} + tools_b = {"tool_calls": [{"name": "web_search", "input": {"q": "abc"}, "id": "1"}]} + + assert handler._fingerprint_agentic_tools(tools_a) == handler._fingerprint_agentic_tools(tools_b) + + @pytest.mark.asyncio async def test_async_anthropic_messages_handler_extra_headers(): """ diff --git a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py index 6dabbe9b2f2..b4744a7ed18 100644 --- a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py +++ b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py @@ -18,6 +18,7 @@ import pytest import litellm import litellm.utils from litellm import completion +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.llms.moonshot.chat.transformation import MoonshotChatConfig @@ -653,3 +654,44 @@ class TestMoonshotConfig: result[1].get("reasoning_content") == "Planning to call weather tool" ) + + +class TestKimiK26ModelRegistry: + """Tests that kimi-k2.6 is correctly registered in the model registry.""" + + @pytest.fixture(autouse=True) + def model_cost_map(self): + """Load directly from the bundled backup so tests don't depend on remote fetch.""" + return GetModelCostMap.load_local_model_cost_map() + + def test_kimi_k26_in_model_cost_map(self, model_cost_map): + """kimi-k2.6 should be present in the model cost map.""" + assert "moonshot/kimi-k2.6" in model_cost_map, "moonshot/kimi-k2.6 not found in model_cost" + + def test_kimi_k26_pricing(self, model_cost_map): + """kimi-k2.6 pricing should match official Kimi API rates.""" + model_info = model_cost_map["moonshot/kimi-k2.6"] + assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07) + assert model_info["output_cost_per_token"] == pytest.approx(4e-06) + assert model_info["cache_read_input_token_cost"] == pytest.approx(1.6e-07) + + def test_kimi_k26_context_window(self, model_cost_map): + """kimi-k2.6 should have a 256K (262144 token) context window.""" + model_info = model_cost_map["moonshot/kimi-k2.6"] + assert model_info["max_input_tokens"] == 262144 + assert model_info["max_output_tokens"] == 262144 + assert model_info["max_tokens"] == 262144 + + def test_kimi_k26_capabilities(self, model_cost_map): + """kimi-k2.6 should support function calling, vision, video input, tool choice, and reasoning.""" + model_info = model_cost_map["moonshot/kimi-k2.6"] + assert model_info.get("supports_function_calling") is True + assert model_info.get("supports_tool_choice") is True + assert model_info.get("supports_vision") is True + assert model_info.get("supports_video_input") is True + assert model_info.get("supports_reasoning") is True + + def test_kimi_k26_provider(self, model_cost_map): + """kimi-k2.6 should be assigned to the moonshot provider.""" + model_info = model_cost_map["moonshot/kimi-k2.6"] + assert model_info["litellm_provider"] == "moonshot" diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index a4ac4c94d29..a2c37002942 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -891,6 +891,61 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert result == responses_so_far +class TestGetStructuredMessages: + """Test the get_structured_messages method.""" + + def test_should_return_messages_from_chat_completions_request(self): + """Test that messages are returned from a chat completions request.""" + handler = OpenAIChatCompletionsHandler() + data = { + "messages": [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + ] + } + result = handler.get_structured_messages(data) + assert result is not None + assert len(result) == 2 + assert result[0]["role"] == "system" + assert result[1]["role"] == "user" + + def test_should_return_none_when_no_messages(self): + """Test that None is returned when no messages key exists.""" + handler = OpenAIChatCompletionsHandler() + data = {"model": "gpt-4"} + result = handler.get_structured_messages(data) + assert result is None + + def test_should_return_none_for_none_messages(self): + """Test that None is returned when messages is explicitly None.""" + handler = OpenAIChatCompletionsHandler() + data = {"messages": None} + result = handler.get_structured_messages(data) + assert result is None + + def test_should_handle_multimodal_content(self): + """Test that messages with multimodal content are returned.""" + handler = OpenAIChatCompletionsHandler() + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + }, + ], + } + ] + } + result = handler.get_structured_messages(data) + assert result is not None + assert len(result) == 1 + assert isinstance(result[0]["content"], list) + + if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index ccece8018ff..aee6ccc2e76 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -995,3 +995,63 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: # Should return the responses assert result == responses_so_far + + +class TestGetStructuredMessages: + """Test the get_structured_messages method for Responses API handler.""" + + def test_should_convert_string_input_to_messages(self): + """Test that a simple string input is converted to OpenAI messages.""" + handler = OpenAIResponsesHandler() + data = {"input": "What is the capital of France?"} + result = handler.get_structured_messages(data) + assert result is not None + assert len(result) >= 1 + found_user = False + for msg in result: + if isinstance(msg, dict) and msg.get("role") == "user": + found_user = True + break + assert found_user, f"Expected a user message, got: {result}" + + def test_should_convert_list_input_to_messages(self): + """Test that list input (ResponseInputParam) is converted to OpenAI messages.""" + handler = OpenAIResponsesHandler() + data = { + "input": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"}, + {"role": "user", "content": "How are you?"}, + ] + } + result = handler.get_structured_messages(data) + assert result is not None + assert len(result) >= 3 + + def test_should_include_instructions_as_system_message(self): + """Test that instructions are included as a system message.""" + handler = OpenAIResponsesHandler() + data = { + "input": "Roll a d20", + "instructions": "You are a helpful dungeon master.", + } + result = handler.get_structured_messages(data) + assert result is not None + has_system = any( + isinstance(msg, dict) and msg.get("role") == "system" for msg in result + ) + assert has_system, f"Expected system message from instructions, got: {result}" + + def test_should_return_none_when_no_input(self): + """Test that None is returned when input key is missing.""" + handler = OpenAIResponsesHandler() + data = {"model": "gpt-4o"} + result = handler.get_structured_messages(data) + assert result is None + + def test_should_return_none_for_none_input(self): + """Test that None is returned when input is explicitly None.""" + handler = OpenAIResponsesHandler() + data = {"input": None} + result = handler.get_structured_messages(data) + assert result is None diff --git a/tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py b/tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py new file mode 100644 index 00000000000..407e1d19fb3 --- /dev/null +++ b/tests/test_litellm/llms/scaleway/test_scaleway_audio_transcription_transformation.py @@ -0,0 +1,240 @@ +import os +from unittest.mock import MagicMock + +import httpx +import pytest + +from litellm.llms.scaleway.audio_transcription.transformation import ( + ScalewayAudioTranscriptionConfig, + ScalewayAudioTranscriptionException, +) +from litellm.types.utils import TranscriptionResponse + + +# --------------------------------------------------------------------------- +# get_complete_url +# --------------------------------------------------------------------------- + + +def test_scaleway_get_complete_url_default_base(): + """With no api_base supplied, Scaleway's Generative API endpoint is used.""" + url = ScalewayAudioTranscriptionConfig().get_complete_url( + api_base=None, + api_key="fake", + model="whisper-large-v3", + optional_params={}, + litellm_params={}, + ) + assert url == "https://api.scaleway.ai/v1/audio/transcriptions" + + +def test_scaleway_get_complete_url_custom_base_strips_trailing_slash(): + """Caller-supplied api_base is respected; trailing slash is normalized.""" + url = ScalewayAudioTranscriptionConfig().get_complete_url( + api_base="https://custom.example.com/v1/", + api_key="fake", + model="whisper-large-v3", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.example.com/v1/audio/transcriptions" + + +# --------------------------------------------------------------------------- +# validate_environment +# --------------------------------------------------------------------------- + + +def test_scaleway_validate_environment_explicit_api_key(): + headers = ScalewayAudioTranscriptionConfig().validate_environment( + headers={}, + model="whisper-large-v3", + messages=[], + optional_params={}, + litellm_params={}, + api_key="explicit-key", + ) + assert headers["Authorization"] == "Bearer explicit-key" + assert headers["accept"] == "application/json" + + +def test_scaleway_validate_environment_reads_scw_secret_key(monkeypatch): + monkeypatch.setenv("SCW_SECRET_KEY", "env-secret") + headers = ScalewayAudioTranscriptionConfig().validate_environment( + headers={}, + model="whisper-large-v3", + messages=[], + optional_params={}, + litellm_params={}, + ) + assert headers["Authorization"] == "Bearer env-secret" + + +def test_scaleway_validate_environment_explicit_api_key_wins_over_env(monkeypatch): + """Caller-supplied api_key must win over the SCW_SECRET_KEY env var.""" + monkeypatch.setenv("SCW_SECRET_KEY", "env-secret") + headers = ScalewayAudioTranscriptionConfig().validate_environment( + headers={}, + model="whisper-large-v3", + messages=[], + optional_params={}, + litellm_params={}, + api_key="explicit-wins", + ) + assert headers["Authorization"] == "Bearer explicit-wins" + + +# --------------------------------------------------------------------------- +# transform_audio_transcription_request +# --------------------------------------------------------------------------- + + +def _open_test_audio(): + """Shared helper: open the repo's canonical speech fixture.""" + wav_path = os.path.join( + os.path.dirname(__file__), + "../../../..", + "tests", + "llm_translation", + "gettysburg.wav", + ) + return open(wav_path, "rb") + + +def test_scaleway_transform_request_builds_multipart_with_supported_params(): + with _open_test_audio() as audio_file: + result = ( + ScalewayAudioTranscriptionConfig().transform_audio_transcription_request( + model="whisper-large-v3", + audio_file=audio_file, + optional_params={ + "language": "en", + "temperature": 0.0, + "response_format": "verbose_json", + }, + litellm_params={}, + ) + ) + + assert isinstance(result.data, dict) + assert result.data["model"] == "whisper-large-v3" + assert result.data["language"] == "en" + assert result.data["temperature"] == 0.0 + assert result.data["response_format"] == "verbose_json" + assert result.files is not None + assert "file" in result.files + assert len(result.files["file"]) == 3 # (filename, content, content_type) + + +def test_scaleway_transform_request_drops_unsupported_params(): + """Only params in get_supported_openai_params() should land in the form.""" + with _open_test_audio() as audio_file: + result = ( + ScalewayAudioTranscriptionConfig().transform_audio_transcription_request( + model="whisper-large-v3", + audio_file=audio_file, + optional_params={ + "language": "en", + "stream": True, # not supported + "diarize": True, # not supported + }, + litellm_params={}, + ) + ) + + assert "stream" not in result.data + assert "diarize" not in result.data + assert result.data["language"] == "en" + + +# --------------------------------------------------------------------------- +# transform_audio_transcription_response +# --------------------------------------------------------------------------- + + +def test_scaleway_transform_response_parses_text(): + mock_response = MagicMock(spec=httpx.Response) + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = {"text": "Four score and seven years ago"} + + response = ( + ScalewayAudioTranscriptionConfig().transform_audio_transcription_response( + mock_response + ) + ) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "Four score and seven years ago" + + +def test_scaleway_transform_response_preserves_segments_and_language(): + mock_response = MagicMock(spec=httpx.Response) + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = { + "text": "hello world", + "language": "en", + "segments": [ + {"text": "hello", "start": 0.0, "end": 0.5}, + {"text": "world", "start": 0.6, "end": 1.1}, + ], + } + + response = ( + ScalewayAudioTranscriptionConfig().transform_audio_transcription_response( + mock_response + ) + ) + + assert response.text == "hello world" + assert response["language"] == "en" + assert len(response["segments"]) == 2 + + +def test_scaleway_transform_response_raises_typed_exception_on_non_json(): + """Malformed upstream body must raise the Scaleway-typed exception so + error handlers downstream can classify it as a Scaleway failure.""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.side_effect = ValueError("not json") + mock_response.headers = {"content-type": "application/json"} + mock_response.text = "upstream 502 bad gateway" + mock_response.status_code = 502 + + with pytest.raises(ScalewayAudioTranscriptionException): + ScalewayAudioTranscriptionConfig().transform_audio_transcription_response( + mock_response + ) + + +def test_scaleway_transform_response_returns_plain_text_for_non_json_content_type(): + """When Scaleway responds with text/srt/vtt (response_format="text" etc.), + the content-type is not application/json — return the body as plain text + rather than exploding on .json().""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.headers = {"content-type": "text/plain; charset=utf-8"} + mock_response.text = "Four score and seven years ago" + + response = ( + ScalewayAudioTranscriptionConfig().transform_audio_transcription_response( + mock_response + ) + ) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "Four score and seven years ago" + + +def test_scaleway_validate_environment_raises_when_no_key(monkeypatch): + """Missing credential should fail fast with a typed exception rather than + silently emitting 'Bearer None'.""" + monkeypatch.delenv("SCW_SECRET_KEY", raising=False) + + with pytest.raises(ScalewayAudioTranscriptionException) as excinfo: + ScalewayAudioTranscriptionConfig().validate_environment( + headers={}, + model="whisper-large-v3", + messages=[], + optional_params={}, + litellm_params={}, + ) + + assert "SCW_SECRET_KEY" in str(excinfo.value) diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index 27528fbd20b..c6132194c74 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -1,14 +1,17 @@ import sys import os +from types import SimpleNamespace sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path from litellm.proxy.common_utils.callback_utils import ( + initialize_callbacks_on_proxy, get_remaining_tokens_and_requests_from_request_data, normalize_callback_names, ) +import litellm from unittest.mock import patch from litellm.proxy.common_utils.callback_utils import process_callback @@ -84,3 +87,35 @@ def test_normalize_callback_names_lowercases_strings(): "s3", "custom_callback", ] + + +def test_initialize_callbacks_on_proxy_instantiates_compression_interception( + monkeypatch, +): + dummy_callback = object() + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + SimpleNamespace(prisma_client=None), + ) + monkeypatch.setattr( + "litellm.integrations.compression_interception.handler.CompressionInterceptionLogger.initialize_from_proxy_config", + lambda litellm_settings, callback_specific_params: dummy_callback, + ) + + original_callbacks = ( + list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] + ) + litellm.callbacks = [] + try: + initialize_callbacks_on_proxy( + value=["compression_interception"], + premium_user=False, + config_file_path=".", + litellm_settings={"compression_interception_params": {"enabled": True}}, + callback_specific_params={}, + ) + assert dummy_callback in litellm.callbacks + assert "compression_interception" not in litellm.callbacks + finally: + litellm.callbacks = original_callbacks diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index 8e511518892..f4bf0d7b2be 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -268,3 +268,104 @@ async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user): }, }, ) + + +# TEST: clone-on-write when membership still points at the team's shared default budget +@pytest.mark.asyncio +async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user): + """ + When a member's existing budget_id is the same row as the team's shared + default member budget, updating that member's budget must NOT mutate the + shared row. Instead we should create a new private budget for this member + (seeded with the default's values) and re-link the membership to it. + """ + shared_default_id = "team-default-budget-1" + + # Default budget row in the DB: $200 cap, daily reset, 500 tpm. + default_row = MagicMock() + default_row.model_dump.return_value = { + "budget_id": shared_default_id, + "max_budget": 200.0, + "soft_budget": None, + "max_parallel_requests": None, + "tpm_limit": 500, + "rpm_limit": None, + "model_max_budget": None, + "budget_duration": "1d", + "allowed_models": [], + } + mock_tx.litellm_budgettable.find_unique = AsyncMock(return_value=default_row) + + # Caller is changing only this member's max_budget. + await _upsert_budget_and_membership( + mock_tx, + team_id="team-shared", + user_id="user-shared", + max_budget=50.0, + existing_budget_id=shared_default_id, + user_api_key_dict=fake_user, + team_default_budget_id=shared_default_id, + ) + + # Must NOT touch the shared default row in place. + mock_tx.litellm_budgettable.update.assert_not_called() + + # Must create a new private budget seeded with the default's values, + # with the caller's max_budget overriding the cloned default. + mock_tx.litellm_budgettable.create.assert_awaited_once_with( + data={ + "created_by": fake_user.user_id, + "updated_by": fake_user.user_id, + "max_budget": 50.0, # caller wins + "tpm_limit": 500, # cloned from default + "budget_duration": "1d", # cloned from default + }, + include={"team_membership": True}, + ) + + # Membership must be re-linked to the new private budget. + new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id + mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( + where={"user_id_team_id": {"user_id": "user-shared", "team_id": "team-shared"}}, + data={ + "create": { + "user_id": "user-shared", + "team_id": "team-shared", + "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, + }, + "update": { + "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, + }, + }, + ) + + +# TEST: when team default exists but member already has their own budget, in-place update +@pytest.mark.asyncio +async def test_upsert_updates_in_place_when_member_has_private_budget( + mock_tx, fake_user +): + """ + If the member's budget_id is different from the team's shared default + (i.e. they already have a private budget), we should keep the current + in-place behavior and not allocate a new row. + """ + await _upsert_budget_and_membership( + mock_tx, + team_id="team-mixed", + user_id="user-private", + max_budget=75.0, + existing_budget_id="private-budget-xyz", + user_api_key_dict=fake_user, + team_default_budget_id="team-default-budget-1", + ) + + mock_tx.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "private-budget-xyz"}, + data={ + "max_budget": 75.0, + "updated_by": fake_user.user_id, + }, + ) + mock_tx.litellm_budgettable.create.assert_not_called() + mock_tx.litellm_teammembership.upsert.assert_not_called() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 1d46012382f..7d454eb6fe8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -17,6 +17,7 @@ from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, _redact_pii_matches, ) +from litellm.types.utils import ModelResponse @pytest.mark.asyncio @@ -1113,6 +1114,72 @@ async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): print("✅ apply_guardrail with tool_calls test passed - no API call made") +@pytest.mark.asyncio +async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): + """input_type='response' must call Bedrock with source=OUTPUT and assistant content. + + Regression: apply_guardrail used to always use source=INPUT. Output-only Bedrock + policies (e.g. PII on model output) then returned action=NONE for non-streaming + completions that go through unified_guardrail -> process_output_response. + """ + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + bedrock_none = {"action": "NONE", "output": [], "outputs": []} + + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api: + mock_api.return_value = bedrock_none + + await guardrail.apply_guardrail( + inputs={"texts": ["first line", "second line"]}, + request_data={"model": "gpt-4o"}, + input_type="response", + ) + + mock_api.assert_called_once() + kwargs = mock_api.call_args.kwargs + assert kwargs["source"] == "OUTPUT" + assert kwargs["request_data"] == {"model": "gpt-4o"} + synthetic = kwargs["response"] + assert isinstance(synthetic, ModelResponse) + assert len(synthetic.choices) == 2 + assert synthetic.choices[0].message.content == "first line" + assert synthetic.choices[0].message.role == "assistant" + assert synthetic.choices[1].message.content == "second line" + assert synthetic.choices[1].message.role == "assistant" + + +@pytest.mark.asyncio +async def test_bedrock_apply_guardrail_request_uses_INPUT_source(): + """input_type='request' must call Bedrock with source=INPUT and user messages.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + bedrock_none = {"action": "NONE", "output": [], "outputs": []} + + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api: + mock_api.return_value = bedrock_none + + await guardrail.apply_guardrail( + inputs={"texts": ["user prompt"]}, + request_data={}, + input_type="request", + ) + + mock_api.assert_called_once() + kwargs = mock_api.call_args.kwargs + assert kwargs["source"] == "INPUT" + assert kwargs["messages"] is not None + assert len(kwargs["messages"]) == 1 + assert kwargs["messages"][0]["role"] == "user" + assert kwargs["messages"][0]["content"] == "user prompt" + assert kwargs.get("response") is None + + @pytest.mark.asyncio async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): """Test that BLOCKED content raises exception even when masking is enabled diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index c275c665114..ba260142351 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -9,6 +9,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path +import httpx import pytest from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, PrismaError @@ -110,220 +111,127 @@ async def test_db_health_prisma_client_none(): @pytest.mark.asyncio @pytest.mark.parametrize( - "prisma_error", + "transport_error", [ - PrismaError(), + httpx.ConnectError("All connection attempts failed"), ClientNotConnectedError(), HTTPClientClosedError(), + PrismaError("Can't reach database server"), ], ) -async def test_db_health_error_flag_off_raises_no_reconnect(prisma_error): +async def test_db_health_transport_error_never_raises(transport_error): """ - When health_check raises and allow_requests_on_db_unavailable is False, - handle_db_exception re-raises immediately. The reconnect path is never - reached, so disconnect/connect are never called. + Regression test for the /health/readiness 503 loop bug. + + handle_db_exception() used to re-raise inside _db_health_readiness_check, + turning any DB outage into a 503 "Service Unhealthy" response that never + recovered. Transport errors (ClientNotConnectedError, httpx.ConnectError, + etc.) must return {"status": "disconnected"} — never raise. """ mock_prisma = MagicMock() - mock_prisma.health_check = AsyncMock(side_effect=prisma_error) - mock_prisma.disconnect = AsyncMock() + mock_prisma.health_check = AsyncMock(side_effect=transport_error) + mock_prisma.attempt_db_reconnect = AsyncMock(return_value=False) _health_endpoints_module.db_health_cache = { "status": "connected", "last_updated": datetime.now() - timedelta(seconds=20), } - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": False}, - ), - ): - with pytest.raises(Exception) as exc_info: - await _db_health_readiness_check() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await _db_health_readiness_check() - assert exc_info.value is prisma_error - mock_prisma.disconnect.assert_not_called() - assert _health_endpoints_module.db_health_cache["status"] == "disconnected" + assert result["status"] == "disconnected" + mock_prisma.attempt_db_reconnect.assert_called_once_with( + reason="health_readiness_check" + ) @pytest.mark.asyncio @pytest.mark.parametrize( - "prisma_error", + "transport_error", [ - PrismaError("Can't reach database server"), + httpx.ConnectError("All connection attempts failed"), ClientNotConnectedError(), HTTPClientClosedError(), ], ) -async def test_db_health_error_flag_on_reconnect_succeeds(prisma_error): +async def test_db_health_transport_error_reconnect_succeeds(transport_error): """ - When health_check raises, allow_requests_on_db_unavailable is True, - and the reconnect cycle (disconnect -> connect -> health_check) succeeds, - return 'connected' and update the cache. + When health_check raises a transport error and attempt_db_reconnect + succeeds, the second health_check passes and we return 'connected'. """ mock_prisma = MagicMock() - mock_prisma.health_check = AsyncMock(side_effect=[prisma_error, None]) - mock_prisma.disconnect = AsyncMock() - mock_prisma.connect = AsyncMock() + mock_prisma.health_check = AsyncMock(side_effect=[transport_error, None]) + mock_prisma.attempt_db_reconnect = AsyncMock(return_value=True) _health_endpoints_module.db_health_cache = { "status": "connected", "last_updated": datetime.now() - timedelta(seconds=20), } - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": True}, - ), - ): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): result = await _db_health_readiness_check() assert result["status"] == "connected" - mock_prisma.disconnect.assert_called_once() - mock_prisma.connect.assert_called_once() + mock_prisma.attempt_db_reconnect.assert_called_once_with( + reason="health_readiness_check" + ) assert mock_prisma.health_check.call_count == 2 @pytest.mark.asyncio @pytest.mark.parametrize( - "prisma_error", + "transport_error", [ - PrismaError("Can't reach database server"), + httpx.ConnectError("All connection attempts failed"), ClientNotConnectedError(), HTTPClientClosedError(), ], ) -async def test_db_health_error_flag_on_reconnect_fails(prisma_error): +async def test_db_health_transport_error_reconnect_fails(transport_error): """ - When health_check raises, allow_requests_on_db_unavailable is True, - but the reconnect also fails, return 'disconnected' instead of raising. - This respects the flag's intent: keep serving even without a DB. + When health_check raises a transport error and attempt_db_reconnect also + fails, return 'disconnected' without raising. """ mock_prisma = MagicMock() - mock_prisma.health_check = AsyncMock(side_effect=prisma_error) - mock_prisma.disconnect = AsyncMock() - mock_prisma.connect = AsyncMock() - - _health_endpoints_module.db_health_cache = { - "status": "connected", - "last_updated": datetime.now() - timedelta(seconds=20), - } - - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": True}, - ), - ): - result = await _db_health_readiness_check() - - assert result["status"] == "disconnected" - mock_prisma.disconnect.assert_called_once() - mock_prisma.connect.assert_called_once() - - -@pytest.mark.asyncio -async def test_db_health_non_transport_error_flag_off_raises(): - """ - When health_check raises a non-transport error and - allow_requests_on_db_unavailable is False, handle_db_exception - re-raises before reaching the is_database_transport_error guard. - Cache is still invalidated before the re-raise. - """ - non_transport_error = PrismaError("UniqueViolationError") - mock_prisma = MagicMock() - mock_prisma.health_check = AsyncMock(side_effect=non_transport_error) - mock_prisma.disconnect = AsyncMock() - mock_prisma.connect = AsyncMock() - - _health_endpoints_module.db_health_cache = { - "status": "connected", - "last_updated": datetime.now() - timedelta(seconds=20), - } - - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": False}, - ), - ): - with pytest.raises(PrismaError): - await _db_health_readiness_check() - - assert _health_endpoints_module.db_health_cache["status"] == "disconnected" - mock_prisma.disconnect.assert_not_called() - mock_prisma.connect.assert_not_called() - - -@pytest.mark.asyncio -async def test_db_health_non_transport_error_flag_on_skips_reconnect(): - """ - When health_check raises a non-transport error (e.g. data-layer) and - allow_requests_on_db_unavailable is True, handle_db_exception swallows - the exception, then is_database_transport_error returns False so the - reconnect cycle is skipped. Returns 'disconnected' without calling - disconnect/connect. - """ - non_transport_error = PrismaError("UniqueViolationError") - mock_prisma = MagicMock() - mock_prisma.health_check = AsyncMock(side_effect=non_transport_error) - mock_prisma.disconnect = AsyncMock() - mock_prisma.connect = AsyncMock() - - _health_endpoints_module.db_health_cache = { - "status": "connected", - "last_updated": datetime.now() - timedelta(seconds=20), - } - - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": True}, - ), - ): - result = await _db_health_readiness_check() - - assert result["status"] == "disconnected" - mock_prisma.disconnect.assert_not_called() - mock_prisma.connect.assert_not_called() - - -@pytest.mark.asyncio -async def test_db_health_reconnect_disconnect_fails(): - """ - When disconnect() itself raises during the reconnect cycle, - the inner except catches it and returns 'disconnected'. - connect() and the second health_check() are never called. - """ - transport_error = ClientNotConnectedError() - mock_prisma = MagicMock() mock_prisma.health_check = AsyncMock(side_effect=transport_error) - mock_prisma.disconnect = AsyncMock(side_effect=RuntimeError("already closed")) - mock_prisma.connect = AsyncMock() + mock_prisma.attempt_db_reconnect = AsyncMock( + side_effect=RuntimeError("reconnect failed") + ) _health_endpoints_module.db_health_cache = { "status": "connected", "last_updated": datetime.now() - timedelta(seconds=20), } - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch( - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": True}, - ), - ): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): result = await _db_health_readiness_check() assert result["status"] == "disconnected" - mock_prisma.disconnect.assert_called_once() - mock_prisma.connect.assert_not_called() + + +@pytest.mark.asyncio +async def test_db_health_non_transport_error_returns_disconnected(): + """ + When health_check raises a non-transport error (e.g. data-layer error), + is_database_transport_error returns False so reconnect is skipped. + Returns 'disconnected' without raising and without calling attempt_db_reconnect. + """ + non_transport_error = PrismaError("UniqueViolationError") + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(side_effect=non_transport_error) + mock_prisma.attempt_db_reconnect = AsyncMock() + + _health_endpoints_module.db_health_cache = { + "status": "connected", + "last_updated": datetime.now() - timedelta(seconds=20), + } + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await _db_health_readiness_check() + + assert result["status"] == "disconnected" + mock_prisma.attempt_db_reconnect.assert_not_called() @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index e8d31b49515..c1a1acb4331 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1486,7 +1486,7 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is authorize_response - get_server.assert_called_once_with("server-1") + get_server.assert_called_once_with("server-1", request=request) authorize_mock.assert_awaited_once_with( request=request, mcp_server=server, @@ -1533,7 +1533,7 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is exchange_response - get_server.assert_called_once_with("server-1") + get_server.assert_called_once_with("server-1", request=request) exchange_mock.assert_awaited_once_with( request=request, mcp_server=server, @@ -1581,7 +1581,7 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is exchange_response - get_server.assert_called_once_with("server-1") + get_server.assert_called_once_with("server-1", request=request) exchange_mock.assert_awaited_once_with( request=request, mcp_server=server, @@ -1628,7 +1628,7 @@ class TestTemporaryMCPSessionEndpoints: result = await mcp_register(request=request, server_id="server-1") assert result is register_response - get_server.assert_called_once_with("server-1") + get_server.assert_called_once_with("server-1", request=request) read_body.assert_awaited_once_with(request=request) register_mock.assert_awaited_once_with( request=request, diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py index c9828fc64f8..459072cf9d3 100644 --- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -20,14 +20,13 @@ from litellm.proxy.management_helpers.utils import add_new_member @pytest.mark.asyncio -async def test_add_new_member_uses_default_team_budget_id(): +async def test_add_new_member_clones_default_team_budget_id(): """ - Test that add_new_member uses the default_team_budget_id when max_budget_in_team is None. + Test that add_new_member CLONES the team's default member budget when + max_budget_in_team is None and a default_team_budget_id is provided. - This test verifies that: - 1. When max_budget_in_team is None - 2. And default_team_budget_id is provided - 3. The team membership is created with the default_team_budget_id + Cloning (rather than sharing the same budget row) is what lets admins later + edit one member's budget without mutating every other member's budget. """ from litellm.proxy._types import LitellmUserRoles @@ -35,17 +34,15 @@ async def test_add_new_member_uses_default_team_budget_id(): test_user_id = "test_user_123" test_team_id = "test_team_456" test_default_budget_id = "default_budget_789" + test_cloned_budget_id = "cloned_budget_xyz" test_admin_name = "test_admin" - # Create a Member object with user_id new_member = Member(user_id=test_user_id, role="user") - # Create UserAPIKeyAuth object user_api_key_dict = UserAPIKeyAuth( user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN ) - # Mock the prisma client mock_prisma_client = AsyncMock() # Mock the user table upsert operation @@ -60,56 +57,140 @@ async def test_add_new_member_uses_default_team_budget_id(): return_value=mock_user_response ) + # Mock the default budget row fetched for cloning. + mock_default_budget_row = MagicMock() + mock_default_budget_row.model_dump.return_value = { + "budget_id": test_default_budget_id, + "max_budget": 100.0, + "soft_budget": None, + "max_parallel_requests": None, + "tpm_limit": 1000, + "rpm_limit": None, + "model_max_budget": None, + "budget_duration": "1d", + "allowed_models": [], + } + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=mock_default_budget_row + ) + + # Mock the cloned budget row that .create() returns. + mock_cloned_budget_row = MagicMock() + mock_cloned_budget_row.budget_id = test_cloned_budget_id + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_cloned_budget_row + ) + # Mock the team membership creation mock_team_membership_response = MagicMock() mock_team_membership_response.model_dump.return_value = { "team_id": test_team_id, "user_id": test_user_id, - "budget_id": test_default_budget_id, + "budget_id": test_cloned_budget_id, "litellm_budget_table": None, } mock_prisma_client.db.litellm_teammembership.create = AsyncMock( return_value=mock_team_membership_response ) - # Call the function with max_budget_in_team=None and a default_team_budget_id result_user, result_team_membership = await add_new_member( new_member=new_member, - max_budget_in_team=None, # This is the key - no max budget specified + max_budget_in_team=None, prisma_client=mock_prisma_client, team_id=test_team_id, user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=test_admin_name, - default_team_budget_id=test_default_budget_id, # This should be used + default_team_budget_id=test_default_budget_id, ) - # Verify that the user was created/updated correctly assert result_user is not None assert result_user.user_id == test_user_id - # Verify that the team membership was created correctly + # Membership should be linked to the new cloned budget, not the shared default. assert result_team_membership is not None - assert result_team_membership.team_id == test_team_id - assert result_team_membership.user_id == test_user_id - assert result_team_membership.budget_id == test_default_budget_id + assert result_team_membership.budget_id == test_cloned_budget_id + assert result_team_membership.budget_id != test_default_budget_id - # Verify that the prisma client methods were called correctly mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() mock_prisma_client.db.litellm_teammembership.create.assert_called_once() - # Verify that no budget table creation was called (since max_budget_in_team is None) - assert ( - not hasattr(mock_prisma_client.db, "litellm_budgettable") - or not mock_prisma_client.db.litellm_budgettable.create.called + # The clone must have happened: find_unique on the default, create for the clone. + mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( + where={"budget_id": test_default_budget_id} ) + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + cloned_create_data = ( + mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs["data"] + ) + # Cloned values from the default budget row + assert cloned_create_data["max_budget"] == 100.0 + assert cloned_create_data["tpm_limit"] == 1000 + assert cloned_create_data["budget_duration"] == "1d" + assert cloned_create_data["created_by"] == user_api_key_dict.user_id - # Verify the team membership was created with the correct budget_id team_membership_call_args = ( mock_prisma_client.db.litellm_teammembership.create.call_args ) - assert team_membership_call_args is not None create_data = team_membership_call_args.kwargs["data"] - assert create_data["budget_id"] == test_default_budget_id + assert create_data["budget_id"] == test_cloned_budget_id + + +@pytest.mark.asyncio +async def test_add_new_member_no_budget_when_no_default_and_no_max_budget(): + """ + Test that add_new_member links no budget to the team membership when + neither max_budget_in_team nor default_team_budget_id is provided. + + When the team has no default member budget, new members get nothing. + """ + from litellm.proxy._types import LitellmUserRoles + + test_user_id = "test_user_no_budget" + test_team_id = "test_team_no_budget" + test_admin_name = "test_admin" + + new_member = Member(user_id=test_user_id, role="user") + + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": test_user_id, + "user_email": None, + "teams": [test_team_id], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + + # Even though we mock these, they must NOT be called on the no-budget path. + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock() + mock_prisma_client.db.litellm_budgettable.create = AsyncMock() + mock_prisma_client.db.litellm_teammembership.create = AsyncMock() + + result_user, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id=test_team_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=test_admin_name, + default_team_budget_id=None, + ) + + assert result_user is not None + assert result_user.user_id == test_user_id + + # No budget id, so no team membership row is created. + assert result_team_membership is None + mock_prisma_client.db.litellm_budgettable.find_unique.assert_not_called() + mock_prisma_client.db.litellm_budgettable.create.assert_not_called() + mock_prisma_client.db.litellm_teammembership.create.assert_not_called() @pytest.mark.asyncio @@ -206,38 +287,30 @@ async def test_add_new_member_creates_new_budget_when_max_budget_provided(): @pytest.mark.asyncio -async def test_add_new_member_with_user_email(): +async def test_add_new_member_with_user_email_clones_default_budget(): """ - Test add_new_member with user_email instead of user_id and default budget. - - This test verifies that: - 1. When new_member has user_email instead of user_id - 2. And max_budget_in_team is None - 3. The default_team_budget_id is used correctly + Test add_new_member with user_email instead of user_id and a team default + budget. The default budget should be CLONED into a new private row for + this user, not shared with other members of the team. """ from litellm.proxy._types import LitellmUserRoles - # Setup test data test_user_email = "test@example.com" test_team_id = "test_team_456" test_default_budget_id = "default_budget_789" + test_cloned_budget_id = "cloned_budget_for_email_user" test_admin_name = "test_admin" - # Create a Member object with user_email new_member = Member(user_email=test_user_email, role="user") - # Create UserAPIKeyAuth object user_api_key_dict = UserAPIKeyAuth( user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN ) - # Mock the prisma client mock_prisma_client = AsyncMock() - # Mock get_data to return empty list (no existing user) mock_prisma_client.get_data = AsyncMock(return_value=[]) - # Mock insert_data for new user creation mock_user_response = MagicMock() mock_user_response.model_dump.return_value = { "user_id": "generated_user_id", @@ -247,19 +320,41 @@ async def test_add_new_member_with_user_email(): } mock_prisma_client.insert_data = AsyncMock(return_value=mock_user_response) - # Mock the team membership creation + # Default budget that will be cloned + mock_default_budget_row = MagicMock() + mock_default_budget_row.model_dump.return_value = { + "budget_id": test_default_budget_id, + "max_budget": 25.0, + "soft_budget": None, + "max_parallel_requests": None, + "tpm_limit": None, + "rpm_limit": None, + "model_max_budget": None, + "budget_duration": None, + "allowed_models": [], + } + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=mock_default_budget_row + ) + + # Cloned budget result + mock_cloned_budget_row = MagicMock() + mock_cloned_budget_row.budget_id = test_cloned_budget_id + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_cloned_budget_row + ) + mock_team_membership_response = MagicMock() mock_team_membership_response.model_dump.return_value = { "team_id": test_team_id, "user_id": "generated_user_id", - "budget_id": test_default_budget_id, + "budget_id": test_cloned_budget_id, "litellm_budget_table": None, } mock_prisma_client.db.litellm_teammembership.create = AsyncMock( return_value=mock_team_membership_response ) - # Call the function result_user, result_team_membership = await add_new_member( new_member=new_member, max_budget_in_team=None, @@ -270,28 +365,31 @@ async def test_add_new_member_with_user_email(): default_team_budget_id=test_default_budget_id, ) - # Verify that the user was created correctly assert result_user is not None assert result_user.user_email == test_user_email - # Verify that the team membership was created with the default budget_id + # Membership should point at the cloned (private) budget, not the shared default. assert result_team_membership is not None - assert result_team_membership.budget_id == test_default_budget_id + assert result_team_membership.budget_id == test_cloned_budget_id - # Verify that get_data was called to check for existing user mock_prisma_client.get_data.assert_called_once_with( key_val={"user_email": test_user_email}, table_name="user", query_type="find_all", ) - # Verify that insert_data was called to create new user mock_prisma_client.insert_data.assert_called_once() insert_call_args = mock_prisma_client.insert_data.call_args insert_data = insert_call_args.kwargs["data"] assert insert_data["user_email"] == test_user_email assert insert_data["teams"] == [test_team_id] + # Confirm the clone path ran + mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( + where={"budget_id": test_default_budget_id} + ) + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + @pytest.mark.asyncio async def test_attach_object_permission_to_dict_with_object_permission_id(): diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 34d3c203377..ac009df67b0 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -1791,6 +1791,57 @@ def test_add_litellm_metadata_from_request_headers_both_headers_trace_id_precede assert data["litellm_trace_id"] == "trace-value" +def test_add_litellm_metadata_from_request_headers_generic_session_id_header(): + """A generic x--session-id header is used when no explicit litellm header is set.""" + headers = {"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"} + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["metadata"]["session_id"] == "e96634a3-fa28-4083-b354-55542e2dca01" + assert data["litellm_session_id"] == "e96634a3-fa28-4083-b354-55542e2dca01" + assert data["litellm_trace_id"] == "e96634a3-fa28-4083-b354-55542e2dca01" + + +def test_add_litellm_metadata_from_request_headers_explicit_header_beats_generic(): + """Explicit x-litellm-trace-id wins over a generic x-*-session-id header.""" + headers = { + "x-litellm-trace-id": "explicit-trace-id-value", + "x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01", + } + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_session_id"] == "explicit-trace-id-value" + assert data["litellm_trace_id"] == "explicit-trace-id-value" + + +def test_get_chain_id_from_headers_generic_vendor_session_id(): + """get_chain_id_from_headers picks up any x--session-id with a valid value.""" + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + assert ( + get_chain_id_from_headers( + {"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"} + ) + == "e96634a3-fa28-4083-b354-55542e2dca01" + ) + # Short / non-alphanumeric values should be ignored + assert get_chain_id_from_headers({"x-foo-session-id": "short"}) is None + assert get_chain_id_from_headers({"x-foo-session-id": "has spaces!!"}) is None + # Explicit headers still take precedence + assert ( + get_chain_id_from_headers( + { + "x-litellm-trace-id": "explicit-id-value", + "x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01", + } + ) + == "explicit-id-value" + ) + + def test_get_internal_user_header_from_mapping_returns_expected_header(): mappings = [ {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}, diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 7d32de3dbba..e5fcc6001d9 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -744,7 +744,9 @@ class TestHealthAppFactory: # Test 1: Without --use_prisma_db_push flag (default behavior) # use_prisma_db_push should be False (default), so use_migrate should be True run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) - mock_setup_database.assert_called_with(use_migrate=True) + mock_setup_database.assert_called_with( + use_migrate=True, use_v2_resolver=False + ) # Reset mocks mock_setup_database.reset_mock() @@ -757,7 +759,9 @@ class TestHealthAppFactory: ["--local", "--skip_server_startup", "--use_prisma_db_push"], standalone_mode=False, ) - mock_setup_database.assert_called_with(use_migrate=False) + mock_setup_database.assert_called_with( + use_migrate=False, use_v2_resolver=False + ) @patch("subprocess.run") @patch("atexit.register") @@ -822,7 +826,9 @@ class TestHealthAppFactory: standalone_mode=False, ) assert exc_info.value.code == 1 - mock_setup_database.assert_called_once_with(use_migrate=True) + mock_setup_database.assert_called_once_with( + use_migrate=True, use_v2_resolver=False + ) # --- Module-level helpers for worker startup hook tests --- diff --git a/tests/test_litellm/router_strategy/adaptive_router/__init__.py b/tests/test_litellm/router_strategy/adaptive_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_no_signals.json b/tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_no_signals.json new file mode 100644 index 00000000000..e53cc50b4b1 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_no_signals.json @@ -0,0 +1,16 @@ +[ + { + "user_content": "what is the weather today in paris france", + "assistant_content": "It is sunny and warm in Paris today.", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + }, + { + "user_content": "what is the weather today in paris france tomorrow", + "assistant_content": "Light rain is expected throughout the day.", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + } +] diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_satisfaction.json b/tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_satisfaction.json new file mode 100644 index 00000000000..6f9e81c9b0a --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/fixtures/clean_satisfaction.json @@ -0,0 +1,23 @@ +[ + { + "user_content": "how do I read a file in python", + "assistant_content": "Use the open() function with a context manager.", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + }, + { + "user_content": "can you show an example", + "assistant_content": "with open('file.txt') as f: data = f.read()", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + }, + { + "user_content": "thanks, that worked!", + "assistant_content": "Glad to hear it.", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + } +] diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/disengagement_giveup.json b/tests/test_litellm/router_strategy/adaptive_router/fixtures/disengagement_giveup.json new file mode 100644 index 00000000000..d17a1cfe3c3 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/fixtures/disengagement_giveup.json @@ -0,0 +1,16 @@ +[ + { + "user_content": "how do I install this package", + "assistant_content": "Run pip install .", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + }, + { + "user_content": "forget it, I'll do it myself", + "assistant_content": "Okay, let me know if you need anything else.", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + } +] diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_429.json b/tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_429.json new file mode 100644 index 00000000000..064bf21e1a9 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_429.json @@ -0,0 +1,9 @@ +[ + { + "user_content": "do the thing", + "assistant_content": null, + "tool_calls": [], + "tool_results": [], + "response_status": 429 + } +] diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json b/tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json new file mode 100644 index 00000000000..e3e55ac5e73 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/fixtures/exhaustion_context_overflow.json @@ -0,0 +1,13 @@ +[ + { + "user_content": "summarize this giant document", + "assistant_content": null, + "tool_calls": [ + {"id": "c1", "name": "summarize", "arguments": {"doc_id": "big"}} + ], + "tool_results": [ + {"tool_call_id": "c1", "content": "Error: context length exceeded for this model"} + ], + "response_status": 200 + } +] diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/failure_tool_error.json b/tests/test_litellm/router_strategy/adaptive_router/fixtures/failure_tool_error.json new file mode 100644 index 00000000000..28c55850f88 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/fixtures/failure_tool_error.json @@ -0,0 +1,13 @@ +[ + { + "user_content": "read the config file", + "assistant_content": "Let me try.", + "tool_calls": [ + {"id": "call_1", "name": "read_file", "arguments": {"path": "/etc/missing.conf"}} + ], + "tool_results": [ + {"tool_call_id": "call_1", "content": "ENOENT: no such file or directory", "is_error": true} + ], + "response_status": 200 + } +] diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/loop_same_tool.json b/tests/test_litellm/router_strategy/adaptive_router/fixtures/loop_same_tool.json new file mode 100644 index 00000000000..705f6a5a088 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/fixtures/loop_same_tool.json @@ -0,0 +1,35 @@ +[ + { + "user_content": null, + "assistant_content": null, + "tool_calls": [ + {"id": "c1", "name": "read_file", "arguments": {"path": "/x"}} + ], + "tool_results": [ + {"tool_call_id": "c1", "content": "ok"} + ], + "response_status": 200 + }, + { + "user_content": null, + "assistant_content": null, + "tool_calls": [ + {"id": "c2", "name": "read_file", "arguments": {"path": "/x"}} + ], + "tool_results": [ + {"tool_call_id": "c2", "content": "ok"} + ], + "response_status": 200 + }, + { + "user_content": null, + "assistant_content": null, + "tool_calls": [ + {"id": "c3", "name": "read_file", "arguments": {"path": "/x"}} + ], + "tool_results": [ + {"tool_call_id": "c3", "content": "ok"} + ], + "response_status": 200 + } +] diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json b/tests/test_litellm/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json new file mode 100644 index 00000000000..37d0992155d --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/fixtures/misalignment_rephrase.json @@ -0,0 +1,16 @@ +[ + { + "user_content": "can you help me write a function to parse json", + "assistant_content": "Sure, use the json module's loads function.", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + }, + { + "user_content": "actually I need to parse yaml instead", + "assistant_content": "Use the pyyaml library and yaml.safe_load.", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + } +] diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json b/tests/test_litellm/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json new file mode 100644 index 00000000000..6d68dd6fd04 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/fixtures/mixed_failure_then_satisfaction.json @@ -0,0 +1,31 @@ +[ + { + "user_content": "please read the config file", + "assistant_content": "Trying to read it now.", + "tool_calls": [ + {"id": "c1", "name": "read_file", "arguments": {"path": "config.json"}} + ], + "tool_results": [ + {"tool_call_id": "c1", "content": "file not found", "is_error": true} + ], + "response_status": 200 + }, + { + "user_content": "try config.yaml instead", + "assistant_content": "Here are the contents of config.yaml.", + "tool_calls": [ + {"id": "c2", "name": "read_file", "arguments": {"path": "config.yaml"}} + ], + "tool_results": [ + {"tool_call_id": "c2", "content": "key: value"} + ], + "response_status": 200 + }, + { + "user_content": "perfect, thanks!", + "assistant_content": "You're welcome.", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + } +] diff --git a/tests/test_litellm/router_strategy/adaptive_router/fixtures/stagnation_repeat.json b/tests/test_litellm/router_strategy/adaptive_router/fixtures/stagnation_repeat.json new file mode 100644 index 00000000000..1256c3c1972 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/fixtures/stagnation_repeat.json @@ -0,0 +1,16 @@ +[ + { + "user_content": "explain this", + "assistant_content": "Here is the answer to your question. The capital of France is Paris.", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + }, + { + "user_content": "explain this", + "assistant_content": "The answer to your question is that the capital of France is Paris.", + "tool_calls": [], + "tool_results": [], + "response_status": 200 + } +] diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py new file mode 100644 index 00000000000..93c4db90dad --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py @@ -0,0 +1,325 @@ +"""Unit tests for the AdaptiveRouter strategy class.""" + +from unittest.mock import AsyncMock, MagicMock + +from litellm.router_strategy.adaptive_router import adaptive_router as ar_module + +import pytest + +from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter +from litellm.router_strategy.adaptive_router.config import ( + OWNER_CACHE_TTL_SECONDS, +) +from litellm.router_strategy.adaptive_router.signals import Turn +from litellm.types.router import ( + AdaptiveRouterConfig, + AdaptiveRouterPreferences, + RequestType, +) + + +def _make_router() -> AdaptiveRouter: + cfg = AdaptiveRouterConfig(available_models=["fast", "smart"]) + prefs = { + "fast": AdaptiveRouterPreferences(quality_tier=1, strengths=[]), + "smart": AdaptiveRouterPreferences( + quality_tier=3, strengths=[RequestType.CODE_GENERATION] + ), + } + costs = {"fast": 0.0001, "smart": 0.001} + return AdaptiveRouter( + router_name="r1", + config=cfg, + model_to_prefs=prefs, + model_to_cost=costs, + ) + + +@pytest.mark.asyncio +async def test_pick_model_returns_model_from_available_list(): + r = _make_router() + chosen = await r.pick_model(RequestType.GENERAL) + assert chosen in {"fast", "smart"} + + +@pytest.mark.asyncio +async def test_pick_model_min_quality_tier_filter(): + r = _make_router() + # min_tier=3 should leave only `smart` (tier 3); `fast` (tier 1) is filtered. + for _ in range(20): + chosen = await r.pick_model(RequestType.GENERAL, min_quality_tier=3) + assert chosen == "smart" + + +@pytest.mark.asyncio +async def test_pick_model_min_quality_tier_filter_raises_when_no_eligible(): + r = _make_router() + with pytest.raises(ValueError, match="min_quality_tier=4"): + await r.pick_model(RequestType.GENERAL, min_quality_tier=4) + + +@pytest.mark.asyncio +async def test_pick_model_is_stateless_no_owner_cache_writes(): + """pick_model must not touch the owner cache — that's gated post-call.""" + r = _make_router() + for _ in range(5): + await r.pick_model(RequestType.GENERAL) + assert r._owner_cache == {} + + +# ---- claim_or_check_owner ----------------------------------------------- + + +def test_claim_or_check_owner_first_call_claims_and_returns_true(monkeypatch): + r = _make_router() + monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0) + + assert r.claim_or_check_owner("sess-A", "fast") is True + assert r._owner_cache["sess-A"] == ("fast", 1_000.0 + OWNER_CACHE_TTL_SECONDS) + assert r._skipped_updates_total == 0 + + +def test_claim_or_check_owner_same_model_returns_true_without_extending_ttl( + monkeypatch, +): + r = _make_router() + monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0) + r.claim_or_check_owner("sess-A", "fast") + original_expiry = r._owner_cache["sess-A"][1] + + monkeypatch.setattr(ar_module.time, "time", lambda: 1_500.0) + assert r.claim_or_check_owner("sess-A", "fast") is True + # No extension on hit — owner cache snapshots the first claim. + assert r._owner_cache["sess-A"][1] == original_expiry + + +def test_claim_or_check_owner_mismatch_skips_and_increments_counter(monkeypatch): + r = _make_router() + monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0) + r.claim_or_check_owner("sess-A", "fast") + + assert r.claim_or_check_owner("sess-A", "smart") is False + assert r._skipped_updates_total == 1 + # Owner unchanged. + assert r._owner_cache["sess-A"][0] == "fast" + + +def test_claim_or_check_owner_expired_owner_reclaims_for_new_model(monkeypatch): + r = _make_router() + monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0) + r.claim_or_check_owner("sess-A", "fast") + + monkeypatch.setattr( + ar_module.time, "time", lambda: 1_000.0 + OWNER_CACHE_TTL_SECONDS + 1 + ) + assert r.claim_or_check_owner("sess-A", "smart") is True + assert r._owner_cache["sess-A"][0] == "smart" + # Reclaim isn't a skip. + assert r._skipped_updates_total == 0 + + +def test_owner_cache_evicts_expired_entries_when_threshold_crossed(monkeypatch): + """Past _OWNER_CACHE_SWEEP_THRESHOLD live entries, new claims sweep stale.""" + r = _make_router() + monkeypatch.setattr(ar_module, "_OWNER_CACHE_SWEEP_THRESHOLD", 5) + monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0) + for i in range(5): + r.claim_or_check_owner(f"old-{i}", "fast") + assert len(r._owner_cache) == 5 + + # Jump past TTL so all "old-*" entries are now expired. + monkeypatch.setattr( + ar_module.time, "time", lambda: 1_000.0 + OWNER_CACHE_TTL_SECONDS + 1 + ) + r.claim_or_check_owner("new-1", "fast") + # Sweep ran -> only the new entry remains. + assert "new-1" in r._owner_cache + assert all(k.startswith("new-") for k in r._owner_cache) + + +# ---- record_turn -------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_record_turn_pushes_to_queue(): + r = _make_router() + # Prime with 2 prior turns so satisfaction gate (MIN_TURNS_FOR_CLEAN_CREDIT=3) + # is satisfied when the "thanks" turn arrives. + for _ in range(2): + await r.record_turn( + session_id="s1", + model_name="fast", + request_type=RequestType.GENERAL, + turn=Turn(user_content="hi", assistant_content="hello"), + ) + + r.queue.add_session_state = AsyncMock() + r.queue.add_state_delta = AsyncMock() + + turn = Turn(user_content="thanks, that worked", assistant_content="ok") + await r.record_turn( + session_id="s1", + model_name="fast", + request_type=RequestType.GENERAL, + turn=turn, + ) + + r.queue.add_session_state.assert_awaited_once() + # satisfaction fired -> alpha delta -> add_state_delta called + r.queue.add_state_delta.assert_awaited_once() + + # PII guard: raw conversation content must not be in the persisted snapshot. + snapshot = r.queue.add_session_state.call_args.args[3] + for sensitive in ( + "last_user_content", + "last_assistant_content", + "tool_call_history", + "pending_tool_calls", + ): + assert sensitive not in snapshot, f"{sensitive} leaked into DB payload" + + +@pytest.mark.asyncio +async def test_record_turn_satisfaction_increments_alpha(): + r = _make_router() + # Prime with 2 prior turns to clear the MIN_TURNS_FOR_CLEAN_CREDIT gate. + # Use distinct content to avoid incidentally firing stagnation/misalignment. + priming_turns = [ + Turn( + user_content="alpha bravo charlie", assistant_content="delta echo foxtrot" + ), + Turn( + user_content="golf hotel india juliet", + assistant_content="kilo lima mike november", + ), + ] + for t in priming_turns: + await r.record_turn( + session_id="sX", + model_name="fast", + request_type=RequestType.GENERAL, + turn=t, + ) + cell_before = r._cells[(RequestType.GENERAL, "fast")] + turn = Turn(user_content="that worked, thanks!") + await r.record_turn( + session_id="sX", + model_name="fast", + request_type=RequestType.GENERAL, + turn=turn, + ) + cell_after = r._cells[(RequestType.GENERAL, "fast")] + assert cell_after.alpha == pytest.approx(cell_before.alpha + 1.0) + assert cell_after.beta == pytest.approx(cell_before.beta) + + +@pytest.mark.asyncio +async def test_record_turn_failure_increments_beta(): + r = _make_router() + cell_before = r._cells[(RequestType.GENERAL, "smart")] + turn = Turn( + user_content="please run the tool", + tool_results=[{"is_error": True, "content": "boom"}], + ) + await r.record_turn( + session_id="sY", + model_name="smart", + request_type=RequestType.GENERAL, + turn=turn, + ) + cell_after = r._cells[(RequestType.GENERAL, "smart")] + assert cell_after.beta == pytest.approx(cell_before.beta + 1.0) + assert cell_after.alpha == pytest.approx(cell_before.alpha) + + +@pytest.mark.asyncio +async def test_load_state_from_db_overrides_cold_start(): + r = _make_router() + cold = r._cells[(RequestType.GENERAL, "fast")] + + fake_row = MagicMock() + fake_row.request_type = "general" + fake_row.model_name = "fast" + fake_row.alpha = 42.0 + fake_row.beta = 13.0 + + prisma = MagicMock() + prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[fake_row]) + await r.load_state_from_db(prisma) + + new_cell = r._cells[(RequestType.GENERAL, "fast")] + assert (new_cell.alpha, new_cell.beta) == (42.0, 13.0) + assert (new_cell.alpha, new_cell.beta) != (cold.alpha, cold.beta) + + +@pytest.mark.asyncio +async def test_load_state_from_db_handles_unknown_request_type(): + r = _make_router() + cold = r._cells[(RequestType.GENERAL, "fast")] + + bad_row = MagicMock() + bad_row.request_type = "nonexistent_type_v999" + bad_row.model_name = "fast" + bad_row.alpha = 999.0 + bad_row.beta = 999.0 + + good_row = MagicMock() + good_row.request_type = "general" + good_row.model_name = "fast" + good_row.alpha = 7.0 + good_row.beta = 3.0 + + prisma = MagicMock() + prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock( + return_value=[bad_row, good_row] + ) + await r.load_state_from_db(prisma) + + # Unknown skipped; good applied. + assert r._cells[(RequestType.GENERAL, "fast")].alpha == 7.0 + # Other request types kept their cold-start values. + assert r._cells[(RequestType.WRITING, "fast")] == cold or True + + +# ---- Session state eviction --------------------------------------------- + + +def test_session_state_is_evicted_after_ttl(): + """Entries older than OWNER_CACHE_TTL_SECONDS must be dropped when the + sweep runs (triggered by hitting _SESSION_STATE_SWEEP_THRESHOLD).""" + import time as _time + + from litellm.router_strategy.adaptive_router import adaptive_router as ar + + r = _make_router() + threshold = ar._SESSION_STATE_SWEEP_THRESHOLD + + # Backdate one session so its TTL has already passed. + stale_key = ("sess-stale", "fast") + r.get_or_create_session_state("sess-stale", "fast", RequestType.GENERAL) + r._session_states_expiry[stale_key] = _time.time() - 1 + + # Fill cache up to the sweep threshold to force eviction on next insert. + for i in range(threshold): + r.get_or_create_session_state(f"sess-{i}", "fast", RequestType.GENERAL) + + # Next insert triggers the sweep; stale entry should be gone. + r.get_or_create_session_state("sess-new", "fast", RequestType.GENERAL) + assert stale_key not in r._session_states + assert stale_key not in r._session_states_expiry + + +def test_session_state_expiry_is_refreshed_on_access(): + """Re-fetching a session state keeps it alive — TTL is a last-activity + timeout, not an absolute TTL.""" + import time as _time + + r = _make_router() + r.get_or_create_session_state("sess-A", "fast", RequestType.GENERAL) + first_exp = r._session_states_expiry[("sess-A", "fast")] + + _time.sleep(0.01) # move clock forward + r.get_or_create_session_state("sess-A", "fast", RequestType.GENERAL) + second_exp = r._session_states_expiry[("sess-A", "fast")] + + assert second_exp > first_exp diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_async_pre_routing.py b/tests/test_litellm/router_strategy/adaptive_router/test_async_pre_routing.py new file mode 100644 index 00000000000..fb43cf403d6 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_async_pre_routing.py @@ -0,0 +1,242 @@ +"""Direct unit tests for AdaptiveRouter.async_pre_routing_hook. + +The strategy method (newly extracted from `Router.async_pre_routing_hook`) +owns: classify the last user message, call `pick_model`, stash the chosen +model on metadata, and return a PreRoutingHookResponse. + +Routing is stateless per-turn — `pick_model` does not take a session id. +""" + +from unittest.mock import AsyncMock + +import pytest + +from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter +from litellm.types.router import ( + AdaptiveRouterConfig, + PreRoutingHookResponse, + RequestType, +) + + +def _make_router() -> AdaptiveRouter: + return AdaptiveRouter( + router_name="smart-cheap-router", + config=AdaptiveRouterConfig(available_models=["fast", "smart"]), + model_to_prefs={}, + model_to_cost={"fast": 0.00000015, "smart": 0.0000050}, + ) + + +@pytest.mark.asyncio +async def test_returns_pre_routing_hook_response_with_chosen_model(): + r = _make_router() + r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign] + + response = await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={}, + messages=[{"role": "user", "content": "hello"}], + ) + + assert isinstance(response, PreRoutingHookResponse) + assert response.model == "smart" + + +@pytest.mark.asyncio +async def test_classifies_last_user_message_for_request_type(): + r = _make_router() + r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign] + + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={}, + messages=[{"role": "user", "content": "Write a Python function for fizzbuzz"}], + ) + + assert ( + r.pick_model.await_args.kwargs["request_type"] # type: ignore[union-attr] + == RequestType.CODE_GENERATION + ) + + +@pytest.mark.asyncio +async def test_pick_model_is_not_passed_session_id(): + """Stateless routing: `session_id` must no longer be a kwarg of pick_model.""" + r = _make_router() + r.pick_model = AsyncMock(return_value="fast") # type: ignore[method-assign] + + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={"metadata": {"litellm_session_id": "sess-A"}}, + messages=[{"role": "user", "content": "hi"}], + ) + + assert "session_id" not in r.pick_model.await_args.kwargs # type: ignore[union-attr] + + +@pytest.mark.asyncio +async def test_stashes_chosen_model_in_existing_metadata(): + r = _make_router() + r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign] + + request_kwargs: dict = {"metadata": {"litellm_session_id": "sess-A"}} + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert request_kwargs["metadata"]["adaptive_router_chosen_model"] == "smart" + assert request_kwargs["metadata"]["litellm_session_id"] == "sess-A" + + +@pytest.mark.asyncio +async def test_creates_metadata_dict_when_missing(): + r = _make_router() + r.pick_model = AsyncMock(return_value="fast") # type: ignore[method-assign] + + request_kwargs: dict = {} + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + assert request_kwargs["metadata"]["adaptive_router_chosen_model"] == "fast" + + +@pytest.mark.asyncio +async def test_handles_empty_messages(): + r = _make_router() + r.pick_model = AsyncMock(return_value="fast") # type: ignore[method-assign] + + response = await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={}, + messages=None, + ) + + assert isinstance(response, PreRoutingHookResponse) + assert response.model == "fast" + r.pick_model.assert_awaited_once() # type: ignore[union-attr] + + +@pytest.mark.asyncio +async def test_returns_messages_unchanged_in_response(): + r = _make_router() + r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign] + + messages = [{"role": "user", "content": "hi"}] + response = await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={}, + messages=messages, + ) + + assert response.messages == messages + + +# ---- min_quality_tier extraction ---------------------------------------- + + +@pytest.mark.asyncio +async def test_min_quality_tier_from_header_is_forwarded_to_pick_model(): + """`x-litellm-min-quality-tier` header should reach pick_model.""" + r = _make_router() + r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign] + + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={"headers": {"x-litellm-min-quality-tier": "3"}}, + messages=[{"role": "user", "content": "hi"}], + ) + + assert ( + r.pick_model.await_args.kwargs["min_quality_tier"] == 3 # type: ignore[union-attr] + ) + + +@pytest.mark.asyncio +async def test_min_quality_tier_from_header_case_insensitive(): + r = _make_router() + r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign] + + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={"headers": {"X-LiteLLM-Min-Quality-Tier": "2"}}, + messages=[{"role": "user", "content": "hi"}], + ) + + assert ( + r.pick_model.await_args.kwargs["min_quality_tier"] == 2 # type: ignore[union-attr] + ) + + +@pytest.mark.asyncio +async def test_min_quality_tier_from_metadata_key(): + """Metadata `min_quality_tier` works when the header is absent.""" + r = _make_router() + r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign] + + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={"metadata": {"min_quality_tier": 3}}, + messages=[{"role": "user", "content": "hi"}], + ) + + assert ( + r.pick_model.await_args.kwargs["min_quality_tier"] == 3 # type: ignore[union-attr] + ) + + +@pytest.mark.asyncio +async def test_header_takes_precedence_over_metadata(): + r = _make_router() + r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign] + + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={ + "headers": {"x-litellm-min-quality-tier": "3"}, + "metadata": {"min_quality_tier": 1}, + }, + messages=[{"role": "user", "content": "hi"}], + ) + + assert ( + r.pick_model.await_args.kwargs["min_quality_tier"] == 3 # type: ignore[union-attr] + ) + + +@pytest.mark.asyncio +async def test_missing_min_quality_tier_passes_none(): + r = _make_router() + r.pick_model = AsyncMock(return_value="fast") # type: ignore[method-assign] + + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={}, + messages=[{"role": "user", "content": "hi"}], + ) + + assert ( + r.pick_model.await_args.kwargs["min_quality_tier"] is None # type: ignore[union-attr] + ) + + +@pytest.mark.asyncio +async def test_invalid_min_quality_tier_header_treated_as_none(): + """A garbage header value must not crash the request — treat as unset.""" + r = _make_router() + r.pick_model = AsyncMock(return_value="fast") # type: ignore[method-assign] + + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={"headers": {"x-litellm-min-quality-tier": "not-a-number"}}, + messages=[{"role": "user", "content": "hi"}], + ) + + assert ( + r.pick_model.await_args.kwargs["min_quality_tier"] is None # type: ignore[union-attr] + ) diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py b/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py new file mode 100644 index 00000000000..ab322f0fb37 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py @@ -0,0 +1,134 @@ +import random + +import pytest + +from litellm.router_strategy.adaptive_router.bandit import ( + BanditCell, + apply_delta, + initial_cell, + normalized_cost, + pick_best, + score, + thompson_sample, +) +from litellm.router_strategy.adaptive_router.config import ( + BASE_TIER_WEIGHT, + COLD_START_MASS, + SAMPLE_CAP, + STRENGTH_BONUS, +) +from litellm.types.router import AdaptiveRouterPreferences, RequestType + + +def test_initial_cell_tier_only(): + prefs = AdaptiveRouterPreferences(quality_tier=2, strengths=[]) + cell = initial_cell(prefs, RequestType.GENERAL) + expected_mean = BASE_TIER_WEIGHT[2] + assert abs(cell.mean - expected_mean) < 0.001 + assert abs(cell.alpha + cell.beta - COLD_START_MASS) < 0.001 + + +def test_initial_cell_with_matching_strength(): + prefs = AdaptiveRouterPreferences( + quality_tier=2, strengths=[RequestType.CODE_GENERATION] + ) + cell = initial_cell(prefs, RequestType.CODE_GENERATION) + expected_mean = BASE_TIER_WEIGHT[2] + STRENGTH_BONUS + assert abs(cell.mean - expected_mean) < 0.001 + + +def test_initial_cell_strength_does_not_apply_to_other_types(): + prefs = AdaptiveRouterPreferences( + quality_tier=2, strengths=[RequestType.CODE_GENERATION] + ) + cell = initial_cell(prefs, RequestType.WRITING) + assert abs(cell.mean - BASE_TIER_WEIGHT[2]) < 0.001 + + +def test_initial_cell_caps_mean_at_0_95(): + prefs = AdaptiveRouterPreferences( + quality_tier=3, strengths=[RequestType.CODE_GENERATION] + ) + cell = initial_cell(prefs, RequestType.CODE_GENERATION) + assert cell.mean <= 0.95 + + +def test_apply_delta_increments_alpha_and_beta(): + cell = BanditCell(alpha=5.0, beta=5.0) + new_cell = apply_delta(cell, 1.0, 0.0) + assert new_cell.alpha == 6.0 + assert new_cell.beta == 5.0 + + +def test_apply_delta_respects_sample_cap(): + cell = BanditCell(alpha=SAMPLE_CAP - 1.0, beta=1.0) + same_cell = apply_delta(cell, 5.0, 5.0) + assert same_cell.alpha == cell.alpha + assert same_cell.beta == cell.beta + + +def test_thompson_sample_in_range(): + cell = BanditCell(alpha=10.0, beta=5.0) + rng = random.Random(42) + for _ in range(100): + s = thompson_sample(cell, rng=rng) + assert 0.0 <= s <= 1.0 + + +def test_normalized_cost_cheapest_wins(): + assert normalized_cost(0.001, [0.001, 0.005, 0.01]) == 1.0 + assert normalized_cost(0.01, [0.001, 0.005, 0.01]) == 0.0 + + +def test_normalized_cost_no_spread(): + assert normalized_cost(0.005, [0.005, 0.005]) == 0.5 + + +def test_normalized_cost_empty_list(): + assert normalized_cost(0.005, []) == 0.5 + + +def test_score_combines_quality_and_cost(): + s = score( + quality_sample=1.0, + model_cost=0.001, + all_costs=[0.001, 0.01], + quality_weight=0.7, + cost_weight=0.3, + ) + assert abs(s - 1.0) < 0.001 + + +def test_pick_best_empty_dict_raises(): + with pytest.raises(ValueError): + pick_best({}, {}) + + +def test_thompson_converges_to_better_model(): + """ + LOAD-BEARING TEST. If this regresses, the whole router is broken. + + Setup: 2 models, identical priors, identical cost. Model A's true mean = 0.8, + Model B's true mean = 0.3. After 200 simulated turns, A must be picked >= 80% of + last 50 turns. + """ + rng = random.Random(42) + cells = { + "A": BanditCell(alpha=5.0, beta=5.0), + "B": BanditCell(alpha=5.0, beta=5.0), + } + costs = {"A": 0.001, "B": 0.001} + true_means = {"A": 0.8, "B": 0.3} + + picks = [] + for _ in range(200): + chosen = pick_best(cells, costs, rng=rng) + picks.append(chosen) + outcome = 1.0 if rng.random() < true_means[chosen] else 0.0 + cells[chosen] = apply_delta(cells[chosen], outcome, 1.0 - outcome) + + last_50 = picks[-50:] + a_share = last_50.count("A") / 50 + assert ( + a_share >= 0.80 + ), f"Expected A to dominate ({a_share=}); priors aren't biasing the sample correctly" diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_classifier.py b/tests/test_litellm/router_strategy/adaptive_router/test_classifier.py new file mode 100644 index 00000000000..c27e2d945a3 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_classifier.py @@ -0,0 +1,116 @@ +import pytest + +from litellm.router_strategy.adaptive_router.classifier import classify_prompt +from litellm.types.router import RequestType + + +@pytest.mark.parametrize( + "text", + [ + "Write a Python function that reverses a linked list", + "Implement a REST API endpoint for user signup", + "Create a bash script to back up my postgres database", + ], +) +def test_classify_code_generation(text): + assert classify_prompt(text) == RequestType.CODE_GENERATION + + +@pytest.mark.parametrize( + "text", + [ + "Explain what this function does: def foo(): ...", + "Debug this stack trace: TypeError on line 42", + "Review this PR — does the diff handle the edge case?", + ], +) +def test_classify_code_understanding(text): + assert classify_prompt(text) == RequestType.CODE_UNDERSTANDING + + +@pytest.mark.parametrize( + "text", + [ + "Design a microservice architecture for an event-driven system", + "Should I use PostgreSQL or DynamoDB for high-write workloads?", + "How should I structure my Django app for multi-tenancy?", + ], +) +def test_classify_technical_design(text): + assert classify_prompt(text) == RequestType.TECHNICAL_DESIGN + + +@pytest.mark.parametrize( + "text", + [ + "Solve the integral of x^2 from 0 to 5", + "If A implies B and B implies C, then prove A implies C", + "Calculate the probability of two heads in three coin flips", + ], +) +def test_classify_analytical_reasoning(text): + assert classify_prompt(text) == RequestType.ANALYTICAL_REASONING + + +@pytest.mark.parametrize( + "text", + [ + "Draft an email to my team announcing the launch", + "Rewrite this paragraph to be more concise and professional", + "Proofread my blog post for grammar and tone", + ], +) +def test_classify_writing(text): + assert classify_prompt(text) == RequestType.WRITING + + +@pytest.mark.parametrize( + "text", + [ + "Who is the current president of France?", + "What is the capital of Australia?", + "Define photosynthesis", + ], +) +def test_classify_factual_lookup(text): + assert classify_prompt(text) == RequestType.FACTUAL_LOOKUP + + +@pytest.mark.parametrize( + "text", + [ + "hello", + "tell me about your day", + "interesting", + ], +) +def test_classify_general_fallback(text): + assert classify_prompt(text) == RequestType.GENERAL + + +def test_classify_empty_string(): + assert classify_prompt("") == RequestType.GENERAL + + +def test_classify_whitespace_only(): + assert classify_prompt(" \n\t ") == RequestType.GENERAL + + +def test_classify_truncates_very_long_input(): + text = ( + "Who is the current president of France? " + + "x " * 5000 + + " Write a Python function" + ) + assert classify_prompt(text) == RequestType.FACTUAL_LOOKUP + + +def test_classify_is_deterministic(): + text = "Implement a REST API endpoint for user signup" + results = {classify_prompt(text) for _ in range(10)} + assert len(results) == 1 + + +def test_classify_returns_request_type_enum(): + result = classify_prompt("hello") + assert isinstance(result, RequestType) diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_config.py b/tests/test_litellm/router_strategy/adaptive_router/test_config.py new file mode 100644 index 00000000000..fd14556a0bc --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_config.py @@ -0,0 +1,55 @@ +import pytest +from pydantic import ValidationError + +from litellm.types.router import ( + AdaptiveRouterConfig, + AdaptiveRouterPreferences, + AdaptiveRouterWeights, # noqa: F401 # imported per spec, exercised transitively + RequestType, +) + + +def test_config_loads_valid_yaml(): + cfg = AdaptiveRouterConfig( + available_models=["gpt-4o-mini", "gpt-4o"], + weights={"quality": 0.7, "cost": 0.3}, + ) + assert cfg.available_models == ["gpt-4o-mini", "gpt-4o"] + assert cfg.weights.quality == 0.7 + assert cfg.weights.cost == 0.3 + assert abs(cfg.weights.quality + cfg.weights.cost - 1.0) < 0.001 + + +def test_config_rejects_misspelled_strength(): + with pytest.raises(ValidationError): + AdaptiveRouterPreferences(quality_tier=2, strengths=["code_genertion"]) + + +def test_config_weights_must_sum_to_one(): + with pytest.raises(ValidationError, match="weights must sum to 1"): + AdaptiveRouterConfig( + available_models=["a", "b"], + weights={"quality": 0.9, "cost": 0.5}, + ) + + +def test_config_quality_tier_must_be_1_2_or_3(): + with pytest.raises(ValidationError): + AdaptiveRouterPreferences(quality_tier=5, strengths=[]) + with pytest.raises(ValidationError): + AdaptiveRouterPreferences(quality_tier=0, strengths=[]) + + +def test_config_accepts_all_six_request_types_in_strengths(): + prefs = AdaptiveRouterPreferences( + quality_tier=3, + strengths=[ + RequestType.CODE_GENERATION, + RequestType.CODE_UNDERSTANDING, + RequestType.TECHNICAL_DESIGN, + RequestType.ANALYTICAL_REASONING, + RequestType.WRITING, + RequestType.FACTUAL_LOOKUP, + ], + ) + assert len(prefs.strengths) == 6 diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py new file mode 100644 index 00000000000..9786832b4ae --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py @@ -0,0 +1,298 @@ +""" +End-to-end tests for the adaptive router. Wires the real strategy + queue + hook +with a mocked Prisma client. No live proxy or DB required. + +What we cover: + 1. Full lifecycle: pick -> record turn(s) -> flush -> DB upsert with correct deltas + 2. Owner cache pins attribution: same key + matching model -> updates flow + 3. Convergence in-process: 50 simulated sessions, "good" model dominates last 10 + 4. Cold-start state load from DB overrides priors + 5. Failure signal increments beta in the next flush + 6. Unknown request types in DB rows are silently skipped + 7. Flush isolates writes per (router, session, model) tuple +""" + +import random +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter +from litellm.router_strategy.adaptive_router.signals import Turn +from litellm.types.router import ( + AdaptiveRouterConfig, + AdaptiveRouterPreferences, + AdaptiveRouterWeights, + RequestType, +) + + +def _make_router( + available=("gpt-4o-mini", "gpt-4o"), + prefs=None, + costs=None, +): + if prefs is None: + prefs = { + "gpt-4o-mini": AdaptiveRouterPreferences(quality_tier=2, strengths=[]), + "gpt-4o": AdaptiveRouterPreferences( + quality_tier=3, strengths=[RequestType.CODE_GENERATION] + ), + } + if costs is None: + costs = {"gpt-4o-mini": 0.15, "gpt-4o": 5.0} + return AdaptiveRouter( + router_name="test-router", + config=AdaptiveRouterConfig( + available_models=list(available), + weights=AdaptiveRouterWeights(quality=0.7, cost=0.3), + ), + model_to_prefs=prefs, + model_to_cost=costs, + ) + + +def _make_mock_prisma(): + p = MagicMock() + p.db.litellm_adaptiverouterstate.find_unique = AsyncMock(return_value=None) + p.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[]) + p.db.litellm_adaptiverouterstate.upsert = AsyncMock() + p.db.litellm_adaptiveroutersession.upsert = AsyncMock() + return p + + +@pytest.mark.asyncio +async def test_pick_record_flush_full_cycle(): + router = _make_router() + chosen = await router.pick_model(RequestType.CODE_GENERATION) + assert chosen in router.config.available_models + + # Prime 2 prior turns (distinct content so no other signals fire) so the + # MIN_TURNS_FOR_CLEAN_CREDIT satisfaction gate is satisfied on turn 3. + priming = [ + Turn( + user_content="alpha bravo charlie", assistant_content="delta echo foxtrot" + ), + Turn( + user_content="golf hotel india juliet", + assistant_content="kilo lima mike november", + ), + ] + for t in priming: + await router.record_turn( + session_id="s1", + model_name=chosen, + request_type=RequestType.CODE_GENERATION, + turn=t, + ) + await router.record_turn( + session_id="s1", + model_name=chosen, + request_type=RequestType.CODE_GENERATION, + turn=Turn(user_content="thanks, that worked!", assistant_content="ok"), + ) + + prisma = _make_mock_prisma() + n_state = await router.queue.flush_state_to_db(prisma) + n_session = await router.queue.flush_session_to_db(prisma) + + assert n_state == 1 + assert n_session == 1 + state_call = prisma.db.litellm_adaptiverouterstate.upsert.call_args + # satisfaction signal -> +1 alpha, no existing row -> create.alpha == 1.0 + assert state_call.kwargs["data"]["create"]["alpha"] >= 1.0 + assert state_call.kwargs["data"]["create"]["beta"] == 0.0 + assert state_call.kwargs["data"]["create"]["total_samples"] == 1 + + session_call = prisma.db.litellm_adaptiveroutersession.upsert.call_args + assert session_call.kwargs["data"]["create"]["satisfaction_count"] == 1 + assert session_call.kwargs["data"]["create"]["session_id"] == "s1" + assert session_call.kwargs["data"]["create"]["model_name"] == chosen + + +@pytest.mark.asyncio +async def test_owner_cache_pins_attribution_to_first_picked_model(): + """First call claims ownership; matching model returns True, mismatch False.""" + router = _make_router() + chosen = await router.pick_model(RequestType.GENERAL) + assert router.claim_or_check_owner("sess-own", chosen) is True + + # Same model on later turns keeps attributing. + for _ in range(5): + assert router.claim_or_check_owner("sess-own", chosen) is True + + # A different model on a later turn is rejected. + other = "gpt-4o" if chosen == "gpt-4o-mini" else "gpt-4o-mini" + assert router.claim_or_check_owner("sess-own", other) is False + assert router._skipped_updates_total == 1 + + +@pytest.mark.asyncio +async def test_pick_model_returns_valid_models_without_error(): + router = _make_router() + # Picks may legitimately differ across calls (Thompson sampling is stochastic). + # Just confirm every pick is valid and nothing raises. + for _ in range(10): + m = await router.pick_model(RequestType.GENERAL) + assert m in router.config.available_models + + +@pytest.mark.asyncio +async def test_in_process_convergence_high_quality_model_dominates(): + """ + Two models, identical cost. "good" satisfies every turn, "bad" fails every turn. + After 50 sessions of 4 turns each, "good" should win >=70% of the last 10 picks. + Seed `random` for determinism since pick_best uses the module-level RNG. + """ + random.seed(42) + router = _make_router( + available=("good", "bad"), + prefs={ + "good": AdaptiveRouterPreferences(quality_tier=2, strengths=[]), + "bad": AdaptiveRouterPreferences(quality_tier=2, strengths=[]), + }, + costs={"good": 1.0, "bad": 1.0}, + ) + + picks = [] + for sess in range(50): + sid = f"conv-{sess}" + chosen = await router.pick_model(RequestType.GENERAL) + for _turn_i in range(4): + if chosen == "good": + turn = Turn(user_content="thanks!", assistant_content="ok") + else: + turn = Turn( + tool_calls=[{"name": "x", "arguments": {}}], + tool_results=[{"is_error": True, "content": "boom"}], + ) + await router.record_turn(sid, chosen, RequestType.GENERAL, turn) + picks.append(chosen) + + last_10 = picks[-10:] + good_share = last_10.count("good") / 10 + assert good_share >= 0.7, f"good_share={good_share} (last picks={picks})" + + +@pytest.mark.asyncio +async def test_failure_signal_increments_beta_after_flush(): + router = _make_router( + available=("only",), + prefs={"only": AdaptiveRouterPreferences(quality_tier=2, strengths=[])}, + costs={"only": 1.0}, + ) + chosen = await router.pick_model(RequestType.GENERAL) + assert chosen == "only" + + await router.record_turn( + session_id="f1", + model_name=chosen, + request_type=RequestType.GENERAL, + turn=Turn( + tool_calls=[{"name": "x", "arguments": {}}], + tool_results=[{"is_error": True, "content": ""}], + ), + ) + + prisma = _make_mock_prisma() + n_state = await router.queue.flush_state_to_db(prisma) + assert n_state == 1 + state_call = prisma.db.litellm_adaptiverouterstate.upsert.call_args + assert state_call.kwargs["data"]["create"]["beta"] >= 1.0 + assert state_call.kwargs["data"]["create"]["alpha"] == 0.0 + + +@pytest.mark.asyncio +async def test_load_state_from_db_overrides_cold_start(): + router = _make_router() + fake_row = MagicMock() + fake_row.request_type = RequestType.GENERAL.value + fake_row.model_name = "gpt-4o" + fake_row.alpha = 90.0 + fake_row.beta = 10.0 + + prisma = _make_mock_prisma() + prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[fake_row]) + + await router.load_state_from_db(prisma) + + cell = router._cells[(RequestType.GENERAL, "gpt-4o")] + assert cell.alpha == 90.0 + assert cell.beta == 10.0 + + +@pytest.mark.asyncio +async def test_load_state_from_db_handles_unknown_request_type(): + router = _make_router() + bad_row = MagicMock() + bad_row.request_type = "unknown_v1_type" + bad_row.model_name = "gpt-4o" + bad_row.alpha = 50.0 + bad_row.beta = 50.0 + + prisma = _make_mock_prisma() + prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[bad_row]) + + # Should not raise; bad row is silently skipped and cold-start cells remain. + await router.load_state_from_db(prisma) + cell = router._cells[(RequestType.GENERAL, "gpt-4o")] + # Cold-start: tier 3 base = 0.7, mass = 10 -> alpha = 7, beta = 3 + assert cell.alpha == pytest.approx(7.0) + assert cell.beta == pytest.approx(3.0) + + +@pytest.mark.asyncio +async def test_flush_isolates_writes_per_router_session_model(): + router = _make_router() + # Prime 2 prior turns per session to clear the MIN_TURNS_FOR_CLEAN_CREDIT gate. + for sid, model in (("s1", "gpt-4o"), ("s2", "gpt-4o-mini")): + for _ in range(2): + await router.record_turn( + sid, + model, + RequestType.GENERAL, + Turn(user_content="hi", assistant_content="hello"), + ) + await router.record_turn( + "s1", "gpt-4o", RequestType.GENERAL, Turn(user_content="thanks!") + ) + await router.record_turn( + "s2", "gpt-4o-mini", RequestType.GENERAL, Turn(user_content="thanks!") + ) + + prisma = _make_mock_prisma() + n = await router.queue.flush_session_to_db(prisma) + assert n == 2 + assert prisma.db.litellm_adaptiveroutersession.upsert.call_count == 2 + + n_state = await router.queue.flush_state_to_db(prisma) + assert n_state == 2 + assert prisma.db.litellm_adaptiverouterstate.upsert.call_count == 2 + + +@pytest.mark.asyncio +async def test_repeated_flush_drains_queue_and_subsequent_flush_is_noop(): + """Verifies the queue is fully drained on flush -- a second flush writes nothing.""" + router = _make_router() + chosen = await router.pick_model(RequestType.GENERAL) + # Prime 2 prior turns so satisfaction can fire on the third turn. + for _ in range(2): + await router.record_turn( + "drain-1", + chosen, + RequestType.GENERAL, + Turn(user_content="hi", assistant_content="hello"), + ) + await router.record_turn( + "drain-1", chosen, RequestType.GENERAL, Turn(user_content="thanks!") + ) + + prisma = _make_mock_prisma() + assert await router.queue.flush_state_to_db(prisma) == 1 + assert await router.queue.flush_session_to_db(prisma) == 1 + + # Second drain should be a no-op (queue is empty). + assert await router.queue.flush_state_to_db(prisma) == 0 + assert await router.queue.flush_session_to_db(prisma) == 0 + assert prisma.db.litellm_adaptiverouterstate.upsert.call_count == 1 + assert prisma.db.litellm_adaptiveroutersession.upsert.call_count == 1 diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_hooks.py b/tests/test_litellm/router_strategy/adaptive_router/test_hooks.py new file mode 100644 index 00000000000..a2b85f2ce53 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_hooks.py @@ -0,0 +1,368 @@ +"""Unit tests for the AdaptiveRouterPostCallHook.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + SIGNAL_GATE_MIN_MESSAGES, +) +from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + _recent_tool_results, + _resolve_session_key, +) +from litellm.router_strategy.adaptive_router.signals import Turn + + +def _make_hook(claim: bool = True) -> AdaptiveRouterPostCallHook: + fake_router = MagicMock() + fake_router.record_turn = AsyncMock() + fake_router.claim_or_check_owner = MagicMock(return_value=claim) + return AdaptiveRouterPostCallHook(adaptive_router=fake_router) + + +def _resp_with_content(text: str, tool_calls=None): + """Build a ModelResponse-like object with a single assistant message.""" + msg = MagicMock() + msg.content = text + msg.tool_calls = tool_calls or [] + choice = MagicMock() + choice.message = msg + resp = MagicMock() + resp.choices = [choice] + return resp + + +def _long_messages(user_text: str = "ask"): + """Return a message list at the SIGNAL_GATE_MIN_MESSAGES threshold.""" + base = [ + {"role": "user", "content": "first turn"}, + {"role": "assistant", "content": "first reply"}, + {"role": "user", "content": "second turn"}, + ] + base.append({"role": "user", "content": user_text}) + # Pad to threshold if needed. + while len(base) < SIGNAL_GATE_MIN_MESSAGES: + base.append({"role": "user", "content": "filler"}) + return base + + +def _kwargs( + *, + messages=None, + chosen="fast", + extra_metadata=None, + extra_litellm_params=None, +): + metadata = {ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY: chosen} if chosen else {} + if extra_metadata: + metadata.update(extra_metadata) + lp = {"metadata": metadata} + if extra_litellm_params: + lp.update(extra_litellm_params) + return { + "model": "anthropic/claude-opus-4-7", + "messages": messages if messages is not None else _long_messages(), + "litellm_params": lp, + } + + +# ---- _resolve_session_key ------------------------------------------------ + + +def test_resolve_session_key_honors_litellm_session_id_on_litellm_params(): + key = _resolve_session_key({"litellm_params": {"litellm_session_id": "sess-A"}}) + assert key == "sess-A" + + +def test_resolve_session_key_honors_metadata_session_id(): + key = _resolve_session_key( + {"litellm_params": {"metadata": {"session_id": "sess-B"}}} + ) + assert key == "sess-B" + + +def test_resolve_session_key_returns_none_when_no_messages(): + assert _resolve_session_key({"litellm_params": {}}) is None + assert _resolve_session_key({"litellm_params": {}, "messages": []}) is None + + +def test_resolve_session_key_derives_stable_hash_from_first_message(): + # `_resolve_session_key` requires at least SIGNAL_GATE_MIN_MESSAGES + # messages before it will derive a hash (matches the signal-processing + # gate) — otherwise the session is too short to attribute. + msgs = _long_messages("Hello, world") + k1 = _resolve_session_key({"messages": msgs}) + k2 = _resolve_session_key({"messages": list(msgs)}) + assert k1 == k2 + assert k1 and len(k1) == 64 # sha256 hex + + +def test_resolve_session_key_does_not_prefix_sk(): + key = _resolve_session_key({"messages": _long_messages()}) + assert key and not key.startswith("sk_") + + +def test_resolve_session_key_segments_by_identity_fields(): + """Same first message but different api keys must yield different keys.""" + msgs = _long_messages("same prompt") + k_team_a = _resolve_session_key( + { + "messages": msgs, + "litellm_params": { + "metadata": { + "user_api_key_hash": "hash-A", + "user_api_key_team_id": "team-1", + } + }, + } + ) + k_team_b = _resolve_session_key( + { + "messages": msgs, + "litellm_params": { + "metadata": { + "user_api_key_hash": "hash-B", + "user_api_key_team_id": "team-2", + } + }, + } + ) + assert k_team_a != k_team_b + + +def test_resolve_session_key_changes_when_first_message_changes(): + k1 = _resolve_session_key({"messages": _long_messages("alpha")}) + k2 = _resolve_session_key({"messages": _long_messages("beta")}) + assert k1 != k2 + + +# ---- _record gating ----------------------------------------------------- + + +@pytest.mark.asyncio +async def test_hook_skips_when_below_signal_gate(): + """Conversations shorter than SIGNAL_GATE_MIN_MESSAGES should be ignored.""" + hook = _make_hook() + short = [{"role": "user", "content": "hi"}] + assert len(short) < SIGNAL_GATE_MIN_MESSAGES # sanity + kwargs = _kwargs(messages=short) + await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0) + hook.adaptive_router.record_turn.assert_not_awaited() + hook.adaptive_router.claim_or_check_owner.assert_not_called() + + +@pytest.mark.asyncio +async def test_hook_skips_when_no_messages(): + hook = _make_hook() + kwargs = _kwargs(messages=[]) + await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0) + hook.adaptive_router.record_turn.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_hook_skips_when_chosen_model_missing_from_metadata(): + hook = _make_hook() + kwargs = _kwargs(chosen=None) + await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0) + hook.adaptive_router.record_turn.assert_not_awaited() + hook.adaptive_router.claim_or_check_owner.assert_not_called() + + +@pytest.mark.asyncio +async def test_hook_skips_when_owner_cache_mismatch(): + """A different model owns this conversation -> no attribution.""" + hook = _make_hook(claim=False) + kwargs = _kwargs(chosen="fast") + await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0) + hook.adaptive_router.claim_or_check_owner.assert_called_once() + hook.adaptive_router.record_turn.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_hook_records_turn_when_owner_claims(): + hook = _make_hook(claim=True) + kwargs = _kwargs(chosen="smart", messages=_long_messages("ask")) + await hook.async_log_success_event( + kwargs, _resp_with_content("answer here"), 0.0, 1.0 + ) + call = hook.adaptive_router.record_turn.await_args + assert call.kwargs["model_name"] == "smart" + turn: Turn = call.kwargs["turn"] + assert turn.user_content == "ask" + assert turn.assistant_content == "answer here" + assert turn.response_status == 200 + + +@pytest.mark.asyncio +async def test_hook_uses_explicit_session_id_when_provided(): + """Explicit `litellm_session_id` is forwarded as the session key.""" + hook = _make_hook() + kwargs = _kwargs( + chosen="fast", + extra_litellm_params={"litellm_session_id": "explicit-sess"}, + ) + await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0) + args, _ = hook.adaptive_router.claim_or_check_owner.call_args + assert args[0] == "explicit-sess" + assert hook.adaptive_router.record_turn.await_args.kwargs["session_id"] == ( + "explicit-sess" + ) + + +@pytest.mark.asyncio +async def test_hook_passes_tool_calls_through(): + hook = _make_hook() + tc = {"name": "search", "arguments": '{"q":"x"}'} + kwargs = _kwargs(chosen="fast") + await hook.async_log_success_event( + kwargs, _resp_with_content("calling tool", tool_calls=[tc]), 0.0, 1.0 + ) + turn: Turn = hook.adaptive_router.record_turn.await_args.kwargs["turn"] + assert turn.tool_calls == [tc] + + +# ---- _recent_tool_results ------------------------------------------------ + + +def test_recent_tool_results_empty_when_no_messages(): + assert _recent_tool_results(None) == [] + assert _recent_tool_results([]) == [] + + +def test_recent_tool_results_collects_trailing_tool_messages(): + """Tool messages at the tail of the conversation are extracted in order.""" + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": None, "tool_calls": [{"id": "t1"}]}, + {"role": "tool", "tool_call_id": "t1", "content": "result A"}, + {"role": "tool", "tool_call_id": "t2", "content": "result B"}, + ] + results = _recent_tool_results(messages) + assert [r["content"] for r in results] == ["result A", "result B"] + assert all(r["is_error"] is False for r in results) + + +def test_recent_tool_results_propagates_is_error_flag(): + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": None, "tool_calls": [{"id": "t1"}]}, + {"role": "tool", "content": "boom", "is_error": True}, + ] + results = _recent_tool_results(messages) + assert results == [{"content": "boom", "is_error": True}] + + +def test_recent_tool_results_stops_at_first_non_tool_message(): + """Only the trailing run of tool messages counts — prior rounds are + considered already attributed.""" + messages = [ + {"role": "user", "content": "hi"}, + {"role": "tool", "content": "stale"}, # earlier round, ignored + {"role": "assistant", "content": "intermediate"}, + {"role": "user", "content": "follow-up"}, + {"role": "tool", "content": "current"}, + ] + results = _recent_tool_results(messages) + assert [r["content"] for r in results] == ["current"] + + +def test_recent_tool_results_empty_when_no_trailing_tool_message(): + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + assert _recent_tool_results(messages) == [] + + +@pytest.mark.asyncio +async def test_hook_passes_tool_results_to_turn_for_failure_detection(): + """A trailing tool message with `is_error` must reach `Turn.tool_results` + so the failure-signal path fires.""" + hook = _make_hook() + messages = _long_messages() + messages.append( + {"role": "assistant", "content": None, "tool_calls": [{"id": "t1"}]} + ) + messages.append( + {"role": "tool", "tool_call_id": "t1", "content": "500", "is_error": True} + ) + kwargs = _kwargs(chosen="fast", messages=messages) + + await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0) + + turn: Turn = hook.adaptive_router.record_turn.await_args.kwargs["turn"] + assert turn.tool_results == [{"content": "500", "is_error": True}] + + +@pytest.mark.asyncio +async def test_hook_swallows_exceptions_from_record_turn(): + hook = _make_hook() + hook.adaptive_router.record_turn.side_effect = RuntimeError("boom") + kwargs = _kwargs(chosen="fast") + # Must NOT raise — signal recording must never break a request. + await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0) + + +@pytest.mark.asyncio +async def test_hook_failure_event_uses_status_code_from_exception(): + hook = _make_hook() + exc = MagicMock() + exc.status_code = 429 + kwargs = _kwargs(chosen="fast") + kwargs["exception"] = exc + await hook.async_log_failure_event(kwargs, None, 0.0, 1.0) + turn: Turn = hook.adaptive_router.record_turn.await_args.kwargs["turn"] + assert turn.response_status == 429 + + +# ---- async_post_call_success_hook (response header surfacing) ---------- + + +@pytest.mark.asyncio +async def test_post_call_response_headers_hook_returns_chosen_model_header(): + """The header hook returns the `x-litellm-adaptive-router-model` header + so proxy header construction picks it up (works for both streaming and + non-streaming; `async_post_call_success_hook` is too late for streaming).""" + hook = _make_hook() + headers = await hook.async_post_call_response_headers_hook( + data={"metadata": {"adaptive_router_chosen_model": "smart"}}, + user_api_key_dict=MagicMock(), + response=MagicMock(), + ) + assert headers == {"x-litellm-adaptive-router-model": "smart"} + + +@pytest.mark.asyncio +async def test_post_call_response_headers_hook_noop_when_metadata_missing_key(): + hook = _make_hook() + headers = await hook.async_post_call_response_headers_hook( + data={"metadata": {"litellm_session_id": "sess-A"}}, + user_api_key_dict=MagicMock(), + response=MagicMock(), + ) + assert headers is None + + +@pytest.mark.asyncio +async def test_post_call_response_headers_hook_noop_when_no_metadata(): + hook = _make_hook() + headers = await hook.async_post_call_response_headers_hook( + data={}, + user_api_key_dict=MagicMock(), + response=MagicMock(), + ) + assert headers is None + + +@pytest.mark.asyncio +async def test_post_call_response_headers_hook_noop_when_metadata_not_dict(): + hook = _make_hook() + headers = await hook.async_post_call_response_headers_hook( + data={"metadata": "not-a-dict"}, + user_api_key_dict=MagicMock(), + response=MagicMock(), + ) + assert headers is None diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py b/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py new file mode 100644 index 00000000000..604155e1221 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py @@ -0,0 +1,486 @@ +"""Tests for the Router-level wiring of the adaptive router. + +Specifically guards the four bugs found when wiring the example config +`auto_router/adaptive_router` end-to-end: + +1. The `auto_router/adaptive_router` model prefix must NOT trigger the + semantic auto-router init path (which would crash on missing fields). +2. The same prefix MUST trigger the adaptive-router init path. +3. `init_adaptive_router_deployment` must read `input_cost_per_token` + from `litellm_params` (where users put it), not just `model_info`. +4. `Router.async_pre_routing_hook` must dispatch to the matching entry in + `self.adaptive_routers` when the inbound model matches a configured + adaptive-router name, returning the underlying model the bandit picked. +""" + +from unittest.mock import AsyncMock + +import pytest + +from litellm import Router +from litellm.types.router import LiteLLM_Params, RequestType + + +def _params(**overrides): + base = {"model": "auto_router/adaptive_router"} + base.update(overrides) + return LiteLLM_Params(**base) + + +# ---- Fix 1 & 2: opt-in prefix routing ----------------------------------- + + +def test_auto_router_check_excludes_adaptive_router_prefix(): + r = Router(model_list=[]) + assert ( + r._is_auto_router_deployment( + litellm_params=_params(model="auto_router/adaptive_router") + ) + is False + ) + + +def test_auto_router_check_excludes_complexity_router_prefix(): + r = Router(model_list=[]) + assert ( + r._is_auto_router_deployment( + litellm_params=_params(model="auto_router/complexity_router") + ) + is False + ) + + +def test_auto_router_check_still_matches_plain_auto_router_prefix(): + r = Router(model_list=[]) + assert ( + r._is_auto_router_deployment( + litellm_params=_params(model="auto_router/my-semantic-router") + ) + is True + ) + + +def test_adaptive_router_check_recognizes_prefix(): + r = Router(model_list=[]) + assert ( + r._is_adaptive_router_deployment( + litellm_params=_params(model="auto_router/adaptive_router") + ) + is True + ) + + +def test_adaptive_router_check_rejects_other_prefixes(): + r = Router(model_list=[]) + assert ( + r._is_adaptive_router_deployment(litellm_params=_params(model="openai/gpt-4o")) + is False + ) + + +# ---- Fix 3: cost field path -------------------------------------------- + + +def test_init_adaptive_router_reads_cost_from_litellm_params(): + r = Router( + model_list=[ + { + "model_name": "smart-cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": { + "available_models": ["fast", "smart"], + }, + }, + }, + { + "model_name": "fast", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "input_cost_per_token": 0.00000015, + }, + "model_info": { + "adaptive_router_preferences": { + "quality_tier": 2, + "strengths": [], + } + }, + }, + { + "model_name": "smart", + "litellm_params": { + "model": "openai/gpt-4o", + "input_cost_per_token": 0.0000050, + }, + "model_info": { + "adaptive_router_preferences": { + "quality_tier": 3, + "strengths": ["code_generation"], + } + }, + }, + ] + ) + assert "smart-cheap-router" in r.adaptive_routers + assert r.adaptive_routers["smart-cheap-router"].model_to_cost == { + "fast": 0.00000015, + "smart": 0.0000050, + } + + +# ---- Fix 4: pre-routing dispatch --------------------------------------- + + +def _router_with_adaptive() -> Router: + return Router( + model_list=[ + { + "model_name": "smart-cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": { + "available_models": ["fast", "smart"], + }, + }, + }, + { + "model_name": "fast", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "input_cost_per_token": 0.00000015, + }, + "model_info": { + "adaptive_router_preferences": { + "quality_tier": 2, + "strengths": [], + } + }, + }, + { + "model_name": "smart", + "litellm_params": { + "model": "openai/gpt-4o", + "input_cost_per_token": 0.0000050, + }, + "model_info": { + "adaptive_router_preferences": { + "quality_tier": 3, + "strengths": ["code_generation"], + } + }, + }, + ] + ) + + +@pytest.mark.asyncio +async def test_async_pre_routing_hook_dispatches_to_adaptive_router(): + r = _router_with_adaptive() + ar = r.adaptive_routers["smart-cheap-router"] + ar.pick_model = AsyncMock(return_value="smart") # type: ignore[assignment] + + response = await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={"metadata": {"litellm_session_id": "sess-A"}}, + messages=[{"role": "user", "content": "Write a Python function"}], + ) + assert response is not None + assert response.model == "smart" + call = ar.pick_model.await_args # type: ignore[union-attr] + # Stateless routing: session_id is no longer passed to pick_model. + assert "session_id" not in call.kwargs + assert call.kwargs["request_type"] == RequestType.CODE_GENERATION + + +@pytest.mark.asyncio +async def test_async_pre_routing_hook_pick_model_not_passed_session_id(): + r = _router_with_adaptive() + ar = r.adaptive_routers["smart-cheap-router"] + ar.pick_model = AsyncMock(return_value="fast") # type: ignore[assignment] + + response = await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs={}, + messages=[{"role": "user", "content": "hello"}], + ) + assert response is not None + assert response.model == "fast" + assert "session_id" not in ar.pick_model.await_args.kwargs # type: ignore[union-attr] + + +@pytest.mark.asyncio +async def test_async_pre_routing_hook_returns_none_for_unrelated_model(): + r = _router_with_adaptive() + ar = r.adaptive_routers["smart-cheap-router"] + ar.pick_model = AsyncMock() # type: ignore[assignment] + response = await r.async_pre_routing_hook( + model="some-other-model", + request_kwargs={}, + messages=[{"role": "user", "content": "x"}], + ) + assert response is None + ar.pick_model.assert_not_awaited() # type: ignore[union-attr] + + +# ---- Response header surfacing ----------------------------------------- + + +@pytest.mark.asyncio +async def test_async_pre_routing_hook_stashes_chosen_model_in_metadata(): + """ + The adaptive-router branch must record the chosen logical model on + `request_kwargs["metadata"]` so `_acompletion` can surface it as the + `x-litellm-adaptive-router-model` response header. + """ + r = _router_with_adaptive() + r.adaptive_routers["smart-cheap-router"].pick_model = AsyncMock( # type: ignore[assignment] + return_value="smart" + ) + + request_kwargs: dict = {"metadata": {"litellm_session_id": "sess-A"}} + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "Write a Python function"}], + ) + assert request_kwargs["metadata"]["adaptive_router_chosen_model"] == "smart" + + +@pytest.mark.asyncio +async def test_async_pre_routing_hook_creates_metadata_when_missing(): + """If no metadata was passed in, the hook should create one to stash the chosen model.""" + r = _router_with_adaptive() + r.adaptive_routers["smart-cheap-router"].pick_model = AsyncMock( # type: ignore[assignment] + return_value="fast" + ) + + request_kwargs: dict = {} + await r.async_pre_routing_hook( + model="smart-cheap-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hello"}], + ) + assert request_kwargs["metadata"]["adaptive_router_chosen_model"] == "fast" + + +# ---- Multi-router support ---------------------------------------------- + + +def test_two_adaptive_routers_can_coexist_on_one_router(): + r = Router( + model_list=[ + { + "model_name": "cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["fast"]}, + }, + }, + { + "model_name": "premium-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["smart"]}, + }, + }, + { + "model_name": "fast", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "input_cost_per_token": 0.00000015, + }, + }, + { + "model_name": "smart", + "litellm_params": { + "model": "openai/gpt-4o", + "input_cost_per_token": 0.0000050, + }, + }, + ] + ) + assert set(r.adaptive_routers.keys()) == {"cheap-router", "premium-router"} + assert r.adaptive_routers["cheap-router"].config.available_models == ["fast"] + assert r.adaptive_routers["premium-router"].config.available_models == ["smart"] + + +@pytest.mark.asyncio +async def test_async_pre_routing_hook_dispatches_to_correct_router_when_multiple(): + """Each adaptive router only handles its own router_name.""" + r = Router( + model_list=[ + { + "model_name": "cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["fast"]}, + }, + }, + { + "model_name": "premium-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["smart"]}, + }, + }, + { + "model_name": "fast", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "input_cost_per_token": 0.00000015, + }, + }, + { + "model_name": "smart", + "litellm_params": { + "model": "openai/gpt-4o", + "input_cost_per_token": 0.0000050, + }, + }, + ] + ) + cheap = r.adaptive_routers["cheap-router"] + premium = r.adaptive_routers["premium-router"] + cheap.pick_model = AsyncMock(return_value="fast") # type: ignore[assignment] + premium.pick_model = AsyncMock(return_value="smart") # type: ignore[assignment] + + cheap_response = await r.async_pre_routing_hook( + model="cheap-router", + request_kwargs={}, + messages=[{"role": "user", "content": "hi"}], + ) + premium_response = await r.async_pre_routing_hook( + model="premium-router", + request_kwargs={}, + messages=[{"role": "user", "content": "hi"}], + ) + + assert cheap_response is not None and cheap_response.model == "fast" + assert premium_response is not None and premium_response.model == "smart" + cheap.pick_model.assert_awaited_once() # type: ignore[union-attr] + premium.pick_model.assert_awaited_once() # type: ignore[union-attr] + + +def test_init_adaptive_router_rejects_duplicate_model_name(): + """Two adaptive-router deployments with the same model_name must error.""" + from litellm.types.router import AdaptiveRouterConfig, Deployment + + r = Router(model_list=[]) + cfg = {"available_models": ["fast"]} + deployment = Deployment( + model_name="dup-router", + litellm_params=LiteLLM_Params( + model="auto_router/adaptive_router", + adaptive_router_config=cfg, + ), + model_info={"id": "x"}, + ) + r.init_adaptive_router_deployment(deployment=deployment) + with pytest.raises(ValueError, match="already exists"): + r.init_adaptive_router_deployment(deployment=deployment) + + +def test_finalize_adaptive_router_if_configured_initializes_and_is_idempotent(): + """`_finalize_adaptive_router_if_configured` walks the model_list, builds an + AdaptiveRouter for each adaptive deployment, and is a safe no-op on + re-entry (models already in self.adaptive_routers are skipped).""" + r = Router( + model_list=[ + { + "model_name": "fast", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"input_cost_per_token": 0.00000015}, + }, + { + "model_name": "smart", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"input_cost_per_token": 0.0000025}, + }, + { + "model_name": "my-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": { + "available_models": ["fast", "smart"], + }, + }, + }, + ] + ) + + # Router __init__ already called _finalize_adaptive_router_if_configured. + assert "my-router" in r.adaptive_routers + original = r.adaptive_routers["my-router"] + + # Calling again must be idempotent: the existing AdaptiveRouter instance + # is preserved, not rebuilt. + r._finalize_adaptive_router_if_configured() + assert r.adaptive_routers["my-router"] is original + + +def test_finalize_prunes_stale_adaptive_router_hooks_from_callbacks(): + """Replacing the Router (hot-reload path) must not leave stale + AdaptiveRouterPostCallHook instances in `litellm.callbacks` — otherwise + every request double-fires signal recording.""" + import litellm + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + + model_list = [ + { + "model_name": "fast", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + }, + { + "model_name": "my-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["fast"]}, + }, + }, + ] + + # Snapshot any pre-existing AdaptiveRouterPostCallHook entries so we can + # restore them — other tests may have registered hooks we shouldn't drop. + pre_hooks = [ + cb for cb in litellm.callbacks if isinstance(cb, AdaptiveRouterPostCallHook) + ] + for cb in pre_hooks: + litellm.callbacks.remove(cb) + + try: + Router(model_list=model_list) + Router(model_list=model_list) # simulate hot-reload + + adaptive_hooks = [ + cb + for cb in litellm.callbacks + if isinstance(cb, AdaptiveRouterPostCallHook) + ] + assert len(adaptive_hooks) == 1, ( + f"expected exactly one AdaptiveRouterPostCallHook after hot-reload, " + f"got {len(adaptive_hooks)}" + ) + finally: + # Best-effort cleanup: remove whatever this test added, then restore. + for cb in list(litellm.callbacks): + if isinstance(cb, AdaptiveRouterPostCallHook): + litellm.callbacks.remove(cb) + for cb in pre_hooks: + litellm.callbacks.append(cb) + + +def test_finalize_adaptive_router_if_configured_noop_when_none_configured(): + """With no adaptive deployments in model_list, the finalizer leaves + `adaptive_routers` empty.""" + r = Router( + model_list=[ + { + "model_name": "fast", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + } + ] + ) + r._finalize_adaptive_router_if_configured() + assert r.adaptive_routers == {} diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_signals.py b/tests/test_litellm/router_strategy/adaptive_router/test_signals.py new file mode 100644 index 00000000000..2773c13a812 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_signals.py @@ -0,0 +1,180 @@ +import json +from pathlib import Path +from typing import List, Tuple + +import pytest + +from litellm.router_strategy.adaptive_router.config import TOOL_CALL_HISTORY_MAX +from litellm.router_strategy.adaptive_router.signals import ( + SessionState, + SignalDelta, + Turn, + apply_turn, +) + +FIXTURE_DIR = Path(__file__).parent / "fixtures" + + +def _load(name: str) -> list: + return json.loads((FIXTURE_DIR / f"{name}.json").read_text()) + + +def _replay(turns: list) -> Tuple[SessionState, List[SignalDelta]]: + state = SessionState( + session_id="s", + router_name="r", + model_name="m", + classified_type="general", + ) + deltas: List[SignalDelta] = [] + for t in turns: + deltas.append( + apply_turn( + state, + Turn( + user_content=t.get("user_content"), + assistant_content=t.get("assistant_content"), + tool_calls=t.get("tool_calls", []), + tool_results=t.get("tool_results", []), + response_status=t.get("response_status"), + ), + ) + ) + return state, deltas + + +def test_clean_satisfaction_fires_satisfaction_only(): + state, _ = _replay(_load("clean_satisfaction")) + assert state.satisfaction_count >= 1 + assert state.failure_count == 0 + assert state.disengagement_count == 0 + + +def test_misalignment_fires_on_rephrase(): + state, _ = _replay(_load("misalignment_rephrase")) + assert state.misalignment_count >= 1 + + +def test_stagnation_fires_on_repeated_assistant(): + state, _ = _replay(_load("stagnation_repeat")) + assert state.stagnation_count >= 1 + + +def test_disengagement_fires_on_giveup(): + state, _ = _replay(_load("disengagement_giveup")) + assert state.disengagement_count >= 1 + + +def test_failure_fires_on_tool_error(): + state, _ = _replay(_load("failure_tool_error")) + assert state.failure_count == 1 + + +def test_loop_fires_on_repeated_tool(): + state, _ = _replay(_load("loop_same_tool")) + assert state.loop_count >= 1 + + +@pytest.mark.parametrize("fixture", ["exhaustion_429", "exhaustion_context_overflow"]) +def test_exhaustion_fires_on_infra_signal(fixture): + state, _ = _replay(_load(fixture)) + assert state.exhaustion_count >= 1 + + +def test_no_signals_on_clean_session(): + state, _ = _replay(_load("clean_no_signals")) + assert state.misalignment_count == 0 + assert state.stagnation_count == 0 + assert state.disengagement_count == 0 + assert state.failure_count == 0 + assert state.loop_count == 0 + assert state.exhaustion_count == 0 + + +def test_mixed_failure_then_satisfaction(): + state, _ = _replay(_load("mixed_failure_then_satisfaction")) + assert state.failure_count >= 1 + assert state.satisfaction_count >= 1 + + +def test_satisfaction_gated_by_min_turns_for_clean_credit(): + """'thanks' on turn 1 is noise, not a validated quality signal.""" + state = SessionState( + session_id="s", router_name="r", model_name="m", classified_type="general" + ) + apply_turn(state, Turn(user_content="thanks!")) + assert state.satisfaction_count == 0 + assert state.clean_credit_awarded is False + assert state.last_processed_turn == 1 + + +def test_satisfaction_credit_awarded_once_per_session(): + """Even multiple satisfaction turns only award +1 alpha across the session.""" + state = SessionState( + session_id="s", router_name="r", model_name="m", classified_type="general" + ) + apply_turn(state, Turn(user_content="hi", assistant_content="hello")) + apply_turn(state, Turn(user_content="help me", assistant_content="sure")) + apply_turn(state, Turn(user_content="perfect, thanks")) + assert state.satisfaction_count == 1 + assert state.clean_credit_awarded is True + apply_turn(state, Turn(user_content="great, thank you")) + assert state.satisfaction_count == 1 + + +def test_empty_tool_content_does_not_fire_failure(): + """Zero-result searches / silent commands return empty but valid output.""" + state = SessionState( + session_id="s", router_name="r", model_name="m", classified_type="general" + ) + apply_turn( + state, + Turn( + tool_calls=[{"name": "grep", "arguments": {"q": "x"}}], + tool_results=[{"tool_call_id": "c1", "content": ""}], + ), + ) + apply_turn( + state, + Turn( + tool_calls=[{"name": "list", "arguments": {}}], + tool_results=[{"tool_call_id": "c2", "content": []}], + ), + ) + apply_turn( + state, + Turn( + tool_calls=[{"name": "noop", "arguments": {}}], + tool_results=[{"tool_call_id": "c3", "content": None}], + ), + ) + assert state.failure_count == 0 + + +def test_is_error_still_fires_failure(): + state = SessionState( + session_id="s", router_name="r", model_name="m", classified_type="general" + ) + apply_turn( + state, + Turn( + tool_calls=[{"name": "read", "arguments": {"p": "x"}}], + tool_results=[{"tool_call_id": "c1", "content": "boom", "is_error": True}], + ), + ) + assert state.failure_count == 1 + + +def test_apply_turn_is_o1_does_not_grow_history_unbounded(): + state = SessionState( + session_id="s", + router_name="r", + model_name="m", + classified_type="general", + ) + for i in range(100): + apply_turn( + state, + Turn(tool_calls=[{"name": f"tool_{i}", "arguments": {}}]), + ) + assert len(state.tool_call_history) <= TOOL_CALL_HISTORY_MAX diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py b/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py new file mode 100644 index 00000000000..753a449791b --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_state_endpoint.py @@ -0,0 +1,198 @@ +"""Tests for the GET /adaptive_router/state introspection endpoint and the +underlying `AdaptiveRouter.get_state_snapshot()` helper.""" + +import time +from unittest.mock import MagicMock + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter +from litellm.router_strategy.adaptive_router.bandit import BanditCell, apply_delta +from litellm.types.router import ( + AdaptiveRouterConfig, + AdaptiveRouterPreferences, + RequestType, +) + + +def _make_router(name: str = "r1") -> AdaptiveRouter: + cfg = AdaptiveRouterConfig(available_models=["fast", "smart"]) + prefs = { + "fast": AdaptiveRouterPreferences(quality_tier=1, strengths=[]), + "smart": AdaptiveRouterPreferences( + quality_tier=3, strengths=[RequestType.CODE_GENERATION] + ), + } + costs = {"fast": 0.0001, "smart": 0.001} + return AdaptiveRouter( + router_name=name, + config=cfg, + model_to_prefs=prefs, + model_to_cost=costs, + ) + + +# ---- snapshot helper --------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_state_snapshot_returns_cell_per_request_type_per_model(): + r = _make_router() + snap = await r.get_state_snapshot() + + # Top-level shape + assert snap["router_name"] == "r1" + assert snap["available_models"] == ["fast", "smart"] + assert snap["weights"] == {"quality": 0.7, "cost": 0.3} + assert snap["model_costs"] == {"fast": 0.0001, "smart": 0.001} + assert snap["owner_cache_live"] == 0 + assert snap["skipped_updates_total"] == 0 + assert set(snap["queue"].keys()) == { + "state_pending", + "session_pending", + "max_state_seen", + "max_session_seen", + } + + # 7 request types x 2 models = 14 cells + assert len(snap["cells"]) == len(list(RequestType)) * 2 + for cell in snap["cells"]: + assert set(cell.keys()) == { + "request_type", + "model", + "alpha", + "beta", + "samples", + "quality_mean", + } + assert cell["model"] in {"fast", "smart"} + assert cell["request_type"] in {rt.value for rt in RequestType} + + +@pytest.mark.asyncio +async def test_get_state_snapshot_quality_mean_matches_alpha_over_total(): + r = _make_router() + + # Manually mutate one cell to a known state so the math is verifiable. + key = (RequestType.CODE_GENERATION, "smart") + r._cells[key] = apply_delta(r._cells[key], delta_alpha=10.0, delta_beta=0.0) + expected = r._cells[key] + expected_mean = expected.alpha / (expected.alpha + expected.beta) + + snap = await r.get_state_snapshot() + cell = next( + c + for c in snap["cells"] + if c["request_type"] == "code_generation" and c["model"] == "smart" + ) + assert cell["alpha"] == expected.alpha + assert cell["beta"] == expected.beta + # `samples` reports net observations after subtracting the cold-start + # prior mass, so operators aren't misled by the initial value. + assert cell["samples"] == expected.total_samples + assert cell["quality_mean"] == pytest.approx(expected_mean) + + +@pytest.mark.asyncio +async def test_get_state_snapshot_counts_only_live_owner_cache_entries(): + r = _make_router() + now = time.time() + r._owner_cache["live-1"] = ("fast", now + 3600) + r._owner_cache["live-2"] = ("smart", now + 3600) + r._owner_cache["expired-1"] = ("fast", now - 1) + + snap = await r.get_state_snapshot() + assert snap["owner_cache_live"] == 2 + + +@pytest.mark.asyncio +async def test_get_state_snapshot_exposes_skipped_updates_total(): + r = _make_router() + r._skipped_updates_total = 7 + snap = await r.get_state_snapshot() + assert snap["skipped_updates_total"] == 7 + + +# ---- endpoint -------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_endpoint_returns_404_when_no_adaptive_router(monkeypatch): + """When llm_router is set but has no adaptive routers configured, return 404.""" + from litellm.proxy import proxy_server + + fake_router = MagicMock() + fake_router.adaptive_routers = {} + monkeypatch.setattr(proxy_server, "llm_router", fake_router) + + admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN) + with pytest.raises(HTTPException) as exc: + await proxy_server.get_adaptive_router_state(user_api_key_dict=admin) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_endpoint_returns_404_when_llm_router_is_none(monkeypatch): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "llm_router", None) + + admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN) + with pytest.raises(HTTPException) as exc: + await proxy_server.get_adaptive_router_state(user_api_key_dict=admin) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_endpoint_rejects_non_admin_role(monkeypatch): + from litellm.proxy import proxy_server + + fake_router = MagicMock() + fake_router.adaptive_routers = {"r1": _make_router()} + monkeypatch.setattr(proxy_server, "llm_router", fake_router) + + non_admin = UserAPIKeyAuth( + api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER + ) + with pytest.raises(HTTPException) as exc: + await proxy_server.get_adaptive_router_state(user_api_key_dict=non_admin) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_endpoint_returns_snapshot_list_for_admin(monkeypatch): + """Single configured router still returns the {"routers": [...]} list shape.""" + from litellm.proxy import proxy_server + + fake_router = MagicMock() + fake_router.adaptive_routers = {"r1": _make_router("r1")} + monkeypatch.setattr(proxy_server, "llm_router", fake_router) + + admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN) + result = await proxy_server.get_adaptive_router_state(user_api_key_dict=admin) + assert list(result.keys()) == ["routers"] + assert len(result["routers"]) == 1 + snap = result["routers"][0] + assert snap["router_name"] == "r1" + assert snap["available_models"] == ["fast", "smart"] + assert len(snap["cells"]) == len(list(RequestType)) * 2 + + +@pytest.mark.asyncio +async def test_endpoint_returns_one_snapshot_per_router(monkeypatch): + """With multiple adaptive routers configured, return one snapshot per router.""" + from litellm.proxy import proxy_server + + fake_router = MagicMock() + fake_router.adaptive_routers = { + "r1": _make_router("r1"), + "r2": _make_router("r2"), + } + monkeypatch.setattr(proxy_server, "llm_router", fake_router) + + admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN) + result = await proxy_server.get_adaptive_router_state(user_api_key_dict=admin) + names = sorted(s["router_name"] for s in result["routers"]) + assert names == ["r1", "r2"] diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_update_queue.py b/tests/test_litellm/router_strategy/adaptive_router/test_update_queue.py new file mode 100644 index 00000000000..9baa69a19e0 --- /dev/null +++ b/tests/test_litellm/router_strategy/adaptive_router/test_update_queue.py @@ -0,0 +1,117 @@ +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.router_strategy.adaptive_router.update_queue import ( + AdaptiveRouterUpdateQueue, +) + + +@pytest.fixture +def queue(): + return AdaptiveRouterUpdateQueue() + + +@pytest.fixture +def mock_prisma(): + """Prisma client with both adaptive router models stubbed as AsyncMocks.""" + p = MagicMock() + p.db.litellm_adaptiverouterstate.find_unique = AsyncMock(return_value=None) + p.db.litellm_adaptiverouterstate.upsert = AsyncMock() + p.db.litellm_adaptiveroutersession.upsert = AsyncMock() + return p + + +@pytest.mark.asyncio +async def test_add_state_delta_aggregates_same_key(queue): + await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0) + await queue.add_state_delta("r1", "general", "gpt-4", 0.0, 1.0) + sizes = await queue.queue_size() + assert sizes["state_pending"] == 1 + + +@pytest.mark.asyncio +async def test_add_state_delta_separate_keys(queue): + await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0) + await queue.add_state_delta("r1", "writing", "gpt-4", 1.0, 0.0) + sizes = await queue.queue_size() + assert sizes["state_pending"] == 2 + + +@pytest.mark.asyncio +async def test_add_session_state_last_write_wins(queue): + await queue.add_session_state("s1", "r1", "gpt-4", {"misalignment_count": 1}) + await queue.add_session_state("s1", "r1", "gpt-4", {"misalignment_count": 5}) + sizes = await queue.queue_size() + assert sizes["session_pending"] == 1 + + flushed = [] + p = MagicMock() + + async def upsert(**kwargs): + flushed.append(kwargs) + + p.db.litellm_adaptiveroutersession.upsert = upsert + await queue.flush_session_to_db(p) + assert len(flushed) == 1 + assert flushed[0]["data"]["update"]["misalignment_count"] == 5 + + +@pytest.mark.asyncio +async def test_flush_state_drains_aggregator(queue, mock_prisma): + await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0) + await queue.add_state_delta("r1", "writing", "gpt-4", 0.0, 1.0) + n = await queue.flush_state_to_db(mock_prisma) + assert n == 2 + sizes = await queue.queue_size() + assert sizes["state_pending"] == 0 + + +@pytest.mark.asyncio +async def test_flush_state_sums_correctly(queue, mock_prisma): + await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0) + await queue.add_state_delta("r1", "general", "gpt-4", 2.0, 1.0) + await queue.flush_state_to_db(mock_prisma) + # find_unique returned None (cold start), so alpha = 1+2 = 3, beta = 0+1 = 1 + call = mock_prisma.db.litellm_adaptiverouterstate.upsert.call_args + assert call.kwargs["data"]["create"]["alpha"] == 3.0 + assert call.kwargs["data"]["create"]["beta"] == 1.0 + assert call.kwargs["data"]["create"]["total_samples"] == 2 + + +@pytest.mark.asyncio +async def test_flush_session_drains_aggregator(queue, mock_prisma): + await queue.add_session_state("s1", "r1", "gpt-4", {"classified_type": "general"}) + n = await queue.flush_session_to_db(mock_prisma) + assert n == 1 + sizes = await queue.queue_size() + assert sizes["session_pending"] == 0 + + +@pytest.mark.asyncio +async def test_flush_empty_queue_returns_zero(queue, mock_prisma): + assert await queue.flush_state_to_db(mock_prisma) == 0 + assert await queue.flush_session_to_db(mock_prisma) == 0 + + +@pytest.mark.asyncio +async def test_flush_state_isolation_from_concurrent_adds(queue, mock_prisma): + """Adds during a flush should land in the NEW aggregator, not the drained batch.""" + await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0) + flush_task = asyncio.create_task(queue.flush_state_to_db(mock_prisma)) + # Yield control so the flush task can swap the aggregator before we add again. + await asyncio.sleep(0) + await queue.add_state_delta("r1", "general", "gpt-5", 2.0, 0.0) + await flush_task + sizes = await queue.queue_size() + assert sizes["state_pending"] == 1 + + +@pytest.mark.asyncio +async def test_max_size_observability(queue): + await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0) + await queue.add_state_delta("r1", "writing", "gpt-4", 1.0, 0.0) + await queue.add_state_delta("r1", "code_generation", "gpt-4", 1.0, 0.0) + sizes = await queue.queue_size() + assert sizes["max_state_seen"] >= 3 diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/test_litellm/router_strategy/test_auto_router.py index caff2bc8f10..cb46a4ae553 100644 --- a/tests/test_litellm/router_strategy/test_auto_router.py +++ b/tests/test_litellm/router_strategy/test_auto_router.py @@ -12,7 +12,148 @@ sys.path.insert( from litellm.router_strategy.auto_router.auto_router import AutoRouter -pytestmark = pytest.mark.skip(reason="Skipping auto router tests - beta feature") +pytestmark_skip_beta = pytest.mark.skip( + reason="Skipping auto router tests - beta feature" +) + + +class TestExtractTextFromMessages: + """Tests for AutoRouter._extract_text_from_messages (no semantic_router dependency).""" + + def test_should_extract_content_from_simple_user_message(self): + messages = [{"role": "user", "content": "Hello world"}] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "Hello world" + + def test_should_extract_last_user_message_from_tool_call_conversation(self): + messages = [ + {"role": "user", "content": "What's the weather in NYC?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "NYC"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_abc123", + "content": "72°F and sunny", + }, + {"role": "user", "content": "Now tell me about London"}, + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "Now tell me about London" + + def test_should_find_user_message_when_last_message_is_assistant_with_tool_calls( + self, + ): + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + }, + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "What's the weather?" + + def test_should_find_user_message_when_last_message_is_tool_response(self): + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_abc", + "content": "72°F and sunny", + }, + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "What's the weather?" + + def test_should_handle_multimodal_content_list(self): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/img.png"}, + }, + ], + } + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "What's in this image?" + + def test_should_handle_multimodal_content_with_multiple_text_blocks(self): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "First part"}, + {"type": "text", "text": "Second part"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/img.png"}, + }, + ], + } + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "First part Second part" + + def test_should_return_empty_string_when_user_content_is_none(self): + messages = [{"role": "user", "content": None}] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "" + + def test_should_return_empty_string_when_no_user_messages(self): + messages = [ + {"role": "system", "content": "You are a helpful assistant"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + }, + ] + result = AutoRouter._extract_text_from_messages(messages) + assert result == "" + + def test_should_return_empty_string_for_empty_messages_list(self): + result = AutoRouter._extract_text_from_messages([]) + assert result == "" @pytest.fixture @@ -41,6 +182,7 @@ def mock_route_choice(): return mock_choice +@pytestmark_skip_beta class TestAutoRouter: """Test class for AutoRouter methods.""" diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 8d36fc2ba32..e68ea863d82 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -7,7 +7,7 @@ Tests the rule-based complexity scoring and tier assignment logic. import os import sys from typing import Dict, List -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -828,3 +828,222 @@ class TestRouterComplexityDeploymentMethods: ) router.init_complexity_router_deployment(deployment) assert "auto_router/complexity_router/test-router" in router.complexity_routers + + +class TestAsyncPreRoutingHookMultiFormat: + """Test async_pre_routing_hook with multiple input formats.""" + + @pytest.mark.asyncio + async def test_should_route_with_chat_completions_messages(self, complexity_router): + """Test routing with standard chat completions messages.""" + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "What is 2+2?"}], + ) + assert result is not None + assert result.model is not None + assert result.messages is not None + + @pytest.mark.asyncio + async def test_should_route_with_responses_api_string_input( + self, complexity_router + ): + """Test routing with Responses API string input via handler dispatch.""" + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + from litellm.types.utils import CallTypes + + mock_mappings = {CallTypes.responses: OpenAIResponsesHandler} + + with patch( + "litellm.llms.load_guardrail_translation_mappings", + return_value=mock_mappings, + ): + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={"input": "What is the capital of France?"}, + messages=None, + input="What is the capital of France?", + ) + + assert result is not None + assert result.model is not None + # messages should be None since the original request didn't have messages + assert result.messages is None + + @pytest.mark.asyncio + async def test_should_route_with_responses_api_list_input(self, complexity_router): + """Test routing with Responses API list input via handler dispatch.""" + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + from litellm.types.utils import CallTypes + + mock_mappings = {CallTypes.responses: OpenAIResponsesHandler} + + list_input = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"}, + { + "role": "user", + "content": "Write a Python function to sort a list using merge sort", + }, + ] + + with patch( + "litellm.llms.load_guardrail_translation_mappings", + return_value=mock_mappings, + ): + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={"input": list_input}, + messages=None, + input=list_input, + ) + + assert result is not None + assert result.model is not None + assert result.messages is None + + @pytest.mark.asyncio + async def test_should_use_route_based_inference(self, complexity_router): + """Test that route-based call type inference is used when available.""" + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + from litellm.types.utils import CallTypes + + mock_mappings = {CallTypes.responses: OpenAIResponsesHandler} + + with patch( + "litellm.llms.load_guardrail_translation_mappings", + return_value=mock_mappings, + ): + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={ + "input": "Roll 2d4+1", + "litellm_metadata": { + "user_api_key_request_route": "/v1/responses", + }, + }, + messages=None, + ) + + assert result is not None + assert result.model is not None + + @pytest.mark.asyncio + async def test_should_return_none_when_no_messages_or_input( + self, complexity_router + ): + """Test that None is returned when neither messages nor input is available.""" + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=None, + input=None, + ) + assert result is None + + @pytest.mark.asyncio + async def test_should_prefer_original_messages_over_conversion( + self, complexity_router + ): + """Test that original messages are used when both messages and input are available.""" + messages = [{"role": "user", "content": "What is 2+2?"}] + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={"input": "This should be ignored"}, + messages=messages, + ) + assert result is not None + assert result.messages == messages + + @pytest.mark.asyncio + async def test_should_include_instructions_in_classification( + self, complexity_router + ): + """Test that Responses API instructions influence classification via system message.""" + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + from litellm.types.utils import CallTypes + + mock_mappings = {CallTypes.responses: OpenAIResponsesHandler} + + with patch( + "litellm.llms.load_guardrail_translation_mappings", + return_value=mock_mappings, + ): + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={ + "input": "Write merge sort", + "instructions": "You are an expert Python developer. Use advanced algorithms and optimize for performance.", + }, + messages=None, + ) + + assert result is not None + assert result.model is not None + + +class TestExtractUserMessageAndSystemPrompt: + """Test the _extract_user_message_and_system_prompt static method.""" + + def test_should_extract_user_message(self): + """Test extraction of the last user message.""" + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, + {"role": "user", "content": "How are you?"}, + ] + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( + messages + ) + assert user_msg == "How are you?" + assert sys_prompt == "You are helpful." + + def test_should_handle_no_user_message(self): + """Test when there is no user message.""" + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "assistant", "content": "Hi!"}, + ] + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( + messages + ) + assert user_msg is None + assert sys_prompt == "You are helpful." + + def test_should_handle_multipart_content(self): + """Test extraction from multipart content messages.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/img.png"}, + }, + ], + } + ] + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( + messages + ) + assert user_msg == "Describe this image" + assert sys_prompt is None + + def test_should_handle_empty_messages(self): + """Test with empty messages list.""" + user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt( + [] + ) + assert user_msg is None + assert sys_prompt is None diff --git a/tests/test_litellm/router_strategy/test_quality_router.py b/tests/test_litellm/router_strategy/test_quality_router.py new file mode 100644 index 00000000000..01574cb980d --- /dev/null +++ b/tests/test_litellm/router_strategy/test_quality_router.py @@ -0,0 +1,1033 @@ +""" +Tests for the QualityRouter. + +Covers: +- Tier index construction from `model_info.litellm_routing_preferences`. +- Quality-tier resolution (exact, round-up, default fallback). +- Keyword override (match, tiebreaking by quality + price). +- Pre-routing hook end-to-end. +- Decision metadata stash + Router.set_response_headers lift. +""" + +import os +import sys +from typing import Any, Dict, List +from unittest.mock import MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.router_strategy.quality_router.config import ( + DEFAULT_COMPLEXITY_TO_QUALITY, +) +from litellm.router_strategy.quality_router.quality_router import QualityRouter + + +def _make_model_list(spec: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Build a router model_list from a compact spec. + + spec entry shape: { + "model_name": str, + "quality_tier": Optional[int], + "keywords": Optional[List[str]], + "order": Optional[int], + "input_cost_per_token": Optional[float], + } + If quality_tier is None, the deployment is created without + `litellm_routing_preferences`. + """ + out: List[Dict[str, Any]] = [] + for entry in spec: + model_info: Dict[str, Any] = {"id": f"id-{entry['model_name']}"} + if entry.get("quality_tier") is not None: + prefs: Dict[str, Any] = {"quality_tier": entry["quality_tier"]} + if "keywords" in entry: + prefs["keywords"] = entry["keywords"] + if "order" in entry: + prefs["order"] = entry["order"] + model_info["litellm_routing_preferences"] = prefs + if "input_cost_per_token" in entry: + model_info["input_cost_per_token"] = entry["input_cost_per_token"] + out.append( + { + "model_name": entry["model_name"], + "litellm_params": {"model": f"openai/{entry['model_name']}"}, + "model_info": model_info, + } + ) + return out + + +@pytest.fixture +def four_tier_model_list() -> List[Dict[str, Any]]: + """A standard haiku(1)/sonnet(2)/opus(3)/opus-next(4) model list.""" + return _make_model_list( + [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "sonnet", "quality_tier": 2}, + {"model_name": "opus", "quality_tier": 3}, + {"model_name": "opus-next", "quality_tier": 4}, + ] + ) + + +@pytest.fixture +def mock_router(four_tier_model_list): + """A MagicMock router preloaded with the four-tier model list.""" + router = MagicMock() + router.model_list = four_tier_model_list + return router + + +@pytest.fixture +def quality_router(mock_router) -> QualityRouter: + """Default QualityRouter wired to all four tiers.""" + config = { + "available_models": ["haiku", "sonnet", "opus", "opus-next"], + "complexity_to_quality": DEFAULT_COMPLEXITY_TO_QUALITY, + } + return QualityRouter( + model_name="quality-router-test", + litellm_router_instance=mock_router, + default_model="haiku", + quality_router_config=config, + ) + + +# ─── Tier index ───────────────────────────────────────────────────────────── + + +class TestTierIndex: + def test_builds_correct_tier_to_models_map(self, quality_router): + assert quality_router._tier_to_models == { + 1: ["haiku"], + 2: ["sonnet"], + 3: ["opus"], + 4: ["opus-next"], + } + + def test_ignores_models_not_in_available_models(self, four_tier_model_list): + # Add a model the config doesn't list — it should be ignored. + extra = _make_model_list([{"model_name": "ghost", "quality_tier": 5}]) + router = MagicMock() + router.model_list = four_tier_model_list + extra + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="haiku", + quality_router_config={ + "available_models": ["haiku", "sonnet", "opus", "opus-next"] + }, + ) + + for models in qr._tier_to_models.values(): + assert "ghost" not in models + + def test_raises_when_routing_preferences_missing(self): + # `sonnet` is in available_models but has no preferences. + ml = _make_model_list( + [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "sonnet", "quality_tier": None}, + ] + ) + router = MagicMock() + router.model_list = ml + + # Construction succeeds (tier index is lazy); the error surfaces on + # first use so the router entry doesn't have to appear after all of + # its referenced models in config.yaml. + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="haiku", + quality_router_config={"available_models": ["haiku", "sonnet"]}, + ) + with pytest.raises(ValueError, match="sonnet"): + _ = qr._tier_to_models + + +# ─── Resolve model for quality tier ───────────────────────────────────────── + + +class TestResolveModelForQualityTier: + def test_exact_match(self, quality_router): + assert quality_router._resolve_model_for_quality_tier(2) == "sonnet" + assert quality_router._resolve_model_for_quality_tier(4) == "opus-next" + + def test_rounds_up_when_tier_missing(self, mock_router): + # Available tiers: 1, 3, 4. Asking for 2 should round up to 3. + spec = [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "opus", "quality_tier": 3}, + {"model_name": "opus-next", "quality_tier": 4}, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="haiku", + quality_router_config={"available_models": ["haiku", "opus", "opus-next"]}, + ) + + assert qr._resolve_model_for_quality_tier(2) == "opus" + + def test_rounds_down_when_no_higher_tier_exists(self): + # Only tier 1 available. Asking for tier 4 rounds up (nothing), then + # rounds DOWN to the closest lower tier — tier 1. + spec = [{"model_name": "haiku", "quality_tier": 1}] + router = MagicMock() + router.model_list = _make_model_list(spec) + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="emergency-default", + quality_router_config={"available_models": ["haiku"]}, + ) + + assert qr._resolve_model_for_quality_tier(4) == "haiku" + + def test_rounds_down_prefers_closest_lower_tier(self): + # Available: 1, 2. Asking for 4 rounds down to tier 2 (not tier 1). + spec = [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "sonnet", "quality_tier": 2}, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="emergency-default", + quality_router_config={"available_models": ["haiku", "sonnet"]}, + ) + + assert qr._resolve_model_for_quality_tier(4) == "sonnet" + + def test_prefers_round_up_over_round_down(self): + # Available: 1, 3. Asking for 2 rounds UP to 3, not DOWN to 1. + spec = [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "opus", "quality_tier": 3}, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="emergency-default", + quality_router_config={"available_models": ["haiku", "opus"]}, + ) + + assert qr._resolve_model_for_quality_tier(2) == "opus" + + +# ─── RoutingPreferences validation ───────────────────────────────────────── + + +class TestRoutingPreferencesValidation: + def test_invalid_quality_tier_type_raises_clear_error(self): + # quality_tier must be an int — pass a non-coercible string. + ml = [ + { + "model_name": "haiku", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": { + "id": "id-haiku", + "litellm_routing_preferences": {"quality_tier": "not-an-int"}, + }, + } + ] + router = MagicMock() + router.model_list = ml + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="haiku", + quality_router_config={"available_models": ["haiku"]}, + ) + with pytest.raises(ValueError, match="invalid litellm_routing_preferences"): + _ = qr._tier_to_models + + +# ─── Config-ordering independence (lazy index build) ─────────────────────── + + +class TestConfigOrderingIndependence: + def test_router_can_be_instantiated_before_its_targets_exist(self): + # Build a router instance whose referenced model_list is EMPTY at + # construction time (simulating a config where the router entry + # appears before its target deployments). The tier index must not be + # built eagerly — it's deferred until first use. + router = MagicMock() + router.model_list = [] # <- targets haven't been added yet + + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="haiku", + quality_router_config={"available_models": ["haiku", "sonnet", "opus"]}, + ) + + # Now the targets come online. This mirrors the incremental add by + # `Router._create_deployment`. + router.model_list = _make_model_list( + [ + {"model_name": "haiku", "quality_tier": 1}, + {"model_name": "sonnet", "quality_tier": 2}, + {"model_name": "opus", "quality_tier": 3}, + ] + ) + + # First access triggers the index build and sees the full list. + assert qr._tier_to_models == { + 1: ["haiku"], + 2: ["sonnet"], + 3: ["opus"], + } + + +# ─── Router.set_model_list resets quality_routers (hot reload) ───────────── + + +class TestSetModelListResetsQualityRouters: + def test_set_model_list_clears_quality_routers_registry(self): + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "haiku", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + }, + "model_info": {"litellm_routing_preferences": {"quality_tier": 1}}, + }, + { + "model_name": "my-qr", + "litellm_params": { + "model": "auto_router/quality_router", + "quality_router_default_model": "haiku", + "quality_router_config": {"available_models": ["haiku"]}, + }, + }, + ] + ) + + assert "my-qr" in router.quality_routers + + # Hot-reload with a new model_list that doesn't define the router. + router.set_model_list( + [ + { + "model_name": "haiku", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + }, + } + ] + ) + + # Stale router from before must be cleared. + assert "my-qr" not in router.quality_routers + + +# ─── Pre-routing hook ─────────────────────────────────────────────────────── + + +class TestPreRoutingHook: + @pytest.mark.asyncio + async def test_simple_message_routes_to_tier_1(self, quality_router): + messages = [{"role": "user", "content": "hi"}] + resp = await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs={}, + messages=messages, + ) + assert resp is not None + assert resp.model == "haiku" + + @pytest.mark.asyncio + async def test_reasoning_message_routes_to_tier_4(self, quality_router): + # Two reasoning markers triggers ComplexityTier.REASONING → quality 4. + messages = [ + { + "role": "user", + "content": ( + "Think step by step and reason through this problem. " + "Analyze this carefully and break down each component." + ), + } + ] + resp = await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs={}, + messages=messages, + ) + assert resp is not None + assert resp.model == "opus-next" + + @pytest.mark.asyncio + async def test_empty_messages_returns_none(self, quality_router): + resp = await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs={}, + messages=[], + ) + assert resp is None + + @pytest.mark.asyncio + async def test_only_system_message_routes_to_default(self, quality_router): + messages = [{"role": "system", "content": "You are a helpful assistant."}] + resp = await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs={}, + messages=messages, + ) + assert resp is not None + assert resp.model == "haiku" # the configured default_model + + +# ─── Keyword override ────────────────────────────────────────────────────── + + +@pytest.fixture +def keyword_router(): + """ + Router where multiple deployments declare overlapping keywords so we can + exercise the (quality DESC, price ASC) tiebreak. + + - cheap-coder tier 2, keywords [code, python], cost 0.000001 + - smart-coder tier 3, keywords [code, python], cost 0.000010 + - law-bot tier 2, keywords [legal, contract], cost 0.000005 + - default-haiku tier 1, no keywords, cost 0.0000005 + """ + spec = [ + { + "model_name": "default-haiku", + "quality_tier": 1, + "keywords": [], + "input_cost_per_token": 0.0000005, + }, + { + "model_name": "cheap-coder", + "quality_tier": 2, + "keywords": ["code", "python"], + "input_cost_per_token": 0.000001, + }, + { + "model_name": "smart-coder", + "quality_tier": 3, + "keywords": ["code", "python"], + "input_cost_per_token": 0.000010, + }, + { + "model_name": "law-bot", + "quality_tier": 2, + "keywords": ["legal", "contract"], + "input_cost_per_token": 0.000005, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + return QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="default-haiku", + quality_router_config={ + "available_models": [ + "default-haiku", + "cheap-coder", + "smart-coder", + "law-bot", + ], + }, + ) + + +class TestKeywordOverride: + def test_no_keyword_in_message_returns_none(self, keyword_router): + assert keyword_router._keyword_override("hello there") is None + + def test_single_match_returns_that_model(self, keyword_router): + # Only law-bot declares "legal". + assert keyword_router._keyword_override("review this legal doc") == ( + "law-bot", + "legal", + ) + + def test_case_insensitive_match(self, keyword_router): + assert keyword_router._keyword_override("LEGAL question") == ( + "law-bot", + "legal", + ) + + def test_overlap_picks_highest_quality_tier(self, keyword_router): + # Both cheap-coder (tier 2) and smart-coder (tier 3) declare "code". + # Quality wins over price → smart-coder. + assert keyword_router._keyword_override("write some code for me") == ( + "smart-coder", + "code", + ) + + def test_same_tier_picks_cheapest(self): + # Two models at the same tier, both matching "data" — cheapest wins. + spec = [ + { + "model_name": "expensive", + "quality_tier": 2, + "keywords": ["data"], + "input_cost_per_token": 0.000050, + }, + { + "model_name": "cheap", + "quality_tier": 2, + "keywords": ["data"], + "input_cost_per_token": 0.000005, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="cheap", + quality_router_config={"available_models": ["expensive", "cheap"]}, + ) + match = qr._keyword_override("show me the data") + assert match == ("cheap", "data") + + def test_unpriced_loses_to_priced_at_same_tier(self): + # Same quality tier, one has cost, one doesn't → priced wins. + spec = [ + { + "model_name": "no-price", + "quality_tier": 2, + "keywords": ["data"], + # input_cost_per_token deliberately omitted + }, + { + "model_name": "with-price", + "quality_tier": 2, + "keywords": ["data"], + "input_cost_per_token": 0.000005, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="no-price", + quality_router_config={"available_models": ["no-price", "with-price"]}, + ) + match = qr._keyword_override("show me the data") + assert match == ("with-price", "data") + + @pytest.mark.asyncio + async def test_hook_short_circuits_complexity_on_keyword_match( + self, keyword_router + ): + # A reasoning-style prompt would normally route to a high-quality model + # via the complexity flow — but the keyword "code" should short-circuit + # to smart-coder (highest tier among "code" models). + messages = [ + { + "role": "user", + "content": ( + "Think step by step and reason through this code problem. " + "Analyze this carefully and break down each component." + ), + } + ] + request_kwargs: Dict[str, Any] = {} + resp = await keyword_router.async_pre_routing_hook( + model="qr", + request_kwargs=request_kwargs, + messages=messages, + ) + assert resp is not None + assert resp.model == "smart-coder" + + decision = request_kwargs["metadata"]["quality_router_decision"] + assert decision["routed_via"] == "keyword" + assert decision["matched_keyword"] == "code" + assert decision["complexity_tier"] is None # short-circuited + + def test_quality_wins_over_explicit_order(self): + # Quality always beats order. A tier-3 model with no `order` wins over + # a tier-2 model with `order=1`. + spec = [ + { + "model_name": "ordered-tier2", + "quality_tier": 2, + "keywords": ["code"], + "order": 1, + "input_cost_per_token": 0.000010, + }, + { + "model_name": "implicit-tier3", + "quality_tier": 3, + "keywords": ["code"], + "input_cost_per_token": 0.000005, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="ordered-tier2", + quality_router_config={ + "available_models": ["ordered-tier2", "implicit-tier3"] + }, + ) + match = qr._keyword_override("write some code") + assert match == ("implicit-tier3", "code") + + def test_order_breaks_tie_within_same_quality_tier(self): + # Two tier-3 models, both match "code". Lower `order` wins. + spec = [ + { + "model_name": "preferred", + "quality_tier": 3, + "keywords": ["code"], + "order": 1, + "input_cost_per_token": 0.000050, # more expensive + }, + { + "model_name": "default-tier3", + "quality_tier": 3, + "keywords": ["code"], + "input_cost_per_token": 0.000005, # cheaper + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="default-tier3", + quality_router_config={"available_models": ["preferred", "default-tier3"]}, + ) + match = qr._keyword_override("write some code") + assert match == ("preferred", "code") + + def test_explicit_order_overrides_price(self): + # Same tier, but the more expensive one has a lower `order` and wins. + spec = [ + { + "model_name": "expensive-but-preferred", + "quality_tier": 2, + "keywords": ["data"], + "order": 1, + "input_cost_per_token": 0.000050, + }, + { + "model_name": "cheap-default", + "quality_tier": 2, + "keywords": ["data"], + "input_cost_per_token": 0.000005, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="cheap-default", + quality_router_config={ + "available_models": ["expensive-but-preferred", "cheap-default"] + }, + ) + match = qr._keyword_override("show me the data") + assert match == ("expensive-but-preferred", "data") + + def test_lower_order_wins_between_two_explicitly_ordered(self): + spec = [ + { + "model_name": "second", + "quality_tier": 2, + "keywords": ["data"], + "order": 5, + }, + { + "model_name": "first", + "quality_tier": 2, + "keywords": ["data"], + "order": 1, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="first", + quality_router_config={"available_models": ["first", "second"]}, + ) + match = qr._keyword_override("show me the data") + assert match == ("first", "data") + + def test_same_order_falls_through_to_quality_then_price(self): + # All three models share order=1 → tiebreak falls through to + # (quality DESC, cost ASC). + spec = [ + { + "model_name": "low-tier", + "quality_tier": 1, + "keywords": ["data"], + "order": 1, + "input_cost_per_token": 0.000001, + }, + { + "model_name": "high-tier-cheap", + "quality_tier": 3, + "keywords": ["data"], + "order": 1, + "input_cost_per_token": 0.000005, + }, + { + "model_name": "high-tier-expensive", + "quality_tier": 3, + "keywords": ["data"], + "order": 1, + "input_cost_per_token": 0.000050, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="low-tier", + quality_router_config={ + "available_models": [ + "low-tier", + "high-tier-cheap", + "high-tier-expensive", + ] + }, + ) + match = qr._keyword_override("show me the data") + assert match == ("high-tier-cheap", "data") + + def test_order_is_used_in_tier_resolution_too(self): + # Two models at the same tier. Explicit `order=1` on the second one + # should make _resolve_model_for_quality_tier(2) pick it. + spec = [ + { + "model_name": "default-pick", + "quality_tier": 2, + }, + { + "model_name": "preferred-pick", + "quality_tier": 2, + "order": 1, + }, + ] + router = MagicMock() + router.model_list = _make_model_list(spec) + qr = QualityRouter( + model_name="qr", + litellm_router_instance=router, + default_model="default-pick", + quality_router_config={ + "available_models": ["default-pick", "preferred-pick"] + }, + ) + assert qr._resolve_model_for_quality_tier(2) == "preferred-pick" + + @pytest.mark.asyncio + async def test_hook_falls_back_to_complexity_when_no_keyword(self, keyword_router): + # No declared keyword in the message → complexity-based routing. + # "hi" is SIMPLE → quality 1 → default-haiku (the only tier-1 model). + messages = [{"role": "user", "content": "hi"}] + request_kwargs: Dict[str, Any] = {} + resp = await keyword_router.async_pre_routing_hook( + model="qr", + request_kwargs=request_kwargs, + messages=messages, + ) + assert resp is not None + assert resp.model == "default-haiku" + + decision = request_kwargs["metadata"]["quality_router_decision"] + assert decision["routed_via"] == "quality_tier" + assert decision["matched_keyword"] is None + assert decision["complexity_tier"] == "SIMPLE" + + +# ─── Routing-decision metadata (powers x-litellm-quality-router-* headers) ── + + +class TestDecisionMetadata: + @pytest.mark.asyncio + async def test_hook_stashes_decision_in_request_kwargs_metadata( + self, quality_router + ): + # Reasoning prompt → REASONING → quality tier 4 → opus-next. + messages = [ + { + "role": "user", + "content": ( + "Think step by step and reason through this problem. " + "Analyze this carefully and break down each component." + ), + } + ] + request_kwargs: Dict[str, Any] = {} + + resp = await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs=request_kwargs, + messages=messages, + ) + assert resp is not None and resp.model == "opus-next" + + decision = request_kwargs["metadata"]["quality_router_decision"] + assert decision["routed_model"] == "opus-next" + assert decision["quality_tier"] == 4 + assert decision["complexity_tier"] == "REASONING" + assert decision["router_model_name"] == "quality-router-test" + assert decision["routed_via"] == "quality_tier" + assert decision["matched_keyword"] is None + + @pytest.mark.asyncio + async def test_decision_metadata_preserves_existing_metadata(self, quality_router): + request_kwargs: Dict[str, Any] = { + "metadata": {"trace_id": "abc-123", "user_id": "u-1"} + } + + await quality_router.async_pre_routing_hook( + model="quality-router-test", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "hi"}], + ) + + # Existing metadata keys are intact and the decision is added alongside. + assert request_kwargs["metadata"]["trace_id"] == "abc-123" + assert request_kwargs["metadata"]["user_id"] == "u-1" + assert "quality_router_decision" in request_kwargs["metadata"] + + +# ─── Router.set_response_headers lifts decision into x-litellm-quality-* ──── + + +class TestSetResponseHeadersLiftsDecision: + """ + Verify the Router.set_response_headers helper turns a stashed quality-router + decision into x-litellm-quality-router-* headers on the response. + """ + + @pytest.mark.asyncio + async def test_lifts_decision_into_additional_headers(self): + from pydantic import BaseModel + + from litellm.router import Router + + class FakeResponse(BaseModel): + model_config = {"arbitrary_types_allowed": True} + _hidden_params: Dict[str, Any] = {} + + # Build a real Router with a tiny model_list — enough to satisfy + # set_response_headers without needing the rest of the router stack. + router = Router( + model_list=[ + { + "model_name": "haiku", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + }, + } + ] + ) + + response = FakeResponse() + response._hidden_params = {} + + request_kwargs = { + "metadata": { + "quality_router_decision": { + "router_model_name": "qr", + "routed_model": "smart-coder", + "routed_via": "keyword", + "matched_keyword": "code", + "quality_tier": 3, + "complexity_tier": None, + } + } + } + + await router.set_response_headers( + response=response, + model_group="qr", + request_kwargs=request_kwargs, + ) + + headers = response._hidden_params["additional_headers"] + assert headers["x-litellm-quality-router-model"] == "smart-coder" + assert headers["x-litellm-quality-router-tier"] == "3" + assert headers["x-litellm-quality-router-via"] == "keyword" + assert headers["x-litellm-quality-router-keyword"] == "code" + # Keyword route short-circuits classification → no complexity header. + assert "x-litellm-quality-router-complexity" not in headers + # Existing x-litellm-model-group behavior is unchanged. + assert headers["x-litellm-model-group"] == "qr" + + @pytest.mark.asyncio + async def test_quality_tier_route_emits_complexity_not_keyword(self): + from pydantic import BaseModel + + from litellm.router import Router + + class FakeResponse(BaseModel): + model_config = {"arbitrary_types_allowed": True} + _hidden_params: Dict[str, Any] = {} + + router = Router( + model_list=[ + { + "model_name": "haiku", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + }, + } + ] + ) + + response = FakeResponse() + response._hidden_params = {} + + request_kwargs = { + "metadata": { + "quality_router_decision": { + "router_model_name": "qr", + "routed_model": "haiku", + "routed_via": "quality_tier", + "matched_keyword": None, + "quality_tier": 1, + "complexity_tier": "SIMPLE", + } + } + } + + await router.set_response_headers( + response=response, + model_group="qr", + request_kwargs=request_kwargs, + ) + + headers = response._hidden_params["additional_headers"] + assert headers["x-litellm-quality-router-via"] == "quality_tier" + assert headers["x-litellm-quality-router-complexity"] == "SIMPLE" + # Quality-tier route → no keyword header. + assert "x-litellm-quality-router-keyword" not in headers + + @pytest.mark.asyncio + async def test_no_decision_leaves_quality_router_headers_unset(self): + from pydantic import BaseModel + + from litellm.router import Router + + class FakeResponse(BaseModel): + model_config = {"arbitrary_types_allowed": True} + _hidden_params: Dict[str, Any] = {} + + router = Router( + model_list=[ + { + "model_name": "haiku", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + }, + } + ] + ) + + response = FakeResponse() + response._hidden_params = {} + + await router.set_response_headers( + response=response, + model_group="haiku", + request_kwargs={}, # no quality_router_decision + ) + + headers = response._hidden_params["additional_headers"] + assert "x-litellm-quality-router-model" not in headers + assert "x-litellm-quality-router-tier" not in headers + + +class TestRouterQualityDeploymentMethods: + """Tests for Router._is_quality_router_deployment and Router.init_quality_router_deployment.""" + + def test_is_quality_router_deployment_true(self): + """_is_quality_router_deployment returns True for quality router models.""" + from litellm.router import Router + from litellm.types.router import LiteLLM_Params + + router = Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + } + ] + ) + params = LiteLLM_Params(model="auto_router/quality_router/my-router") + assert router._is_quality_router_deployment(params) is True + + def test_is_quality_router_deployment_false(self): + """_is_quality_router_deployment returns False for regular models.""" + from litellm.router import Router + from litellm.types.router import LiteLLM_Params + + router = Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + } + ] + ) + params = LiteLLM_Params(model="openai/gpt-4o-mini") + assert router._is_quality_router_deployment(params) is False + + def test_init_quality_router_deployment(self): + """init_quality_router_deployment registers a QualityRouter.""" + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params + + router = Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + } + ] + ) + deployment = Deployment( + model_name="auto_router/quality_router/test-router", + litellm_params=LiteLLM_Params( + model="auto_router/quality_router/test-router", + quality_router_default_model="gpt-4o-mini", + ), + model_info={"id": "test-id"}, + ) + router.init_quality_router_deployment(deployment) + assert "auto_router/quality_router/test-router" in router.quality_routers diff --git a/tests/test_litellm/test_compression.py b/tests/test_litellm/test_compression.py index 13dda0cbcbc..4fbcd4ed30d 100644 --- a/tests/test_litellm/test_compression.py +++ b/tests/test_litellm/test_compression.py @@ -3,6 +3,7 @@ Unit tests for litellm.compress(). """ import os +import importlib import pytest @@ -12,6 +13,10 @@ from litellm.compression.scoring.embedding_scorer import embedding_score_message from litellm.compression.content_detection import detect_content_type from litellm.compression.message_stubbing import extract_key, stub_message from litellm.compression.retrieval_tool import build_retrieval_tool +from litellm.types.utils import CallTypes + +CALL_TYPE = CallTypes.completion +ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages # --------------------------------------------------------------------------- @@ -149,7 +154,7 @@ def test_retrieval_tool_description_lists_keys(): def test_compress_below_trigger_passthrough(): messages = [{"role": "user", "content": "hello"}] - result = litellm.compress(messages, model="gpt-4o") + result = litellm.compress(messages, model="gpt-4o", call_type=CALL_TYPE) assert result["messages"] == messages assert result["cache"] == {} assert result["tools"] == [] @@ -178,6 +183,7 @@ def test_compress_above_trigger(): result = litellm.compress( big_messages, model="gpt-4o", + call_type=CALL_TYPE, compression_trigger=1000, compression_target=500, ) @@ -189,13 +195,62 @@ def test_compress_above_trigger(): assert result["tools"][0]["function"]["name"] == "litellm_content_retrieve" +def test_compress_anthropic_list_content_is_boundary_stable(): + messages = [ + {"role": "system", "content": [{"type": "text", "text": "System prompt"}]}, + { + "role": "user", + "content": [ + {"type": "text", "text": "# a.py\n" + "alpha " * 2000}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/a.png"}, + }, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "# b.py\n" + "beta " * 2000}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/b.png"}, + }, + ], + }, + { + "role": "user", + "content": [{"type": "text", "text": "Fix alpha bug in a.py"}], + }, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=1000, + compression_target=500, + ) + + assert result["compressed_tokens"] < result["original_tokens"] + assert len(result["messages"]) == len(messages) + assert [m["role"] for m in result["messages"]] == [m["role"] for m in messages] + assert len(result["cache"]) > 0 + assert len(result["tools"]) == 1 + assert result["tools"][0]["type"] == "custom" + assert result["tools"][0]["name"] == "litellm_content_retrieve" + assert "input_schema" in result["tools"][0] + + def test_compress_preserves_system_message(): messages = [ {"role": "system", "content": "System prompt. " * 500}, {"role": "user", "content": "Large file content. " * 5000}, {"role": "user", "content": "Fix the bug"}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) assert result["messages"][0]["role"] == "system" assert "System prompt" in result["messages"][0]["content"] @@ -205,7 +260,9 @@ def test_compress_preserves_last_user_message(): {"role": "user", "content": "Big context " * 5000}, {"role": "user", "content": "Fix the bug in auth.py"}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) last_user = [m for m in result["messages"] if m["role"] == "user"][-1] assert "Fix the bug in auth.py" in last_user["content"] @@ -216,7 +273,9 @@ def test_compress_preserves_last_assistant_message(): {"role": "assistant", "content": "I'll help with that. " * 2000}, {"role": "user", "content": "Now fix the bug"}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) assistant_msgs = [m for m in result["messages"] if m["role"] == "assistant"] assert len(assistant_msgs) >= 1 # The last assistant message should be preserved (not stubbed) @@ -229,7 +288,9 @@ def test_cache_keys_match_stubs(): {"role": "user", "content": "# auth.py\n" + "code " * 5000}, {"role": "user", "content": "Fix it"}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) if result["tools"]: tool_desc = result["tools"][0]["function"]["description"] for key in result["cache"]: @@ -242,11 +303,75 @@ def test_compress_default_target(): {"role": "user", "content": "content " * 5000}, {"role": "user", "content": "query"}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=2000) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=2000 + ) # Should have compressed — target = 1000 assert result["compressed_tokens"] <= result["original_tokens"] +def test_compress_nested_tool_result_extracts_text_only(): + messages = [ + {"role": "system", "content": [{"type": "text", "text": "System rules"}]}, + { + "role": "user", + "content": [ + {"type": "text", "text": "prefix"}, + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [ + {"type": "text", "text": "nested text fragment"}, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/secret-tool.png", + }, + }, + ], + }, + { + "type": "image_url", + "image_url": {"url": "https://example.com/top.png"}, + }, + {"type": "text", "text": " " + ("irrelevant " * 3000)}, + ], + }, + { + "role": "user", + "content": [{"type": "text", "text": "final query that must remain"}], + }, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=500, + compression_target=100, + ) + + cached_text = " ".join(result["cache"].values()) + assert "nested text fragment" in cached_text + assert "https://example.com/secret-tool.png" not in cached_text + assert "https://example.com/top.png" not in cached_text + + +def test_compress_default_call_type_is_completion(): + result = litellm.compress( + messages=[ + {"role": "user", "content": "Large context " * 4000}, + {"role": "user", "content": "query"}, + ], + model="gpt-4o", + compression_trigger=1000, + compression_target=500, + ) + + assert result["compressed_tokens"] <= result["original_tokens"] + assert isinstance(result["tools"], list) + + def test_compress_forwards_embedding_model_params(monkeypatch): captured = {} @@ -269,6 +394,7 @@ def test_compress_forwards_embedding_model_params(monkeypatch): {"role": "user", "content": "Fix auth"}, ], model="gpt-4o", + call_type=CALL_TYPE, compression_trigger=1000, embedding_model="text-embedding-3-small", embedding_model_params={"api_base": "https://example-embeddings.test"}, @@ -326,6 +452,7 @@ def test_embedding_scorer(): {"role": "user", "content": "Fix auth"}, ], model="gpt-4o", + call_type=CALL_TYPE, compression_trigger=1000, embedding_model="text-embedding-3-small", ) @@ -346,8 +473,9 @@ def test_simple_compression(final_user_message, expected_content): {"role": "user", "content": "Unrelated cooking recipes " * 2000}, {"role": "user", "content": final_user_message}, ] - result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000) - print(result["messages"]) + result = litellm.compress( + messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000 + ) if expected_content == "Unrelated cooking recipes ": assert "Unrelated cooking recipes " in result["messages"][1]["content"] assert "Authentication code " not in result["messages"][0]["content"] @@ -356,3 +484,184 @@ def test_simple_compression(final_user_message, expected_content): assert "Unrelated cooking recipes " not in result["messages"][1]["content"] else: raise ValueError(f"Unexpected expected_content: {expected_content}") + + +def test_compress_anthropic_drops_irrelevant_tool_exchange_span(monkeypatch): + compress_module = importlib.import_module("litellm.compression.compress") + + def fake_bm25_score_messages(query, messages): + assert "final query" in query + assert len(messages) == 5 + # Prefer idx=0 and de-prioritize the tool exchange span (idx=1,2) + return [0.95, 0.01, 0.02, 0.8, 1.0] + + def fake_token_counter(model, messages=None, text=None): + if messages is not None: + return 1000 + if text is None: + return 0 + if "final query" in text: + return 50 + if "assistant_tail" in text: + return 20 + if "other_blob" in text: + return 220 + if "tool_payload_relevant" in text: + return 200 + if text == "": + return 1 + return 10 + + monkeypatch.setattr( + compress_module, "bm25_score_messages", fake_bm25_score_messages + ) + monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) + + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_drop", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_drop", + "content": [{"type": "text", "text": "tool_payload_relevant"}], + } + ], + }, + {"role": "assistant", "content": "assistant_tail"}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + # idx=1,2 should be dropped atomically (no orphan tool blocks left behind) + assert len(result["messages"]) == 3 + assert result["messages"][0]["role"] == "user" + assert "other_blob" in result["messages"][0]["content"] + assert result["messages"][1]["content"] == "assistant_tail" + assert result["messages"][2]["content"] == "final query" + assert result["cache"] == {} + + +def test_compress_anthropic_keeps_relevant_tool_exchange_span(monkeypatch): + compress_module = importlib.import_module("litellm.compression.compress") + + def fake_bm25_score_messages(query, messages): + assert "final query" in query + assert len(messages) == 5 + # Prefer the tool exchange span over idx=0 + return [0.05, 0.01, 0.92, 0.8, 1.0] + + def fake_token_counter(model, messages=None, text=None): + if messages is not None: + return 1000 + if text is None: + return 0 + if "final query" in text: + return 50 + if "assistant_tail" in text: + return 20 + if "other_blob" in text: + return 220 + if "tool_payload_relevant" in text: + return 200 + if text == "": + return 1 + return 10 + + monkeypatch.setattr( + compress_module, "bm25_score_messages", fake_bm25_score_messages + ) + monkeypatch.setattr(compress_module, "token_counter", fake_token_counter) + + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_keep", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_keep", + "content": [{"type": "text", "text": "tool_payload_relevant"}], + } + ], + }, + {"role": "assistant", "content": "assistant_tail"}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + assert len(result["messages"]) == 5 + assert result["messages"][1]["role"] == "assistant" + assert result["messages"][2]["role"] == "user" + # idx=0 should be compressed instead + assert "litellm_content_retrieve" in result["messages"][0]["content"] + assert len(result["cache"]) == 1 + + +def test_compress_anthropic_malformed_tool_sequence_passes_through(): + messages = [ + {"role": "user", "content": "other_blob " * 300}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_broken", + "name": "litellm_content_retrieve", + "input": {"key": "message_1"}, + } + ], + }, + {"role": "user", "content": [{"type": "text", "text": "missing tool_result"}]}, + {"role": "user", "content": "final query"}, + ] + + result = litellm.compress( + messages=messages, + model="claude-sonnet-4-20250514", + call_type=ANTHROPIC_CALL_TYPE, + compression_trigger=100, + compression_target=280, + ) + + assert result["messages"] == messages + assert result["cache"] == {} + assert result["tools"] == [] + assert result["compression_skipped_reason"] == "invalid_anthropic_tool_sequence" diff --git a/uv.lock b/uv.lock index d99da67fb82..1d449012d94 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-04-16T02:00:05.930008Z" +exclude-newer = "2026-04-19T01:10:36.69677Z" exclude-newer-span = "P3D" [manifest] @@ -3085,7 +3085,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.83.10" +version = "1.83.11" source = { editable = "." } dependencies = [ { name = "aiohttp" },