mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_batch_enqueued_token_limit
This commit is contained in:
commit
d5ac49588a
388 changed files with 24525 additions and 11732 deletions
|
|
@ -1,15 +1,17 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
category="${1:?usage: classify_changes.sh <backend|client>}"
|
||||
category="${1:?usage: classify_changes.sh <backend|client|ui>}"
|
||||
|
||||
has_client=false
|
||||
has_backend=false
|
||||
has_ci=false
|
||||
while IFS= read -r file || [ -n "$file" ]; do
|
||||
[ -n "$file" ] || continue
|
||||
case "$file" in
|
||||
ui/* | tests/e2e/ui/*) has_client=true ;;
|
||||
docs/* | *.md | *.mdx) : ;;
|
||||
.github/* | .circleci/*) has_ci=true; has_backend=true ;;
|
||||
*) has_backend=true ;;
|
||||
esac
|
||||
done
|
||||
|
|
@ -21,6 +23,9 @@ case "$category" in
|
|||
client)
|
||||
{ [ "$has_client" = true ] || [ "$has_backend" = true ]; } && echo run || echo skip
|
||||
;;
|
||||
ui)
|
||||
{ [ "$has_client" = true ] || [ "$has_ci" = true ]; } && echo run || echo skip
|
||||
;;
|
||||
*)
|
||||
echo run
|
||||
;;
|
||||
|
|
|
|||
2
.github/CODEOWNERS
vendored
2
.github/CODEOWNERS
vendored
|
|
@ -1,3 +1,5 @@
|
|||
/ui/ @yuneng-jiang @ryan-crabbe-berri
|
||||
/litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri
|
||||
/ui/litellm-dashboard/src/lib/http/schema.d.ts
|
||||
/model_prices_and_context_window.json @mateo-berri
|
||||
/litellm/model_prices_and_context_window_backup.json @mateo-berri
|
||||
|
|
|
|||
|
|
@ -1,48 +0,0 @@
|
|||
name: "Detect backend-relevant changes"
|
||||
description: >-
|
||||
Classify the pull request's changed files with .circleci/scripts/classify_changes.sh
|
||||
and expose decision=run|skip. decision=skip means only ui/**, **.md or **.mdx files
|
||||
changed, so callers can short-circuit expensive steps while the job still completes
|
||||
successfully and satisfies its required status check. The decision defaults to run for
|
||||
any non pull_request event or whenever the changed set cannot be resolved, so tests are
|
||||
never skipped when the classification is uncertain.
|
||||
|
||||
outputs:
|
||||
decision:
|
||||
description: "run when backend-relevant files changed, otherwise skip"
|
||||
value: ${{ steps.classify.outputs.decision }}
|
||||
|
||||
runs:
|
||||
using: composite
|
||||
steps:
|
||||
- id: classify
|
||||
shell: bash
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
set -uo pipefail
|
||||
if [ -z "${BASE_SHA:-}" ]; then
|
||||
echo "detect-backend-changes: not a pull_request event; running job"
|
||||
echo "decision=run" >> "${GITHUB_OUTPUT}"
|
||||
exit 0
|
||||
fi
|
||||
if ! git fetch --no-tags --depth=1 origin "${BASE_SHA}" >/dev/null 2>&1; then
|
||||
echo "detect-backend-changes: could not fetch base ${BASE_SHA}; running job"
|
||||
echo "decision=run" >> "${GITHUB_OUTPUT}"
|
||||
exit 0
|
||||
fi
|
||||
changed="$(git diff --name-only "${BASE_SHA}" HEAD 2>/dev/null)" || {
|
||||
echo "detect-backend-changes: git diff failed; running job"
|
||||
echo "decision=run" >> "${GITHUB_OUTPUT}"
|
||||
exit 0
|
||||
}
|
||||
if [ -z "${changed}" ]; then
|
||||
echo "detect-backend-changes: no changed files vs ${BASE_SHA}; skipping job"
|
||||
echo "decision=skip" >> "${GITHUB_OUTPUT}"
|
||||
exit 0
|
||||
fi
|
||||
echo "detect-backend-changes: changed files vs ${BASE_SHA}:"
|
||||
printf '%s\n' "${changed}" | sed 's/^/ /'
|
||||
decision="$(printf '%s\n' "${changed}" | bash .circleci/scripts/classify_changes.sh backend)" || decision="run"
|
||||
echo "detect-backend-changes: decision=${decision}"
|
||||
echo "decision=${decision}" >> "${GITHUB_OUTPUT}"
|
||||
41
.github/actions/detect-changes/action.yml
vendored
Normal file
41
.github/actions/detect-changes/action.yml
vendored
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
name: "Detect relevant changes"
|
||||
description: >-
|
||||
Classify the pull request's changed files with .circleci/scripts/classify_changes.sh
|
||||
and expose decision=run|skip for one category. backend means anything outside ui/,
|
||||
docs/ and markdown; ui means the dashboard sources alone. decision=skip lets callers
|
||||
short-circuit expensive steps while the job still completes successfully and satisfies
|
||||
its required status check, which a paths: filter cannot do because a workflow that
|
||||
never starts never reports. The file list comes from the pull request itself rather
|
||||
than from a git diff, because the checked-out merge ref is recomputed as the base
|
||||
branch advances and would otherwise attribute the base branch's own commits to the
|
||||
pull request. The decision defaults to run for any non pull_request event or whenever
|
||||
the changed set cannot be resolved, so jobs are never skipped when the classification
|
||||
is uncertain.
|
||||
|
||||
inputs:
|
||||
category:
|
||||
description: "Which classification to apply: backend, client or ui"
|
||||
required: false
|
||||
default: backend
|
||||
github-token:
|
||||
description: "Token used to list the pull request's files; needs pull-requests: read"
|
||||
required: false
|
||||
default: ${{ github.token }}
|
||||
|
||||
outputs:
|
||||
decision:
|
||||
description: "run when category-relevant files changed, otherwise skip"
|
||||
value: ${{ steps.classify.outputs.decision }}
|
||||
|
||||
runs:
|
||||
using: composite
|
||||
steps:
|
||||
- id: classify
|
||||
shell: bash
|
||||
env:
|
||||
GH_TOKEN: ${{ inputs.github-token }}
|
||||
CATEGORY: ${{ inputs.category }}
|
||||
REPO: ${{ github.repository }}
|
||||
PR_NUMBER: ${{ github.event.pull_request.number }}
|
||||
CHANGED_FILE_COUNT: ${{ github.event.pull_request.changed_files }}
|
||||
run: bash "${GITHUB_ACTION_PATH}/../../scripts/detect_changes.sh"
|
||||
3
.github/pull_request_template.md
vendored
3
.github/pull_request_template.md
vendored
|
|
@ -53,7 +53,8 @@ After: the same request comes back with real token counts, so the dashboard show
|
|||
**Please complete all items before asking a LiteLLM maintainer to review your PR**
|
||||
|
||||
- [ ] I have added meaningful tests
|
||||
- [ ] My PR passes all CI/CD checks (e.g., lint, format, unit tests)
|
||||
- [ ] The handful of test files covering my change pass locally, e.g. `uv run pytest tests/test_litellm/<your_test_file>.py -v`. Leave the suites (`make test-unit-*`, `make test-unit`) to CI: it finishes in ~15 minutes where a laptop takes an hour or more
|
||||
- [ ] My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
|
||||
- [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem
|
||||
- [ ] I have received a Greptile **Confidence Score of at least 4/5** before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment `@greptileai` to re-request a review after pushing changes)
|
||||
|
||||
|
|
|
|||
42
.github/scripts/detect_changes.sh
vendored
Executable file
42
.github/scripts/detect_changes.sh
vendored
Executable file
|
|
@ -0,0 +1,42 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
readonly API_FILE_CEILING=3000
|
||||
readonly CATEGORY="${CATEGORY:-backend}"
|
||||
|
||||
decide() {
|
||||
echo "detect-changes[${CATEGORY}]: decision=$1"
|
||||
[ -z "${GITHUB_OUTPUT:-}" ] || echo "decision=$1" >>"${GITHUB_OUTPUT}"
|
||||
exit 0
|
||||
}
|
||||
|
||||
run_full() {
|
||||
echo "detect-changes[${CATEGORY}]: $1; running job"
|
||||
decide run
|
||||
}
|
||||
|
||||
here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
classify="${here}/../../.circleci/scripts/classify_changes.sh"
|
||||
|
||||
[ -n "${PR_NUMBER:-}" ] || run_full "not a pull_request event"
|
||||
[ -n "${REPO:-}" ] || run_full "no repository in the environment"
|
||||
|
||||
case "${CHANGED_FILE_COUNT:-}" in
|
||||
'' | *[!0-9]*) run_full "the event payload carries no changed_files count" ;;
|
||||
esac
|
||||
[ "${CHANGED_FILE_COUNT}" -le "${API_FILE_CEILING}" ] ||
|
||||
run_full "PR #${PR_NUMBER} changes ${CHANGED_FILE_COUNT} files, past the ${API_FILE_CEILING}-file listing ceiling"
|
||||
|
||||
changed="$(gh api "repos/${REPO}/pulls/${PR_NUMBER}/files" --paginate --jq '.[].filename')" ||
|
||||
run_full "could not list the files on PR #${PR_NUMBER}"
|
||||
[ -n "${changed}" ] || run_full "the API listed no files on PR #${PR_NUMBER}"
|
||||
|
||||
echo "detect-changes[${CATEGORY}]: files changed by PR #${PR_NUMBER}:"
|
||||
printf '%s\n' "${changed}" | sed 's/^/ /'
|
||||
|
||||
decision="$(printf '%s\n' "${changed}" | bash "${classify}" "${CATEGORY}")" ||
|
||||
run_full "classify_changes.sh failed"
|
||||
case "${decision}" in
|
||||
run | skip) decide "${decision}" ;;
|
||||
*) run_full "classify_changes.sh printed an unexpected decision: ${decision}" ;;
|
||||
esac
|
||||
10
.github/workflows/_test-unit-base.yml
vendored
10
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -60,6 +60,9 @@ jobs:
|
|||
name: Run tests
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: ${{ inputs.job-timeout-minutes }}
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
outputs:
|
||||
decision: ${{ steps.changes.outputs.decision }}
|
||||
|
||||
|
|
@ -69,24 +72,27 @@ jobs:
|
|||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect backend-relevant changes
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
timeout-minutes: 2
|
||||
uses: ./.github/actions/detect-backend-changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 3
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 3
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 5
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
|
|
|
|||
23
.github/workflows/test-linting.yml
vendored
23
.github/workflows/test-linting.yml
vendored
|
|
@ -24,6 +24,7 @@ jobs:
|
|||
# re-running basedpyright over the merge-base tree.
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
actions: read
|
||||
|
||||
steps:
|
||||
|
|
@ -37,7 +38,12 @@ jobs:
|
|||
clean: true
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
|
||||
- name: Fetch gate base (merge-base with target branch)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
|
|
@ -50,39 +56,47 @@ jobs:
|
|||
echo "GATE_BASE_SHA=$MERGE_BASE" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Clean Python cache
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
find . -type d -name "__pycache__" -exec rm -rf {} + || true
|
||||
find . -name "*.pyc" -delete || true
|
||||
|
||||
- name: Check uv.lock is up to date
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv lock --check || (echo "❌ uv.lock is out of sync with pyproject.toml. Run 'uv lock' locally and commit the result." && exit 1)
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv sync --frozen --group proxy-dev --group e2e-dev
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
# basedpyright resolves Prisma's generated client (litellm/proxy/schema.prisma)
|
||||
# only after `prisma generate` writes prisma/client.py et al. Without this the
|
||||
# DB wrappers typed against the generated client would degrade to Unknown.
|
||||
- name: Generate Prisma client
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Check ruff format
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
|
||||
if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then
|
||||
|
|
@ -92,6 +106,7 @@ jobs:
|
|||
xargs uv run --no-sync ruff format --check --exclude '/enterprise/' < "$RUNNER_TEMP/ruff_format_files.txt"
|
||||
|
||||
- name: Debug - Check file state
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
echo "Current branch:"
|
||||
git branch --show-current
|
||||
|
|
@ -101,30 +116,36 @@ jobs:
|
|||
head -50 litellm/litellm_core_utils/custom_logger_registry.py | tail -10
|
||||
|
||||
- name: Run Ruff linting
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
cd litellm
|
||||
uv run --no-sync ruff check .
|
||||
cd ..
|
||||
|
||||
- name: Check strict-rule budget (delta vs base)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync python scripts/ruff_strict_gate.py --base "$GATE_BASE_SHA"
|
||||
|
||||
- name: Check type-discipline budget (mutable collections / casts / type guards / kwargs / unexplained suppressions, delta vs base)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync python scripts/type_discipline_gate.py --base "$GATE_BASE_SHA"
|
||||
|
||||
- name: Print OpenAI version
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync python -c "import openai; print(f'OpenAI version: {openai.__version__}')"
|
||||
|
||||
- name: Check basedpyright budget (delta vs base)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
run: |
|
||||
uv run --no-sync python scripts/type_check_gate.py --base "$GATE_BASE_SHA"
|
||||
|
||||
- name: Check tests/e2e basedpyright (zero errors)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- 'tests/e2e/**/*.py' | grep -q .; then
|
||||
uv run --no-sync basedpyright tests/e2e
|
||||
|
|
@ -133,12 +154,14 @@ jobs:
|
|||
fi
|
||||
|
||||
- name: Check for circular imports
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
cd litellm
|
||||
uv run --no-sync python ../tests/documentation_tests/test_circular_imports.py
|
||||
cd ..
|
||||
|
||||
- name: Check import safety
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync python -c "from litellm import *" || (echo '🚨 import failed, this means you introduced unprotected imports! 🚨'; exit 1)
|
||||
|
||||
|
|
|
|||
10
.github/workflows/test-litellm-ui-build.yml
vendored
10
.github/workflows/test-litellm-ui-build.yml
vendored
|
|
@ -1,6 +1,7 @@
|
|||
name: UI Build Check
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
|
|
@ -28,7 +29,14 @@ jobs:
|
|||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
with:
|
||||
category: ui
|
||||
|
||||
- name: Setup Node.js
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version-file: ui/litellm-dashboard/.nvmrc
|
||||
|
|
@ -36,7 +44,9 @@ jobs:
|
|||
cache-dependency-path: ui/litellm-dashboard/package-lock.json
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: npm ci
|
||||
|
||||
- name: Build
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: npm run build
|
||||
|
|
|
|||
11
.github/workflows/test-litellm-ui-unit.yml
vendored
11
.github/workflows/test-litellm-ui-unit.yml
vendored
|
|
@ -1,6 +1,7 @@
|
|||
name: UI Unit Tests
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
|
|
@ -32,7 +33,14 @@ jobs:
|
|||
fetch-depth: 1
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
with:
|
||||
category: ui
|
||||
|
||||
- name: Setup Node.js
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version-file: ui/litellm-dashboard/.nvmrc
|
||||
|
|
@ -40,14 +48,17 @@ jobs:
|
|||
cache-dependency-path: ui/litellm-dashboard/package-lock.json
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: npm ci
|
||||
|
||||
- name: Run UI type tests (Vitest)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
CI: "true"
|
||||
run: npm run test:types
|
||||
|
||||
- name: Run UI unit tests (Vitest)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
CI: "true"
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
|
|
|
|||
9
.github/workflows/test-mcp.yml
vendored
9
.github/workflows/test-mcp.yml
vendored
|
|
@ -10,6 +10,7 @@ on:
|
|||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
|
|
@ -25,26 +26,34 @@ jobs:
|
|||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
|
||||
- name: Thank You Message
|
||||
run: |
|
||||
echo "### 🙏 Thank you for contributing to LiteLLM!" >> $GITHUB_STEP_SUMMARY
|
||||
echo "Your PR is being tested now. We appreciate your help in making LiteLLM better!" >> $GITHUB_STEP_SUMMARY
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv lock --check
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group proxy-dev --extra proxy --extra semantic-router
|
||||
|
||||
- name: Run MCP tests
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync pytest tests/mcp_tests -x -vv -n 4 --cov=./litellm --cov-report=xml --durations=5
|
||||
|
|
|
|||
15
.github/workflows/test-unit-documentation.yml
vendored
15
.github/workflows/test-unit-documentation.yml
vendored
|
|
@ -23,34 +23,41 @@ jobs:
|
|||
documentation:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-changes
|
||||
|
||||
- name: Checkout litellm-docs into docs/my-website (for documentation_tests)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
repository: BerriAI/litellm-docs
|
||||
path: docs/my-website
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect backend-relevant changes
|
||||
id: changes
|
||||
uses: ./.github/actions/detect-backend-changes
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
|
|
|
|||
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -1,9 +1,11 @@
|
|||
.python-version
|
||||
.venv
|
||||
tests/e2e/.fixtures/
|
||||
.venv-typecheck
|
||||
.venv_policy_test
|
||||
.env
|
||||
.claude
|
||||
CLAUDE.local.md
|
||||
.newenv
|
||||
newenv/*
|
||||
litellm/proxy/myenv/*
|
||||
|
|
|
|||
|
|
@ -13,8 +13,8 @@ Here are the core requirements for any PR submitted to LiteLLM:
|
|||
|
||||
- [ ] **Add testing** - Adding at least 1 test is a hard requirement - [see details](#adding-testing)
|
||||
- [ ] **Ensure your PR passes all checks**:
|
||||
- [ ] [Unit Tests](#running-unit-tests) - `make test-unit`
|
||||
- [ ] [Linting / Formatting](#running-linting-and-formatting-checks) - `make lint`
|
||||
- [ ] [The tests covering your change](#running-unit-tests) pass, e.g. `uv run pytest tests/test_litellm/<your_test_file>.py -v`. CI runs the full unit test matrix, so you don't need to run the whole suite locally
|
||||
|
||||
#### UI PRs
|
||||
|
||||
|
|
@ -71,8 +71,8 @@ make format
|
|||
# Run all linting checks (matches CI exactly)
|
||||
make lint
|
||||
|
||||
# Run unit tests to ensure nothing is broken
|
||||
make test-unit
|
||||
# Run the tests covering your change (CI runs the full suite)
|
||||
uv run pytest tests/test_litellm/<your_test_file>.py -v
|
||||
|
||||
# Commit your changes (must follow Conventional Commits — see above)
|
||||
git add .
|
||||
|
|
@ -123,12 +123,13 @@ def test_your_feature():
|
|||
|
||||
### Running Unit Tests
|
||||
|
||||
Run all unit tests (uses parallel execution for speed):
|
||||
|
||||
Run the tests covering your change:
|
||||
```bash
|
||||
make test-unit
|
||||
uv run pytest tests/test_litellm/test_your_file.py -v
|
||||
```
|
||||
|
||||
`tests/test_litellm` holds thousands of tests, so running all of it locally takes a long time. CI runs it as a parallel matrix (`make test-unit-llms`, `make test-unit-proxy-core`, and the other `test-unit-*` targets) on beefier boxes, so if, for whatever reason, you must run the whole suite, it's better to rely on CI to do that.
|
||||
|
||||
If you're running broader test suites, proxy tests, or anything that touches PostgreSQL-backed fixtures/plugins, install the full local test environment first:
|
||||
|
||||
```bash
|
||||
|
|
@ -137,11 +138,6 @@ make install-test-deps
|
|||
|
||||
This syncs the locked test environment used across the repo, including `psycopg` v3 plus `psycopg-binary` (used by `pytest-postgresql`), `psycopg2-binary` (used by some proxy E2E tests), and a generated Prisma client for DB-backed proxy tests, so pytest startup matches CI without manual package installs.
|
||||
|
||||
Run specific test files:
|
||||
```bash
|
||||
uv run pytest tests/test_litellm/test_your_file.py -v
|
||||
```
|
||||
|
||||
### Running Linting and Formatting Checks
|
||||
|
||||
Run all linting checks (matches CI exactly):
|
||||
|
|
|
|||
|
|
@ -57,7 +57,7 @@
|
|||
"limit": 5663
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15557
|
||||
"limit": 15555
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -105,13 +105,13 @@
|
|||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 39043
|
||||
"limit": 39017
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19887
|
||||
"limit": 19885
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 30574
|
||||
"limit": 30572
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 117
|
||||
|
|
|
|||
|
|
@ -145,6 +145,11 @@ NUMBER_KEYS: dict[str, JsonSchema] = {
|
|||
"minimum": 1,
|
||||
"description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).",
|
||||
},
|
||||
"regional_endpoint_uplift_multiplier": {
|
||||
"type": "number",
|
||||
"minimum": 1,
|
||||
"description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).",
|
||||
},
|
||||
}
|
||||
|
||||
COST_DESCRIPTIONS: dict[str, str] = {
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ type: application
|
|||
# This is the chart version. This version number should be incremented each time you make changes
|
||||
# to the chart and its templates, including the app version.
|
||||
# Versions are expected to follow Semantic Versioning (https://semver.org/)
|
||||
version: 1.1.1
|
||||
version: 1.1.2
|
||||
|
||||
# This is the version number of the application being deployed. This version number should be
|
||||
# incremented each time you make changes to the application. Versions are not expected to
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ If `db.useStackgresOperator` is used (not yet implemented):
|
|||
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A |
|
||||
| `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
|
||||
| `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
|
||||
| `image.repository` | LiteLLM Proxy image repository | `docker.litellm.ai/berriai/litellm` |
|
||||
| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` |
|
||||
| `image.pullPolicy` | LiteLLM Proxy image pull policy | `IfNotPresent` |
|
||||
| `image.tag` | Overrides the image tag whose default the latest version of LiteLLM at the time this chart was published. | `""` |
|
||||
| `imagePullSecrets` | Registry credentials for the LiteLLM and initContainer images. | `[]` |
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ tests:
|
|||
pattern: -litellm$
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: ghcr.io/berriai/litellm-database:test
|
||||
value: ghcr.io/berriai/litellm:test
|
||||
- it: should work with tolerations
|
||||
template: deployment.yaml
|
||||
set:
|
||||
|
|
@ -337,7 +337,7 @@ tests:
|
|||
template: deployment.yaml
|
||||
set:
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
repository: ghcr.io/berriai/litellm
|
||||
tag: test
|
||||
extraInitContainers:
|
||||
- name: init-tpl
|
||||
|
|
@ -348,7 +348,7 @@ tests:
|
|||
path: spec.template.spec.initContainers
|
||||
content:
|
||||
name: init-tpl
|
||||
image: "ghcr.io/berriai/litellm-database:test"
|
||||
image: "ghcr.io/berriai/litellm:test"
|
||||
command: ["echo", "hello"]
|
||||
- it: should work with extraContainers
|
||||
template: deployment.yaml
|
||||
|
|
@ -366,7 +366,7 @@ tests:
|
|||
template: deployment.yaml
|
||||
set:
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
repository: ghcr.io/berriai/litellm
|
||||
tag: test
|
||||
extraContainers:
|
||||
- name: sidecar-tpl
|
||||
|
|
@ -376,12 +376,12 @@ tests:
|
|||
path: spec.template.spec.containers
|
||||
content:
|
||||
name: sidecar-tpl
|
||||
image: "ghcr.io/berriai/litellm-database:test"
|
||||
image: "ghcr.io/berriai/litellm:test"
|
||||
- it: should support tpl in podAnnotations
|
||||
template: deployment.yaml
|
||||
set:
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
repository: ghcr.io/berriai/litellm
|
||||
tag: test
|
||||
# Mirrors the real-world scenario this feature unblocks:
|
||||
# user disables the built-in ConfigMap (and its built-in checksum/config
|
||||
|
|
@ -398,7 +398,7 @@ tests:
|
|||
value: "test"
|
||||
- equal:
|
||||
path: spec.template.metadata.annotations["example.com/some-key"]
|
||||
value: "ghcr.io/berriai/litellm-database"
|
||||
value: "ghcr.io/berriai/litellm"
|
||||
- equal:
|
||||
path: spec.template.metadata.annotations["example.com/literal"]
|
||||
value: "plain-string-value"
|
||||
|
|
|
|||
|
|
@ -208,7 +208,7 @@ tests:
|
|||
template: migrations-job.yaml
|
||||
set:
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
repository: ghcr.io/berriai/litellm
|
||||
tag: test
|
||||
migrationJob:
|
||||
enabled: true
|
||||
|
|
@ -221,7 +221,7 @@ tests:
|
|||
path: spec.template.spec.initContainers
|
||||
content:
|
||||
name: init-tpl
|
||||
image: "ghcr.io/berriai/litellm-database:test"
|
||||
image: "ghcr.io/berriai/litellm:test"
|
||||
command: ["echo", "hello"]
|
||||
- it: should work with extraContainers
|
||||
template: migrations-job.yaml
|
||||
|
|
@ -241,7 +241,7 @@ tests:
|
|||
template: migrations-job.yaml
|
||||
set:
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
repository: ghcr.io/berriai/litellm
|
||||
tag: test
|
||||
migrationJob:
|
||||
enabled: true
|
||||
|
|
@ -253,7 +253,7 @@ tests:
|
|||
path: spec.template.spec.containers
|
||||
content:
|
||||
name: sidecar-tpl
|
||||
image: "ghcr.io/berriai/litellm-database:test"
|
||||
image: "ghcr.io/berriai/litellm:test"
|
||||
- it: should render the pod-level securityContext from podSecurityContext
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
|
|
|
|||
|
|
@ -6,8 +6,9 @@ replicaCount: 1
|
|||
# numWorkers: 2
|
||||
|
||||
image:
|
||||
# Use "ghcr.io/berriai/litellm-database" for optimized image with database
|
||||
repository: ghcr.io/berriai/litellm-database
|
||||
# Bundles the prisma CLI and engines, which is what lets the migrations job
|
||||
# and the proxy's own schema check run without network access.
|
||||
repository: ghcr.io/berriai/litellm
|
||||
pullPolicy: Always
|
||||
# Overrides the image tag whose default is the chart appVersion.
|
||||
# tag: "latest"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,9 @@
|
|||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_ProxyWorkerHeartbeat" (
|
||||
"worker_id" TEXT NOT NULL,
|
||||
"hostname" TEXT NOT NULL,
|
||||
"started_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"last_heartbeat_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_ProxyWorkerHeartbeat_pkey" PRIMARY KEY ("worker_id")
|
||||
);
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "group_id" TEXT;
|
||||
|
||||
UPDATE "LiteLLM_ShadowEvalJob" SET "group_id" = "id" WHERE "group_id" IS NULL;
|
||||
|
||||
ALTER TABLE "LiteLLM_ShadowEvalJob" ALTER COLUMN "group_id" SET NOT NULL;
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_group_id_idx" ON "LiteLLM_ShadowEvalJob"("group_id");
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN "stopped_by" TEXT;
|
||||
|
||||
UPDATE "LiteLLM_ShadowEvalJob" SET stopped_by = 'unknown'
|
||||
WHERE stopped_at IS NOT NULL AND ends_at > (NOW() AT TIME ZONE 'utc');
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
UPDATE "LiteLLM_SpendLogs"
|
||||
SET "created_at" = "endTime",
|
||||
"updated_at" = "endTime"
|
||||
WHERE "created_at" > "endTime" + interval '1 hour';
|
||||
|
|
@ -947,6 +947,17 @@ model LiteLLM_DailyTagSpend {
|
|||
}
|
||||
|
||||
|
||||
// One row per live proxy worker process. Workers upsert their row on a fixed
|
||||
// heartbeat; counting rows with a recent heartbeat tells how many workers share
|
||||
// this database, which lets the Admin UI hide its "no Redis" warning for
|
||||
// deployments that are provably a single worker.
|
||||
model LiteLLM_ProxyWorkerHeartbeat {
|
||||
worker_id String @id
|
||||
hostname String
|
||||
started_at DateTime @default(now())
|
||||
last_heartbeat_at DateTime @default(now())
|
||||
}
|
||||
|
||||
// Track the status of cron jobs running. Only allow one pod to run the job at a time
|
||||
model LiteLLM_CronJob {
|
||||
cronjob_id String @id @default(cuid()) // Unique ID for the record
|
||||
|
|
@ -1467,28 +1478,38 @@ model LiteLLM_AutoRouterSession {
|
|||
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
|
||||
}
|
||||
|
||||
// Shadow eval: evaluation of an auto-router against a key's live traffic, in either
|
||||
// direction. forward duplicates the requests the key did not route through the router
|
||||
// through it, answering whether the key should adopt it; reverse duplicates the requests
|
||||
// the router did serve against a fixed baseline model, answering whether a key already on
|
||||
// it still benefits. Either way a sampled slice runs in a detached task and an LLM judge
|
||||
// compares real vs shadow responses blind. The job row is immutable config plus
|
||||
// stopped_at; every count, status, and spend figure is derived from the append-only
|
||||
// attempt rows, so nothing can disagree across pods or stop races.
|
||||
// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in
|
||||
// either direction. forward duplicates the requests the keys did not route through the
|
||||
// router through it, answering whether they should adopt it; reverse duplicates the
|
||||
// requests the router did serve against a fixed baseline model, answering whether a key
|
||||
// already on it still benefits. Either way a sampled slice runs in a detached task and an
|
||||
// LLM judge compares real vs shadow responses blind. Each row is ONE key's leg of a job:
|
||||
// immutable config plus that key's own turn budget and stop state, so one key exhausting
|
||||
// its budget never ends a sibling's sampling. A job is the set of legs sharing group_id
|
||||
// (the id the API reports), written together by one atomic create_many with identical
|
||||
// config; single-key jobs predating group_id were backfilled group_id = id. "One active
|
||||
// job per (key, direction)" is a partial unique index on (api_key_id, direction) WHERE
|
||||
// stopped_at IS NULL, expressed only in the migration because schema.prisma cannot state
|
||||
// partial indexes; it is what makes a concurrent start on another pod race-safe rather
|
||||
// than read-then-create. Every count, status, and spend figure is derived from the
|
||||
// append-only attempt rows, so nothing can disagree across pods or stop races.
|
||||
model LiteLLM_ShadowEvalJob {
|
||||
id String @id @default(cuid())
|
||||
api_key_id String // hashed virtual key whose traffic is shadowed
|
||||
group_id String // legs of one job share this; the API's job id
|
||||
api_key_id String // hashed virtual key whose traffic this leg shadows
|
||||
router_name String // the auto-router under evaluation, in either direction
|
||||
direction String @default("forward") // forward | reverse
|
||||
baseline_model String? // reverse only: the fixed model the router is judged against
|
||||
judge_model String
|
||||
shadow_percentage Float
|
||||
max_turns Int // sample budget: judge at most this many turns
|
||||
max_turns Int // this key's sample budget: judge at most this many turns
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
ends_at DateTime
|
||||
stopped_at DateTime?
|
||||
stopped_by String? // operator who stopped it early; null when it ended on its own
|
||||
|
||||
@@index([group_id])
|
||||
@@index([api_key_id])
|
||||
@@index([created_at])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -221,7 +221,7 @@ overwrite_user_with_key_hash: bool = (
|
|||
bedrock_request_metadata_fields: Optional[Sequence[str]] = (
|
||||
None # allow-list of `user_api_key_*` fields (+ `spend_logs_metadata`) sent as Bedrock `requestMetadata`
|
||||
)
|
||||
store_audit_logs = False # Enterprise feature, allow users to see audit logs
|
||||
store_audit_logs: bool | None = None
|
||||
skip_system_message_in_guardrail: bool = False
|
||||
skip_tool_message_in_guardrail: bool = False
|
||||
### end of callbacks #############
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import json
|
||||
from collections.abc import Iterable, Iterator
|
||||
from collections.abc import Iterable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
|
|
@ -87,7 +87,7 @@ async def _handle_completed_batch(
|
|||
return batch_cost, batch_usage, [model_name]
|
||||
|
||||
return _aggregate_batch_cost_usage_models(
|
||||
entries=_iter_batch_input_entries(file_content),
|
||||
entries=_iter_batch_output_entries(file_content),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
model_info=model_info,
|
||||
|
|
@ -111,43 +111,91 @@ def _iter_successful_output_line_stats(
|
|||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> Iterator[_BatchOutputLineStats]:
|
||||
for entry in entries:
|
||||
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info)
|
||||
if stats is not None:
|
||||
yield stats
|
||||
|
||||
|
||||
def _safe_output_line_stats(
|
||||
entry: Mapping[str, Any],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> _BatchOutputLineStats | None:
|
||||
"""Return the stats for one batch output line, or None for a line that is
|
||||
unsuccessful or cannot be costed, so a single bad line never aborts the
|
||||
whole batch's cost accounting."""
|
||||
custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None
|
||||
try:
|
||||
if not _batch_response_was_successful(entry, custom_llm_provider):
|
||||
return None
|
||||
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info)
|
||||
except Exception as e: # noqa: BLE001 # any single line's costing failure must not abort the whole batch
|
||||
verbose_logger.warning(
|
||||
"batch output line could not be costed, so it is billed at $0 and the rest of the batch "
|
||||
"is still billed. custom_id=%s error=%s",
|
||||
custom_id,
|
||||
str(e),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _compute_output_line_stats(
|
||||
entry: Mapping[str, Any],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> _BatchOutputLineStats:
|
||||
response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
|
||||
usage: Final = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
|
||||
prompt_details: Final = parse_prompt_tokens_details(usage)
|
||||
raw_model: Final = response_body.get("model")
|
||||
response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None
|
||||
return _BatchOutputLineStats(
|
||||
cost=_output_line_cost(
|
||||
response_body=response_body,
|
||||
usage=usage,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
response_model=response_model,
|
||||
model_info=model_info,
|
||||
),
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
completion_tokens=usage.completion_tokens,
|
||||
total_tokens=usage.total_tokens,
|
||||
cache_read_tokens=prompt_details["cache_hit_tokens"],
|
||||
cache_creation_tokens=prompt_details["cache_creation_tokens"],
|
||||
model=response_model,
|
||||
)
|
||||
|
||||
|
||||
def _output_line_cost(
|
||||
response_body: Mapping[str, Any],
|
||||
usage: Usage,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None,
|
||||
response_model: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> float:
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
|
||||
for entry in entries:
|
||||
if not _batch_response_was_successful(entry, custom_llm_provider):
|
||||
continue
|
||||
response_body = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
|
||||
usage = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
|
||||
prompt_details = parse_prompt_tokens_details(usage)
|
||||
raw_model = response_body.get("model")
|
||||
response_model = raw_model if isinstance(raw_model, str) and raw_model else None
|
||||
if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"):
|
||||
if custom_llm_provider == "bedrock" and model_name:
|
||||
cost_model = model_name
|
||||
else:
|
||||
cost_model = response_model or model_name or ""
|
||||
prompt_cost, completion_cost = batch_cost_calculator(
|
||||
usage=usage,
|
||||
model=cost_model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_info=model_info,
|
||||
)
|
||||
line_cost = prompt_cost + completion_cost
|
||||
else:
|
||||
line_cost = litellm.completion_cost(
|
||||
completion_response=response_body,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
call_type=CallTypes.aretrieve_batch.value,
|
||||
)
|
||||
yield _BatchOutputLineStats(
|
||||
cost=line_cost,
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
completion_tokens=usage.completion_tokens,
|
||||
total_tokens=usage.total_tokens,
|
||||
cache_read_tokens=prompt_details["cache_hit_tokens"],
|
||||
cache_creation_tokens=prompt_details["cache_creation_tokens"],
|
||||
model=response_model,
|
||||
if model_info is None and custom_llm_provider not in ("anthropic", "bedrock"):
|
||||
return litellm.completion_cost(
|
||||
completion_response=response_body,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
call_type=CallTypes.aretrieve_batch.value,
|
||||
)
|
||||
cost_model: Final = (
|
||||
model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or ""
|
||||
)
|
||||
prompt_cost, completion_cost = batch_cost_calculator(
|
||||
usage=usage,
|
||||
model=cost_model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_info=model_info,
|
||||
)
|
||||
return prompt_cost + completion_cost
|
||||
|
||||
|
||||
def _aggregate_batch_cost_usage_models(
|
||||
|
|
@ -338,9 +386,10 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict:
|
|||
|
||||
def _get_file_content_as_dictionary(file_content: bytes) -> list[dict]:
|
||||
"""
|
||||
Get the file content as a list of dictionaries from JSON Lines format
|
||||
Get the file content as a list of dictionaries from JSON Lines format,
|
||||
skipping malformed lines
|
||||
"""
|
||||
return list(_iter_batch_input_entries(file_content))
|
||||
return list(_iter_batch_output_entries(file_content))
|
||||
|
||||
|
||||
def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]:
|
||||
|
|
@ -361,15 +410,29 @@ def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]:
|
|||
yield line
|
||||
|
||||
|
||||
def _iter_batch_input_entries(file_content: bytes) -> Iterator[dict]:
|
||||
def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict]:
|
||||
"""
|
||||
Yield parsed batch input JSONL entries one at a time without materializing the
|
||||
whole file as a list, so peak memory stays bounded. Raises on a malformed line;
|
||||
callers that must survive bad rows should iterate ``_iter_batch_input_lines``
|
||||
and parse per-row instead.
|
||||
Yield parsed batch output JSONL entries one at a time without materializing
|
||||
the whole file as a list, so peak memory stays bounded. A malformed or
|
||||
non-object line is skipped with a warning so one bad line never aborts the
|
||||
whole batch's cost accounting.
|
||||
"""
|
||||
for line in _iter_batch_input_lines(file_content):
|
||||
yield json.loads(line)
|
||||
entry = _parse_batch_output_line(line)
|
||||
if entry is not None:
|
||||
yield entry
|
||||
|
||||
|
||||
def _parse_batch_output_line(line: bytes) -> dict | None:
|
||||
try:
|
||||
parsed: Final = json.loads(line)
|
||||
except ValueError as e:
|
||||
verbose_logger.warning("skipping malformed batch output line: %s", str(e))
|
||||
return None
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
verbose_logger.warning("skipping non-object batch output line of type %s", type(parsed).__name__)
|
||||
return None
|
||||
|
||||
|
||||
# A batch request's input tokens scale roughly with its serialized size, so this
|
||||
|
|
@ -440,7 +503,9 @@ def _count_prompt_or_input_tokens(model: str, value: Any) -> int:
|
|||
return 0
|
||||
|
||||
|
||||
def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_provider: str = "openai") -> Usage:
|
||||
def _get_batch_job_usage_from_response_body(
|
||||
response_body: Mapping[str, Any], custom_llm_provider: str = "openai"
|
||||
) -> Usage:
|
||||
"""
|
||||
Get the tokens of a batch job from the response body
|
||||
"""
|
||||
|
|
@ -472,7 +537,7 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov
|
|||
return usage
|
||||
|
||||
|
||||
def _get_anthropic_result_from_batch_results_line(batch_results_line: dict) -> dict:
|
||||
def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> dict:
|
||||
"""
|
||||
Get the ``result`` object from a line of an Anthropic message batch results JSONL file.
|
||||
|
||||
|
|
@ -482,7 +547,9 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: dict) -> d
|
|||
return batch_results_line.get("result", None) or {}
|
||||
|
||||
|
||||
def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom_llm_provider: str = "openai") -> Any:
|
||||
def _get_response_from_batch_job_output_file(
|
||||
batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
|
||||
) -> Any:
|
||||
"""
|
||||
Get the response from the batch job output file
|
||||
"""
|
||||
|
|
@ -495,7 +562,9 @@ def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom
|
|||
return _response_body
|
||||
|
||||
|
||||
def _batch_response_was_successful(batch_job_output_file: dict, custom_llm_provider: str = "openai") -> bool:
|
||||
def _batch_response_was_successful(
|
||||
batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the batch job response was successful
|
||||
|
||||
|
|
|
|||
|
|
@ -327,6 +327,8 @@ def cost_per_token(
|
|||
service_tier: str | None = None, # for OpenAI service tier pricing
|
||||
### DATA RESIDENCY ###
|
||||
data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us")
|
||||
### VERTEX LOCATION ###
|
||||
vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global")
|
||||
response: Any | None = None,
|
||||
### REQUEST MODEL ###
|
||||
request_model: str | None = None, # original request model for router detection
|
||||
|
|
@ -587,6 +589,7 @@ def cost_per_token(
|
|||
prompt_characters=prompt_characters,
|
||||
completion_characters=completion_characters,
|
||||
usage=usage_block,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
elif cost_router == "cost_per_token":
|
||||
return google_cost_per_token(
|
||||
|
|
@ -594,6 +597,7 @@ def cost_per_token(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
usage=usage_block,
|
||||
service_tier=service_tier,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
elif custom_llm_provider == "anthropic":
|
||||
return anthropic_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
|
||||
|
|
@ -1071,6 +1075,7 @@ def _store_cost_breakdown_in_logging_obj(
|
|||
reasoning_cost: float | None = None,
|
||||
service_tier: str | None = None,
|
||||
data_residency: str | None = None,
|
||||
vertex_location: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Helper function to store cost breakdown in the logging object.
|
||||
|
|
@ -1090,6 +1095,7 @@ def _store_cost_breakdown_in_logging_obj(
|
|||
margin_total_amount: Total margin added in USD
|
||||
service_tier: Tier the costs above were priced on, already resolved
|
||||
data_residency: Region uplift the costs above were priced on, already resolved
|
||||
vertex_location: Vertex AI location the costs above were priced on, already resolved
|
||||
"""
|
||||
if litellm_logging_obj is None:
|
||||
return
|
||||
|
|
@ -1113,6 +1119,7 @@ def _store_cost_breakdown_in_logging_obj(
|
|||
reasoning_cost=reasoning_cost,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
except Exception as breakdown_error:
|
||||
|
|
@ -1149,6 +1156,8 @@ def completion_cost(
|
|||
service_tier: str | None = None, # for OpenAI service tier pricing
|
||||
### DATA RESIDENCY ###
|
||||
data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us")
|
||||
### VERTEX LOCATION ###
|
||||
vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global")
|
||||
) -> float:
|
||||
"""
|
||||
Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm.
|
||||
|
|
@ -1577,6 +1586,7 @@ def completion_cost(
|
|||
rerank_billed_units=rerank_billed_units,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
response=completion_response,
|
||||
request_model=request_model_for_cost,
|
||||
)
|
||||
|
|
@ -1664,6 +1674,7 @@ def completion_cost(
|
|||
usage=cost_per_token_usage_object,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
_reasoning_cost = _token_type_breakdown.reasoning_cost
|
||||
_cache_read_cost = _token_type_breakdown.cache_read_cost
|
||||
|
|
@ -1686,6 +1697,7 @@ def completion_cost(
|
|||
reasoning_cost=_reasoning_cost,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
return _final_cost
|
||||
|
|
@ -1765,6 +1777,8 @@ def response_cost_calculator(
|
|||
service_tier: str | None = None, # for OpenAI service tier pricing
|
||||
### DATA RESIDENCY ###
|
||||
data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us")
|
||||
### VERTEX LOCATION ###
|
||||
vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global")
|
||||
) -> float:
|
||||
"""
|
||||
Returns
|
||||
|
|
@ -1797,6 +1811,7 @@ def response_cost_calculator(
|
|||
litellm_logging_obj=litellm_logging_obj,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
return response_cost
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ from litellm.types.utils import (
|
|||
# OpenTelemetry imports moved to individual functions to avoid import errors when not installed
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.sdk.resources import Resource as _Resource
|
||||
from opentelemetry.sdk.trace import TracerProvider as _SDKTracerProvider
|
||||
from opentelemetry.sdk.trace.export import SpanExporter as _SpanExporter
|
||||
from opentelemetry.trace import Context as _Context
|
||||
|
|
@ -389,6 +390,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
self._tracer_provider_cache: OrderedDict[str, _CachedTracerProvider] = OrderedDict()
|
||||
self._tracer_provider_cache_lock: Final = threading.Lock()
|
||||
self._max_dynamic_tracer_providers: Final = max(1, max_dynamic_tracer_providers)
|
||||
self._litellm_resource_memo: _Resource | None = None
|
||||
self._init_tracing(tracer_provider)
|
||||
|
||||
_debug_otel: Final = str(os.getenv("DEBUG_OTEL", "False")).lower()
|
||||
|
|
@ -414,7 +416,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
self._init_otel_logger_on_litellm_proxy()
|
||||
|
||||
@staticmethod
|
||||
def _get_litellm_resource(config: OpenTelemetryConfig):
|
||||
def _get_litellm_resource(config: OpenTelemetryConfig) -> "_Resource":
|
||||
"""Create an OpenTelemetry Resource using config-driven defaults."""
|
||||
from opentelemetry.sdk.resources import OTELResourceDetector, Resource
|
||||
|
||||
|
|
@ -429,6 +431,21 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
env_resource: Final = otel_resource_detector.detect()
|
||||
return base_resource.merge(env_resource)
|
||||
|
||||
def _litellm_resource(self) -> "_Resource":
|
||||
"""The Resource every provider on this logger is built with, frozen at first use.
|
||||
|
||||
``Resource.create`` scans every installed distribution's entry points, roughly 3ms and
|
||||
200 file opens, and the dynamic providers reach it from the async logging path. Freezing
|
||||
also keeps them consistent with whatever this logger built at startup. ``cached_property``
|
||||
locks class-wide before 3.12, which this file still supports.
|
||||
"""
|
||||
memo: Final = self._litellm_resource_memo
|
||||
if memo is not None:
|
||||
return memo
|
||||
built: Final = self._get_litellm_resource(self.config)
|
||||
self._litellm_resource_memo = built
|
||||
return built
|
||||
|
||||
def _init_otel_logger_on_litellm_proxy(self):
|
||||
"""
|
||||
Initializes OpenTelemetry for litellm proxy server
|
||||
|
|
@ -596,7 +613,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
from opentelemetry.trace import SpanKind
|
||||
|
||||
def create_tracer_provider():
|
||||
provider: Final = TracerProvider(resource=self._get_litellm_resource(self.config))
|
||||
provider: Final = TracerProvider(resource=self._litellm_resource())
|
||||
provider.add_span_processor(self._get_span_processor())
|
||||
return provider
|
||||
|
||||
|
|
@ -634,7 +651,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
metric_reader: Final = self._get_metric_reader()
|
||||
return MeterProvider(
|
||||
metric_readers=[metric_reader],
|
||||
resource=self._get_litellm_resource(self.config),
|
||||
resource=self._litellm_resource(),
|
||||
)
|
||||
|
||||
meter_provider = self._get_or_create_provider(
|
||||
|
|
@ -692,7 +709,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
|
||||
|
||||
def create_logger_provider():
|
||||
provider: Final = OTLoggerProvider(resource=self._get_litellm_resource(self.config))
|
||||
provider: Final = OTLoggerProvider(resource=self._litellm_resource())
|
||||
log_exporter: Final = self._get_log_exporter()
|
||||
provider.add_log_record_processor(BatchLogRecordProcessor(log_exporter))
|
||||
return provider
|
||||
|
|
@ -1144,9 +1161,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
owns_exporter: Final = _provider_owns_exporter(dynamic_config.exporter)
|
||||
|
||||
def _build() -> "_SDKTracerProvider":
|
||||
provider: Final = TracerProvider(
|
||||
resource=self._get_litellm_resource(self.config), shutdown_on_exit=owns_exporter
|
||||
)
|
||||
provider: Final = TracerProvider(resource=self._litellm_resource(), shutdown_on_exit=owns_exporter)
|
||||
provider.add_span_processor(self._get_span_processor(config_override=dynamic_config))
|
||||
return provider
|
||||
|
||||
|
|
@ -1162,9 +1177,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
owns_exporter: Final = _provider_owns_exporter(self.OTEL_EXPORTER)
|
||||
|
||||
def _build() -> "_SDKTracerProvider":
|
||||
provider: Final = TracerProvider(
|
||||
resource=self._get_litellm_resource(self.config), shutdown_on_exit=owns_exporter
|
||||
)
|
||||
provider: Final = TracerProvider(resource=self._litellm_resource(), shutdown_on_exit=owns_exporter)
|
||||
provider.add_span_processor(self._get_span_processor(dynamic_headers=dynamic_headers))
|
||||
return provider
|
||||
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ import traceback
|
|||
from collections.abc import Callable, Mapping, Sequence
|
||||
from datetime import datetime as dt_object
|
||||
from functools import lru_cache
|
||||
from types import TracebackType
|
||||
from types import MappingProxyType, TracebackType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast
|
||||
|
||||
from httpx import Response
|
||||
|
|
@ -372,6 +372,35 @@ def _published_pricing(deployment_model: str | None) -> ModelInfo | None:
|
|||
return None
|
||||
|
||||
|
||||
def _resolve_vertex_location_for_cost(
|
||||
custom_llm_provider: str | None,
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
optional_params: Mapping[str, object] | None,
|
||||
model: str,
|
||||
) -> str | None:
|
||||
"""
|
||||
The Vertex AI location a request was served from, resolved the same way
|
||||
dispatch resolves it, so regional deployments price with the
|
||||
regional-endpoint uplift. None for non-Vertex providers.
|
||||
|
||||
Chat dispatch reads the location from request kwargs, which reach this
|
||||
logging object through optional_params: on the proxy the logging object is
|
||||
created before the router picks a deployment, so the deployment's location
|
||||
never lands in litellm_params.
|
||||
"""
|
||||
if custom_llm_provider is None or not custom_llm_provider.startswith("vertex_ai"):
|
||||
return None
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
||||
empty: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
configured_location: Final = (
|
||||
VertexBase.explicit_vertex_ai_location(optional_params or empty)
|
||||
or VertexBase.explicit_vertex_ai_location(litellm_params or empty)
|
||||
or VertexBase.safe_get_vertex_ai_location(empty)
|
||||
)
|
||||
return VertexBase.get_vertex_region(configured_location, model)
|
||||
|
||||
|
||||
class Logging(LiteLLMLoggingBaseClass):
|
||||
global \
|
||||
supabaseClient, \
|
||||
|
|
@ -1432,6 +1461,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
reasoning_cost: float | None = None,
|
||||
service_tier: str | None = None,
|
||||
data_residency: str | None = None,
|
||||
vertex_location: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Helper method to store cost breakdown in the logging object.
|
||||
|
|
@ -1450,6 +1480,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
margin_total_amount: Total margin added in USD
|
||||
service_tier: Tier the costs above were priced on, already resolved
|
||||
data_residency: Region uplift the costs above were priced on, already resolved
|
||||
vertex_location: Vertex AI location the costs above were priced on, already resolved
|
||||
"""
|
||||
|
||||
self.cost_breakdown = CostBreakdown(
|
||||
|
|
@ -1459,6 +1490,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
tool_usage_cost=cost_for_built_in_tools_cost_usd_dollar,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
if cache_read_cost is not None and cache_read_cost > 0:
|
||||
self.cost_breakdown["cache_read_cost"] = cache_read_cost
|
||||
|
|
@ -1574,6 +1606,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if hasattr(self, "litellm_params") and self.litellm_params
|
||||
else None
|
||||
),
|
||||
"vertex_location": _resolve_vertex_location_for_cost(
|
||||
custom_llm_provider=self.model_call_details.get("custom_llm_provider", None),
|
||||
litellm_params=(self.litellm_params if hasattr(self, "litellm_params") else None),
|
||||
optional_params=self.optional_params,
|
||||
model=litellm_model_name or self.model,
|
||||
),
|
||||
}
|
||||
except Exception as e: # error creating kwargs for cost calculation
|
||||
debug_info = StandardLoggingModelCostFailureDebugInformation(
|
||||
|
|
|
|||
|
|
@ -757,6 +757,33 @@ def _get_regional_uplift_multiplier(model_info: ModelInfo, data_residency: str |
|
|||
return 1.0
|
||||
|
||||
|
||||
def get_vertex_regional_endpoint_uplift(model_info: ModelInfo, vertex_location: str | None) -> float:
|
||||
"""
|
||||
Resolve the per-model uplift multiplier for Vertex AI non-global (regional and
|
||||
multi-region) endpoints.
|
||||
|
||||
Google prices every non-global endpoint at a flat premium over the global
|
||||
endpoint (e.g. 1.10 = +10%) on all token types for the models that carry
|
||||
regional pricing. The multiplier is stored on the model entry as
|
||||
``regional_endpoint_uplift_multiplier``.
|
||||
|
||||
Returns 1.0 (no uplift) when ``vertex_location`` is ``None`` or ``"global"``,
|
||||
or when the model has no multiplier configured.
|
||||
"""
|
||||
if vertex_location is None or vertex_location.lower() == "global":
|
||||
return 1.0
|
||||
multiplier: Final = model_info.get("regional_endpoint_uplift_multiplier")
|
||||
if multiplier is None:
|
||||
return 1.0
|
||||
try:
|
||||
return float(cast(float, multiplier))
|
||||
except (TypeError, ValueError):
|
||||
verbose_logger.exception(
|
||||
"Invalid regional_endpoint_uplift_multiplier for model; defaulting to 1.0",
|
||||
)
|
||||
return 1.0
|
||||
|
||||
|
||||
def get_provider_specific_geo_multiplier(model_info: ModelInfo, usage: Usage) -> float:
|
||||
"""
|
||||
Resolve the provider-specific regional pricing multiplier for the geo the
|
||||
|
|
@ -798,6 +825,7 @@ def generic_cost_per_token(
|
|||
service_tier: str | None = None,
|
||||
data_residency: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
vertex_location: str | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
|
||||
|
|
@ -809,6 +837,9 @@ def generic_cost_per_token(
|
|||
- usage: LiteLLM Usage block, containing anthropic caching information
|
||||
- data_residency: optional OpenAI data-residency region (e.g. "eu", "us"),
|
||||
used to apply the per-model regional-processing uplift multiplier.
|
||||
- vertex_location: optional Vertex AI location the request was served from
|
||||
(e.g. "us-east5", "global"), used to apply the per-model
|
||||
regional-endpoint uplift multiplier when non-global.
|
||||
|
||||
Returns:
|
||||
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
|
||||
|
|
@ -968,6 +999,11 @@ def generic_cost_per_token(
|
|||
prompt_cost *= uplift
|
||||
completion_cost *= uplift
|
||||
|
||||
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)
|
||||
if vertex_uplift != 1.0:
|
||||
prompt_cost *= vertex_uplift
|
||||
completion_cost *= vertex_uplift
|
||||
|
||||
return prompt_cost, completion_cost
|
||||
|
||||
|
||||
|
|
@ -988,6 +1024,7 @@ def get_token_type_cost_breakdown(
|
|||
usage: Usage,
|
||||
service_tier: str | None = None,
|
||||
data_residency: str | None = None,
|
||||
vertex_location: str | None = None,
|
||||
) -> TokenTypeCostBreakdown:
|
||||
"""
|
||||
Provider-agnostic cost of reasoning and cache tokens, derived from the usage
|
||||
|
|
@ -1069,6 +1106,12 @@ def get_token_type_cost_breakdown(
|
|||
cache_read_cost *= uplift
|
||||
cache_creation_cost *= uplift
|
||||
|
||||
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)
|
||||
if vertex_uplift != 1.0:
|
||||
reasoning_cost *= vertex_uplift
|
||||
cache_read_cost *= vertex_uplift
|
||||
cache_creation_cost *= vertex_uplift
|
||||
|
||||
# Mirror the provider-specific geo uplift (e.g. Anthropic us: 1.1) the totals
|
||||
# apply, so cache and reasoning line items stay reconciled with them.
|
||||
geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage)
|
||||
|
|
|
|||
149
litellm/litellm_core_utils/ptu_pricing.py
Normal file
149
litellm/litellm_core_utils/ptu_pricing.py
Normal file
|
|
@ -0,0 +1,149 @@
|
|||
"""Which deployments accrue PTU flat cost, and what that costs them per token.
|
||||
|
||||
Reserved provisioned throughput is billed by the hour whether or not requests are sent, so
|
||||
a deployment that accrues flat cost must not also bill per token. The two halves live here
|
||||
together because they have to agree: a deployment the rollup declines to charge but the
|
||||
router prices at zero serves its traffic for free.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
from litellm.types.router import ModelInfo
|
||||
from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams
|
||||
|
||||
PTU_COST_ATTRIBUTION_ENV_VAR: Final = "LITELLM_ENABLE_PTU_COST_ATTRIBUTION"
|
||||
|
||||
|
||||
def is_ptu_cost_attribution_enabled() -> bool:
|
||||
"""Whether PTU flat-cost attribution is turned on for this process."""
|
||||
return get_secret_bool(PTU_COST_ATTRIBUTION_ENV_VAR, False) is True
|
||||
|
||||
|
||||
PTU_ZEROED_PRICING_FIELDS: Final = tuple(f for f in MirroredPricingParams.model_fields if f != "tiered_pricing") + (
|
||||
"cache_creation_input_token_cost_above_1hr",
|
||||
"cache_creation_input_token_cost_above_200k_tokens",
|
||||
"cache_read_input_token_cost_above_200k_tokens",
|
||||
)
|
||||
# tiered_pricing is emptied rather than zeroed: its tiers outrank the zeros written beside
|
||||
# them, so a zero here would leave the cost map's tiers billing the traffic the reserved
|
||||
# capacity already covers.
|
||||
PTU_EMPTIED_PRICING_FIELDS: Final = frozenset(("tiered_pricing",))
|
||||
# search_context_cost_per_query holds its rates in a table keyed by context size, and an
|
||||
# absent table means the provider's own default rather than free, so it is zeroed in place
|
||||
# and written on every PTU deployment rather than only where a table is already stored.
|
||||
PTU_ZEROED_TABLE_FIELDS: Final = frozenset(("search_context_cost_per_query",))
|
||||
SEARCH_CONTEXT_SIZES: Final = ("search_context_size_low", "search_context_size_medium", "search_context_size_high")
|
||||
# Rate fields only. CustomPricingLiteLLMParams also carries settings that are not charges,
|
||||
# and zeroing one of those would destroy the deployment's configuration rather than stop a
|
||||
# charge.
|
||||
CUSTOM_PRICING_FIELDS: Final = frozenset(f for f in CustomPricingLiteLLMParams.model_fields if "cost" in f)
|
||||
PTU_ZEROED_PRICING: Final[Mapping[str, float | tuple[()] | Mapping[str, float]]] = MappingProxyType(
|
||||
{
|
||||
**dict.fromkeys(PTU_ZEROED_PRICING_FIELDS, 0.0),
|
||||
**dict.fromkeys(PTU_EMPTIED_PRICING_FIELDS, ()),
|
||||
**dict.fromkeys(PTU_ZEROED_TABLE_FIELDS, MappingProxyType(dict.fromkeys(SEARCH_CONTEXT_SIZES, 0.0))),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PTUTerms:
|
||||
"""The reservation a deployment declares, once every field has been validated."""
|
||||
|
||||
team_id: str
|
||||
ptu_count: int
|
||||
cost_per_ptu_per_hour: float
|
||||
effective_from: datetime
|
||||
effective_to: datetime | None
|
||||
|
||||
|
||||
def _to_utc(parsed: datetime) -> datetime:
|
||||
"""``parsed`` as UTC, reading a naive value as UTC rather than local time."""
|
||||
return parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def _as_utc(value: object) -> datetime | None:
|
||||
"""A model_info datetime as UTC, parsing an ISO string, else None."""
|
||||
if isinstance(value, datetime):
|
||||
return _to_utc(value)
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
try:
|
||||
return _to_utc(datetime.fromisoformat(value.replace("Z", "+00:00")))
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def ptu_terms(model_info: Mapping[str, object]) -> PTUTerms | None:
|
||||
"""The reservation this deployment accrues flat cost for, else None.
|
||||
|
||||
A start is required rather than inferred because flat cost accrues from it, and a
|
||||
present but unparseable bound would read as no bound and widen the window to the whole
|
||||
day, so either one leaves the deployment unpriced until the config is fixed.
|
||||
"""
|
||||
ptu_count: Final = model_info.get("ptu_count")
|
||||
cost_per_hour: Final = model_info.get("cost_per_ptu_per_hour")
|
||||
team_id: Final = model_info.get("team_id")
|
||||
if ptu_count is None or cost_per_hour is None or not team_id:
|
||||
return None
|
||||
try:
|
||||
ptu_count_int: Final = int(ptu_count)
|
||||
cost_per_hour_float: Final = float(cost_per_hour)
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return None
|
||||
if not 0 < ptu_count_int <= ModelInfo.MAX_PTU_COUNT:
|
||||
return None
|
||||
if not 0 <= cost_per_hour_float <= ModelInfo.MAX_COST_PER_PTU_PER_HOUR:
|
||||
return None
|
||||
|
||||
raw_from: Final = model_info.get("ptu_effective_from")
|
||||
raw_to: Final = model_info.get("ptu_effective_to")
|
||||
effective_from: Final = _as_utc(raw_from)
|
||||
effective_to: Final = _as_utc(raw_to)
|
||||
if effective_from is None or (raw_to is not None and effective_to is None):
|
||||
return None
|
||||
if effective_to is not None and effective_to <= effective_from:
|
||||
return None
|
||||
return PTUTerms(
|
||||
team_id=str(team_id),
|
||||
ptu_count=ptu_count_int,
|
||||
cost_per_ptu_per_hour=cost_per_hour_float,
|
||||
effective_from=effective_from,
|
||||
effective_to=effective_to,
|
||||
)
|
||||
|
||||
|
||||
def zeroed_ptu_pricing(
|
||||
model_info: Mapping[str, object], declared: Mapping[str, object]
|
||||
) -> Mapping[str, float | tuple[()] | Mapping[str, float]] | None:
|
||||
"""The pricing a deployment accruing flat cost must carry, else None.
|
||||
|
||||
Both conditions hold or nothing is zeroed. Without the flag no flat cost accrues, so
|
||||
zeroing would leave the deployment serving for free with nothing charged in its place,
|
||||
which is what an SDK user who happens to carry ptu_count would otherwise get. The terms
|
||||
are checked first only because they are a few dict reads, while the flag can resolve
|
||||
through a configured secret manager, and this runs for every deployment registered.
|
||||
|
||||
Any further rate the deployment itself declares is zeroed alongside the standing set,
|
||||
since one left standing bills the traffic the reserved capacity already paid for.
|
||||
"""
|
||||
if ptu_terms(model_info) is None:
|
||||
return None
|
||||
if not is_ptu_cost_attribution_enabled():
|
||||
return None
|
||||
return MappingProxyType(
|
||||
{
|
||||
**PTU_ZEROED_PRICING,
|
||||
**dict.fromkeys(
|
||||
CUSTOM_PRICING_FIELDS.intersection(declared)
|
||||
.difference(PTU_ZEROED_TABLE_FIELDS)
|
||||
.difference(PTU_EMPTIED_PRICING_FIELDS),
|
||||
0.0,
|
||||
),
|
||||
}
|
||||
)
|
||||
|
|
@ -723,7 +723,7 @@ class ChunkProcessor:
|
|||
for choice in response.choices:
|
||||
if (
|
||||
hasattr(cast(Choices, choice).message, "reasoning_content")
|
||||
and cast(Choices, choice).message.reasoning_content is not None
|
||||
and cast(Choices, choice).message.reasoning_content
|
||||
):
|
||||
if reasoning_tokens is None:
|
||||
reasoning_tokens = 0
|
||||
|
|
@ -987,7 +987,12 @@ class ChunkProcessor:
|
|||
returned_usage.completion_tokens_details is not None
|
||||
and returned_usage.completion_tokens_details.reasoning_tokens is None
|
||||
):
|
||||
returned_usage.completion_tokens_details.reasoning_tokens = reasoning_tokens
|
||||
capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
|
||||
returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens
|
||||
if returned_usage.completion_tokens_details.text_tokens is None:
|
||||
returned_usage.completion_tokens_details.text_tokens = (
|
||||
returned_usage.completion_tokens - capped_reasoning_tokens
|
||||
)
|
||||
if prompt_tokens_details is not None:
|
||||
returned_usage.prompt_tokens_details = prompt_tokens_details
|
||||
|
||||
|
|
|
|||
|
|
@ -1830,6 +1830,20 @@ class CustomStreamWrapper:
|
|||
return
|
||||
self.chunks.append(model_response.model_copy(update={"choices": []}))
|
||||
|
||||
@staticmethod
|
||||
def _resolve_provider_reported_cost(usage_cost: object) -> float | None:
|
||||
"""
|
||||
Providers report usage.cost either as a number or, for Perplexity, as a
|
||||
breakdown object whose total lives under ``total_cost``.
|
||||
"""
|
||||
if isinstance(usage_cost, bool):
|
||||
return None
|
||||
if isinstance(usage_cost, (int, float)):
|
||||
return float(usage_cost)
|
||||
if isinstance(usage_cost, dict):
|
||||
return CustomStreamWrapper._resolve_provider_reported_cost(usage_cost.get("total_cost"))
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _propagate_usage_cost_to_hidden_params(
|
||||
response: "ModelResponse",
|
||||
|
|
@ -1840,10 +1854,11 @@ class CustomStreamWrapper:
|
|||
calculator uses it instead of a token-based estimate.
|
||||
"""
|
||||
_usage: Final[Usage | None] = getattr(response, "usage", None)
|
||||
if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None:
|
||||
_cost: Final = CustomStreamWrapper._resolve_provider_reported_cost(getattr(_usage, "cost", None))
|
||||
if _cost is not None:
|
||||
if "additional_headers" not in response._hidden_params:
|
||||
response._hidden_params["additional_headers"] = {}
|
||||
response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(_usage.cost)
|
||||
response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = _cost
|
||||
|
||||
def __next__(self) -> "ModelResponseStream":
|
||||
cache_hit = False
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from collections.abc import Mapping, Sequence
|
|||
from typing import TYPE_CHECKING, Any, Final, NoReturn, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
|
|
@ -39,6 +40,7 @@ from litellm.types.llms.anthropic import (
|
|||
AnthropicMessagesTool,
|
||||
AnthropicMessagesToolChoice,
|
||||
AnthropicOutputSchema,
|
||||
AnthropicOutputTokensDetails,
|
||||
AnthropicSystemMessageContent,
|
||||
AnthropicThinkingParam,
|
||||
AnthropicWebSearchTool,
|
||||
|
|
@ -2104,6 +2106,68 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
compaction_blocks,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None:
|
||||
details: Final = usage_object.get("output_tokens_details")
|
||||
if not isinstance(details, Mapping):
|
||||
return None
|
||||
try:
|
||||
return AnthropicOutputTokensDetails.model_validate(details).thinking_tokens
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _response_has_thinking_block(completion_response: Mapping[str, object] | None) -> bool:
|
||||
if completion_response is None:
|
||||
return False
|
||||
content: Final = completion_response.get("content")
|
||||
if not isinstance(content, list):
|
||||
return False
|
||||
return any(
|
||||
isinstance(block, Mapping) and block.get("type") in ("thinking", "redacted_thinking") for block in content
|
||||
)
|
||||
|
||||
def _build_completion_token_details(
|
||||
self,
|
||||
usage_object: Mapping[str, object],
|
||||
iterations: Sequence[object] | None,
|
||||
completion_tokens: int,
|
||||
reasoning_content: str | None,
|
||||
completion_response: Mapping[str, object] | None,
|
||||
) -> CompletionTokensDetailsWrapper:
|
||||
iteration_thinking_tokens: Final = self._sum_iteration_thinking_tokens(iterations) if iterations else None
|
||||
reported_thinking_tokens: Final = (
|
||||
iteration_thinking_tokens
|
||||
if iteration_thinking_tokens is not None
|
||||
else self._thinking_tokens_from_usage(usage_object)
|
||||
)
|
||||
if reported_thinking_tokens is not None:
|
||||
capped_reported: Final = min(max(0, reported_thinking_tokens), completion_tokens)
|
||||
return CompletionTokensDetailsWrapper(
|
||||
reasoning_tokens=capped_reported,
|
||||
text_tokens=completion_tokens - capped_reported,
|
||||
)
|
||||
if reasoning_content:
|
||||
estimated: Final = min(
|
||||
token_counter(text=reasoning_content, count_response_tokens=True),
|
||||
completion_tokens,
|
||||
)
|
||||
return CompletionTokensDetailsWrapper(
|
||||
reasoning_tokens=max(0, estimated),
|
||||
text_tokens=completion_tokens - max(0, estimated),
|
||||
)
|
||||
if self._response_has_thinking_block(completion_response):
|
||||
return CompletionTokensDetailsWrapper(reasoning_tokens=None, text_tokens=None)
|
||||
return CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=completion_tokens)
|
||||
|
||||
def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None:
|
||||
per_iteration: Final = tuple(
|
||||
self._thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None
|
||||
for iteration in iterations
|
||||
)
|
||||
reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None)
|
||||
return sum(reported) if len(reported) == len(per_iteration) else None
|
||||
|
||||
@staticmethod
|
||||
def is_anthropic_usage_object(usage_object: dict) -> bool:
|
||||
"""Anthropic reports prompt cache tokens as top-level ``cache_read_input_tokens`` /
|
||||
|
|
@ -2222,14 +2286,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
cache_creation_token_details=cache_creation_token_details,
|
||||
text_tokens=raw_input_tokens,
|
||||
)
|
||||
# Always populate completion_token_details, not just when there's reasoning_content
|
||||
estimated_reasoning_tokens: Final = (
|
||||
token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0
|
||||
)
|
||||
reasoning_tokens: Final = min(estimated_reasoning_tokens, completion_tokens)
|
||||
completion_token_details: Final = CompletionTokensDetailsWrapper(
|
||||
reasoning_tokens=max(0, reasoning_tokens),
|
||||
text_tokens=(completion_tokens - reasoning_tokens if reasoning_tokens > 0 else completion_tokens),
|
||||
completion_token_details: Final = self._build_completion_token_details(
|
||||
usage_object=_usage,
|
||||
iterations=iterations,
|
||||
completion_tokens=completion_tokens,
|
||||
reasoning_content=reasoning_content,
|
||||
completion_response=completion_response,
|
||||
)
|
||||
total_tokens: Final = prompt_tokens + completion_tokens
|
||||
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from typing_extensions import TypedDict
|
|||
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.proxy.pass_through_endpoints.success_handler import (
|
||||
PassThroughEndpointLogging,
|
||||
)
|
||||
|
|
@ -134,8 +135,11 @@ class BaseAnthropicMessagesStreamingIterator:
|
|||
if self.completion_start_time is not None:
|
||||
self.litellm_logging_obj.completion_start_time = self.completion_start_time
|
||||
self.litellm_logging_obj.model_call_details["completion_start_time"] = self.completion_start_time
|
||||
asyncio.create_task(
|
||||
PassThroughStreamingHandler._route_streaming_logging_to_handler(
|
||||
# Enqueue on the rooted logging worker rather than asyncio.create_task:
|
||||
# this also runs during generator teardown after a client disconnect,
|
||||
# where an unrooted task could be garbage-collected before it bills.
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
async_coroutine=PassThroughStreamingHandler._route_streaming_logging_to_handler(
|
||||
litellm_logging_obj=self.litellm_logging_obj,
|
||||
passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
|
||||
url_route="/v1/messages",
|
||||
|
|
@ -197,13 +201,21 @@ class BaseAnthropicMessagesStreamingIterator:
|
|||
collected_chunks: Final = []
|
||||
saw_terminal_event = False
|
||||
|
||||
async for chunk in completion_stream:
|
||||
if self.completion_start_time is None:
|
||||
self.completion_start_time = datetime.now()
|
||||
saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk)
|
||||
encoded_chunk = self._convert_chunk_to_sse_format(chunk)
|
||||
collected_chunks.append(encoded_chunk)
|
||||
yield encoded_chunk
|
||||
try:
|
||||
async for chunk in completion_stream:
|
||||
if self.completion_start_time is None:
|
||||
self.completion_start_time = datetime.now()
|
||||
saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk)
|
||||
encoded_chunk = self._convert_chunk_to_sse_format(chunk)
|
||||
collected_chunks.append(encoded_chunk)
|
||||
yield encoded_chunk
|
||||
except (GeneratorExit, asyncio.CancelledError):
|
||||
# A client disconnect tears the generator down at the yield, so the
|
||||
# post-loop logging below never runs and the tokens already streamed
|
||||
# (and billed by the provider) would never reach spend tracking. See LIT-5839.
|
||||
if collected_chunks:
|
||||
await self._handle_streaming_logging(collected_chunks)
|
||||
raise
|
||||
|
||||
if not saw_terminal_event:
|
||||
yield _incomplete_stream_error_sse_event()
|
||||
|
|
|
|||
|
|
@ -183,6 +183,29 @@ class BaseSearchConfig:
|
|||
"""
|
||||
return headers
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: dict[str, str], # mutable-ok: matches the request header dict every other hook on this base takes
|
||||
optional_params: dict[str, object], # mutable-ok: matches every other hook on this base
|
||||
request_data: dict[str, object] | list[dict[str, object]], # mutable-ok: transform_search_request's body
|
||||
api_base: str,
|
||||
api_key: str | None = None,
|
||||
) -> tuple[dict[str, str], bytes | None]: # mutable-ok: the handler passes these headers straight to httpx
|
||||
"""
|
||||
OPTIONAL
|
||||
|
||||
Sign the request. Providers like Bedrock AgentCore need to SigV4-sign
|
||||
the request before sending it to the API.
|
||||
|
||||
For all other providers, this is a no-op and we just return the headers.
|
||||
|
||||
Returns:
|
||||
Tuple of (headers, signed_json_body). When signed_json_body is not
|
||||
None, the handler MUST send it verbatim as the request body —
|
||||
re-serializing the payload would invalidate the signature.
|
||||
"""
|
||||
return headers, None
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
|
|
|
|||
|
|
@ -1834,6 +1834,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
self,
|
||||
usage: ConverseTokenUsageBlock,
|
||||
reasoning_content: str | None = None,
|
||||
thinking_ran: bool = False,
|
||||
) -> Usage:
|
||||
input_tokens = usage["inputTokens"]
|
||||
output_tokens: Final = usage["outputTokens"]
|
||||
|
|
@ -1854,10 +1855,19 @@ class AmazonConverseConfig(BaseConfig):
|
|||
cache_creation_tokens=cache_creation_input_tokens,
|
||||
text_tokens=raw_input_tokens,
|
||||
)
|
||||
reasoning_tokens = token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0
|
||||
completion_tokens_details: Final = CompletionTokensDetailsWrapper(
|
||||
reasoning_tokens=reasoning_tokens,
|
||||
text_tokens=(output_tokens - reasoning_tokens if reasoning_tokens > 0 else output_tokens),
|
||||
reasoning_tokens: Final = (
|
||||
token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0
|
||||
)
|
||||
completion_tokens_details: Final = (
|
||||
CompletionTokensDetailsWrapper(
|
||||
reasoning_tokens=reasoning_tokens,
|
||||
text_tokens=output_tokens - reasoning_tokens,
|
||||
)
|
||||
if reasoning_tokens > 0
|
||||
else CompletionTokensDetailsWrapper(
|
||||
reasoning_tokens=None if thinking_ran else 0,
|
||||
text_tokens=None if thinking_ran else output_tokens,
|
||||
)
|
||||
)
|
||||
openai_usage: Final = Usage(
|
||||
prompt_tokens=input_tokens,
|
||||
|
|
@ -2254,6 +2264,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
usage: Final = self.transform_usage(
|
||||
completion_response["usage"],
|
||||
reasoning_content=chat_completion_message.get("reasoning_content"),
|
||||
thinking_ran=reasoningContentBlocks is not None,
|
||||
)
|
||||
|
||||
## HANDLE TOOL CALLS
|
||||
|
|
|
|||
|
|
@ -330,6 +330,7 @@ class AWSEventStreamDecoder:
|
|||
self.response_id: str | None = None
|
||||
self.json_mode = json_mode
|
||||
self._current_tool_name: str | None = None
|
||||
self._thinking_ran = False
|
||||
|
||||
def check_empty_tool_call_args(self) -> bool:
|
||||
"""
|
||||
|
|
@ -559,7 +560,12 @@ class AWSEventStreamDecoder:
|
|||
elif "stopReason" in chunk_data:
|
||||
finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop"))
|
||||
elif "usage" in chunk_data:
|
||||
usage = converse_config.transform_usage(chunk_data.get("usage", {}))
|
||||
usage = converse_config.transform_usage(
|
||||
chunk_data.get("usage", {}),
|
||||
thinking_ran=self._thinking_ran,
|
||||
)
|
||||
if thinking_blocks:
|
||||
self._thinking_ran = True
|
||||
|
||||
model_response_provider_specific_fields: Final = {}
|
||||
if "trace" in chunk_data:
|
||||
|
|
|
|||
|
|
@ -842,8 +842,15 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
|
||||
patched_stream: Final = self._promote_message_stop_usage(completion_stream)
|
||||
|
||||
async for chunk in handler.async_sse_wrapper(patched_stream):
|
||||
yield chunk
|
||||
sse_stream: Final = handler.async_sse_wrapper(patched_stream)
|
||||
try:
|
||||
async for chunk in sse_stream:
|
||||
yield chunk
|
||||
finally:
|
||||
# Close the inner generator deterministically so a client disconnect
|
||||
# (GeneratorExit here) reaches async_sse_wrapper's partial-spend logging
|
||||
# now instead of at garbage collection. See LIT-5839.
|
||||
await sse_stream.aclose()
|
||||
|
||||
@staticmethod
|
||||
def _merge_message_start_cache_into_delta_usage(
|
||||
|
|
|
|||
0
litellm/llms/bedrock/search/__init__.py
Normal file
0
litellm/llms/bedrock/search/__init__.py
Normal file
455
litellm/llms/bedrock/search/transformation.py
Normal file
455
litellm/llms/bedrock/search/transformation.py
Normal file
|
|
@ -0,0 +1,455 @@
|
|||
"""
|
||||
Calls an Amazon Bedrock AgentCore Gateway web-search target (MCP protocol) to search the web.
|
||||
|
||||
Web Search on Amazon Bedrock AgentCore exposes Amazon's managed web index through
|
||||
an AgentCore Gateway MCP endpoint.
|
||||
|
||||
AWS docs: https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-target-connector-web-search-tool.html
|
||||
|
||||
Authentication (matches the gateway's inbound authorizer type):
|
||||
- AWS_IAM gateway: the request is SigV4-signed. Credentials come from explicit
|
||||
params (aws_access_key_id / aws_secret_access_key / aws_session_token /
|
||||
aws_region_name, also settable in a proxy search_tools entry) or the
|
||||
standard AWS credential chain (env / profile / IRSA / assumed role)
|
||||
- CUSTOM_JWT gateway: pass the OAuth2 bearer token (e.g. Cognito
|
||||
client_credentials) as api_key, or set AGENTCORE_GATEWAY_TOKEN
|
||||
|
||||
Setup:
|
||||
1. Create an AgentCore Gateway with a web-search connector target
|
||||
2. Set AGENTCORE_GATEWAY_URL (or pass api_base) to the gateway MCP endpoint, e.g.
|
||||
https://<gateway-id>.gateway.bedrock-agentcore.<region>.amazonaws.com/mcp
|
||||
3. AWS_IAM: ensure the credentials allow bedrock-agentcore:InvokeGateway
|
||||
CUSTOM_JWT: set AGENTCORE_GATEWAY_TOKEN (or pass api_key)
|
||||
|
||||
Usage:
|
||||
response = litellm.search(
|
||||
query="latest AI developments",
|
||||
search_provider="agentcore",
|
||||
max_results=5,
|
||||
aws_access_key_id="...", # optional, omit to use the default chain
|
||||
aws_secret_access_key="...",
|
||||
)
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.search.transformation import (
|
||||
BaseSearchConfig,
|
||||
SearchResponse,
|
||||
SearchResult,
|
||||
)
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
# AgentCore web-search rejects queries longer than 200 characters
|
||||
AGENTCORE_MAX_QUERY_LENGTH: Final = 200
|
||||
|
||||
# The provider contract documents a default of 10 results, send it explicitly
|
||||
# so the gateway can't silently apply a different default.
|
||||
AGENTCORE_DEFAULT_MAX_RESULTS: Final = 10
|
||||
|
||||
# Default MCP tool name for a gateway web-search connector target:
|
||||
# "<target-name>___<tool-name>". Override with AGENTCORE_SEARCH_TOOL_NAME
|
||||
# or optional_params["tool_name"] when the target uses a custom name.
|
||||
AGENTCORE_DEFAULT_TOOL_NAME: Final = "web-search-tool___WebSearch"
|
||||
|
||||
# All web-search connector tools share this suffix; rejecting other names keeps
|
||||
# a caller-supplied tool_name from invoking unrelated tools on the same gateway
|
||||
# with the proxy's credentials.
|
||||
AGENTCORE_TOOL_NAME_SUFFIX: Final = "___WebSearch"
|
||||
|
||||
# MCP revision this provider speaks. Sent on every request because the gateway is
|
||||
# called statelessly, without an initialize handshake to negotiate a version.
|
||||
# AgentCore gateways whose protocolConfiguration leaves supportedVersions unset
|
||||
# accept only 2025-03-26 and reject anything newer with a -32600 error, so that
|
||||
# is the default; a gateway pinned to another version needs
|
||||
# AGENTCORE_MCP_PROTOCOL_VERSION set to match.
|
||||
AGENTCORE_DEFAULT_MCP_PROTOCOL_VERSION: Final = "2025-03-26"
|
||||
|
||||
# Matched against the URL host so a crafted path or query string can't pass for
|
||||
# a gateway hostname.
|
||||
_GATEWAY_HOST_PATTERN: Final = re.compile(r"[a-z0-9-]+\.gateway\.bedrock-agentcore\.([a-z0-9-]+)\.amazonaws\.com")
|
||||
|
||||
_SSE_EVENT_SEPARATOR: Final = re.compile(r"\r?\n[ \t]*\r?\n")
|
||||
|
||||
_SSE_LINE_PREFIXES: Final = ("event:", "data:", ":", "id:", "retry:")
|
||||
|
||||
|
||||
def _gateway_host_match(api_base: str) -> re.Match[str] | None:
|
||||
return _GATEWAY_HOST_PATTERN.fullmatch(httpx.URL(api_base).host)
|
||||
|
||||
|
||||
_LOOPBACK_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"})
|
||||
|
||||
|
||||
def _credential_safe_transport(api_base: str) -> bool:
|
||||
url: Final = httpx.URL(api_base)
|
||||
return url.scheme == "https" or url.host in _LOOPBACK_HOSTS
|
||||
|
||||
|
||||
def _string_field(item: Mapping[str, object], *keys: str) -> str | None:
|
||||
return next(
|
||||
(value for key in keys if isinstance(value := item.get(key), str) and value),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _to_search_result(item: Mapping[str, object]) -> SearchResult:
|
||||
return SearchResult(
|
||||
title=_string_field(item, "title") or "",
|
||||
url=_string_field(item, "url") or "",
|
||||
snippet=_string_field(item, "text", "snippet") or "",
|
||||
date=_string_field(item, "publishedDate", "date"),
|
||||
last_updated=None,
|
||||
)
|
||||
|
||||
|
||||
def _result_items(parsed: object) -> tuple[Mapping[str, object], ...]:
|
||||
items: Final = parsed.get("results", ()) if isinstance(parsed, Mapping) else parsed
|
||||
if not isinstance(items, Sequence) or isinstance(items, (str, bytes)):
|
||||
return ()
|
||||
return tuple(item for item in items if isinstance(item, Mapping))
|
||||
|
||||
|
||||
def _parse_result_items(raw_text: object) -> tuple[Mapping[str, object], ...]:
|
||||
"""
|
||||
Parse one MCP text block into the search result objects it carries.
|
||||
|
||||
A block holds either a JSON list of results or a {"results": [...]} object;
|
||||
anything unparseable is skipped rather than failing the whole response.
|
||||
"""
|
||||
if not isinstance(raw_text, str):
|
||||
return ()
|
||||
try:
|
||||
parsed: Final = json.loads(raw_text)
|
||||
except json.JSONDecodeError:
|
||||
return ()
|
||||
return _result_items(parsed)
|
||||
|
||||
|
||||
def _iter_sse_events(text: str) -> Iterator[Mapping[str, object]]:
|
||||
"""
|
||||
Yield the JSON payload of each SSE event in a Streamable HTTP MCP response.
|
||||
|
||||
Per the SSE spec an event's data is the concatenation of all its ``data:``
|
||||
lines (joined with newlines), and a stream may carry several events, e.g.
|
||||
progress notifications before the JSON-RPC response.
|
||||
"""
|
||||
for chunk in _SSE_EVENT_SEPARATOR.split(text):
|
||||
payload = "\n".join(line[len("data:") :].lstrip() for line in chunk.splitlines() if line.startswith("data:"))
|
||||
if not payload:
|
||||
continue
|
||||
try:
|
||||
parsed = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(parsed, dict):
|
||||
yield parsed
|
||||
|
||||
|
||||
class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
|
||||
def __init__(self) -> None:
|
||||
BaseSearchConfig.__init__(self)
|
||||
BaseAWSLLM.__init__(self)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "Web Search on Amazon Bedrock"
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict, # mutable-ok: BaseSearchConfig hands providers the mutable request header dict
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
**kwargs: object, # kwargs-ok: BaseSearchConfig.validate_environment forwards provider-specific extras
|
||||
) -> dict: # mutable-ok: the handler passes these headers straight to httpx, which wants a dict
|
||||
"""
|
||||
Set MCP transport headers. Per the MCP Streamable HTTP transport spec,
|
||||
the client MUST accept both application/json and text/event-stream, and
|
||||
declare its protocol revision with MCP-Protocol-Version.
|
||||
|
||||
Authentication itself happens in sign_request(): bearer token for
|
||||
CUSTOM_JWT gateways, AWS SigV4 for AWS_IAM gateways.
|
||||
"""
|
||||
return { # mutable-ok: httpx request headers are a dict
|
||||
**headers,
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json, text/event-stream",
|
||||
"MCP-Protocol-Version": get_secret_str("AGENTCORE_MCP_PROTOCOL_VERSION")
|
||||
or AGENTCORE_DEFAULT_MCP_PROTOCOL_VERSION,
|
||||
}
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
optional_params: dict, # mutable-ok: BaseSearchConfig passes optional params as a dict
|
||||
data: dict | list[dict] | None = None, # mutable-ok: BaseSearchConfig request bodies are JSON dicts
|
||||
**kwargs: object, # kwargs-ok: BaseSearchConfig.get_complete_url forwards provider-specific extras
|
||||
) -> str:
|
||||
gateway_url: Final = api_base or get_secret_str("AGENTCORE_GATEWAY_URL")
|
||||
if not gateway_url:
|
||||
raise ValueError(
|
||||
"AGENTCORE_GATEWAY_URL is not set. Set it to your AgentCore Gateway MCP "
|
||||
"endpoint (https://<gateway-id>.gateway.bedrock-agentcore.<region>"
|
||||
".amazonaws.com/mcp) or pass api_base."
|
||||
)
|
||||
return gateway_url
|
||||
|
||||
def transform_search_request(
|
||||
self,
|
||||
query: str | list[str], # mutable-ok: BaseSearchConfig accepts a list of queries
|
||||
optional_params: dict, # mutable-ok: BaseSearchConfig passes optional params as a dict
|
||||
**kwargs: object, # kwargs-ok: BaseSearchConfig.transform_search_request forwards provider-specific extras
|
||||
) -> dict: # mutable-ok: the JSON-RPC body is serialized as a JSON object
|
||||
"""
|
||||
Transform Search request to an MCP tools/call request.
|
||||
|
||||
Args:
|
||||
query: Search query (string or list of strings). AgentCore only
|
||||
supports single string queries; lists are joined with spaces.
|
||||
optional_params: Optional parameters for the request
|
||||
- max_results: Maximum number of results (1-25), default 10
|
||||
- tool_name: Override the MCP tool name of the gateway target
|
||||
|
||||
Returns:
|
||||
Dict with the JSON-RPC 2.0 request body
|
||||
"""
|
||||
joined_query: Final = " ".join(query) if isinstance(query, list) else query
|
||||
tool_name: Final = (
|
||||
optional_params.get("tool_name")
|
||||
or get_secret_str("AGENTCORE_SEARCH_TOOL_NAME")
|
||||
or AGENTCORE_DEFAULT_TOOL_NAME
|
||||
)
|
||||
if not tool_name.endswith(AGENTCORE_TOOL_NAME_SUFFIX):
|
||||
raise ValueError(
|
||||
f"Invalid AgentCore search tool_name '{tool_name}': must end with "
|
||||
f"'{AGENTCORE_TOOL_NAME_SUFFIX}' (a web-search connector tool). "
|
||||
"Other gateway tools cannot be invoked through this provider."
|
||||
)
|
||||
|
||||
return { # mutable-ok: JSON-RPC request bodies are JSON objects
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/call",
|
||||
"params": { # mutable-ok: JSON-RPC request bodies are JSON objects
|
||||
"name": tool_name,
|
||||
"arguments": { # mutable-ok: JSON-RPC request bodies are JSON objects
|
||||
"query": joined_query[:AGENTCORE_MAX_QUERY_LENGTH],
|
||||
"maxResults": optional_params.get("max_results", AGENTCORE_DEFAULT_MAX_RESULTS),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: dict[str, str], # mutable-ok: BaseSearchConfig hands providers the mutable request header dict
|
||||
optional_params: dict[str, object], # mutable-ok: BaseSearchConfig passes optional params as a dict
|
||||
request_data: dict[str, object] | list[dict[str, object]], # mutable-ok: request bodies are JSON dicts
|
||||
api_base: str,
|
||||
api_key: str | None = None,
|
||||
) -> tuple[dict[str, str], bytes | None]: # mutable-ok: BaseSearchConfig.sign_request returns httpx headers
|
||||
"""
|
||||
Authenticate the MCP request.
|
||||
|
||||
CUSTOM_JWT gateways: attach the caller's OAuth2 bearer token (api_key
|
||||
or AGENTCORE_GATEWAY_TOKEN), no AWS credentials involved.
|
||||
|
||||
AWS_IAM gateways: SigV4-sign with the bedrock-agentcore service name.
|
||||
"""
|
||||
if not isinstance(request_data, dict):
|
||||
raise TypeError("AgentCore search expects a single dict request body")
|
||||
|
||||
if not _credential_safe_transport(api_base):
|
||||
raise ValueError(
|
||||
f"Refusing to send AgentCore credentials over plaintext HTTP to '{api_base}': a bearer "
|
||||
"token or SigV4 signature would be readable in transit. Use an https gateway URL "
|
||||
"(plain http is allowed only for localhost)."
|
||||
)
|
||||
|
||||
# Server-managed credentials only go to a trusted host, otherwise an
|
||||
# authenticated caller could point api_base at their own server (e.g. via
|
||||
# /search_tools/test_connection) and collect AGENTCORE_GATEWAY_TOKEN or a
|
||||
# SigV4 signature with the proxy's credential scope and session token.
|
||||
gateway_host_match: Final = _gateway_host_match(api_base)
|
||||
bearer_token: Final = self.resolve_server_api_key(
|
||||
caller_api_key=api_key,
|
||||
caller_api_base=api_base,
|
||||
key_env_vars=("AGENTCORE_GATEWAY_TOKEN",),
|
||||
base_env_var="AGENTCORE_GATEWAY_URL",
|
||||
default_api_base=api_base if gateway_host_match else None,
|
||||
)
|
||||
if bearer_token:
|
||||
bearer_headers: Final = { # mutable-ok: httpx request headers are a dict
|
||||
**headers,
|
||||
"Authorization": f"Bearer {bearer_token}",
|
||||
}
|
||||
return bearer_headers, json.dumps(request_data).encode()
|
||||
|
||||
if gateway_host_match is None and not self._is_configured_gateway(api_base):
|
||||
raise ValueError(
|
||||
f"Refusing to send SigV4-signed AgentCore requests to '{api_base}': it is neither an "
|
||||
"AgentCore gateway hostname nor the host in AGENTCORE_GATEWAY_URL. Set "
|
||||
"AGENTCORE_GATEWAY_URL to authorize a custom gateway hostname."
|
||||
)
|
||||
|
||||
signing_params: Final = (
|
||||
optional_params
|
||||
if optional_params.get("aws_region_name") is not None
|
||||
else { # mutable-ok: BaseAWSLLM._sign_request takes optional params as a dict
|
||||
**optional_params,
|
||||
"aws_region_name": self._signing_region(api_base),
|
||||
}
|
||||
)
|
||||
|
||||
# api_key="" (not None, but falsy) disables BaseAWSLLM's fallback to the
|
||||
# AWS_BEARER_TOKEN_BEDROCK env var: that token is a Bedrock Runtime
|
||||
# credential and must not be sent to an AgentCore gateway.
|
||||
return self._sign_request(
|
||||
service_name="bedrock-agentcore",
|
||||
headers=headers,
|
||||
optional_params=signing_params,
|
||||
request_data=request_data,
|
||||
api_base=api_base,
|
||||
api_key="",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_configured_gateway(api_base: str) -> bool:
|
||||
configured: Final = get_secret_str("AGENTCORE_GATEWAY_URL")
|
||||
if not configured:
|
||||
return False
|
||||
return httpx.URL(configured).host == httpx.URL(api_base).host
|
||||
|
||||
@staticmethod
|
||||
def _signing_region(api_base: str) -> str:
|
||||
"""
|
||||
Resolve the SigV4 signing region, which must match the gateway's region.
|
||||
|
||||
Standard gateway hostnames carry it, so callers don't have to set
|
||||
aws_region_name to a region different from their default. For custom or
|
||||
private hostnames, defer to the AWS configuration chain (env vars and
|
||||
the shared config / profile region), and error out when that yields
|
||||
nothing rather than silently signing for a guessed region the gateway
|
||||
would reject with a confusing auth error.
|
||||
"""
|
||||
match: Final = _gateway_host_match(api_base)
|
||||
if match:
|
||||
return match.group(1)
|
||||
|
||||
# boto3's session resolution covers env vars AND the AWS shared config
|
||||
# (profile region), unlike BaseAWSLLM's helper, which silently defaults
|
||||
# to us-west-2 when nothing is configured.
|
||||
import boto3
|
||||
|
||||
configured_region: Final = boto3.Session().region_name
|
||||
if configured_region:
|
||||
return configured_region
|
||||
raise ValueError(
|
||||
f"Cannot derive the SigV4 signing region from api_base '{api_base}' "
|
||||
"or the AWS configuration chain. Set aws_region_name (or AWS_DEFAULT_REGION / "
|
||||
"a profile region) to the gateway's region when using a custom hostname."
|
||||
)
|
||||
|
||||
def transform_search_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
**kwargs: object, # kwargs-ok: BaseSearchConfig.transform_search_response forwards provider-specific extras
|
||||
) -> SearchResponse:
|
||||
"""
|
||||
Transform an MCP tools/call response to LiteLLM unified SearchResponse.
|
||||
|
||||
The gateway returns JSON-RPC (as plain JSON or a single-message SSE
|
||||
stream) whose result.content[] text blocks contain a JSON list of
|
||||
{title, url, date/publishedDate, text} entries. Web-search connector
|
||||
1.1.0 and later repeat that list in result.structuredContent, which is
|
||||
the only machine-readable copy when the text block holds prose instead.
|
||||
"""
|
||||
response_json: Final = self._parse_mcp_body(raw_response)
|
||||
|
||||
error: Final = response_json.get("error")
|
||||
if error is not None:
|
||||
raise BedrockError(
|
||||
status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
|
||||
message=f"AgentCore gateway MCP error: {error}",
|
||||
)
|
||||
|
||||
# A failed tools/call is reported in-band, as HTTP 200 with result.isError
|
||||
# and the failure text where the results would be.
|
||||
result: Final = response_json.get("result")
|
||||
if isinstance(result, dict) and result.get("isError"):
|
||||
raise BedrockError(
|
||||
status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
|
||||
message=f"AgentCore web search tool error: {self._tool_error_message(response_json)}",
|
||||
)
|
||||
|
||||
text_items: Final = tuple(
|
||||
item for block in self._text_blocks(response_json) for item in _parse_result_items(block.get("text"))
|
||||
)
|
||||
structured: Final = result.get("structuredContent") if isinstance(result, Mapping) else None
|
||||
items: Final = text_items or _result_items(structured)
|
||||
|
||||
results: Final = [_to_search_result(item) for item in items] # mutable-ok: pydantic list field
|
||||
|
||||
return SearchResponse(results=results, object="search")
|
||||
|
||||
def _tool_error_message(self, response_json: Mapping[str, object]) -> str:
|
||||
texts: Final = tuple(
|
||||
text for block in self._text_blocks(response_json) if isinstance(text := block.get("text"), str)
|
||||
)
|
||||
return " ".join(texts) if texts else json.dumps(response_json.get("result"))[:500]
|
||||
|
||||
@staticmethod
|
||||
def _text_blocks(response_json: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
|
||||
result: Final = response_json.get("result")
|
||||
content: Final = result.get("content") if isinstance(result, dict) else None
|
||||
if not isinstance(content, Sequence) or isinstance(content, (str, bytes)):
|
||||
return ()
|
||||
return tuple(block for block in content if isinstance(block, dict) and block.get("type") == "text")
|
||||
|
||||
@staticmethod
|
||||
def _parse_mcp_body(raw_response: httpx.Response) -> Mapping[str, object]:
|
||||
"""
|
||||
Parse a JSON or SSE-framed (Streamable HTTP transport) MCP response.
|
||||
|
||||
Return the event whose payload carries the JSON-RPC response, i.e. one
|
||||
containing ``result`` or ``error``, falling back to the last event when
|
||||
the stream carries only notifications.
|
||||
"""
|
||||
text: Final = raw_response.text
|
||||
if not text.lstrip().startswith(_SSE_LINE_PREFIXES):
|
||||
return raw_response.json()
|
||||
|
||||
events: Final = tuple(_iter_sse_events(text))
|
||||
response_event: Final = next(
|
||||
(event for event in events if "result" in event or "error" in event),
|
||||
None,
|
||||
)
|
||||
if response_event is not None:
|
||||
return response_event
|
||||
if events:
|
||||
return events[-1]
|
||||
raise BedrockError(
|
||||
status_code=502,
|
||||
message=f"AgentCore gateway returned SSE without a JSON data frame: {text[:200]}",
|
||||
)
|
||||
|
||||
def get_error_class(
|
||||
self,
|
||||
error_message: str,
|
||||
status_code: int,
|
||||
headers: dict, # mutable-ok: BaseSearchConfig.get_error_class takes the response headers as a dict
|
||||
) -> Exception:
|
||||
return BaseLLMException(
|
||||
status_code=status_code,
|
||||
message=error_message,
|
||||
headers=headers,
|
||||
)
|
||||
|
|
@ -5,7 +5,7 @@ import ssl
|
|||
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from functools import lru_cache
|
||||
from types import ModuleType
|
||||
from types import MappingProxyType, ModuleType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, Union, cast, get_type_hints
|
||||
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
|
||||
|
||||
|
|
@ -1788,6 +1788,14 @@ class BaseLLMHTTPHandler:
|
|||
api_key=api_key,
|
||||
)
|
||||
|
||||
signed_headers, signed_json_body = provider_config.sign_request(
|
||||
headers=headers,
|
||||
optional_params=optional_params,
|
||||
request_data=data,
|
||||
api_base=complete_url,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=query if isinstance(query, str) else str(query),
|
||||
|
|
@ -1811,14 +1819,15 @@ class BaseLLMHTTPHandler:
|
|||
# Note: timeout is set on the client itself, not per-request for GET
|
||||
response = client.get(
|
||||
url=complete_url,
|
||||
headers=headers,
|
||||
headers=signed_headers,
|
||||
)
|
||||
else:
|
||||
# Make POST request with JSON data
|
||||
# A signed body must be sent verbatim, re-serializing it would break the signature
|
||||
response = client.post(
|
||||
url=complete_url,
|
||||
headers=headers,
|
||||
json=data,
|
||||
headers=signed_headers,
|
||||
data=signed_json_body,
|
||||
json=data if signed_json_body is None else None,
|
||||
timeout=timeout,
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -1872,6 +1881,14 @@ class BaseLLMHTTPHandler:
|
|||
api_key=api_key,
|
||||
)
|
||||
|
||||
signed_headers, signed_json_body = provider_config.sign_request(
|
||||
headers=headers,
|
||||
optional_params=optional_params,
|
||||
request_data=data,
|
||||
api_base=complete_url,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=query if isinstance(query, str) else str(query),
|
||||
|
|
@ -1900,14 +1917,15 @@ class BaseLLMHTTPHandler:
|
|||
# Note: timeout is set on the client itself, not per-request for GET
|
||||
response = await async_httpx_client.get(
|
||||
url=complete_url,
|
||||
headers=headers,
|
||||
headers=signed_headers,
|
||||
)
|
||||
else:
|
||||
# Make async POST request with JSON data
|
||||
# A signed body must be sent verbatim, re-serializing it would break the signature
|
||||
response = await async_httpx_client.post(
|
||||
url=complete_url,
|
||||
headers=headers,
|
||||
json=data,
|
||||
headers=signed_headers,
|
||||
data=signed_json_body,
|
||||
json=data if signed_json_body is None else None,
|
||||
timeout=timeout,
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -2069,6 +2087,14 @@ class BaseLLMHTTPHandler:
|
|||
if anthropic_messages_provider_config.should_filter_anthropic_beta_headers():
|
||||
headers = update_headers_with_filtered_beta(headers=headers, provider=custom_llm_provider)
|
||||
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
||||
explicit_vertex_location: Final = VertexBase.explicit_vertex_ai_location(MappingProxyType(dict(litellm_params)))
|
||||
vertex_location_params: Final = (
|
||||
MappingProxyType({"vertex_location": explicit_vertex_location})
|
||||
if explicit_vertex_location
|
||||
else MappingProxyType({})
|
||||
)
|
||||
logging_obj.update_from_kwargs(
|
||||
kwargs=kwargs,
|
||||
model=model,
|
||||
|
|
@ -2077,6 +2103,7 @@ class BaseLLMHTTPHandler:
|
|||
"preset_cache_key": None,
|
||||
"stream_response": {},
|
||||
"model_info": kwargs.get("model_info"),
|
||||
**vertex_location_params,
|
||||
**anthropic_messages_optional_request_params,
|
||||
},
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from litellm import verbose_logger
|
|||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
_is_above_128k,
|
||||
generic_cost_per_token,
|
||||
get_vertex_regional_endpoint_uplift,
|
||||
)
|
||||
from litellm.types.utils import ModelInfo, Usage
|
||||
|
||||
|
|
@ -63,6 +64,7 @@ def cost_per_character(
|
|||
usage: Usage,
|
||||
prompt_characters: float | None = None,
|
||||
completion_characters: float | None = None,
|
||||
vertex_location: str | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculates the cost per character for a given VertexAI model, input messages, and response object.
|
||||
|
|
@ -72,6 +74,8 @@ def cost_per_character(
|
|||
- custom_llm_provider: str, "vertex_ai-*"
|
||||
- prompt_characters: float, the number of input characters
|
||||
- completion_characters: float, the number of output characters
|
||||
- vertex_location: the Vertex AI location serving the request; non-global
|
||||
locations apply the model's regional-endpoint uplift multiplier
|
||||
|
||||
Returns:
|
||||
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
|
||||
|
|
@ -79,8 +83,6 @@ def cost_per_character(
|
|||
Raises:
|
||||
Exception if model requires >128k pricing, but model cost not mapped
|
||||
"""
|
||||
model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
|
||||
## GET MODEL INFO
|
||||
model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
|
||||
|
|
@ -162,7 +164,8 @@ def cost_per_character(
|
|||
usage=usage,
|
||||
)
|
||||
|
||||
return prompt_cost, completion_cost
|
||||
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)
|
||||
return prompt_cost * vertex_uplift, completion_cost * vertex_uplift
|
||||
|
||||
|
||||
def _handle_128k_pricing(
|
||||
|
|
@ -196,6 +199,7 @@ def cost_per_token(
|
|||
custom_llm_provider: str,
|
||||
usage: Usage,
|
||||
service_tier: str | None = None,
|
||||
vertex_location: str | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
|
||||
|
|
@ -207,6 +211,8 @@ def cost_per_token(
|
|||
- completion_tokens: float, the number of output tokens
|
||||
- service_tier: optional tier derived from Gemini trafficType
|
||||
("priority" for ON_DEMAND_PRIORITY, "flex" for FLEX/batch).
|
||||
- vertex_location: the Vertex AI location serving the request; non-global
|
||||
locations apply the model's regional-endpoint uplift multiplier
|
||||
|
||||
Returns:
|
||||
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
|
||||
|
|
@ -222,14 +228,17 @@ def cost_per_token(
|
|||
input_cost_per_token_above_128k_tokens: Final = model_info.get("input_cost_per_token_above_128k_tokens")
|
||||
output_cost_per_token_above_128k_tokens: Final = model_info.get("output_cost_per_token_above_128k_tokens")
|
||||
if input_cost_per_token_above_128k_tokens is not None or output_cost_per_token_above_128k_tokens is not None:
|
||||
return _handle_128k_pricing(
|
||||
prompt_cost_128k, completion_cost_128k = _handle_128k_pricing(
|
||||
model_info=model_info,
|
||||
usage=usage,
|
||||
)
|
||||
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)
|
||||
return prompt_cost_128k * vertex_uplift, completion_cost_128k * vertex_uplift
|
||||
|
||||
return generic_cost_per_token(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
usage=usage,
|
||||
service_tier=service_tier,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import asyncio
|
|||
import json
|
||||
import os
|
||||
import threading
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
|
@ -68,7 +69,8 @@ class VertexBase:
|
|||
# re-acquire it without deadlocking the current thread.
|
||||
self._sync_refresh_lock = threading.RLock()
|
||||
|
||||
def get_vertex_region(self, vertex_region: str | None, model: str) -> str:
|
||||
@staticmethod
|
||||
def get_vertex_region(vertex_region: str | None, model: str) -> str:
|
||||
import litellm
|
||||
|
||||
# Try to get supported_regions directly from model_cost
|
||||
|
|
@ -1191,7 +1193,18 @@ class VertexBase:
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def safe_get_vertex_ai_location(litellm_params: dict) -> str | None:
|
||||
def explicit_vertex_ai_location(params: Mapping[str, object]) -> str | None:
|
||||
"""
|
||||
The location explicitly configured in the given params, without any
|
||||
module-level or environment fallback. None when not configured.
|
||||
"""
|
||||
for configured in (params.get("vertex_location"), params.get("vertex_ai_location")):
|
||||
if isinstance(configured, str) and configured:
|
||||
return configured
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def safe_get_vertex_ai_location(litellm_params: Mapping[str, object]) -> str | None:
|
||||
"""
|
||||
Safely get Vertex AI location without mutating the litellm_params dict.
|
||||
|
||||
|
|
@ -1205,8 +1218,7 @@ class VertexBase:
|
|||
Vertex AI location/region or None
|
||||
"""
|
||||
return (
|
||||
litellm_params.get("vertex_location")
|
||||
or litellm_params.get("vertex_ai_location")
|
||||
VertexBase.explicit_vertex_ai_location(litellm_params)
|
||||
or litellm.vertex_location
|
||||
or get_secret_str("VERTEXAI_LOCATION")
|
||||
or get_secret_str("VERTEX_LOCATION")
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -75,6 +75,24 @@ class MCPUpstreamAuthError(Exception):
|
|||
)
|
||||
|
||||
|
||||
class MCPOpenApiUpstreamError(Exception):
|
||||
"""An OpenAPI-backed MCP tool's upstream answered with a non-2xx that is not a 401.
|
||||
|
||||
Carries the status only. The upstream's response body is deliberately dropped rather than served
|
||||
as tool content: it crosses a trust boundary and may hold prose, urls, or an error document that
|
||||
reads as data, which is how these failures came to be reported as successful tool output. This
|
||||
matches ``outcome_wire_value``'s contract for listing faults, category and status and nothing
|
||||
else. A 401 is raised as ``MCPUpstreamAuthError`` instead, so the caller learns to
|
||||
re-authenticate; every other status stays here, mirroring the regular MCP path where a 403
|
||||
deliberately does not produce a challenge.
|
||||
"""
|
||||
|
||||
def __init__(self, status_code: int, server_name: str) -> None:
|
||||
self.status_code = status_code
|
||||
self.server_name = server_name
|
||||
super().__init__(f"upstream returned HTTP {status_code}")
|
||||
|
||||
|
||||
class MCPToolResultError(Exception):
|
||||
"""An MCP tool call completed with ``isError=True`` in its result.
|
||||
|
||||
|
|
|
|||
|
|
@ -2330,7 +2330,15 @@ class MCPServerManager:
|
|||
input_schema = build_input_schema(resolved_operation)
|
||||
|
||||
# Create tool function with headers using imported function
|
||||
tool_func = create_tool_function(path, method, resolved_operation, base_url, headers=headers)
|
||||
tool_func = create_tool_function(
|
||||
path,
|
||||
method,
|
||||
resolved_operation,
|
||||
base_url,
|
||||
headers=headers,
|
||||
server_label=server.name or server.server_name or server.alias or server.server_id,
|
||||
relays_upstream_auth=server.is_client_forwarded_token,
|
||||
)
|
||||
tool_func.__name__ = prefixed_tool_name
|
||||
tool_func.__doc__ = description
|
||||
|
||||
|
|
@ -4979,6 +4987,12 @@ class MCPServerManager:
|
|||
|
||||
return result
|
||||
|
||||
except MCPUpstreamAuthError:
|
||||
# The caller must re-authenticate upstream, so this keeps its type all the way to the
|
||||
# renderers: the streamable path turns it into an isError result naming the status, and
|
||||
# the REST path relays a real 401 with the upstream's WWW-Authenticate. Flattening it
|
||||
# into the generic message below would lose both.
|
||||
raise
|
||||
except Exception as e:
|
||||
error_msg = f"Error calling OpenAPI tool {tool_name}: {e}"
|
||||
verbose_logger.error(error_msg)
|
||||
|
|
|
|||
|
|
@ -15,6 +15,12 @@ from urllib.parse import quote
|
|||
import httpx
|
||||
from typing_extensions import ReadOnly, Required
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import (
|
||||
MCPOpenApiUpstreamError,
|
||||
MCPUpstreamAuthError,
|
||||
)
|
||||
|
||||
# Tool names emitted from OpenAPI specs must work across all major LLM providers.
|
||||
# OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to
|
||||
# ^[a-zA-Z0-9_-]+$ on tool names. Many specs (notably GitHub's REST API) use
|
||||
|
|
@ -392,12 +398,40 @@ def _merge_openapi_tool_request_headers(
|
|||
return effective_headers
|
||||
|
||||
|
||||
def _raise_for_upstream_failure(
|
||||
response: httpx.Response,
|
||||
upstream: str,
|
||||
relays_upstream_auth: bool,
|
||||
) -> None:
|
||||
"""Turn a non-2xx upstream response into the right typed failure, or return for a 2xx.
|
||||
|
||||
Both call sites feed this: ``get`` hands back the response for a 4xx, while post/put/patch/delete
|
||||
raise ``MaskedHTTPStatusError`` from inside the HTTP handler, so without one classifier the
|
||||
non-GET tools would keep serving an error body as tool output.
|
||||
|
||||
Only the client-forwarded modes carry the caller's own upstream token, so only they can act on a
|
||||
401 by re-authenticating; ``_call_regular_mcp_tool`` gates its re-auth signal the same way. Every
|
||||
other status carries the code alone, never the upstream's body, which crosses a trust boundary.
|
||||
"""
|
||||
if response.status_code < 400:
|
||||
return
|
||||
if response.status_code == 401 and relays_upstream_auth:
|
||||
raise MCPUpstreamAuthError(
|
||||
status_code=response.status_code,
|
||||
www_authenticate=response.headers.get("www-authenticate"),
|
||||
server_name=upstream,
|
||||
)
|
||||
raise MCPOpenApiUpstreamError(response.status_code, upstream)
|
||||
|
||||
|
||||
def create_tool_function(
|
||||
path: str,
|
||||
method: str,
|
||||
operation: _OpenAPIOperation,
|
||||
base_url: str,
|
||||
headers: dict[str, str] | None = None,
|
||||
server_label: str | None = None,
|
||||
relays_upstream_auth: bool = False,
|
||||
):
|
||||
"""Create a tool function for an OpenAPI operation.
|
||||
|
||||
|
|
@ -477,20 +511,26 @@ def create_tool_function(
|
|||
json_body = {"data": body_value}
|
||||
|
||||
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
|
||||
upstream: Final = server_label or f"{original_method.upper()} {path}"
|
||||
|
||||
if original_method == "get":
|
||||
response = await client.get(url, params=params, headers=effective_headers)
|
||||
elif original_method == "post":
|
||||
response = await client.post(url, params=params, json=json_body, headers=effective_headers)
|
||||
elif original_method == "put":
|
||||
response = await client.put(url, params=params, json=json_body, headers=effective_headers)
|
||||
elif original_method == "delete":
|
||||
response = await client.delete(url, params=params, headers=effective_headers)
|
||||
elif original_method == "patch":
|
||||
response = await client.patch(url, params=params, json=json_body, headers=effective_headers)
|
||||
else:
|
||||
return f"Unsupported HTTP method: {original_method}"
|
||||
try:
|
||||
if original_method == "get":
|
||||
response = await client.get(url, params=params, headers=effective_headers)
|
||||
elif original_method == "post":
|
||||
response = await client.post(url, params=params, json=json_body, headers=effective_headers)
|
||||
elif original_method == "put":
|
||||
response = await client.put(url, params=params, json=json_body, headers=effective_headers)
|
||||
elif original_method == "delete":
|
||||
response = await client.delete(url, params=params, headers=effective_headers)
|
||||
elif original_method == "patch":
|
||||
response = await client.patch(url, params=params, json=json_body, headers=effective_headers)
|
||||
else:
|
||||
return f"Unsupported HTTP method: {original_method}"
|
||||
except MaskedHTTPStatusError as e:
|
||||
_raise_for_upstream_failure(e.response, upstream, relays_upstream_auth)
|
||||
raise
|
||||
|
||||
_raise_for_upstream_failure(response, upstream, relays_upstream_auth)
|
||||
return response.text
|
||||
|
||||
return tool_function
|
||||
|
|
|
|||
|
|
@ -407,8 +407,6 @@ if MCP_AVAILABLE:
|
|||
StreamableHTTPSessionManager = None
|
||||
from mcp.types import (
|
||||
CallToolResult,
|
||||
EmbeddedResource,
|
||||
ImageContent,
|
||||
ListToolsResult,
|
||||
Prompt,
|
||||
TextContent,
|
||||
|
|
@ -2861,12 +2859,11 @@ if MCP_AVAILABLE:
|
|||
_extra_token: Final = _request_extra_headers.set(forwarded_headers)
|
||||
_resolved_token: Final = _request_resolved_auth_headers.set(resolved_auth_headers)
|
||||
try:
|
||||
local_content = await _handle_local_mcp_tool(name, arguments)
|
||||
response = await _handle_local_mcp_tool(name, arguments)
|
||||
finally:
|
||||
_request_auth_header.reset(_auth_token)
|
||||
_request_extra_headers.reset(_extra_token)
|
||||
_request_resolved_auth_headers.reset(_resolved_token)
|
||||
response = CallToolResult(content=local_content, isError=False)
|
||||
|
||||
# Try managed MCP server tool (the name is bare; the prefix boundary was
|
||||
# already resolved above against this server's registered prefixes)
|
||||
|
|
@ -2940,8 +2937,7 @@ if MCP_AVAILABLE:
|
|||
if "arguments" in hook_result:
|
||||
arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args
|
||||
|
||||
local_content = await _handle_local_mcp_tool(original_tool_name, arguments)
|
||||
response = CallToolResult(content=local_content, isError=False)
|
||||
response = await _handle_local_mcp_tool(original_tool_name, arguments)
|
||||
|
||||
return await _run_post_mcp_call_guardrails(
|
||||
result=response,
|
||||
|
|
@ -3319,11 +3315,18 @@ if MCP_AVAILABLE:
|
|||
verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result)
|
||||
return call_tool_result
|
||||
|
||||
async def _handle_local_mcp_tool(
|
||||
name: str, arguments: dict[str, object]
|
||||
) -> list[TextContent | ImageContent | EmbeddedResource]:
|
||||
"""
|
||||
Handle tool execution for local registry tools
|
||||
async def _handle_local_mcp_tool(name: str, arguments: dict[str, object]) -> CallToolResult:
|
||||
"""Execute a local-registry tool and report whether it succeeded.
|
||||
|
||||
Returns the result rather than bare content because the verdict is part of it: the content
|
||||
alone cannot say whether the handler failed, so callers used to stamp isError=False on every
|
||||
outcome and an upstream rejection was served as tool output.
|
||||
|
||||
A failure is reported as ``isError=True`` here rather than raised, because the REST surface
|
||||
turns an unrecognized exception into a 500 and an upstream 403 or 429 is not a gateway crash.
|
||||
``MCPUpstreamAuthError`` is the exception: it propagates so the caller is told to
|
||||
re-authenticate, which both renderers already know how to say.
|
||||
|
||||
Note: Local tools don't use prefixes, so we use the original name
|
||||
"""
|
||||
import inspect
|
||||
|
|
@ -3333,15 +3336,16 @@ if MCP_AVAILABLE:
|
|||
raise HTTPException(status_code=404, detail=f"Tool '{name}' not found")
|
||||
|
||||
try:
|
||||
# Check if handler is async or sync
|
||||
if inspect.iscoroutinefunction(tool.handler):
|
||||
result = await tool.handler(**arguments)
|
||||
else:
|
||||
result = tool.handler(**arguments)
|
||||
return [TextContent(text=str(result), type="text")]
|
||||
except MCPUpstreamAuthError:
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error executing local tool %s: %s", name, e)
|
||||
return [TextContent(text=f"Error: {e}", type="text")]
|
||||
return CallToolResult(content=[TextContent(text=f"Error: {e}", type="text")], isError=True)
|
||||
return CallToolResult(content=[TextContent(text=str(result), type="text")], isError=False)
|
||||
|
||||
def _get_mcp_servers_in_path(path: str) -> list[str] | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -2439,6 +2439,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
None,
|
||||
description="max request size in MB, if a request is larger than this size it will be rejected",
|
||||
)
|
||||
max_batch_file_size_mb: int | None = Field(
|
||||
None,
|
||||
description="max batch input file size in MB for /v1/files uploads with purpose=batch, if a file is larger than this size it will be rejected before being forwarded to the provider",
|
||||
)
|
||||
max_response_size_mb: int | None = Field(
|
||||
None,
|
||||
description="max response size in MB, if a response is larger than this size it will be rejected",
|
||||
|
|
|
|||
|
|
@ -521,6 +521,20 @@ This is a one-time file patch and restore, not a live traffic interceptor. A Cla
|
|||
|
||||
Cursor is not supported: it has no equivalent file-based config to hot-patch this way, since its model routing lives in its own app storage and is configured through its GUI.
|
||||
|
||||
#### Making It Permanent at Login
|
||||
|
||||
`lite up` holds the patch only for as long as it runs. To wire Claude Code up once and leave it that way, pass `--config-claude` to `lite login`:
|
||||
|
||||
```bash
|
||||
lite --base-url https://your-proxy.example.com login --config-claude
|
||||
```
|
||||
|
||||
It writes the same two settings `lite up` does, `env.ANTHROPIC_BASE_URL` and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
|
||||
|
||||
Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`.
|
||||
|
||||
Run it again to point Claude Code at a different proxy; the base URL and the helper are both rewritten. `lite up` and `--config-claude` manage the same file, so the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first, rather than writing settings that `lite up` would silently revert when it stops.
|
||||
|
||||
### QA Complexity-Based Auto-Routing Against Your Real Proxy
|
||||
|
||||
`lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session.
|
||||
|
|
|
|||
|
|
@ -16,6 +16,12 @@ from typing_extensions import NotRequired, TypedDict
|
|||
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
|
||||
from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh
|
||||
|
||||
from .claude_settings import (
|
||||
CLAUDE_SETTINGS_PATH,
|
||||
SETTINGS_FILE_OWNERS,
|
||||
ClaudeSettingsError,
|
||||
write_claude_settings,
|
||||
)
|
||||
from .private_json import write_private_json
|
||||
|
||||
|
||||
|
|
@ -629,9 +635,28 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
def _configure_claude_code(base_url: str) -> None:
|
||||
"""Point Claude Code at base_url by patching ~/.claude/settings.json."""
|
||||
try:
|
||||
write_claude_settings(base_url, CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS)
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}")
|
||||
click.echo(f"\nConfigured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url.rstrip('/')}.")
|
||||
click.echo("Your other Claude Code settings were left untouched. Restart Claude Code to pick this up.")
|
||||
|
||||
|
||||
@click.command(name="login")
|
||||
@click.option(
|
||||
"--config-claude",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help=(
|
||||
"After logging in, update ~/.claude/settings.json so Claude Code routes through this proxy. "
|
||||
"Unrelated settings are preserved."
|
||||
),
|
||||
)
|
||||
@click.pass_context
|
||||
def login(ctx: click.Context):
|
||||
def login(ctx: click.Context, config_claude: bool):
|
||||
"""Login to LiteLLM proxy using SSO authentication"""
|
||||
from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
|
||||
from litellm.proxy.client.cli.interface import show_commands
|
||||
|
|
@ -683,6 +708,9 @@ def login(ctx: click.Context):
|
|||
click.echo(f"JWT Token: {api_key[:20]}...")
|
||||
click.echo("You can now use the CLI without specifying --api-key")
|
||||
|
||||
if config_claude:
|
||||
_configure_claude_code(base_url)
|
||||
|
||||
# Show available commands after successful login
|
||||
click.echo("\n" + "=" * 60)
|
||||
show_commands()
|
||||
|
|
@ -698,6 +726,10 @@ def login(ctx: click.Context):
|
|||
except KeyboardInterrupt:
|
||||
click.echo("\nAuthentication cancelled by user.")
|
||||
return
|
||||
except click.ClickException:
|
||||
# Login itself already succeeded; only the post-login step failed, so this
|
||||
# must not be relabelled as an authentication failure by the handler below.
|
||||
raise
|
||||
except Exception as e:
|
||||
click.echo(f"Authentication failed: {e}")
|
||||
return
|
||||
|
|
|
|||
|
|
@ -10,11 +10,16 @@ import click
|
|||
import yaml
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from ..up import CLAUDE_SETTINGS_PATH, UpError, load_json_or_empty, restore_claude_settings, write_backup
|
||||
from ..claude_settings import (
|
||||
AUTOROUTE_BACKUP_PATH,
|
||||
CLAUDE_SETTINGS_PATH,
|
||||
ClaudeSettingsError,
|
||||
load_json_or_empty,
|
||||
)
|
||||
from ..up import BackupRecord as ClaudeBackupRecord
|
||||
from ..up import restore_claude_settings, write_backup
|
||||
from .config import master_key_from_config
|
||||
from .process import (
|
||||
AUTOROUTE_DIR,
|
||||
CONFIG_PATH,
|
||||
DEFAULT_AUTOROUTE_PORT,
|
||||
LOG_PATH,
|
||||
|
|
@ -35,8 +40,6 @@ from .process import (
|
|||
from .settings import merge_claude_settings_static_token
|
||||
from .wizard import run_configure_wizard
|
||||
|
||||
AUTOROUTE_BACKUP_PATH: Final = AUTOROUTE_DIR / "claude_settings_backup.json"
|
||||
|
||||
_GENERATED_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
|
|
@ -108,7 +111,7 @@ def up(port: int) -> None:
|
|||
|
||||
try:
|
||||
existing_pid: Final = read_pid_record()
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(str(e))
|
||||
if existing_pid is not None and is_running(existing_pid.pid):
|
||||
raise click.ClickException(
|
||||
|
|
@ -157,7 +160,7 @@ def up(port: int) -> None:
|
|||
CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
with secure_create(CLAUDE_SETTINGS_PATH) as f:
|
||||
json.dump(merged, f, indent=2)
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
terminate(process.pid)
|
||||
clear_pid_record()
|
||||
raise click.ClickException(str(e))
|
||||
|
|
@ -175,7 +178,7 @@ def up(port: int) -> None:
|
|||
clear_pid_record()
|
||||
try:
|
||||
restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH)
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
# Runs from atexit/a signal handler too, outside Click's own exception
|
||||
# handling -- raising here would only produce an unhandled-exception
|
||||
# warning on stderr, not a clean message.
|
||||
|
|
@ -207,7 +210,7 @@ def down() -> None:
|
|||
"""Restore Claude Code settings and stop a leftover ephemeral proxy, if any"""
|
||||
try:
|
||||
record: PidRecord | None = read_pid_record()
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
# down is the crash-recovery path -- a corrupt pid record must not block it; clear the
|
||||
# unusable record and keep going rather than leaving the user with no way to clean up.
|
||||
click.echo(f"{e} Clearing it and continuing cleanup.", err=True)
|
||||
|
|
@ -219,7 +222,7 @@ def down() -> None:
|
|||
|
||||
try:
|
||||
restored: Final = restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH)
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(str(e))
|
||||
if restored is None:
|
||||
click.echo("Nothing to restore.")
|
||||
|
|
|
|||
155
litellm/proxy/client/cli/commands/claude_settings.py
Normal file
155
litellm/proxy/client/cli/commands/claude_settings.py
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
"""Shared handling of Claude Code's ~/.claude/settings.json.
|
||||
|
||||
`lite up` patches this file temporarily and restores it on exit; `lite login
|
||||
--config-claude` patches it persistently. Both need the same merge and the same
|
||||
apiKeyHelper command, and `up` already imports from `auth`, so the shared parts
|
||||
live here rather than in either command module.
|
||||
"""
|
||||
|
||||
import shlex
|
||||
import shutil
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from .private_json import write_private_json
|
||||
|
||||
ENV_KEY: Final = "env"
|
||||
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
|
||||
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
|
||||
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
|
||||
|
||||
CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json"
|
||||
BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json"
|
||||
AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SettingsFileOwner:
|
||||
"""A command that takes temporary ownership of CLAUDE_SETTINGS_PATH and restores it later."""
|
||||
|
||||
backup_path: Path
|
||||
start_command: str
|
||||
stop_command: str
|
||||
|
||||
|
||||
SETTINGS_FILE_OWNERS: Final = (
|
||||
SettingsFileOwner(BACKUP_PATH, "lite up", "lite down"),
|
||||
SettingsFileOwner(AUTOROUTE_BACKUP_PATH, "lite autoroute up", "lite autoroute down"),
|
||||
)
|
||||
|
||||
_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
class ClaudeSettingsError(Exception):
|
||||
"""Raised for any user-actionable failure while reading or writing Claude Code settings."""
|
||||
|
||||
|
||||
def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
|
||||
try:
|
||||
content: Final = path.read_bytes() if path.exists() else b""
|
||||
except OSError as e:
|
||||
raise ClaudeSettingsError(f"Could not read {path}: {e}") from e
|
||||
if not content.strip():
|
||||
return {}
|
||||
try:
|
||||
return _SETTINGS_ADAPTER.validate_json(content)
|
||||
except ValidationError:
|
||||
raise ClaudeSettingsError(
|
||||
f"{path} contains invalid JSON (or its root is not an object); cannot proceed safely."
|
||||
)
|
||||
|
||||
|
||||
def merge_claude_settings(
|
||||
settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str
|
||||
) -> dict[str, JsonValue]:
|
||||
"""Return a new settings dict wired to route Claude Code through the proxy.
|
||||
|
||||
Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a
|
||||
stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
|
||||
token (same reasoning as build_agent_env in agents.py). Every other key is
|
||||
preserved untouched.
|
||||
"""
|
||||
raw_env: Final = settings.get(ENV_KEY, {})
|
||||
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
|
||||
env: Final = {
|
||||
**{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY},
|
||||
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
|
||||
}
|
||||
return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
|
||||
|
||||
|
||||
def resolve_api_key_helper(base_url: str) -> str:
|
||||
"""Build the shell command Claude Code should run for its apiKeyHelper.
|
||||
|
||||
Resolves `lite` to an absolute path so the helper works regardless of the
|
||||
PATH visible to whatever subprocess Claude Code spawns it from. Passing
|
||||
--base-url explicitly (rather than relying on the bare invocation Claude
|
||||
Code would otherwise use) makes `print-token` enforce that the cached
|
||||
token was actually issued for this proxy -- without it, a token minted
|
||||
for a different, previously-logged-into proxy would be handed to
|
||||
whichever server the settings currently point at.
|
||||
|
||||
--base-url belongs to the top-level `lite` group, so it has to precede the
|
||||
subcommand; click rejects it outright after `print-token`.
|
||||
"""
|
||||
lite_path: Final = shutil.which("lite")
|
||||
if lite_path is None:
|
||||
raise ClaudeSettingsError(
|
||||
"Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs an absolute path to it."
|
||||
)
|
||||
return f"{shlex.quote(lite_path)} --base-url {shlex.quote(base_url)} auth print-token"
|
||||
|
||||
|
||||
def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
|
||||
"""Persistently point Claude Code at base_url, preserving every unrelated setting.
|
||||
|
||||
Refuses while any owner holds a backup: each restores its backup when it
|
||||
stops, which would silently undo this write.
|
||||
"""
|
||||
for owner in owners:
|
||||
if owner.backup_path.exists():
|
||||
raise ClaudeSettingsError(
|
||||
f"`{owner.start_command}` is currently managing {settings_path} (backup at "
|
||||
f"{owner.backup_path}) and will restore it when it stops. "
|
||||
f"Run `{owner.stop_command}` first, then retry."
|
||||
)
|
||||
normalized_base_url: Final = base_url.rstrip("/")
|
||||
api_key_helper: Final = resolve_api_key_helper(normalized_base_url)
|
||||
existing: Final = load_json_or_empty(settings_path)
|
||||
raw_env: Final = existing.get(ENV_KEY)
|
||||
if raw_env is not None and not isinstance(raw_env, dict):
|
||||
raise ClaudeSettingsError(
|
||||
f'{settings_path} has a non-object "{ENV_KEY}" value, which this would discard. '
|
||||
"Fix or remove it, then retry."
|
||||
)
|
||||
merged: Final = merge_claude_settings(existing, normalized_base_url, api_key_helper)
|
||||
# os.replace() swaps the symlink itself for a regular file, silently detaching a
|
||||
# settings.json that is symlinked into a dotfiles repo. There is no backup to undo
|
||||
# that here, unlike `lite up`, so write through to the link's target instead.
|
||||
target: Final = settings_path.resolve() if settings_path.is_symlink() else settings_path
|
||||
try:
|
||||
write_private_json(str(target), merged)
|
||||
except OSError as e:
|
||||
raise ClaudeSettingsError(f"Could not write {target}: {e}") from e
|
||||
|
||||
|
||||
__all__ = (
|
||||
"ANTHROPIC_API_KEY_KEY",
|
||||
"ANTHROPIC_BASE_URL_KEY",
|
||||
"API_KEY_HELPER_KEY",
|
||||
"AUTOROUTE_BACKUP_PATH",
|
||||
"BACKUP_PATH",
|
||||
"CLAUDE_SETTINGS_PATH",
|
||||
"ENV_KEY",
|
||||
"SETTINGS_FILE_OWNERS",
|
||||
"ClaudeSettingsError",
|
||||
"SettingsFileOwner",
|
||||
"load_json_or_empty",
|
||||
"merge_claude_settings",
|
||||
"resolve_api_key_helper",
|
||||
"write_claude_settings",
|
||||
)
|
||||
|
|
@ -2,12 +2,10 @@ import atexit
|
|||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import shutil
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Iterator, Mapping
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import FrameType
|
||||
|
|
@ -20,17 +18,17 @@ from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh
|
|||
|
||||
from .agents import AgentRunError, resolve_api_key, verify_proxy_key
|
||||
from .auth import load_token, login
|
||||
|
||||
ENV_KEY: Final = "env"
|
||||
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
|
||||
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
|
||||
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
|
||||
|
||||
CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json"
|
||||
BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json"
|
||||
from .claude_settings import (
|
||||
BACKUP_PATH,
|
||||
CLAUDE_SETTINGS_PATH,
|
||||
ClaudeSettingsError,
|
||||
load_json_or_empty,
|
||||
merge_claude_settings,
|
||||
resolve_api_key_helper,
|
||||
)
|
||||
|
||||
|
||||
class UpError(Exception):
|
||||
class UpError(ClaudeSettingsError):
|
||||
"""Raised for any user-actionable failure while starting/stopping interception."""
|
||||
|
||||
|
||||
|
|
@ -42,40 +40,9 @@ class BackupRecord:
|
|||
content: dict[str, JsonValue] | None
|
||||
|
||||
|
||||
_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_BACKUP_RECORD_ADAPTER: Final = TypeAdapter(BackupRecord)
|
||||
|
||||
|
||||
def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
|
||||
if not path.exists():
|
||||
return {}
|
||||
with open(path, "r") as f:
|
||||
content: Final = f.read()
|
||||
if not content.strip():
|
||||
return {}
|
||||
try:
|
||||
return _SETTINGS_ADAPTER.validate_json(content)
|
||||
except ValidationError:
|
||||
raise UpError(f"{path} contains invalid JSON (or its root is not an object); cannot proceed safely.")
|
||||
|
||||
|
||||
def merge_claude_settings(
|
||||
settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str
|
||||
) -> dict[str, JsonValue]:
|
||||
"""Return a new settings dict wired to route Claude Code through the proxy.
|
||||
|
||||
Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a
|
||||
stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
|
||||
token (same reasoning as build_agent_env in agents.py). Every other key is
|
||||
preserved untouched.
|
||||
"""
|
||||
raw_env: Final = settings.get(ENV_KEY, {})
|
||||
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
|
||||
env: Final = {**base_env, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/")}
|
||||
env.pop(ANTHROPIC_API_KEY_KEY, None)
|
||||
return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def secure_create(path: Path) -> Iterator[IO[str]]:
|
||||
"""Open path for writing with mode 0600 fixed up before any content is written.
|
||||
|
|
@ -136,26 +103,6 @@ def restore_claude_settings(settings_path: Path | None = None, backup_path: Path
|
|||
return record
|
||||
|
||||
|
||||
def resolve_api_key_helper(base_url: str) -> str:
|
||||
"""Build the shell command Claude Code should run for its apiKeyHelper.
|
||||
|
||||
Resolves `lite` to an absolute path so the helper works regardless of the
|
||||
PATH visible to whatever subprocess Claude Code spawns it from. Passing
|
||||
--base-url explicitly (rather than relying on the bare invocation Claude
|
||||
Code would otherwise use) makes `print-token` enforce that the cached
|
||||
token was actually issued for this proxy -- without it, a token minted
|
||||
for a different, previously-logged-into proxy would be handed to
|
||||
whichever server `up` currently points at.
|
||||
"""
|
||||
lite_path: Final = shutil.which("lite")
|
||||
if lite_path is None:
|
||||
raise UpError(
|
||||
"Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs "
|
||||
"an absolute path to it, so `lite up` cannot continue."
|
||||
)
|
||||
return f"{shlex.quote(lite_path)} auth print-token --base-url {shlex.quote(base_url)}"
|
||||
|
||||
|
||||
def _ensure_fresh_login(ctx: click.Context) -> None:
|
||||
base_url: Final = ctx.obj["base_url"].rstrip("/")
|
||||
token_data = load_token()
|
||||
|
|
@ -224,7 +171,7 @@ def up(ctx: click.Context) -> None:
|
|||
merged: Final = merge_claude_settings(original_settings, base_url, api_key_helper)
|
||||
with open(CLAUDE_SETTINGS_PATH, "w") as f:
|
||||
json.dump(merged, f, indent=2)
|
||||
except (AgentRunError, UpError) as e:
|
||||
except (AgentRunError, ClaudeSettingsError) as e:
|
||||
raise click.ClickException(str(e))
|
||||
|
||||
click.echo(f"litellm: routing Claude Code through proxy at {base_url.rstrip('/')}")
|
||||
|
|
@ -241,7 +188,7 @@ def up(ctx: click.Context) -> None:
|
|||
return
|
||||
try:
|
||||
_restore_and_report()
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
# Runs from atexit/a signal handler, outside Click's own exception
|
||||
# handling -- raising here would only produce an unhandled-exception
|
||||
# warning on stderr, not a clean message.
|
||||
|
|
@ -264,7 +211,7 @@ def down() -> None:
|
|||
"""
|
||||
try:
|
||||
_restore_and_report()
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(str(e))
|
||||
|
||||
|
||||
|
|
@ -272,6 +219,7 @@ __all__ = [
|
|||
"BACKUP_PATH",
|
||||
"CLAUDE_SETTINGS_PATH",
|
||||
"BackupRecord",
|
||||
"ClaudeSettingsError",
|
||||
"UpError",
|
||||
"down",
|
||||
"load_json_or_empty",
|
||||
|
|
|
|||
|
|
@ -1204,6 +1204,29 @@ class DBSpendUpdateWriter:
|
|||
except Exception as e:
|
||||
verbose_proxy_logger.debug("_flush_tool_discovery_queue error (non-blocking): %s", e)
|
||||
|
||||
@staticmethod
|
||||
async def _handle_spend_update_failure(
|
||||
e: Exception,
|
||||
attempt: int,
|
||||
n_retry_times: int,
|
||||
start_time: float,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> None:
|
||||
"""Retry a failed spend-update transaction on connection errors or deadlocks, else re-raise."""
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.utils import _raise_failed_update_spend_exception
|
||||
|
||||
is_retryable = isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) or PrismaDBExceptionHandler.is_deadlock_error(e)
|
||||
if not is_retryable or attempt >= n_retry_times:
|
||||
_raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj)
|
||||
verbose_proxy_logger.warning(
|
||||
"Retrying spend update after retryable DB error (attempt %s/%s): %s",
|
||||
attempt + 1,
|
||||
n_retry_times,
|
||||
e,
|
||||
)
|
||||
await asyncio.sleep(random.uniform(2**attempt, 2 ** (attempt + 1)))
|
||||
|
||||
async def _commit_spend_updates_to_db(
|
||||
self,
|
||||
prisma_client: PrismaClient,
|
||||
|
|
@ -1215,10 +1238,7 @@ class DBSpendUpdateWriter:
|
|||
Commits all the spend `UPDATE` transactions to the Database
|
||||
|
||||
"""
|
||||
from litellm.proxy.utils import (
|
||||
ProxyUpdateSpend,
|
||||
_raise_failed_update_spend_exception,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyUpdateSpend
|
||||
|
||||
### UPDATE USER TABLE ###
|
||||
user_list_transactions: Final = db_spend_update_transactions["user_list_transactions"]
|
||||
|
|
@ -1238,18 +1258,13 @@ class DBSpendUpdateWriter:
|
|||
data={"spend": {"increment": response_cost}},
|
||||
)
|
||||
break
|
||||
except DB_RETRY_SAFE_ERROR_TYPES as e:
|
||||
if i >= n_retry_times: # If we've reached the maximum number of retries
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
# Optionally, sleep for a bit before retrying
|
||||
await asyncio.sleep(2**i) # Exponential backoff
|
||||
except Exception as e:
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
|
||||
await self._handle_spend_update_failure(
|
||||
e=e,
|
||||
attempt=i,
|
||||
n_retry_times=n_retry_times,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
### UPDATE END-USER TABLE ###
|
||||
|
|
@ -1281,18 +1296,13 @@ class DBSpendUpdateWriter:
|
|||
},
|
||||
)
|
||||
break
|
||||
except DB_RETRY_SAFE_ERROR_TYPES as e:
|
||||
if i >= n_retry_times: # If we've reached the maximum number of retries
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
# Optionally, sleep for a bit before retrying
|
||||
await asyncio.sleep(2**i) # Exponential backoff
|
||||
except Exception as e:
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
|
||||
await self._handle_spend_update_failure(
|
||||
e=e,
|
||||
attempt=i,
|
||||
n_retry_times=n_retry_times,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
### UPDATE TEAM TABLE ###
|
||||
|
|
@ -1314,18 +1324,13 @@ class DBSpendUpdateWriter:
|
|||
data={"spend": {"increment": response_cost}},
|
||||
)
|
||||
break
|
||||
except DB_RETRY_SAFE_ERROR_TYPES as e:
|
||||
if i >= n_retry_times: # If we've reached the maximum number of retries
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
# Optionally, sleep for a bit before retrying
|
||||
await asyncio.sleep(2**i) # Exponential backoff
|
||||
except Exception as e:
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
|
||||
await self._handle_spend_update_failure(
|
||||
e=e,
|
||||
attempt=i,
|
||||
n_retry_times=n_retry_times,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
### UPDATE TEAM Membership TABLE with spend ###
|
||||
|
|
@ -1361,18 +1366,13 @@ class DBSpendUpdateWriter:
|
|||
)
|
||||
# Transaction succeeded, break out of retry loop
|
||||
break
|
||||
except DB_RETRY_SAFE_ERROR_TYPES as e:
|
||||
if i >= n_retry_times: # If we've reached the maximum number of retries
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
# Optionally, sleep for a bit before retrying
|
||||
await asyncio.sleep(2**i) # Exponential backoff
|
||||
except Exception as e:
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
|
||||
await self._handle_spend_update_failure(
|
||||
e=e,
|
||||
attempt=i,
|
||||
n_retry_times=n_retry_times,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# Invalidate cache for updated team memberships
|
||||
|
|
@ -1403,25 +1403,13 @@ class DBSpendUpdateWriter:
|
|||
data={"spend": {"increment": response_cost}},
|
||||
)
|
||||
break
|
||||
except DB_RETRY_SAFE_ERROR_TYPES as e:
|
||||
if i >= n_retry_times: # If we've reached the maximum number of retries
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
# Optionally, sleep for a bit before retrying
|
||||
await asyncio.sleep(
|
||||
# Sleep a random amount to avoid retrying and deadlocking again: when two transactions deadlock they are
|
||||
# cancelled basically at the same time, so if they wait the same time they will also retry at the same time
|
||||
# and thus they are more likely to deadlock again.
|
||||
# Instead, we sleep a random amount so that they retry at slightly different times, lowering the chance of
|
||||
# repeated deadlocks, and therefore of exceeding the retry limit.
|
||||
random.uniform(2**i, 2 ** (i + 1))
|
||||
)
|
||||
except Exception as e:
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
|
||||
await self._handle_spend_update_failure(
|
||||
e=e,
|
||||
attempt=i,
|
||||
n_retry_times=n_retry_times,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
### UPDATE TAG TABLE ###
|
||||
|
|
@ -1470,8 +1458,6 @@ class DBSpendUpdateWriter:
|
|||
prisma_client: Prisma client instance
|
||||
proxy_logging_obj: Proxy logging object
|
||||
"""
|
||||
from litellm.proxy.utils import _raise_failed_update_spend_exception
|
||||
|
||||
verbose_proxy_logger.debug("%s Spend transactions: %s", entity_name, transactions)
|
||||
if transactions is not None and len(transactions.keys()) > 0:
|
||||
for i in range(n_retry_times + 1):
|
||||
|
|
@ -1493,17 +1479,13 @@ class DBSpendUpdateWriter:
|
|||
data={"spend": {"increment": response_cost}},
|
||||
)
|
||||
break
|
||||
except DB_RETRY_SAFE_ERROR_TYPES as e:
|
||||
if i >= n_retry_times:
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
await asyncio.sleep(2**i) # Exponential backoff
|
||||
except Exception as e:
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
|
||||
await DBSpendUpdateWriter._handle_spend_update_failure(
|
||||
e=e,
|
||||
attempt=i,
|
||||
n_retry_times=n_retry_times,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# fmt: off
|
||||
|
|
@ -1672,7 +1654,16 @@ class DBSpendUpdateWriter:
|
|||
|
||||
break
|
||||
|
||||
except DB_RETRY_SAFE_ERROR_TYPES as e:
|
||||
except Exception as e:
|
||||
from litellm.proxy.db.exception_handler import (
|
||||
PrismaDBExceptionHandler,
|
||||
)
|
||||
|
||||
is_retryable = isinstance(
|
||||
e, DB_RETRY_SAFE_ERROR_TYPES
|
||||
) or PrismaDBExceptionHandler.is_deadlock_error(e)
|
||||
if not is_retryable:
|
||||
raise
|
||||
if i >= n_retry_times:
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e,
|
||||
|
|
|
|||
|
|
@ -166,6 +166,22 @@ class PrismaDBExceptionHandler:
|
|||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def is_deadlock_error(e: Exception) -> bool:
|
||||
"""True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma."""
|
||||
import prisma
|
||||
|
||||
if not isinstance(e, prisma.errors.PrismaError):
|
||||
return False
|
||||
if getattr(e, "code", None) == "P2034":
|
||||
return True
|
||||
error_message = str(e).lower()
|
||||
return (
|
||||
"deadlock detected" in error_message
|
||||
or "40p01" in error_message
|
||||
or "write conflict or a deadlock" in error_message
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def is_prisma_engine_internal_error(e: Exception) -> bool:
|
||||
"""True iff ``e`` is a non-``PrismaError`` exception raised from inside
|
||||
|
|
|
|||
93
litellm/proxy/db/proxy_worker_heartbeat.py
Normal file
93
litellm/proxy/db/proxy_worker_heartbeat.py
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
"""
|
||||
Live proxy worker census, one row per worker process.
|
||||
|
||||
Every uvicorn worker upserts its own row on a fixed heartbeat, so counting
|
||||
rows with a recent heartbeat answers "how many workers share this database?"
|
||||
without any coordination. The Admin UI's "no Redis" banner uses that count to
|
||||
hide itself for deployments that are provably a single worker, where per-worker
|
||||
rate limits, budgets, and router state are already global. All timestamps are
|
||||
written and compared with the database's own clock, so pods with skewed clocks
|
||||
still agree.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import socket
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS: Final = 60
|
||||
PROXY_WORKER_LIVENESS_WINDOW_SECONDS: Final = 3 * PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS
|
||||
STALE_ROW_RETENTION_SECONDS: Final = 3600
|
||||
|
||||
BEAT_SQL: Final = """
|
||||
INSERT INTO "LiteLLM_ProxyWorkerHeartbeat" (worker_id, hostname, last_heartbeat_at)
|
||||
VALUES ($1, $2, NOW())
|
||||
ON CONFLICT (worker_id) DO UPDATE SET last_heartbeat_at = NOW()
|
||||
"""
|
||||
|
||||
PRUNE_SQL: Final = """
|
||||
DELETE FROM "LiteLLM_ProxyWorkerHeartbeat"
|
||||
WHERE last_heartbeat_at < NOW() - make_interval(secs => $1)
|
||||
"""
|
||||
|
||||
COUNT_SQL: Final = """
|
||||
SELECT COUNT(*)::int AS live_workers FROM "LiteLLM_ProxyWorkerHeartbeat"
|
||||
WHERE last_heartbeat_at > NOW() - make_interval(secs => $1)
|
||||
"""
|
||||
|
||||
DEREGISTER_SQL: Final = """
|
||||
DELETE FROM "LiteLLM_ProxyWorkerHeartbeat" WHERE worker_id = $1
|
||||
"""
|
||||
|
||||
|
||||
class _LiveWorkerCountRow(TypedDict):
|
||||
live_workers: ReadOnly[int]
|
||||
|
||||
|
||||
_COUNT_ROWS_ADAPTER: Final = TypeAdapter(tuple[_LiveWorkerCountRow, ...])
|
||||
|
||||
|
||||
class ProxyWorkerHeartbeat:
|
||||
def __init__(self, prisma_client: PrismaClient, worker_id: str | None = None) -> None:
|
||||
self.prisma_client: Final = prisma_client
|
||||
self.worker_id: Final[str] = worker_id or str(uuid.uuid4())
|
||||
self.hostname: Final = socket.gethostname()
|
||||
|
||||
async def beat(self) -> None:
|
||||
try:
|
||||
await self.prisma_client.db.execute_raw(BEAT_SQL, self.worker_id, self.hostname)
|
||||
await self.prisma_client.db.execute_raw(PRUNE_SQL, STALE_ROW_RETENTION_SECONDS)
|
||||
except Exception as beat_err: # noqa: BLE001 # a missed heartbeat must never take down the worker
|
||||
verbose_proxy_logger.debug("Proxy worker heartbeat write failed: %s", beat_err)
|
||||
|
||||
async def deregister(self) -> None:
|
||||
try:
|
||||
await self.prisma_client.db.execute_raw(DEREGISTER_SQL, self.worker_id)
|
||||
except Exception as deregister_err: # noqa: BLE001 # best-effort cleanup; the liveness window ages the row out anyway
|
||||
verbose_proxy_logger.debug("Proxy worker heartbeat deregister failed: %s", deregister_err)
|
||||
|
||||
|
||||
async def count_live_proxy_workers(prisma_client: PrismaClient) -> int | None:
|
||||
"""
|
||||
The number of workers with a recent heartbeat, or None when the database
|
||||
cannot answer. Callers must treat None as "unknown", not as zero. Always
|
||||
counts on the primary: a lagging read replica must never undercount.
|
||||
"""
|
||||
try:
|
||||
db: Final = prisma_client.db
|
||||
primary_db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) else db
|
||||
rows: Final = await primary_db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS)
|
||||
return _COUNT_ROWS_ADAPTER.validate_python(rows)[0]["live_workers"]
|
||||
except Exception as count_err: # noqa: BLE001 # an unknown count must degrade to "warn", never to a 503
|
||||
verbose_proxy_logger.debug("Live proxy worker count unavailable: %s", count_err)
|
||||
return None
|
||||
|
|
@ -0,0 +1,40 @@
|
|||
# Claude Code / Anthropic-native web search on Bedrock, backed by
|
||||
# Amazon Bedrock AgentCore Web Search (AWS-managed web index, no third-party
|
||||
# search API). See litellm/llms/bedrock/search/transformation.py for details.
|
||||
|
||||
model_list:
|
||||
- model_name: claude-sonnet
|
||||
litellm_params:
|
||||
model: bedrock/us.anthropic.claude-sonnet-5
|
||||
aws_region_name: us-east-1
|
||||
|
||||
search_tools:
|
||||
- search_tool_name: agentcore-search
|
||||
litellm_params:
|
||||
search_provider: agentcore
|
||||
# Your AgentCore Gateway MCP endpoint (gateway must have a `web-search`
|
||||
# connector target). Alternatively set the AGENTCORE_GATEWAY_URL env var.
|
||||
api_base: https://<gateway-id>.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp
|
||||
|
||||
# The gateway exposes the connector as "<target-name>___WebSearch".
|
||||
# Default is "web-search-tool___WebSearch", matching the target name used
|
||||
# in the AWS docs' boto3/CLI setup examples. If your target was created
|
||||
# with a different name (misconfiguration surfaces as an MCP "tool not
|
||||
# found" error), set the AGENTCORE_SEARCH_TOOL_NAME env var or pass
|
||||
# tool_name in the request body. The search router forwards only
|
||||
# search_provider / api_key / api_base from this litellm_params block,
|
||||
# so a tool_name set here would be silently ignored.
|
||||
|
||||
# AWS_IAM gateway (default): SigV4-signed using the standard AWS
|
||||
# credential chain (env / profile / IRSA / instance role). Explicit
|
||||
# aws_access_key_id / aws_secret_access_key set here would be silently
|
||||
# ignored for the same reason; pass them per request instead.
|
||||
|
||||
# CUSTOM_JWT gateway alternative — OAuth2 bearer token instead of SigV4:
|
||||
# api_key: os.environ/AGENTCORE_GATEWAY_TOKEN
|
||||
|
||||
litellm_settings:
|
||||
callbacks: ["websearch_interception"]
|
||||
websearch_interception_params:
|
||||
enabled_providers: ["bedrock"]
|
||||
search_tool_name: agentcore-search
|
||||
|
|
@ -34,6 +34,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers
|
||||
from litellm.proxy.health_check import (
|
||||
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS,
|
||||
_clean_endpoint_data,
|
||||
|
|
@ -1451,7 +1452,7 @@ def callback_name(callback):
|
|||
DISABLE_NO_REDIS_WARNING_ENV_VAR: Final = "LITELLM_DISABLE_NO_REDIS_WARNING"
|
||||
|
||||
|
||||
def _show_no_redis_warning() -> bool:
|
||||
async def _show_no_redis_warning() -> bool:
|
||||
"""
|
||||
Whether the UI should warn that no Redis is configured.
|
||||
|
||||
|
|
@ -1461,16 +1462,22 @@ def _show_no_redis_warning() -> bool:
|
|||
coordination cache (from a Redis response cache, general_settings.
|
||||
coordination_redis, or the REDIS_* env fallback) and the router's own
|
||||
Redis (router_settings.redis_host), which backs cooldowns and usage-based
|
||||
routing on its own. Operators who know they run one worker can silence the
|
||||
warning with LITELLM_DISABLE_NO_REDIS_WARNING=true.
|
||||
routing on its own. A deployment whose worker-heartbeat census proves it
|
||||
is exactly one worker needs no cross-worker coordination, so it never
|
||||
warns; when the census is unavailable or shows more than one worker, the
|
||||
warning stands unless LITELLM_DISABLE_NO_REDIS_WARNING=true silences it.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router, redis_usage_cache
|
||||
from litellm.proxy.proxy_server import llm_router, prisma_client, redis_usage_cache
|
||||
|
||||
if redis_usage_cache is not None:
|
||||
return False
|
||||
if llm_router is not None and llm_router.cache.redis_cache is not None:
|
||||
return False
|
||||
return get_secret_bool(DISABLE_NO_REDIS_WARNING_ENV_VAR, False) is not True
|
||||
if get_secret_bool(DISABLE_NO_REDIS_WARNING_ENV_VAR, False) is True:
|
||||
return False
|
||||
if prisma_client is None:
|
||||
return True
|
||||
return await count_live_proxy_workers(prisma_client) != 1
|
||||
|
||||
|
||||
async def _get_health_readiness_details(
|
||||
|
|
@ -1513,7 +1520,7 @@ async def _get_health_readiness_details(
|
|||
# check log level
|
||||
log_level_name: Final = logging.getLevelName(verbose_logger.getEffectiveLevel())
|
||||
is_detailed_debug: Final = verbose_logger.isEnabledFor(logging.DEBUG)
|
||||
show_no_redis_warning: Final = _show_no_redis_warning()
|
||||
show_no_redis_warning: Final = await _show_no_redis_warning()
|
||||
|
||||
# check DB
|
||||
if prisma_client is not None: # if db passed in, check if it's connected
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ class KeyManagementEventHooks:
|
|||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
get_audit_log_changed_by,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
|
|
@ -53,8 +54,7 @@ class KeyManagementEventHooks:
|
|||
except Exception as e:
|
||||
verbose_proxy_logger.warning("Failed to send key created email: %s", e)
|
||||
|
||||
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
|
||||
if litellm.store_audit_logs is True:
|
||||
if is_audit_logging_enabled():
|
||||
_updated_values: Final = response.model_dump_json(exclude_none=True)
|
||||
asyncio.create_task(
|
||||
create_audit_log_for_update(
|
||||
|
|
@ -103,11 +103,11 @@ class KeyManagementEventHooks:
|
|||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
get_audit_log_changed_by,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
|
||||
if litellm.store_audit_logs is True:
|
||||
if is_audit_logging_enabled():
|
||||
_updated_values: Final = json.dumps(data.json(exclude_none=True), default=str)
|
||||
|
||||
_before_value = existing_key_row.json(exclude_none=True)
|
||||
|
|
@ -144,6 +144,7 @@ class KeyManagementEventHooks:
|
|||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
get_audit_log_changed_by,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
|
|
@ -180,7 +181,7 @@ class KeyManagementEventHooks:
|
|||
verbose_proxy_logger.warning("Failed to send key rotated email: %s", e)
|
||||
|
||||
# store the audit log
|
||||
if litellm.store_audit_logs is True and existing_key_row.token is not None:
|
||||
if is_audit_logging_enabled() and existing_key_row.token is not None:
|
||||
asyncio.create_task(
|
||||
create_audit_log_for_update(
|
||||
request_data=LiteLLM_AuditLogs(
|
||||
|
|
@ -218,12 +219,12 @@ class KeyManagementEventHooks:
|
|||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
get_audit_log_changed_by,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
|
||||
# we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes
|
||||
if litellm.store_audit_logs is True and data.keys is not None:
|
||||
if is_audit_logging_enabled() and data.keys is not None:
|
||||
# make an audit log for each key deleted
|
||||
for key in keys_being_deleted:
|
||||
if key.token is None:
|
||||
|
|
|
|||
|
|
@ -20,7 +20,10 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
WebhookEvent,
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
|
||||
|
||||
|
|
@ -203,7 +206,7 @@ class UserManagementEventHooks:
|
|||
- user_api_key_dict: UserAPIKeyAuth - The user api key dictionary.
|
||||
- litellm_proxy_admin_name: Optional[str] - The name of the proxy admin.
|
||||
"""
|
||||
if not litellm.store_audit_logs:
|
||||
if not is_audit_logging_enabled():
|
||||
return
|
||||
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
|
|
|
|||
|
|
@ -6,10 +6,13 @@ POST /auto_router/test_routing - Route one prompt through an unsaved complexity-
|
|||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from itertools import groupby
|
||||
from operator import attrgetter
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Protocol
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, field_validator
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.exceptions import BudgetExceededError
|
||||
|
|
@ -41,6 +44,8 @@ from litellm.types.management_endpoints.auto_router_endpoints import (
|
|||
AutoRouterRoutingTestRequest,
|
||||
AutoRouterRoutingTestResponse,
|
||||
RequestComplexityRouterConfig,
|
||||
ShadowEvalDirection,
|
||||
ShadowEvalJobKeyResponse,
|
||||
ShadowEvalJobResponse,
|
||||
ShadowEvalResult,
|
||||
ShadowEvalSlice,
|
||||
|
|
@ -89,17 +94,9 @@ class _ShadowEvalJobRow(Protocol):
|
|||
|
||||
|
||||
class _ShadowEvalJobTable(Protocol):
|
||||
async def find_unique(self, *, where: Mapping[str, object]) -> _ShadowEvalJobRow | None: ...
|
||||
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_ShadowEvalJobRow]: ...
|
||||
|
||||
async def find_first(self, *, where: Mapping[str, object]) -> _ShadowEvalJobRow | None: ...
|
||||
|
||||
async def find_many(
|
||||
self, *, where: Mapping[str, object], order: Mapping[str, str], take: int
|
||||
) -> Sequence[_ShadowEvalJobRow]: ...
|
||||
|
||||
async def create(self, data: Mapping[str, object]) -> _ShadowEvalJobRow: ...
|
||||
|
||||
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> _ShadowEvalJobRow | None: ...
|
||||
async def create_many(self, data: Sequence[Mapping[str, object]]) -> int: ...
|
||||
|
||||
|
||||
class _ShadowEvalAttemptRow(Protocol):
|
||||
|
|
@ -606,18 +603,19 @@ _ATTEMPT_AGG_SELECT: Final = """
|
|||
COUNT(*) FILTER (WHERE outcome = 'tie')::int AS ties,
|
||||
AVG(confidence)::float AS avg_confidence
|
||||
FROM "LiteLLM_ShadowEvalAttempt"
|
||||
WHERE job_id = $1 AND outcome != 'error'
|
||||
WHERE job_id = ANY($1::text[]) AND outcome != 'error'
|
||||
GROUP BY 1
|
||||
"""
|
||||
|
||||
_ATTEMPT_AGG_BY_TIER_SQL: Final = "SELECT COALESCE(tier, 'UNCLASSIFIED') AS grp," + _ATTEMPT_AGG_SELECT
|
||||
_ATTEMPT_AGG_BY_MODEL_SQL: Final = "SELECT COALESCE(real_model, 'unknown') AS grp," + _ATTEMPT_AGG_SELECT
|
||||
_ATTEMPT_AGG_BY_LEG_SQL: Final = "SELECT job_id AS grp," + _ATTEMPT_AGG_SELECT
|
||||
|
||||
_SWEEP_FINISHED_JOBS_SQL: Final = """
|
||||
UPDATE "LiteLLM_ShadowEvalJob" j SET stopped_at = NOW()
|
||||
WHERE j.api_key_id = $1 AND j.stopped_at IS NULL
|
||||
UPDATE "LiteLLM_ShadowEvalJob" j SET stopped_at = (NOW() AT TIME ZONE 'utc')
|
||||
WHERE j.api_key_id = ANY($1::text[]) AND j.stopped_at IS NULL
|
||||
AND (
|
||||
j.ends_at <= NOW()
|
||||
j.ends_at <= (NOW() AT TIME ZONE 'utc')
|
||||
OR (SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a WHERE a.job_id = j.id) >= j.max_turns
|
||||
)
|
||||
"""
|
||||
|
|
@ -628,7 +626,52 @@ SELECT
|
|||
COUNT(*) FILTER (WHERE outcome = 'error')::int AS error_count,
|
||||
COALESCE(SUM(judge_cost), 0)::float AS judge_spend
|
||||
FROM "LiteLLM_ShadowEvalAttempt"
|
||||
WHERE job_id = $1
|
||||
WHERE job_id = ANY($1::text[])
|
||||
"""
|
||||
|
||||
_ATTEMPT_COUNTS_SQL: Final = """
|
||||
SELECT a.job_id, COUNT(*)::int AS attempt_count
|
||||
FROM "LiteLLM_ShadowEvalAttempt" a
|
||||
JOIN "LiteLLM_ShadowEvalJob" j ON j.id = a.job_id
|
||||
WHERE a.job_id = ANY($1::text[]) AND (j.stopped_at IS NULL OR a.created_at <= j.stopped_at)
|
||||
GROUP BY a.job_id
|
||||
"""
|
||||
|
||||
_STOP_JOB_SQL: Final = """
|
||||
UPDATE "LiteLLM_ShadowEvalJob"
|
||||
SET stopped_by = $2, stopped_at = COALESCE(stopped_at, $3::timestamp)
|
||||
WHERE group_id = $1 AND stopped_by IS NULL
|
||||
AND ends_at > (NOW() AT TIME ZONE 'utc')
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM "LiteLLM_ShadowEvalJob" k
|
||||
WHERE k.group_id = $1 AND k.stopped_at IS NULL
|
||||
AND (SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a WHERE a.job_id = k.id) < k.max_turns
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
class _AttemptCountRow(BaseModel):
|
||||
job_id: str
|
||||
attempt_count: int
|
||||
|
||||
|
||||
_ATTEMPT_COUNT_ROWS: Final = TypeAdapter(list[_AttemptCountRow])
|
||||
|
||||
|
||||
_LIST_LEGS_SQL: Final = """
|
||||
SELECT * FROM "LiteLLM_ShadowEvalJob"
|
||||
WHERE group_id IN (
|
||||
SELECT group_id FROM "LiteLLM_ShadowEvalJob"
|
||||
GROUP BY group_id ORDER BY MAX(created_at) DESC LIMIT $1::int
|
||||
)
|
||||
"""
|
||||
|
||||
_LIST_LEGS_BY_KEY_SQL: Final = """
|
||||
SELECT * FROM "LiteLLM_ShadowEvalJob"
|
||||
WHERE group_id IN (
|
||||
SELECT group_id FROM "LiteLLM_ShadowEvalJob" WHERE api_key_id = $2
|
||||
GROUP BY group_id ORDER BY MAX(created_at) DESC LIMIT $1::int
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
|
|
@ -659,18 +702,98 @@ def _slices(rows: Sequence[_AttemptAggRow]) -> tuple[ShadowEvalSlice, ...]:
|
|||
)
|
||||
|
||||
|
||||
class _LegRow(BaseModel):
|
||||
"""One LiteLLM_ShadowEvalJob row, validated off the untyped prisma record. A row is
|
||||
one key's leg of a job; the legs of a job share group_id and identical config, written
|
||||
together by one create_many. The API's job id is the group id, so leg ids never leave
|
||||
the server (attempts reference them internally)."""
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
id: str
|
||||
group_id: str
|
||||
api_key_id: str
|
||||
router_name: str
|
||||
direction: ShadowEvalDirection
|
||||
baseline_model: str | None = None
|
||||
judge_model: str
|
||||
shadow_percentage: float
|
||||
max_turns: int
|
||||
created_at: datetime
|
||||
ends_at: datetime
|
||||
stopped_at: datetime | None = None
|
||||
stopped_by: str | None = None
|
||||
|
||||
@field_validator("created_at", "ends_at", "stopped_at")
|
||||
@classmethod
|
||||
def _as_aware_utc(cls, value: datetime | None) -> datetime | None:
|
||||
"""The columns store naive UTC wall time (prisma's convention); prisma reads hand
|
||||
back aware datetimes while raw SQL reads hand back naive ones, so this boundary
|
||||
makes every read aware UTC before anything compares or serializes them."""
|
||||
if value is None or value.tzinfo is not None:
|
||||
return value
|
||||
return value.replace(tzinfo=timezone.utc)
|
||||
|
||||
|
||||
_LEG_ROWS: Final = TypeAdapter(list[_LegRow])
|
||||
|
||||
|
||||
async def _leg_attempt_counts(prisma_client: "PrismaClient", legs: Sequence[_LegRow]) -> Mapping[str, int]:
|
||||
"""Each leg's attempt count by leg id, judged and errored alike, in one grouped read.
|
||||
It is the same count the sampler budgets against max_turns, so the derived status
|
||||
flips to completed exactly when sampling actually ends. A stamped leg's count freezes
|
||||
at its stopped_at: in-flight attempts that land after the stamp are excluded, so they
|
||||
can never reclassify a leg that was stopped under budget as budget-spent."""
|
||||
if not legs:
|
||||
return MappingProxyType({})
|
||||
rows: Final = _ATTEMPT_COUNT_ROWS.validate_python(
|
||||
await _query_raw(prisma_client, _ATTEMPT_COUNTS_SQL, [leg.id for leg in legs]) # mutable-ok: query param
|
||||
or ()
|
||||
)
|
||||
return MappingProxyType({row.job_id: row.attempt_count for row in rows})
|
||||
|
||||
|
||||
def _group_response(group_id: str, legs: Sequence[_LegRow], attempt_counts: Mapping[str, int]) -> ShadowEvalJobResponse:
|
||||
"""The one constructor of a job response: the caller names the group and passes that
|
||||
group's legs. Config is read off the first leg because every leg carries the same copy,
|
||||
written by one create_many. No caller may serialize a raw row (that would leak a leg id
|
||||
as the job id)."""
|
||||
first: Final = legs[0]
|
||||
return ShadowEvalJobResponse(
|
||||
job_id=group_id,
|
||||
keys=tuple(
|
||||
ShadowEvalJobKeyResponse(
|
||||
api_key_id=leg.api_key_id,
|
||||
max_turns=leg.max_turns,
|
||||
stopped_at=leg.stopped_at,
|
||||
attempt_count=attempt_counts.get(leg.id, 0),
|
||||
)
|
||||
for leg in sorted(legs, key=lambda leg: leg.api_key_id)
|
||||
),
|
||||
router_name=first.router_name,
|
||||
direction=first.direction,
|
||||
baseline_model=first.baseline_model,
|
||||
judge_model=first.judge_model,
|
||||
shadow_percentage=first.shadow_percentage,
|
||||
created_at=first.created_at,
|
||||
ends_at=first.ends_at,
|
||||
stopped_by=next((leg.stopped_by for leg in legs if leg.stopped_by is not None), None),
|
||||
)
|
||||
|
||||
|
||||
_NO_KEY_LABELS: Final[tuple[str | None, str | None]] = (None, None)
|
||||
|
||||
|
||||
async def _with_key_labels(
|
||||
prisma_client: "PrismaClient", responses: Sequence[ShadowEvalJobResponse]
|
||||
) -> tuple[ShadowEvalJobResponse, ...]:
|
||||
"""Resolve each job's key hash to the key's alias and masked name in one batched read,
|
||||
"""Resolve every scoped key's hash to its alias and masked name in one batched read,
|
||||
so the UI can say whose traffic a job shadows. Deleted keys resolve to None."""
|
||||
if not responses:
|
||||
return ()
|
||||
tokens: Final = sorted(frozenset(key.api_key_id for response in responses for key in response.keys))
|
||||
key_rows: Final = await _verification_tokens(prisma_client).find_many(
|
||||
where={"token": {"in": sorted({response.api_key_id for response in responses})}} # mutable-ok: Prisma filter
|
||||
where={"token": {"in": tokens}} # mutable-ok: Prisma filter
|
||||
)
|
||||
labels: Final[Mapping[str, tuple[str | None, str | None]]] = {
|
||||
row.token: (row.key_alias, row.key_name) for row in key_rows or ()
|
||||
|
|
@ -678,32 +801,50 @@ async def _with_key_labels(
|
|||
return tuple(
|
||||
response.model_copy(
|
||||
update={ # mutable-ok: pydantic update payload
|
||||
"key_alias": labels.get(response.api_key_id, _NO_KEY_LABELS)[0],
|
||||
"key_name": labels.get(response.api_key_id, _NO_KEY_LABELS)[1],
|
||||
"keys": tuple(
|
||||
key.model_copy(
|
||||
update={ # mutable-ok: pydantic update payload
|
||||
"key_alias": labels.get(key.api_key_id, _NO_KEY_LABELS)[0],
|
||||
"key_name": labels.get(key.api_key_id, _NO_KEY_LABELS)[1],
|
||||
}
|
||||
)
|
||||
for key in response.keys
|
||||
)
|
||||
}
|
||||
)
|
||||
for response in responses
|
||||
)
|
||||
|
||||
|
||||
async def _shadow_eval_results(prisma_client: "PrismaClient", job_id: str) -> ShadowEvalResult | None:
|
||||
"""Both stratifications of one job's verdicts. Tier answers "where does the router do
|
||||
well"; the model stratification groups by whichever model served the real arm, so it
|
||||
answers "which of the models this key uses today would the router beat" forward, and
|
||||
"for the turns the router sent to X, did X beat the baseline" in reverse. Reads are
|
||||
bounded by the job's own attempts (<= max_turns) via the job_id index."""
|
||||
async def _shadow_eval_results(prisma_client: "PrismaClient", legs: Sequence[_LegRow]) -> ShadowEvalResult | None:
|
||||
"""All three stratifications of one job's verdicts. Tier answers "where does the router
|
||||
do well"; the model stratification groups by whichever model served the real arm, so it
|
||||
answers "which of the models these keys use today would the router beat" forward, and
|
||||
"for the turns the router sent to X, did X beat the baseline" in reverse; key answers
|
||||
"which key's traffic does the router suit". Reads are bounded by the job's own attempts
|
||||
(<= the sum of its keys' max_turns) via the job_id index."""
|
||||
leg_ids: Final = [leg.id for leg in legs] # mutable-ok: query param
|
||||
by_tier: Final = _ATTEMPT_AGG_ROWS.validate_python(
|
||||
await _query_raw(prisma_client, _ATTEMPT_AGG_BY_TIER_SQL, job_id) or ()
|
||||
await _query_raw(prisma_client, _ATTEMPT_AGG_BY_TIER_SQL, leg_ids) or ()
|
||||
)
|
||||
if not by_tier:
|
||||
return None
|
||||
by_model: Final = _ATTEMPT_AGG_ROWS.validate_python(
|
||||
await _query_raw(prisma_client, _ATTEMPT_AGG_BY_MODEL_SQL, job_id) or ()
|
||||
await _query_raw(prisma_client, _ATTEMPT_AGG_BY_MODEL_SQL, leg_ids) or ()
|
||||
)
|
||||
key_by_leg: Final = MappingProxyType({leg.id: leg.api_key_id for leg in legs})
|
||||
by_leg: Final = _ATTEMPT_AGG_ROWS.validate_python(
|
||||
await _query_raw(prisma_client, _ATTEMPT_AGG_BY_LEG_SQL, leg_ids) or ()
|
||||
)
|
||||
by_key: Final = tuple(
|
||||
row.model_copy(update={"grp": key_by_leg[row.grp]}) # mutable-ok: pydantic update payload
|
||||
for row in by_leg
|
||||
)
|
||||
total_turns: Final = sum(r.turn_count for r in by_tier)
|
||||
return ShadowEvalResult(
|
||||
by_tier=_slices(by_tier),
|
||||
by_current_model=_slices(by_model),
|
||||
by_key=_slices(by_key),
|
||||
overall_shadow_win_rate_pct=_pct_of(sum(r.shadow_wins for r in by_tier), total_turns),
|
||||
overall_tie_rate_pct=_pct_of(sum(r.ties for r in by_tier), total_turns),
|
||||
)
|
||||
|
|
@ -721,20 +862,21 @@ async def start_shadow_eval(
|
|||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> ShadowEvalJobResponse:
|
||||
"""
|
||||
Start a shadow eval: duplicate a sampled slice of a key's live traffic against a second
|
||||
arm, judge the two responses blind, and stratify win rates by tier and by the model that
|
||||
served the real arm.
|
||||
Start a shadow eval: duplicate a sampled slice of one or more keys' live traffic against
|
||||
a second arm, judge the two responses blind, and stratify win rates by tier, by the model
|
||||
that served the real arm, and by key.
|
||||
|
||||
A forward job answers whether the key should adopt router_name: it samples the requests
|
||||
A forward job answers whether the keys should adopt router_name: it samples the requests
|
||||
the router did not serve and duplicates them through it. A reverse job answers whether a
|
||||
key already on the router still gains from it: it samples the requests the router did
|
||||
serve and duplicates them against baseline_model. A key can hold one active job per
|
||||
direction, so both questions can run at once.
|
||||
|
||||
Shadow responses are never served to users. The job samples until it has judged
|
||||
max_turns turns, reaches the end of its window, or is stopped; sampling changes
|
||||
propagate to pods within about 10 seconds. Shadow and judge calls bill to the
|
||||
shadowed key but are excluded from request counts and auto-router adoption metrics.
|
||||
Shadow responses are never served to users. Each key samples until it has judged
|
||||
max_turns turns of its own traffic, the job's window ends, or the job is stopped, so one
|
||||
key running out of budget does not end sampling for the others; sampling changes
|
||||
propagate to pods within about 10 seconds. Shadow and judge calls bill to the shadowed
|
||||
key but are excluded from request counts and auto-router adoption metrics.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router, prisma_client
|
||||
|
||||
|
|
@ -746,48 +888,58 @@ async def start_shadow_eval(
|
|||
_validate_plain_model(llm_router, data.judge_model, "judge_model")
|
||||
if data.baseline_model is not None:
|
||||
_validate_plain_model(llm_router, data.baseline_model, "baseline_model")
|
||||
key_row: Final = await _verification_tokens(prisma_client).find_unique(
|
||||
where={"token": data.api_key_id} # mutable-ok: Prisma filter
|
||||
token_rows: Final = await _verification_tokens(prisma_client).find_many(
|
||||
where={"token": {"in": list(data.api_key_ids)}} # mutable-ok: Prisma filter
|
||||
)
|
||||
if key_row is None:
|
||||
unknown: Final = tuple(sorted(frozenset(data.api_key_ids) - frozenset(row.token for row in token_rows or ())))
|
||||
if unknown:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
f"api_key_id '{data.api_key_id}' is not a key on this proxy; pass the key's token hash, "
|
||||
f"api_key_ids not on this proxy: {', '.join(unknown)}; pass each key's token hash, "
|
||||
"the value the key list and key info endpoints report"
|
||||
),
|
||||
)
|
||||
|
||||
# A job that expired or exhausted its turn budget stopped sampling on its own, but
|
||||
# still holds its slot in the per-key, per-direction partial unique index until
|
||||
# stamped; free it so a new eval can start. Sweeping both directions is deliberate.
|
||||
await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, data.api_key_id)
|
||||
active: Final = await _shadow_eval_jobs(prisma_client).find_first(
|
||||
# A job whose window passed or whose turn budget ran out stopped sampling on its own,
|
||||
# but its legs still hold their slots in the per-key, per-direction partial unique index
|
||||
# until stamped; free them so a new eval can start. Sweeping both directions is deliberate.
|
||||
requested: Final = list(data.api_key_ids) # mutable-ok: query param
|
||||
await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, requested)
|
||||
claimed: Final = await _shadow_eval_jobs(prisma_client).find_many(
|
||||
where={ # mutable-ok: Prisma filter
|
||||
"api_key_id": data.api_key_id,
|
||||
"api_key_id": {"in": requested}, # mutable-ok: Prisma filter
|
||||
"direction": data.direction,
|
||||
"stopped_at": None,
|
||||
},
|
||||
)
|
||||
if active is not None:
|
||||
if claimed:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=f"Key already has an active {data.direction} shadow eval job ({active.id}). Stop it first.",
|
||||
detail=(
|
||||
f"Already in an active {data.direction} shadow eval job: "
|
||||
+ ", ".join(sorted(f"{row.api_key_id} (job {row.group_id})" for row in claimed))
|
||||
+ ". Stop it first."
|
||||
),
|
||||
)
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
group_id: Final = str(uuid4())
|
||||
ends_at: Final = now + timedelta(days=data.duration_days)
|
||||
shared_config: Final = { # mutable-ok: Prisma payload
|
||||
"group_id": group_id,
|
||||
"router_name": data.router_name,
|
||||
"direction": data.direction,
|
||||
"baseline_model": data.baseline_model,
|
||||
"judge_model": data.judge_model,
|
||||
"shadow_percentage": data.shadow_percentage,
|
||||
"max_turns": data.max_turns,
|
||||
"created_by": user_api_key_dict.user_id,
|
||||
"created_at": now,
|
||||
"ends_at": ends_at,
|
||||
}
|
||||
try:
|
||||
job: Final = await _shadow_eval_jobs(prisma_client).create(
|
||||
data={ # mutable-ok: Prisma payload
|
||||
"api_key_id": data.api_key_id,
|
||||
"router_name": data.router_name,
|
||||
"direction": data.direction,
|
||||
"baseline_model": data.baseline_model,
|
||||
"judge_model": data.judge_model,
|
||||
"shadow_percentage": data.shadow_percentage,
|
||||
"max_turns": data.max_turns,
|
||||
"created_by": user_api_key_dict.user_id,
|
||||
"ends_at": now + timedelta(days=data.duration_days),
|
||||
}
|
||||
await _shadow_eval_jobs(prisma_client).create_many(
|
||||
data=[{**shared_config, "api_key_id": key} for key in data.api_key_ids] # mutable-ok: Prisma payload
|
||||
)
|
||||
except Exception as e:
|
||||
if not _is_unique_violation(e):
|
||||
|
|
@ -795,11 +947,28 @@ async def start_shadow_eval(
|
|||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=(
|
||||
f"Key already has an active {data.direction} shadow eval job (started concurrently). Stop it first."
|
||||
f"A requested key was claimed by another {data.direction} shadow eval job concurrently. Stop it first."
|
||||
),
|
||||
) from e
|
||||
return ShadowEvalJobResponse.model_validate(job, from_attributes=True).model_copy(
|
||||
update={"key_alias": key_row.key_alias, "key_name": key_row.key_name} # mutable-ok: pydantic update payload
|
||||
labels: Final = MappingProxyType({row.token: row for row in token_rows})
|
||||
return ShadowEvalJobResponse(
|
||||
job_id=group_id,
|
||||
keys=tuple(
|
||||
ShadowEvalJobKeyResponse(
|
||||
api_key_id=api_key_id,
|
||||
max_turns=data.max_turns,
|
||||
key_alias=labels[api_key_id].key_alias,
|
||||
key_name=labels[api_key_id].key_name,
|
||||
)
|
||||
for api_key_id in sorted(data.api_key_ids)
|
||||
),
|
||||
router_name=data.router_name,
|
||||
direction=data.direction,
|
||||
baseline_model=data.baseline_model,
|
||||
judge_model=data.judge_model,
|
||||
shadow_percentage=data.shadow_percentage,
|
||||
created_at=now,
|
||||
ends_at=ends_at,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -811,23 +980,38 @@ async def start_shadow_eval(
|
|||
)
|
||||
async def list_shadow_eval_jobs(
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
api_key_id: Annotated[str | None, Query(description="Filter to jobs shadowing this key")] = None,
|
||||
api_key_id: Annotated[
|
||||
str | None, Query(description="Filter to jobs that shadow this key, alone or alongside others")
|
||||
] = None,
|
||||
limit: Annotated[int, Query(ge=1, le=200, description="Newest jobs to return")] = 50,
|
||||
) -> tuple[ShadowEvalJobResponse, ...]:
|
||||
"""List shadow eval jobs, newest first. Counts and results ride the detail endpoint only."""
|
||||
"""List shadow eval jobs, newest first, each key with its attempt count so status is
|
||||
accurate. Judged counts, spend, and results ride the detail endpoint only."""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
_require_admin_viewer(user_api_key_dict, "view shadow evals")
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
records: Final = await _shadow_eval_jobs(prisma_client).find_many(
|
||||
where={"api_key_id": api_key_id} if api_key_id else {}, # mutable-ok: Prisma filter
|
||||
order={"created_at": "desc"}, # mutable-ok: Prisma order
|
||||
take=limit,
|
||||
legs: Final = _LEG_ROWS.validate_python(
|
||||
(
|
||||
await _query_raw(prisma_client, _LIST_LEGS_BY_KEY_SQL, limit, api_key_id)
|
||||
if api_key_id
|
||||
else await _query_raw(prisma_client, _LIST_LEGS_SQL, limit)
|
||||
)
|
||||
or ()
|
||||
)
|
||||
by_group: Final[Mapping[str, tuple[_LegRow, ...]]] = MappingProxyType(
|
||||
{
|
||||
group_id: tuple(group)
|
||||
for group_id, group in groupby(sorted(legs, key=attrgetter("group_id")), key=attrgetter("group_id"))
|
||||
}
|
||||
)
|
||||
newest_first: Final = sorted(
|
||||
by_group, key=lambda group_id: max(leg.created_at for leg in by_group[group_id]), reverse=True
|
||||
)
|
||||
counts: Final = await _leg_attempt_counts(prisma_client, legs)
|
||||
return await _with_key_labels(
|
||||
prisma_client,
|
||||
tuple(ShadowEvalJobResponse.model_validate(record, from_attributes=True) for record in records or ()),
|
||||
prisma_client, tuple(_group_response(group_id, by_group[group_id], counts) for group_id in newest_first)
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -847,20 +1031,24 @@ async def get_shadow_eval_job(
|
|||
_require_admin_viewer(user_api_key_dict, "view shadow evals")
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
record: Final = await _shadow_eval_jobs(prisma_client).find_unique(
|
||||
where={"id": job_id} # mutable-ok: Prisma filter
|
||||
legs: Final = _LEG_ROWS.validate_python(
|
||||
await _shadow_eval_jobs(prisma_client).find_many(
|
||||
where={"group_id": job_id} # mutable-ok: Prisma filter
|
||||
)
|
||||
or ()
|
||||
)
|
||||
if record is None:
|
||||
if not legs:
|
||||
raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}")
|
||||
leg_ids: Final = [leg.id for leg in legs] # mutable-ok: query param
|
||||
totals: Final = _ATTEMPT_TOTALS_ROWS.validate_python(
|
||||
await _query_raw(prisma_client, _ATTEMPT_TOTALS_SQL, job_id) or ()
|
||||
await _query_raw(prisma_client, _ATTEMPT_TOTALS_SQL, leg_ids) or ()
|
||||
)
|
||||
latest_error: Final = await _shadow_eval_attempts(prisma_client).find_first(
|
||||
where={"job_id": job_id, "outcome": "error"}, # mutable-ok: Prisma filter
|
||||
where={"job_id": {"in": leg_ids}, "outcome": "error"}, # mutable-ok: Prisma filter
|
||||
order={"created_at": "desc"}, # mutable-ok: Prisma order
|
||||
)
|
||||
labeled: Final = await _with_key_labels(
|
||||
prisma_client, (ShadowEvalJobResponse.model_validate(record, from_attributes=True),)
|
||||
prisma_client, (_group_response(job_id, legs, await _leg_attempt_counts(prisma_client, legs)),)
|
||||
)
|
||||
return labeled[0].model_copy(
|
||||
update={ # mutable-ok: pydantic update payload
|
||||
|
|
@ -868,7 +1056,7 @@ async def get_shadow_eval_job(
|
|||
"error_count": totals[0].error_count if totals else 0,
|
||||
"judge_spend": round(totals[0].judge_spend, 6) if totals else 0.0,
|
||||
"last_error": latest_error.error if latest_error else None,
|
||||
"results": await _shadow_eval_results(prisma_client, job_id),
|
||||
"results": await _shadow_eval_results(prisma_client, legs),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -883,25 +1071,33 @@ async def stop_shadow_eval_job(
|
|||
job_id: str,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> ShadowEvalJobResponse:
|
||||
"""Stop an active shadow eval job. Attempts are kept; sampling halts within ~10s."""
|
||||
"""Stop an active shadow eval job, every key it scopes at once. Attempts are kept;
|
||||
sampling halts within ~10s. Keys that already stopped on their own budget keep the
|
||||
stopped_at they earned. The statement is the whole state machine: it claims the job
|
||||
only while a leg still samples inside the window with no stop recorded, so a racing
|
||||
operator, a same-instant budget spend, and a repeat stop all read the same 400 with
|
||||
the status the job actually holds."""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
_require_admin_writer(user_api_key_dict, "stop a shadow eval")
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
record: Final = await _shadow_eval_jobs(prisma_client).find_unique(
|
||||
where={"id": job_id} # mutable-ok: Prisma filter
|
||||
stamp: Final = datetime.now(timezone.utc)
|
||||
operator: Final = user_api_key_dict.user_id or "operator"
|
||||
claimed: Final = await prisma_client.db.execute_raw(
|
||||
_STOP_JOB_SQL, job_id, operator, stamp.replace(tzinfo=None).isoformat()
|
||||
)
|
||||
if record is None:
|
||||
legs: Final = _LEG_ROWS.validate_python(
|
||||
await _shadow_eval_jobs(prisma_client).find_many(
|
||||
where={"group_id": job_id} # mutable-ok: Prisma filter
|
||||
)
|
||||
or ()
|
||||
)
|
||||
if not legs:
|
||||
raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}")
|
||||
current: Final = ShadowEvalJobResponse.model_validate(record, from_attributes=True)
|
||||
if current.status != "running":
|
||||
counts: Final = await _leg_attempt_counts(prisma_client, legs)
|
||||
current: Final = _group_response(job_id, legs, counts)
|
||||
if claimed == 0:
|
||||
raise HTTPException(status_code=400, detail=f"Job {job_id} is already {current.status}")
|
||||
updated: Final = await _shadow_eval_jobs(prisma_client).update(
|
||||
where={"id": job_id}, # mutable-ok: Prisma filter
|
||||
data={"stopped_at": datetime.now(timezone.utc)}, # mutable-ok: Prisma payload
|
||||
)
|
||||
labeled: Final = await _with_key_labels(
|
||||
prisma_client, (ShadowEvalJobResponse.model_validate(updated, from_attributes=True),)
|
||||
)
|
||||
labeled: Final = await _with_key_labels(prisma_client, (current,))
|
||||
return labeled[0]
|
||||
|
|
|
|||
|
|
@ -17,7 +17,6 @@ from typing import TYPE_CHECKING, Any, Final, Protocol
|
|||
from fastapi import APIRouter, Depends, Header, HTTPException
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._redis import _redis_kwargs_from_environment
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -299,14 +298,15 @@ async def _emit_cache_settings_audit_log(
|
|||
exception. Captured under ``LiteLLM_CacheConfig`` so the row
|
||||
co-locates with the table it mutates.
|
||||
"""
|
||||
if litellm.store_audit_logs is not True:
|
||||
return
|
||||
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
if not is_audit_logging_enabled():
|
||||
return
|
||||
|
||||
task: Final = asyncio.create_task(
|
||||
create_audit_log_for_update(
|
||||
request_data=LiteLLM_AuditLogs(
|
||||
|
|
|
|||
|
|
@ -100,16 +100,15 @@ async def _emit_hashicorp_vault_audit_log(
|
|||
``LiteLLM_ConfigOverrides`` so the row co-locates with the table it
|
||||
mutates.
|
||||
"""
|
||||
import litellm
|
||||
|
||||
if litellm.store_audit_logs is not True:
|
||||
return
|
||||
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
if not is_audit_logging_enabled():
|
||||
return
|
||||
|
||||
task: Final = asyncio.create_task(
|
||||
create_audit_log_for_update(
|
||||
request_data=LiteLLM_AuditLogs(
|
||||
|
|
|
|||
|
|
@ -243,12 +243,15 @@ async def _emit_coordination_redis_audit_log(
|
|||
litellm_changed_by: str | None,
|
||||
) -> None:
|
||||
"""Emit an audit-log row for a /coordination_redis/settings mutation."""
|
||||
if litellm.store_audit_logs is not True:
|
||||
return
|
||||
|
||||
from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
if not is_audit_logging_enabled():
|
||||
return
|
||||
|
||||
task: Final = asyncio.create_task(
|
||||
create_audit_log_for_update(
|
||||
request_data=LiteLLM_AuditLogs(
|
||||
|
|
|
|||
|
|
@ -2220,6 +2220,7 @@ async def delete_user(
|
|||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import (
|
||||
create_audit_log_for_update,
|
||||
|
|
@ -2298,9 +2299,8 @@ async def delete_user(
|
|||
},
|
||||
)
|
||||
|
||||
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
|
||||
# we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes
|
||||
if litellm.store_audit_logs is True:
|
||||
if is_audit_logging_enabled():
|
||||
# make an audit log for each team deleted
|
||||
_user_row = user_row.json(exclude_none=True)
|
||||
|
||||
|
|
|
|||
|
|
@ -1423,6 +1423,11 @@ async def _check_team_key_limits(
|
|||
)
|
||||
|
||||
|
||||
_INHERITED_MODEL_SENTINELS: Final = frozenset(
|
||||
{SpecialModelNames.all_team_models.value, SpecialModelNames.all_proxy_models.value}
|
||||
)
|
||||
|
||||
|
||||
async def _check_project_key_limits(
|
||||
project_id: str,
|
||||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||||
|
|
@ -1432,7 +1437,8 @@ async def _check_project_key_limits(
|
|||
"""
|
||||
Validate that key's models and budget respect its project's limits.
|
||||
|
||||
- Key models must be a subset of project models
|
||||
- Key models must be a subset of project models, except the all-team-models / all-proxy-models
|
||||
sentinels, which inherit a parent scope and are narrowed by the project at request time
|
||||
- Key max_budget must be <= project max_budget
|
||||
"""
|
||||
project_obj: Final = await get_project_object(
|
||||
|
|
@ -1450,7 +1456,7 @@ async def _check_project_key_limits(
|
|||
# Validate key models are a subset of project models
|
||||
if data.models and len(project_obj.models) > 0:
|
||||
for m in data.models:
|
||||
if m not in project_obj.models:
|
||||
if m not in project_obj.models and m not in _INHERITED_MODEL_SENTINELS:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
|
|
@ -6266,6 +6272,7 @@ async def block_key(
|
|||
"""
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import (
|
||||
create_audit_log_for_update,
|
||||
|
|
@ -6312,7 +6319,7 @@ async def block_key(
|
|||
code=status.HTTP_404_NOT_FOUND,
|
||||
)
|
||||
|
||||
if litellm.store_audit_logs is True:
|
||||
if is_audit_logging_enabled():
|
||||
asyncio.create_task(
|
||||
create_audit_log_for_update(
|
||||
request_data=LiteLLM_AuditLogs(
|
||||
|
|
@ -6379,6 +6386,7 @@ async def unblock_key(
|
|||
"""
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import (
|
||||
create_audit_log_for_update,
|
||||
|
|
@ -6425,7 +6433,7 @@ async def unblock_key(
|
|||
code=status.HTTP_404_NOT_FOUND,
|
||||
)
|
||||
|
||||
if litellm.store_audit_logs is True:
|
||||
if is_audit_logging_enabled():
|
||||
asyncio.create_task(
|
||||
create_audit_log_for_update(
|
||||
request_data=LiteLLM_AuditLogs(
|
||||
|
|
|
|||
|
|
@ -64,7 +64,10 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import get_audit_log_changed_by
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.repositories.table_repositories import (
|
||||
MCPServerRepository,
|
||||
MCPUserCredentialsRepository,
|
||||
|
|
@ -2018,7 +2021,7 @@ if MCP_AVAILABLE:
|
|||
await global_mcp_server_manager.reload_servers_from_database()
|
||||
|
||||
# TODO: Enterprise: Finish audit log trail
|
||||
if litellm.store_audit_logs:
|
||||
if is_audit_logging_enabled():
|
||||
pass
|
||||
|
||||
# TODO: Delete from virtual keys
|
||||
|
|
@ -2613,7 +2616,7 @@ if MCP_AVAILABLE:
|
|||
)
|
||||
|
||||
# TODO: Enterprise: Finish audit log trail
|
||||
if litellm.store_audit_logs:
|
||||
if is_audit_logging_enabled():
|
||||
pass
|
||||
|
||||
return _redact_mcp_credentials(mcp_server_record_updated)
|
||||
|
|
|
|||
|
|
@ -24,6 +24,13 @@ from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
|
||||
from litellm.litellm_core_utils.ptu_pricing import (
|
||||
CUSTOM_PRICING_FIELDS,
|
||||
PTU_EMPTIED_PRICING_FIELDS,
|
||||
PTU_ZEROED_PRICING_FIELDS,
|
||||
PTU_ZEROED_TABLE_FIELDS,
|
||||
SEARCH_CONTEXT_SIZES,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
BlockModelRequest,
|
||||
CommonProxyErrors,
|
||||
|
|
@ -89,7 +96,6 @@ from litellm.types.router import (
|
|||
ModelInfo,
|
||||
updateDeployment,
|
||||
)
|
||||
from litellm.types.utils import CustomPricingLiteLLMParams
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
|
@ -346,12 +352,8 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None:
|
|||
# tiered_pricing is the one mirrored field that is a table of ranges, not a rate, so it is stored
|
||||
# empty (see _PTU_EMPTIED_PRICING_FIELDS): its tiers outrank the zeros written beside them, so
|
||||
# dropping it would leave the cost map's tiers billing the traffic the reserved capacity covers.
|
||||
_PTU_ZEROED_PRICING_FIELDS: Final = tuple(f for f in SPECIAL_MODEL_INFO_PARAMS if f != "tiered_pricing") + (
|
||||
"cache_creation_input_token_cost_above_1hr",
|
||||
"cache_creation_input_token_cost_above_200k_tokens",
|
||||
"cache_read_input_token_cost_above_200k_tokens",
|
||||
)
|
||||
_PTU_EMPTIED_PRICING_FIELDS: Final = frozenset({"tiered_pricing"})
|
||||
_PTU_ZEROED_PRICING_FIELDS: Final = PTU_ZEROED_PRICING_FIELDS
|
||||
_PTU_EMPTIED_PRICING_FIELDS: Final = PTU_EMPTIED_PRICING_FIELDS
|
||||
_PTU_ZEROED_PRICING: Final[Mapping[str, float | tuple[()]]] = MappingProxyType(
|
||||
{
|
||||
**dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0),
|
||||
|
|
@ -363,13 +365,13 @@ _EMPTY_MODEL_INFO: Final[Mapping[str, object]] = _NO_PRICING_OVERRIDE
|
|||
# Rate fields only. CustomPricingLiteLLMParams also carries settings that are not charges
|
||||
# (an embedding's output_vector_size, the regional uplift multipliers), and zeroing one of
|
||||
# those would destroy the deployment's configuration rather than stop a charge.
|
||||
_CUSTOM_PRICING_FIELDS: Final = frozenset(f for f in CustomPricingLiteLLMParams.model_fields if "cost" in f)
|
||||
_CUSTOM_PRICING_FIELDS: Final = CUSTOM_PRICING_FIELDS
|
||||
# search_context_cost_per_query holds its rates in a table keyed by context size, and an absent
|
||||
# table means the provider's own default rate rather than free (litellm/llms/gemini/cost_calculator
|
||||
# falls back to $0.035), so it is zeroed in place rather than emptied like tiered_pricing, and
|
||||
# written on every PTU deployment rather than only where a table is already stored.
|
||||
_PTU_ZEROED_TABLE_FIELDS: Final = frozenset({"search_context_cost_per_query"})
|
||||
_SEARCH_CONTEXT_SIZES: Final = ("search_context_size_low", "search_context_size_medium", "search_context_size_high")
|
||||
_PTU_ZEROED_TABLE_FIELDS: Final = PTU_ZEROED_TABLE_FIELDS
|
||||
_SEARCH_CONTEXT_SIZES: Final = SEARCH_CONTEXT_SIZES
|
||||
|
||||
|
||||
def _is_nonzero_rate(value: object) -> bool:
|
||||
|
|
|
|||
|
|
@ -13,7 +13,6 @@ from typing import Annotated, Any, Final
|
|||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.proxy._types import (
|
||||
|
|
@ -182,14 +181,15 @@ async def _emit_team_callback_audit_log(
|
|||
Callback secrets are redacted before serialization so the audit table
|
||||
cannot itself become a credential-harvest sink.
|
||||
"""
|
||||
if litellm.store_audit_logs is not True:
|
||||
return
|
||||
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_audit_log_for_update,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name
|
||||
|
||||
if not is_audit_logging_enabled():
|
||||
return
|
||||
|
||||
redacted_before: Final = _redact_callback_secrets(before_metadata)
|
||||
redacted_after: Final = _redact_callback_secrets(after_metadata)
|
||||
|
||||
|
|
|
|||
|
|
@ -1252,6 +1252,7 @@ async def new_team(
|
|||
try:
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import (
|
||||
_license_check,
|
||||
|
|
@ -1560,8 +1561,7 @@ async def new_team(
|
|||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
|
||||
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
|
||||
if litellm.store_audit_logs is True:
|
||||
if is_audit_logging_enabled():
|
||||
_updated_values = complete_team_data.json(exclude_none=True)
|
||||
|
||||
_updated_values = json.dumps(_updated_values, default=str)
|
||||
|
|
@ -1953,6 +1953,7 @@ async def update_team(
|
|||
```
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.management_helpers.audit_logs import is_audit_logging_enabled
|
||||
from litellm.proxy.proxy_server import (
|
||||
litellm_proxy_admin_name,
|
||||
llm_router,
|
||||
|
|
@ -2261,8 +2262,7 @@ async def update_team(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
|
||||
if litellm.store_audit_logs is True:
|
||||
if is_audit_logging_enabled():
|
||||
await _create_team_update_audit_log(
|
||||
existing_team_row=existing_team_row,
|
||||
updated_kv=updated_kv,
|
||||
|
|
@ -3727,6 +3727,7 @@ async def delete_team(
|
|||
"""
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
get_audit_log_changed_by,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
from litellm.proxy.proxy_server import (
|
||||
create_audit_log_for_update,
|
||||
|
|
@ -3771,9 +3772,8 @@ async def delete_team(
|
|||
litellm_changed_by=litellm_changed_by,
|
||||
)
|
||||
|
||||
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
|
||||
# we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes
|
||||
if litellm.store_audit_logs is True:
|
||||
if is_audit_logging_enabled():
|
||||
# make an audit log for each team deleted
|
||||
for team_id in data.team_ids:
|
||||
team_row: LiteLLM_TeamTable | None = await prisma_client.get_data(
|
||||
|
|
|
|||
|
|
@ -24,6 +24,22 @@ _audit_log_callback_cache: Final[dict[str, CustomLogger]] = {}
|
|||
ALLOW_LITELLM_CHANGED_BY_HEADER_METADATA_KEY: Final = "allow_litellm_changed_by_header"
|
||||
|
||||
|
||||
def is_audit_logging_enabled(store_audit_logs: bool | None = None) -> bool:
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
configured_value: Final[bool | None] = litellm.store_audit_logs if store_audit_logs is None else store_audit_logs
|
||||
if configured_value is not None:
|
||||
return configured_value
|
||||
|
||||
environment_value: Final[bool | None] = get_secret_bool("LITELLM_STORE_AUDIT_LOGS")
|
||||
if environment_value is not None:
|
||||
return environment_value
|
||||
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
||||
return premium_user is True
|
||||
|
||||
|
||||
def _allows_litellm_changed_by_header(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
for admin_metadata in (user_api_key_dict.metadata, user_api_key_dict.team_metadata):
|
||||
if (
|
||||
|
|
@ -164,11 +180,7 @@ async def create_object_audit_log(
|
|||
- user_api_key_dict: UserAPIKeyAuth - The user api key dictionary.
|
||||
- litellm_proxy_admin_name: Optional[str] - The name of the proxy admin.
|
||||
"""
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
_store_audit_logs: Final[bool | None] = litellm.store_audit_logs or get_secret_bool("LITELLM_STORE_AUDIT_LOGS")
|
||||
|
||||
if _store_audit_logs is not True:
|
||||
if not is_audit_logging_enabled():
|
||||
return
|
||||
|
||||
_changed_by: Final = get_audit_log_changed_by(
|
||||
|
|
@ -196,10 +208,7 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
|
|||
"""
|
||||
Create an audit log for an object.
|
||||
"""
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
_store_audit_logs: Final[bool | None] = litellm.store_audit_logs or get_secret_bool("LITELLM_STORE_AUDIT_LOGS")
|
||||
if _store_audit_logs is not True:
|
||||
if not is_audit_logging_enabled():
|
||||
return
|
||||
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client
|
||||
|
|
|
|||
182
litellm/proxy/openai_files_endpoints/batch_file_validation.py
Normal file
182
litellm/proxy/openai_files_endpoints/batch_file_validation.py
Normal file
|
|
@ -0,0 +1,182 @@
|
|||
import json
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from typing import BinaryIO, Final, NoReturn, assert_never
|
||||
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
BATCH_LINE_REQUIRED_KEYS: Final = ("custom_id", "method", "url", "body")
|
||||
_MB: Final = 1024 * 1024
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BatchFileTooLarge:
|
||||
size_bytes: int
|
||||
limit_mb: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BatchFileWrongExtension:
|
||||
filename: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BatchFileEmpty:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BatchFileInvalidJsonLine:
|
||||
line_number: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BatchFileLineNotObject:
|
||||
line_number: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BatchFileMissingLineKey:
|
||||
line_number: int
|
||||
key: str
|
||||
|
||||
|
||||
BatchFileValidationFailure = (
|
||||
BatchFileTooLarge
|
||||
| BatchFileWrongExtension
|
||||
| BatchFileEmpty
|
||||
| BatchFileInvalidJsonLine
|
||||
| BatchFileLineNotObject
|
||||
| BatchFileMissingLineKey
|
||||
)
|
||||
|
||||
|
||||
def _file_size_bytes(file_source: bytes | BinaryIO) -> int:
|
||||
if isinstance(file_source, bytes):
|
||||
return len(file_source)
|
||||
file_source.seek(0, 2)
|
||||
size: Final = file_source.tell()
|
||||
file_source.seek(0)
|
||||
return size
|
||||
|
||||
|
||||
def _iter_lines(file_source: bytes | BinaryIO) -> Iterator[bytes]:
|
||||
if isinstance(file_source, bytes):
|
||||
return iter(file_source.splitlines())
|
||||
file_source.seek(0)
|
||||
return iter(file_source)
|
||||
|
||||
|
||||
def _check_line(line_number: int, raw_line: bytes) -> BatchFileValidationFailure | None:
|
||||
try:
|
||||
parsed: Final = json.loads(raw_line)
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
return BatchFileInvalidJsonLine(line_number=line_number)
|
||||
if not isinstance(parsed, dict):
|
||||
return BatchFileLineNotObject(line_number=line_number)
|
||||
missing: Final = next((key for key in BATCH_LINE_REQUIRED_KEYS if key not in parsed), None)
|
||||
if missing is None:
|
||||
return None
|
||||
return BatchFileMissingLineKey(line_number=line_number, key=missing)
|
||||
|
||||
|
||||
def _scan_lines(file_source: bytes | BinaryIO) -> BatchFileValidationFailure | None:
|
||||
content_lines: Final = (
|
||||
(line_number, raw_line)
|
||||
for line_number, raw_line in enumerate(_iter_lines(file_source), start=1)
|
||||
if raw_line.strip()
|
||||
)
|
||||
first_line: Final = next(content_lines, None)
|
||||
if first_line is None:
|
||||
return BatchFileEmpty()
|
||||
return next(
|
||||
(
|
||||
failure
|
||||
for line_number, raw_line in chain((first_line,), content_lines)
|
||||
for failure in (_check_line(line_number, raw_line),)
|
||||
if failure is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def check_batch_file_upload(
|
||||
filename: str | None,
|
||||
file_source: bytes | BinaryIO,
|
||||
max_batch_file_size_mb: int | None,
|
||||
) -> BatchFileValidationFailure | None:
|
||||
if filename is None or not filename.lower().endswith(".jsonl"):
|
||||
return BatchFileWrongExtension(filename=filename or "")
|
||||
if max_batch_file_size_mb is not None and max_batch_file_size_mb > 0:
|
||||
size_bytes: Final = _file_size_bytes(file_source)
|
||||
if size_bytes > max_batch_file_size_mb * _MB:
|
||||
return BatchFileTooLarge(size_bytes=size_bytes, limit_mb=max_batch_file_size_mb)
|
||||
scan_failure: Final = _scan_lines(file_source)
|
||||
if not isinstance(file_source, bytes):
|
||||
file_source.seek(0)
|
||||
return scan_failure
|
||||
|
||||
|
||||
def raise_batch_file_validation_failure(failure: BatchFileValidationFailure) -> NoReturn:
|
||||
match failure:
|
||||
case BatchFileTooLarge(size_bytes=size_bytes, limit_mb=limit_mb):
|
||||
raise ProxyException(
|
||||
message=(
|
||||
f"Batch input file is {size_bytes / _MB:.1f} MB, which exceeds the configured "
|
||||
f"max_batch_file_size_mb of {limit_mb} MB. The file was not forwarded to the provider."
|
||||
),
|
||||
type="invalid_request_error",
|
||||
param="file",
|
||||
code=413,
|
||||
)
|
||||
case BatchFileWrongExtension(filename=filename):
|
||||
raise ProxyException(
|
||||
message=(
|
||||
f"Invalid file format for Batch API: '{filename}'. "
|
||||
"Batch input files must be .jsonl files. The file was not forwarded to the provider."
|
||||
),
|
||||
type="invalid_request_error",
|
||||
param="file",
|
||||
code=400,
|
||||
)
|
||||
case BatchFileEmpty():
|
||||
raise ProxyException(
|
||||
message="Batch input file has no request lines. The file was not forwarded to the provider.",
|
||||
type="invalid_request_error",
|
||||
param="file",
|
||||
code=400,
|
||||
)
|
||||
case BatchFileInvalidJsonLine(line_number=line_number):
|
||||
raise ProxyException(
|
||||
message=(
|
||||
f"Batch input file line {line_number} is not valid JSON. "
|
||||
"The file was not forwarded to the provider."
|
||||
),
|
||||
type="invalid_request_error",
|
||||
param="file",
|
||||
code=400,
|
||||
)
|
||||
case BatchFileLineNotObject(line_number=line_number):
|
||||
raise ProxyException(
|
||||
message=(
|
||||
f"Batch input file line {line_number} must be a JSON object. "
|
||||
"The file was not forwarded to the provider."
|
||||
),
|
||||
type="invalid_request_error",
|
||||
param="file",
|
||||
code=400,
|
||||
)
|
||||
case BatchFileMissingLineKey(line_number=line_number, key=key):
|
||||
raise ProxyException(
|
||||
message=(
|
||||
f"Missing required parameter: '{key}' (batch input file line {line_number}). "
|
||||
f"Each line must be a JSON object with keys {', '.join(BATCH_LINE_REQUIRED_KEYS)}. "
|
||||
"The file was not forwarded to the provider."
|
||||
),
|
||||
type="invalid_request_error",
|
||||
param=key,
|
||||
code=400,
|
||||
)
|
||||
case _:
|
||||
assert_never(failure)
|
||||
|
|
@ -21,6 +21,7 @@ from fastapi import (
|
|||
UploadFile,
|
||||
status,
|
||||
)
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm import CreateFileRequest, get_secret_str
|
||||
|
|
@ -41,6 +42,10 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
|
|||
get_custom_llm_provider_from_request_headers,
|
||||
get_custom_llm_provider_from_request_query,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.batch_file_validation import (
|
||||
check_batch_file_upload,
|
||||
raise_batch_file_validation_failure,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_is_base64_encoded_unified_file_id,
|
||||
add_internal_model_credentials,
|
||||
|
|
@ -65,6 +70,8 @@ from litellm.types.llms.openai import (
|
|||
|
||||
router: Final = APIRouter()
|
||||
|
||||
_MAX_BATCH_FILE_SIZE_MB_ADAPTER: Final = TypeAdapter(int | None)
|
||||
|
||||
files_config = None
|
||||
|
||||
|
||||
|
|
@ -361,18 +368,27 @@ async def create_file(
|
|||
|
||||
# Prepare the data for forwarding
|
||||
|
||||
# Replace with:
|
||||
valid_purposes: Final = get_args(OpenAIFilesPurpose)
|
||||
if purpose not in valid_purposes:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"Invalid purpose: {purpose}. Must be one of: {valid_purposes}",
|
||||
},
|
||||
raise ProxyException(
|
||||
message=f"Invalid purpose: {purpose}. Must be one of: {valid_purposes}",
|
||||
type="invalid_request_error",
|
||||
param="purpose",
|
||||
code=400,
|
||||
)
|
||||
# Cast purpose to OpenAIFilesPurpose type
|
||||
purpose = cast(OpenAIFilesPurpose, purpose)
|
||||
|
||||
if purpose == "batch":
|
||||
batch_file_failure: Final = await asyncio.to_thread(
|
||||
check_batch_file_upload,
|
||||
file.filename,
|
||||
file_source,
|
||||
_MAX_BATCH_FILE_SIZE_MB_ADAPTER.validate_python(general_settings.get("max_batch_file_size_mb")),
|
||||
)
|
||||
if batch_file_failure is not None:
|
||||
raise_batch_file_validation_failure(batch_file_failure)
|
||||
|
||||
data = {}
|
||||
|
||||
# Parse expires_after if provided
|
||||
|
|
@ -552,6 +568,8 @@ async def create_file(
|
|||
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
|
||||
)
|
||||
verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_file(): Exception occured - %s", e)
|
||||
if isinstance(e, ProxyException):
|
||||
raise e
|
||||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e.detail)),
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import VERTEX_BATCH_PREDICTION_JOBS_ROUTE
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.vertex_ai.common_utils import get_vertex_location_from_url
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator as VertexModelResponseIterator,
|
||||
)
|
||||
|
|
@ -60,6 +61,9 @@ class VertexPassthroughLoggingHandler:
|
|||
request_body: dict | None = None,
|
||||
**kwargs,
|
||||
) -> PassThroughEndpointLoggingTypedDict:
|
||||
vertex_location: Final = get_vertex_location_from_url(url_route)
|
||||
if vertex_location is not None:
|
||||
logging_obj.optional_params["vertex_location"] = vertex_location
|
||||
if "predictLongRunning" in url_route:
|
||||
model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route)
|
||||
|
||||
|
|
@ -82,6 +86,7 @@ class VertexPassthroughLoggingHandler:
|
|||
model=model,
|
||||
custom_llm_provider="vertex_ai",
|
||||
call_type="create_video",
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
# Set response_cost in _hidden_params to prevent recalculation
|
||||
|
|
@ -123,6 +128,7 @@ class VertexPassthroughLoggingHandler:
|
|||
end_time=end_time,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route),
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
return {
|
||||
|
|
@ -190,6 +196,7 @@ class VertexPassthroughLoggingHandler:
|
|||
end_time=end_time,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
return {
|
||||
|
|
@ -206,6 +213,7 @@ class VertexPassthroughLoggingHandler:
|
|||
model="vertex_ai/search_api",
|
||||
custom_llm_provider="vertex_ai",
|
||||
call_type="vector_store_search",
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
standard_pass_through_response_object: Final[StandardPassThroughResponseObject] = {
|
||||
|
|
@ -302,6 +310,7 @@ class VertexPassthroughLoggingHandler:
|
|||
completion_response=litellm_prediction_response,
|
||||
model=model,
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_location=get_vertex_location_from_url(url_route),
|
||||
)
|
||||
|
||||
kwargs["response_cost"] = response_cost
|
||||
|
|
@ -381,6 +390,7 @@ class VertexPassthroughLoggingHandler:
|
|||
completion_response=litellm_embedding_response,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_location=get_vertex_location_from_url(url_route),
|
||||
)
|
||||
|
||||
kwargs["response_cost"] = response_cost
|
||||
|
|
@ -413,6 +423,9 @@ class VertexPassthroughLoggingHandler:
|
|||
- Logs in litellm callbacks
|
||||
"""
|
||||
kwargs: dict[str, Any] = {}
|
||||
vertex_location: Final = get_vertex_location_from_url(url_route)
|
||||
if vertex_location is not None:
|
||||
litellm_logging_obj.optional_params["vertex_location"] = vertex_location
|
||||
model = model or VertexPassthroughLoggingHandler.extract_model_from_url(url_route)
|
||||
complete_streaming_response: Final = VertexPassthroughLoggingHandler._build_complete_streaming_response(
|
||||
all_chunks=all_chunks,
|
||||
|
|
@ -438,6 +451,7 @@ class VertexPassthroughLoggingHandler:
|
|||
end_time=end_time,
|
||||
logging_obj=litellm_logging_obj,
|
||||
custom_llm_provider=VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route),
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
return {
|
||||
|
|
@ -591,6 +605,7 @@ class VertexPassthroughLoggingHandler:
|
|||
end_time: datetime,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
custom_llm_provider: str,
|
||||
vertex_location: str | None,
|
||||
) -> dict:
|
||||
"""
|
||||
Create the standard logging object for Vertex passthrough generateContent (streaming and non-streaming)
|
||||
|
|
@ -601,6 +616,7 @@ class VertexPassthroughLoggingHandler:
|
|||
completion_response=litellm_model_response,
|
||||
model=model,
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
kwargs["response_cost"] = response_cost
|
||||
|
|
|
|||
|
|
@ -384,6 +384,10 @@ from litellm.proxy.db.gateway_request_tracking import (
|
|||
GatewayRequestAccumulator,
|
||||
flush_gateway_requests,
|
||||
)
|
||||
from litellm.proxy.db.proxy_worker_heartbeat import (
|
||||
PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS,
|
||||
ProxyWorkerHeartbeat,
|
||||
)
|
||||
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
|
||||
from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router
|
||||
from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router
|
||||
|
|
@ -635,6 +639,7 @@ from litellm.secret_managers.main import (
|
|||
get_secret_bool,
|
||||
get_secret_str,
|
||||
normalize_nonempty_secret_str,
|
||||
secret_manager_would_be_consulted,
|
||||
str_to_bool,
|
||||
)
|
||||
from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs
|
||||
|
|
@ -874,9 +879,11 @@ async def _flush_spend_logs_queue_on_shutdown() -> None:
|
|||
verbose_proxy_logger.exception("Error flushing spend logs queue on shutdown: %s", e)
|
||||
|
||||
|
||||
async def proxy_shutdown_event() -> None:
|
||||
async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = None) -> None:
|
||||
global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update
|
||||
verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server")
|
||||
if worker_heartbeat is not None and prisma_client:
|
||||
await worker_heartbeat.deregister()
|
||||
if prisma_client:
|
||||
# Drain the SGR fold first: it lives in memory, so an un-drained interval
|
||||
# is lost, and a write attempted after disconnect raises
|
||||
|
|
@ -1210,7 +1217,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
|
|||
)
|
||||
|
||||
### START BATCH WRITING DB + CHECKING NEW MODELS###
|
||||
if prisma_client is not None:
|
||||
worker_heartbeat: Final = (
|
||||
await ProxyStartupEvent.initialize_scheduled_background_jobs(
|
||||
general_settings=general_settings,
|
||||
prisma_client=prisma_client,
|
||||
|
|
@ -1219,7 +1226,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
|
|||
proxy_batch_write_at=proxy_batch_write_at,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
if prisma_client is not None
|
||||
else None
|
||||
)
|
||||
if prisma_client is not None:
|
||||
await ProxyStartupEvent._update_default_team_member_budget()
|
||||
|
||||
## SYNC UI SETTINGS ##
|
||||
|
|
@ -1290,7 +1300,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
|
|||
|
||||
await proxy_config.stop_auth_cache_invalidation_subscriber()
|
||||
|
||||
await proxy_shutdown_event()
|
||||
await proxy_shutdown_event(worker_heartbeat=worker_heartbeat)
|
||||
|
||||
|
||||
def _generate_stable_operation_id(route: "APIRoute") -> str:
|
||||
|
|
@ -4371,9 +4381,55 @@ class ProxyConfig:
|
|||
item = self._check_for_os_environ_vars(config=item, depth=depth + 1, max_depth=max_depth)
|
||||
# if the value is a string and starts with "os.environ/" - then it's an environment variable
|
||||
elif isinstance(value, str) and value.startswith("os.environ/"):
|
||||
config[key] = get_secret(value)
|
||||
resolved = get_secret(value)
|
||||
if resolved is None and secret_manager_would_be_consulted(value):
|
||||
verbose_proxy_logger.warning("%s is absent from the configured secret manager", value)
|
||||
config[key] = resolved
|
||||
return config
|
||||
|
||||
def _initialize_secret_manager_from_raw_config(
|
||||
self, config: Mapping[str, object], config_file_path: str | None
|
||||
) -> None:
|
||||
"""
|
||||
Bring the secret manager up before `os.environ/<KEY>` references are resolved.
|
||||
|
||||
`_check_for_os_environ_vars` writes whatever it resolves back into the config, so a key
|
||||
held only by the secret manager would otherwise become a permanent `None` that the later
|
||||
fallbacks in `load_config` can no longer recover from.
|
||||
|
||||
`get_config` also runs on management-endpoint request paths, so this returns early once a
|
||||
manager exists rather than rebuilding the client on every request.
|
||||
|
||||
The manager's own settings can only come from real environment variables, so they are
|
||||
resolved against a throwaway copy and the config is left untouched for the main pass.
|
||||
"""
|
||||
if litellm.secret_manager_client is not None:
|
||||
return
|
||||
|
||||
general_settings: Final = config.get("general_settings")
|
||||
if not isinstance(general_settings, dict):
|
||||
return
|
||||
|
||||
raw_system: Final = general_settings.get("key_management_system")
|
||||
key_management_system: Final = (
|
||||
get_secret(raw_system)
|
||||
if isinstance(raw_system, str) and raw_system.startswith("os.environ/")
|
||||
else raw_system
|
||||
)
|
||||
if not isinstance(key_management_system, str):
|
||||
return
|
||||
|
||||
raw_settings: Final = general_settings.get("key_management_settings")
|
||||
if isinstance(raw_settings, dict):
|
||||
litellm._key_management_settings = KeyManagementSettings(
|
||||
**self._check_for_os_environ_vars(config=copy.deepcopy(raw_settings))
|
||||
)
|
||||
|
||||
self.initialize_secret_manager(
|
||||
key_management_system=key_management_system,
|
||||
config_file_path=config_file_path,
|
||||
)
|
||||
|
||||
def _get_team_config(self, team_id: str, all_teams_config: list[dict]) -> dict:
|
||||
team_config: dict = {}
|
||||
for team in all_teams_config:
|
||||
|
|
@ -4544,6 +4600,8 @@ class ProxyConfig:
|
|||
printed_yaml: Final = copy.deepcopy(config)
|
||||
printed_yaml.pop("environment_variables", None)
|
||||
|
||||
self._initialize_secret_manager_from_raw_config(config=config, config_file_path=config_file_path)
|
||||
|
||||
config = self._check_for_os_environ_vars(config=config)
|
||||
|
||||
self.update_config_state(config=config)
|
||||
|
|
@ -4977,6 +5035,7 @@ class ProxyConfig:
|
|||
)
|
||||
elif key == "audit_log_callbacks":
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
is_audit_logging_enabled,
|
||||
reset_audit_log_callback_cache,
|
||||
)
|
||||
|
||||
|
|
@ -4995,14 +5054,14 @@ class ProxyConfig:
|
|||
litellm.audit_log_callbacks.append(callback)
|
||||
|
||||
_store_audit_logs = litellm_settings.get("store_audit_logs", litellm.store_audit_logs)
|
||||
if _store_audit_logs:
|
||||
if is_audit_logging_enabled(store_audit_logs=_store_audit_logs):
|
||||
print( # noqa: T201
|
||||
f"{blue_color_code} Initialized Audit Log Callbacks - {litellm.audit_log_callbacks} {reset_color_code}"
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.warning(
|
||||
"'audit_log_callbacks' is configured but 'store_audit_logs' is not enabled. "
|
||||
"Audit log callbacks will not fire until 'store_audit_logs: true' is added to litellm_settings."
|
||||
"'audit_log_callbacks' is configured but audit logging is not enabled. "
|
||||
"Audit log callbacks will not fire."
|
||||
)
|
||||
elif key == "cache_params":
|
||||
# this is set in the cache branch
|
||||
|
|
@ -5114,17 +5173,14 @@ class ProxyConfig:
|
|||
key: general_settings[key] for key in SPEND_LOG_CLEANUP_BOUND_SETTINGS if key in general_settings
|
||||
}
|
||||
|
||||
### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ###
|
||||
### LOAD KEY MANAGEMENT SETTINGS ###
|
||||
# The secret manager itself is brought up by get_config(), which runs before the
|
||||
# `os.environ/` references in this config were resolved. Re-reading the settings here
|
||||
# picks up any of them that were themselves secret-manager backed.
|
||||
key_management_settings: Final = general_settings.get("key_management_settings", None)
|
||||
if key_management_settings is not None:
|
||||
litellm._key_management_settings = KeyManagementSettings(**key_management_settings)
|
||||
|
||||
### LOAD SECRET MANAGER ###
|
||||
key_management_system: Final = general_settings.get("key_management_system", None)
|
||||
self.initialize_secret_manager(
|
||||
key_management_system=key_management_system,
|
||||
config_file_path=config_file_path,
|
||||
)
|
||||
### [DEPRECATED] LOAD FROM GOOGLE KMS ### old way of loading from google kms
|
||||
use_google_kms: Final = general_settings.get("use_google_kms", False)
|
||||
load_google_kms(use_google_kms=use_google_kms)
|
||||
|
|
@ -6316,6 +6372,9 @@ class ProxyConfig:
|
|||
if "global_max_parallel_requests" in _general_settings:
|
||||
general_settings["global_max_parallel_requests"] = _general_settings["global_max_parallel_requests"]
|
||||
|
||||
if "max_batch_file_size_mb" not in self._yaml_general_settings_keys:
|
||||
general_settings["max_batch_file_size_mb"] = _general_settings.get("max_batch_file_size_mb")
|
||||
|
||||
## ALERTING ARGS ##
|
||||
if "alerting_args" in _general_settings:
|
||||
general_settings["alerting_args"] = _general_settings["alerting_args"]
|
||||
|
|
@ -8733,7 +8792,7 @@ class ProxyStartupEvent:
|
|||
proxy_budget_rescheduler_max_time: int,
|
||||
proxy_batch_write_at: int,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
):
|
||||
) -> ProxyWorkerHeartbeat:
|
||||
"""Initializes scheduled background jobs"""
|
||||
global store_model_in_db, scheduler
|
||||
|
||||
|
|
@ -8778,6 +8837,18 @@ class ProxyStartupEvent:
|
|||
# Ensure minimum interval of 30 seconds for batch writing to prevent memory issues
|
||||
batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, 5)
|
||||
|
||||
### PROXY WORKER HEARTBEAT ###
|
||||
worker_heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma_client)
|
||||
await worker_heartbeat.beat()
|
||||
scheduler.add_job(
|
||||
worker_heartbeat.beat,
|
||||
"interval",
|
||||
seconds=PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS,
|
||||
id="proxy_worker_heartbeat_job",
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
|
||||
### RESET BUDGET ###
|
||||
if general_settings.get("disable_reset_budget", False) is False:
|
||||
budget_reset_job: Final = ResetBudgetJob(
|
||||
|
|
@ -9117,6 +9188,7 @@ class ProxyStartupEvent:
|
|||
"APScheduler started with memory leak prevention settings: removed jitter, increased intervals, misfire_grace_time=%s",
|
||||
APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
return worker_heartbeat
|
||||
|
||||
@classmethod
|
||||
async def _initialize_spend_tracking_background_jobs(cls, scheduler: AsyncIOScheduler):
|
||||
|
|
@ -15689,6 +15761,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
|
|||
"max_parallel_requests": "Integer",
|
||||
"global_max_parallel_requests": "Integer",
|
||||
"max_request_size_mb": "Integer",
|
||||
"max_batch_file_size_mb": "Integer",
|
||||
"max_response_size_mb": "Integer",
|
||||
"proxy_config_reload_interval_seconds": "Integer",
|
||||
"pass_through_endpoints": "PydanticModel",
|
||||
|
|
|
|||
|
|
@ -9,7 +9,8 @@ effects.
|
|||
Instead we reuse ``ProxyConfig.get_config`` — the actual config reader — so the
|
||||
gateway inherits the same heavy lifting the proxy does: ``include:`` merging,
|
||||
``os.environ/`` + secret-manager resolution, and DB-stored models (when a DB is
|
||||
configured). It has no proxy-setup side effects. Returns the resolved
|
||||
configured). Its only proxy-setup side effect is bringing up the configured
|
||||
secret manager, which is what makes that resolution work. Returns the resolved
|
||||
``model_list``; the Rust side deserializes each entry into its ``Deployment``.
|
||||
"""
|
||||
|
||||
|
|
|
|||
|
|
@ -947,6 +947,17 @@ model LiteLLM_DailyTagSpend {
|
|||
}
|
||||
|
||||
|
||||
// One row per live proxy worker process. Workers upsert their row on a fixed
|
||||
// heartbeat; counting rows with a recent heartbeat tells how many workers share
|
||||
// this database, which lets the Admin UI hide its "no Redis" warning for
|
||||
// deployments that are provably a single worker.
|
||||
model LiteLLM_ProxyWorkerHeartbeat {
|
||||
worker_id String @id
|
||||
hostname String
|
||||
started_at DateTime @default(now())
|
||||
last_heartbeat_at DateTime @default(now())
|
||||
}
|
||||
|
||||
// Track the status of cron jobs running. Only allow one pod to run the job at a time
|
||||
model LiteLLM_CronJob {
|
||||
cronjob_id String @id @default(cuid()) // Unique ID for the record
|
||||
|
|
@ -1467,28 +1478,38 @@ model LiteLLM_AutoRouterSession {
|
|||
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
|
||||
}
|
||||
|
||||
// Shadow eval: evaluation of an auto-router against a key's live traffic, in either
|
||||
// direction. forward duplicates the requests the key did not route through the router
|
||||
// through it, answering whether the key should adopt it; reverse duplicates the requests
|
||||
// the router did serve against a fixed baseline model, answering whether a key already on
|
||||
// it still benefits. Either way a sampled slice runs in a detached task and an LLM judge
|
||||
// compares real vs shadow responses blind. The job row is immutable config plus
|
||||
// stopped_at; every count, status, and spend figure is derived from the append-only
|
||||
// attempt rows, so nothing can disagree across pods or stop races.
|
||||
// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in
|
||||
// either direction. forward duplicates the requests the keys did not route through the
|
||||
// router through it, answering whether they should adopt it; reverse duplicates the
|
||||
// requests the router did serve against a fixed baseline model, answering whether a key
|
||||
// already on it still benefits. Either way a sampled slice runs in a detached task and an
|
||||
// LLM judge compares real vs shadow responses blind. Each row is ONE key's leg of a job:
|
||||
// immutable config plus that key's own turn budget and stop state, so one key exhausting
|
||||
// its budget never ends a sibling's sampling. A job is the set of legs sharing group_id
|
||||
// (the id the API reports), written together by one atomic create_many with identical
|
||||
// config; single-key jobs predating group_id were backfilled group_id = id. "One active
|
||||
// job per (key, direction)" is a partial unique index on (api_key_id, direction) WHERE
|
||||
// stopped_at IS NULL, expressed only in the migration because schema.prisma cannot state
|
||||
// partial indexes; it is what makes a concurrent start on another pod race-safe rather
|
||||
// than read-then-create. Every count, status, and spend figure is derived from the
|
||||
// append-only attempt rows, so nothing can disagree across pods or stop races.
|
||||
model LiteLLM_ShadowEvalJob {
|
||||
id String @id @default(cuid())
|
||||
api_key_id String // hashed virtual key whose traffic is shadowed
|
||||
group_id String // legs of one job share this; the API's job id
|
||||
api_key_id String // hashed virtual key whose traffic this leg shadows
|
||||
router_name String // the auto-router under evaluation, in either direction
|
||||
direction String @default("forward") // forward | reverse
|
||||
baseline_model String? // reverse only: the fixed model the router is judged against
|
||||
judge_model String
|
||||
shadow_percentage Float
|
||||
max_turns Int // sample budget: judge at most this many turns
|
||||
max_turns Int // this key's sample budget: judge at most this many turns
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
ends_at DateTime
|
||||
stopped_at DateTime?
|
||||
stopped_by String? // operator who stopped it early; null when it ended on its own
|
||||
|
||||
@@index([group_id])
|
||||
@@index([api_key_id])
|
||||
@@index([created_at])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,18 +1,12 @@
|
|||
"""Opt-in flag for PTU (provisioned throughput unit) flat-cost attribution.
|
||||
"""Re-exported from ``litellm.litellm_core_utils.ptu_pricing``.
|
||||
|
||||
The whole feature is inert unless an operator sets
|
||||
``LITELLM_ENABLE_PTU_COST_ATTRIBUTION``: the daily rollup is not scheduled, the
|
||||
model endpoints reject PTU config, the daily activity read path reports zero flat
|
||||
cost, and the model form hides the PTU inputs.
|
||||
The flag lives in core because the router reads it while registering a deployment, and
|
||||
router code cannot import from the proxy.
|
||||
"""
|
||||
|
||||
from typing import Final
|
||||
from litellm.litellm_core_utils.ptu_pricing import (
|
||||
PTU_COST_ATTRIBUTION_ENV_VAR,
|
||||
is_ptu_cost_attribution_enabled,
|
||||
)
|
||||
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
PTU_COST_ATTRIBUTION_ENV_VAR: Final = "LITELLM_ENABLE_PTU_COST_ATTRIBUTION"
|
||||
|
||||
|
||||
def is_ptu_cost_attribution_enabled() -> bool:
|
||||
"""Report whether this deployment opted into PTU flat-cost attribution."""
|
||||
return get_secret_bool(PTU_COST_ATTRIBUTION_ENV_VAR, False) is True
|
||||
__all__ = ("PTU_COST_ATTRIBUTION_ENV_VAR", "is_ptu_cost_attribution_enabled")
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ and share the existing unique constraint.
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, datetime, time, timedelta, timezone
|
||||
|
|
@ -29,14 +30,15 @@ from litellm.constants import (
|
|||
PTU_ROLLUP_MAX_BACKFILL_DAYS,
|
||||
PTU_SENTINEL_API_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.ptu_pricing import ptu_terms
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
_HOURS_PER_DAY: Final = 24
|
||||
_PRUNE_ID_CHUNK_SIZE: Final = 5_000
|
||||
_UPSERT_ATTEMPTS: Final = 3
|
||||
_UPSERT_RETRY_BACKOFF_SECONDS: Final = 0.5
|
||||
|
||||
|
|
@ -72,28 +74,6 @@ class PTUModel:
|
|||
effective_to: datetime | None = None
|
||||
|
||||
|
||||
def _parse_utc_datetime(value: object) -> datetime | None:
|
||||
"""Parse a model_info datetime (ISO string or datetime) into a UTC-aware datetime, else None."""
|
||||
parsed: Final = _coerce_datetime(value)
|
||||
if parsed is None:
|
||||
return None
|
||||
if parsed.tzinfo is None:
|
||||
return parsed.replace(tzinfo=timezone.utc)
|
||||
return parsed.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def _coerce_datetime(value: object) -> datetime | None:
|
||||
"""``value`` as a datetime, parsing an ISO string, else None."""
|
||||
if isinstance(value, datetime):
|
||||
return value
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
try:
|
||||
return datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def _public_model_name(row: object, model_info: Mapping[str, object]) -> str:
|
||||
"""The name an operator recognises for this deployment.
|
||||
|
||||
|
|
@ -167,46 +147,20 @@ def _parse_ptu_model(row: object) -> PTUModel | None:
|
|||
Valid means model_info has a positive ptu_count, a non-negative
|
||||
cost_per_ptu_per_hour, and a team_id (1 model -> 1 team).
|
||||
"""
|
||||
raw_model_info: Final = getattr(row, "model_info", None)
|
||||
model_info: Final = _decode_model_info(raw_model_info)
|
||||
model_info: Final = _decode_model_info(getattr(row, "model_info", None))
|
||||
if model_info is None:
|
||||
return None
|
||||
ptu_count: Final = model_info.get("ptu_count")
|
||||
cost_per_hour: Final = model_info.get("cost_per_ptu_per_hour")
|
||||
team_id: Final = model_info.get("team_id")
|
||||
if ptu_count is None or cost_per_hour is None or not team_id:
|
||||
return None
|
||||
try:
|
||||
ptu_count_int: Final = int(ptu_count)
|
||||
cost_per_hour_float: Final = float(cost_per_hour)
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return None
|
||||
if not 0 < ptu_count_int <= ModelInfo.MAX_PTU_COUNT:
|
||||
return None
|
||||
if not 0 <= cost_per_hour_float <= ModelInfo.MAX_COST_PER_PTU_PER_HOUR:
|
||||
return None
|
||||
if model_info.get("ptu_effective_from") is None:
|
||||
# The endpoints require a start; a row without one predates that rule or was
|
||||
# written around them, and inferring one would bill days the deployment did not exist
|
||||
return None
|
||||
raw_from: Final = model_info.get("ptu_effective_from")
|
||||
raw_to: Final = model_info.get("ptu_effective_to")
|
||||
effective_from: Final = _parse_utc_datetime(raw_from)
|
||||
effective_to: Final = _parse_utc_datetime(raw_to)
|
||||
# A present-but-unparseable bound would read as "no bound" and silently widen the
|
||||
# window to the whole day, so the deployment is skipped until the config is fixed
|
||||
if (raw_from is not None and effective_from is None) or (raw_to is not None and effective_to is None):
|
||||
return None
|
||||
if effective_from is not None and effective_to is not None and effective_to <= effective_from:
|
||||
terms: Final = ptu_terms(model_info)
|
||||
if terms is None:
|
||||
return None
|
||||
return PTUModel(
|
||||
model_id=str(getattr(row, "model_id", "") or ""),
|
||||
model_name=_public_model_name(row, model_info),
|
||||
team_id=str(team_id),
|
||||
ptu_count=ptu_count_int,
|
||||
cost_per_ptu_per_hour=cost_per_hour_float,
|
||||
effective_from=effective_from,
|
||||
effective_to=effective_to,
|
||||
team_id=terms.team_id,
|
||||
ptu_count=terms.ptu_count,
|
||||
cost_per_ptu_per_hour=terms.cost_per_ptu_per_hour,
|
||||
effective_from=terms.effective_from,
|
||||
effective_to=terms.effective_to,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -358,10 +312,70 @@ async def _upsert_charge_with_retry(
|
|||
return False
|
||||
|
||||
|
||||
async def _load_ptu_models(prisma_client: "PrismaClient") -> tuple[PTUModel, ...]:
|
||||
"""Every model deployment currently carrying valid manual PTU config."""
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _LoadedDeployments:
|
||||
"""The deployments a run will price, and every deployment id it looked at.
|
||||
|
||||
The id set is deliberately wider than the priced set. A deployment whose PTU config
|
||||
was removed produces no charge and still has to be prunable, so bounding the prune on
|
||||
what priced would strand its old rows forever. It is also a guaranteed superset of the
|
||||
priced set, or a run could write a charge that falls outside its own delete filter.
|
||||
"""
|
||||
|
||||
models: tuple[PTUModel, ...]
|
||||
scanned_ids: frozenset[str]
|
||||
config_sourced: bool
|
||||
|
||||
|
||||
def _running_router() -> object | None:
|
||||
"""The proxy's router, or None outside a running proxy.
|
||||
|
||||
Read out of ``sys.modules`` rather than imported, so a rollup driven from a test or a
|
||||
script does not pull the whole proxy server in behind it.
|
||||
"""
|
||||
proxy_server: Final = sys.modules.get("litellm.proxy.proxy_server")
|
||||
return getattr(proxy_server, "llm_router", None) if proxy_server is not None else None
|
||||
|
||||
|
||||
def _config_deployments(router: object | None, *, owned_by_db: frozenset[str]) -> tuple[_PTUDeployment, ...]:
|
||||
"""Deployments the router holds that no ``LiteLLM_ProxyModelTable`` row owns.
|
||||
|
||||
``db_model`` is forced True on every deployment loaded from that table and defaults to
|
||||
False on ModelInfo, so the complement is what config.yaml declared. A per-request
|
||||
credential clone carries ``original_model_id`` and reuses its source's PTU config under
|
||||
a fresh id, so pricing it would bill one reservation once per distinct client key.
|
||||
"""
|
||||
entries: Final = tuple(getattr(router, "model_list", None) or ())
|
||||
records: Final = tuple(_router_deployment(entry) for entry in entries)
|
||||
return tuple(
|
||||
record
|
||||
for record in records
|
||||
if record is not None
|
||||
and record.model_info.get("db_model") is not True
|
||||
and record.model_info.get("original_model_id") is None
|
||||
and record.model_id not in owned_by_db
|
||||
)
|
||||
|
||||
|
||||
async def _load_ptu_models(prisma_client: "PrismaClient") -> _LoadedDeployments:
|
||||
"""Every deployment carrying valid manual PTU config, and every id the scan saw.
|
||||
|
||||
Reserved capacity is billed by the provider whichever file declared it, so a
|
||||
deployment the proxy only knows from config.yaml accrues alongside the stored ones.
|
||||
"""
|
||||
rows: Final = await prisma_client.db.litellm_proxymodeltable.find_many()
|
||||
return tuple(parsed for parsed in (_parse_ptu_model(row) for row in rows) if parsed is not None)
|
||||
db_ids: Final = frozenset(model_id for row in rows if (model_id := str(getattr(row, "model_id", "") or "")))
|
||||
config_records: Final = _config_deployments(_running_router(), owned_by_db=db_ids)
|
||||
models: Final = tuple(
|
||||
parsed for parsed in (_parse_ptu_model(row) for row in (*rows, *config_records)) if parsed is not None
|
||||
)
|
||||
return _LoadedDeployments(
|
||||
models=models,
|
||||
config_sourced=bool(config_records),
|
||||
scanned_ids=db_ids
|
||||
| frozenset(record.model_id for record in config_records)
|
||||
| frozenset(model.model_id for model in models),
|
||||
)
|
||||
|
||||
|
||||
async def run_ptu_flat_cost_rollup(
|
||||
|
|
@ -378,8 +392,10 @@ async def run_ptu_flat_cost_rollup(
|
|||
The prune predicate is ``updated_at < run_started`` rather than "not in the charge
|
||||
set I computed", which matters under concurrency: whether a row is garbage becomes a
|
||||
property of the row instead of one run's in-memory config snapshot, so a run can
|
||||
never delete a row a concurrent run just wrote. It is still skipped when any charge
|
||||
failed to write, since a row whose replacement never landed would look unrefreshed.
|
||||
never delete a row a concurrent run just wrote. It is bounded to the deployments this
|
||||
run looked at, so a row it cannot account for is out of reach either way. It is still
|
||||
skipped when any charge failed to write, since a row whose replacement never landed
|
||||
would look unrefreshed.
|
||||
"""
|
||||
day: Final = target_date or (datetime.now(timezone.utc).date() - timedelta(days=1))
|
||||
|
||||
|
|
@ -390,7 +406,8 @@ async def run_ptu_flat_cost_rollup(
|
|||
date_str: Final = day.isoformat()
|
||||
run_started: Final = datetime.now(timezone.utc)
|
||||
|
||||
ptu_models: Final = await _load_ptu_models(prisma_client)
|
||||
loaded: Final = await _load_ptu_models(prisma_client)
|
||||
ptu_models: Final = loaded.models
|
||||
charges: Final = _aggregate_charges(ptu_models, day)
|
||||
|
||||
landed: Final = tuple(
|
||||
|
|
@ -415,7 +432,12 @@ async def run_ptu_flat_cost_rollup(
|
|||
date_str,
|
||||
)
|
||||
else:
|
||||
await _prune_unrefreshed_sentinel_rows(prisma_client, date_str=date_str, run_started=run_started)
|
||||
await _prune_unrefreshed_sentinel_rows(
|
||||
prisma_client,
|
||||
date_str=date_str,
|
||||
run_started=run_started,
|
||||
scanned_ids=loaded.scanned_ids if loaded.config_sourced else None,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"PTU rollup for %s: %d PTU models processed, %d rows written, %d rows failed",
|
||||
|
|
@ -524,7 +546,7 @@ async def run_ptu_flat_cost_backfill(
|
|||
verbose_proxy_logger.warning("PTU backfill: prisma_client is None, skipping")
|
||||
return BackfillResult(start=end, end=end, days_scanned=0, rows_written=0)
|
||||
|
||||
ptu_models: Final = await _load_ptu_models(prisma_client)
|
||||
ptu_models: Final = (await _load_ptu_models(prisma_client)).models
|
||||
days: Final = _backfill_window(ptu_models, end)
|
||||
|
||||
if not days:
|
||||
|
|
@ -707,26 +729,61 @@ async def _prune_unrefreshed_sentinel_rows(
|
|||
*,
|
||||
date_str: str,
|
||||
run_started: datetime,
|
||||
scanned_ids: frozenset[str] | None,
|
||||
) -> None:
|
||||
"""Delete the day's PTU sentinel rows this run did not refresh.
|
||||
"""Delete the day's PTU sentinel rows this run looked at and did not refresh.
|
||||
|
||||
Every charge the run wrote bumps ``updated_at`` past ``run_started``, so anything
|
||||
left below that mark is a (team, model) the current config no longer prices. The mark
|
||||
is pulled back by ``PTU_PRUNE_SKEW_GRACE_SECONDS`` because the two timestamps come
|
||||
from different hosts: a stale row is hours old, a concurrently written one is seconds
|
||||
old, and the grace separates them without waiting on clocks agreeing. The
|
||||
predicate reads only the row, never the caller's config snapshot, which is what
|
||||
makes it safe to run twice, out of order, or beside another pod: a row written
|
||||
after this run began is out of reach of its delete. Mirrors the retention predicate
|
||||
``SpendLogCleanup`` deletes by."""
|
||||
Two conditions, and a row survives unless it meets both. It must be stale: every
|
||||
charge the run wrote bumps ``updated_at`` past ``run_started``, so anything left below
|
||||
that mark is a (team, model) the current config no longer prices. The mark is pulled
|
||||
back by ``PTU_PRUNE_SKEW_GRACE_SECONDS`` because the two timestamps come from
|
||||
different hosts, and the grace separates a row that is hours old from one written
|
||||
seconds ago without waiting on clocks agreeing.
|
||||
|
||||
A run that priced a deployment only its own host declares must also name the
|
||||
deployments it scanned. Staleness alone is sufficient while every run derives its
|
||||
charges from the same table, because then any two runs compute the same set, so a
|
||||
database-only run still sweeps by timestamp exactly as it always has. Once one host's
|
||||
charges come from a file the others cannot read, a row it never considered is not
|
||||
evidence of anything, and deleting it drops a charge that host is responsible for.
|
||||
|
||||
Where the bound applies the ids go out in chunks, because each is one bind variable and
|
||||
the server rejects a statement carrying more than 32767 of them, which a proxy holding
|
||||
that many deployments would otherwise hit every night with no handler above here.
|
||||
"""
|
||||
cutoff: Final = run_started - timedelta(seconds=PTU_PRUNE_SKEW_GRACE_SECONDS)
|
||||
await prisma_client.db.litellm_dailyteamspend.delete_many(
|
||||
where={ # mutable-ok: prisma delete filter
|
||||
"date": date_str,
|
||||
"api_key": PTU_SENTINEL_API_KEY,
|
||||
"updated_at": {"lt": cutoff}, # mutable-ok: prisma comparison filter
|
||||
}
|
||||
unbounded: Final = { # mutable-ok: prisma delete filter
|
||||
"date": date_str,
|
||||
"api_key": PTU_SENTINEL_API_KEY,
|
||||
"updated_at": {"lt": cutoff}, # mutable-ok: prisma comparison filter
|
||||
}
|
||||
ordered: Final = () if scanned_ids is None else tuple(sorted(scanned_ids))
|
||||
filters: Final = (
|
||||
(unbounded,)
|
||||
if scanned_ids is None
|
||||
else tuple(
|
||||
MappingProxyType(
|
||||
{
|
||||
**unbounded,
|
||||
"model": { # mutable-ok: prisma membership filter
|
||||
"in": ordered[start : start + _PRUNE_ID_CHUNK_SIZE]
|
||||
},
|
||||
}
|
||||
)
|
||||
for start in range(0, len(ordered), _PRUNE_ID_CHUNK_SIZE)
|
||||
)
|
||||
)
|
||||
deletions: Final = tuple(
|
||||
[await prisma_client.db.litellm_dailyteamspend.delete_many(where=where) for where in filters]
|
||||
)
|
||||
deleted: Final = sum(deletions)
|
||||
if deleted:
|
||||
verbose_proxy_logger.info(
|
||||
"PTU rollup for %s: pruned %s stale sentinel row(s) across %s deployment(s)",
|
||||
date_str,
|
||||
deleted,
|
||||
"every" if scanned_ids is None else len(scanned_ids),
|
||||
)
|
||||
|
||||
|
||||
__all__ = (
|
||||
|
|
|
|||
|
|
@ -130,6 +130,7 @@ class PricingBasis(NamedTuple):
|
|||
|
||||
service_tier: str | None = None
|
||||
data_residency: str | None = None
|
||||
vertex_location: str | None = None
|
||||
|
||||
|
||||
_STANDARD_RATES: Final = PricingBasis()
|
||||
|
|
@ -141,8 +142,8 @@ def _pricing_basis(cost_breakdown: Mapping[str, object] | None) -> PricingBasis:
|
|||
Rows written before this field shipped carry neither key, and there is no backfill:
|
||||
they price at standard rates, which is what they already did.
|
||||
|
||||
Both values survive a JSON round trip on the way here, so neither is guaranteed to be
|
||||
a string. `generic_cost_per_token` calls `.lower()` on both without a type check, and
|
||||
These values survive a JSON round trip on the way here, so none is guaranteed to be
|
||||
a string. `generic_cost_per_token` calls `.lower()` on them without a type check, and
|
||||
the resulting `AttributeError` would be swallowed into a silent zero by the caller's
|
||||
`except`, so anything that is not a string is dropped here instead.
|
||||
"""
|
||||
|
|
@ -150,9 +151,11 @@ def _pricing_basis(cost_breakdown: Mapping[str, object] | None) -> PricingBasis:
|
|||
return _STANDARD_RATES
|
||||
service_tier: Final = cost_breakdown.get("service_tier")
|
||||
data_residency: Final = cost_breakdown.get("data_residency")
|
||||
vertex_location: Final = cost_breakdown.get("vertex_location")
|
||||
return PricingBasis(
|
||||
service_tier=service_tier if isinstance(service_tier, str) else None,
|
||||
data_residency=data_residency if isinstance(data_residency, str) else None,
|
||||
vertex_location=vertex_location if isinstance(vertex_location, str) else None,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -193,6 +196,7 @@ def _cost_of_usage(
|
|||
service_tier=basis.service_tier,
|
||||
data_residency=basis.data_residency,
|
||||
model_info=model_info,
|
||||
vertex_location=basis.vertex_location,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models; degrade to zero savings
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
|
|||
|
|
@ -216,6 +216,15 @@ def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> d
|
|||
return {}
|
||||
|
||||
|
||||
def _sl_attribution_fallback(
|
||||
standard_logging_payload: StandardLoggingPayload | None,
|
||||
field: Literal["model_id", "model_group", "api_base", "custom_llm_provider"],
|
||||
) -> str:
|
||||
if standard_logging_payload is None:
|
||||
return ""
|
||||
return standard_logging_payload.get(field) or ""
|
||||
|
||||
|
||||
def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogsPayload:
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
|
@ -288,8 +297,15 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
|
|||
): # use 'tags' from standard logging payload instead
|
||||
request_tags = safe_dumps(standard_logging_payload["request_tags"])
|
||||
|
||||
_model_id: Final = metadata.get("model_info", {}).get("id", "")
|
||||
_model_group: Final = metadata.get("model_group", "")
|
||||
_model_id: Final = metadata.get("model_info", {}).get("id", "") or _sl_attribution_fallback(
|
||||
standard_logging_payload, "model_id"
|
||||
)
|
||||
_model_group: Final = metadata.get("model_group", "") or _sl_attribution_fallback(
|
||||
standard_logging_payload, "model_group"
|
||||
)
|
||||
_api_base: Final = litellm_params.get("api_base", "") or _sl_attribution_fallback(
|
||||
standard_logging_payload, "api_base"
|
||||
)
|
||||
|
||||
# Extract overhead from hidden_params if available
|
||||
litellm_overhead_time_ms = None
|
||||
|
|
@ -389,7 +405,11 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
|
|||
|
||||
# Extract agent_id for A2A requests (set directly on model_call_details)
|
||||
agent_id: Final[str | None] = kwargs.get("agent_id") or metadata.get("agent_id")
|
||||
custom_llm_provider: Final = kwargs.get("custom_llm_provider")
|
||||
custom_llm_provider: Final = (
|
||||
kwargs.get("custom_llm_provider")
|
||||
or _sl_attribution_fallback(standard_logging_payload, "custom_llm_provider")
|
||||
or None
|
||||
)
|
||||
raw_model: Final = cast(str, kwargs.get("model") or "")
|
||||
model_name: Final = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})
|
||||
|
||||
|
|
@ -414,13 +434,13 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
|
|||
completion_tokens=usage.get("completion_tokens", standard_logging_completion_tokens),
|
||||
request_tags=request_tags,
|
||||
end_user=end_user_id or "",
|
||||
api_base=litellm_params.get("api_base", ""),
|
||||
api_base=_api_base,
|
||||
model_group=_model_group,
|
||||
model_id=_model_id,
|
||||
mcp_namespaced_tool_name=mcp_namespaced_tool_name,
|
||||
agent_id=agent_id,
|
||||
requester_ip_address=clean_metadata.get("requester_ip_address", None),
|
||||
custom_llm_provider=kwargs.get("custom_llm_provider", ""),
|
||||
custom_llm_provider=custom_llm_provider or "",
|
||||
messages=_get_messages_for_spend_logs_payload(
|
||||
standard_logging_payload=standard_logging_payload, metadata=metadata
|
||||
),
|
||||
|
|
|
|||
|
|
@ -30,7 +30,6 @@ from litellm.constants import (
|
|||
SPEND_LOG_WRITE_BATCH_MAX_BYTES,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
DB_RETRY_SAFE_ERROR_TYPES,
|
||||
CommonProxyErrors,
|
||||
ProxyErrorTypes,
|
||||
ProxyException,
|
||||
|
|
@ -517,6 +516,34 @@ def _failure_usage_to_lift(
|
|||
return estimated_usage, 0.0
|
||||
|
||||
|
||||
_EMPTY_LIFT: Final = MappingProxyType({})
|
||||
|
||||
|
||||
def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""Failure-path callbacks run after ``litellm_logging_obj`` is popped from
|
||||
request_data (it is not serialisable), so the caller merges these fields
|
||||
onto request_data first: the first-handoff instant for preprocessing
|
||||
latency, recovered or estimated usage for token counts, and the standard
|
||||
logging object for deployment attribution on failed-request spend logs."""
|
||||
_logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
if _logging_obj is None:
|
||||
return _EMPTY_LIFT
|
||||
_model_call_details: Final = getattr(_logging_obj, "model_call_details", {})
|
||||
_first_handoff: Final = _model_call_details.get("first_api_call_start_time")
|
||||
_usage_to_lift: Final = _failure_usage_to_lift(
|
||||
model_call_details=_model_call_details,
|
||||
request_body=request_data,
|
||||
dispatched=_first_handoff is not None,
|
||||
)
|
||||
_entries: Final = (
|
||||
("first_api_call_start_time", _first_handoff),
|
||||
("combined_usage_object", None if _usage_to_lift is None else _usage_to_lift[0]),
|
||||
("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)),
|
||||
("standard_logging_object", _model_call_details.get("standard_logging_object")),
|
||||
)
|
||||
return MappingProxyType({key: value for key, value in _entries if value is not None})
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _CallbackCapabilities:
|
||||
"""Cached per-hook capability flags derived from ``litellm.callbacks``.
|
||||
|
|
@ -2281,6 +2308,11 @@ class ProxyLogging:
|
|||
)
|
||||
)
|
||||
|
||||
# Auth and pass-through failure bodies are unstripped client input, and
|
||||
# the logging handler below flattens body keys into model_call_details,
|
||||
# so drop the key before it can masquerade as the built payload.
|
||||
request_data.pop("standard_logging_object", None)
|
||||
|
||||
### LOGGING ###
|
||||
if self._is_proxy_only_llm_api_error(
|
||||
original_exception=original_exception,
|
||||
|
|
@ -2294,29 +2326,7 @@ class ProxyLogging:
|
|||
original_exception=original_exception,
|
||||
)
|
||||
|
||||
# Lift the first-handoff instant onto request_data (top-level
|
||||
# internal key, not metadata) so failure-path callbacks can still
|
||||
# compute preprocessing latency after the logging object is popped.
|
||||
_logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
if _logging_obj is not None:
|
||||
_model_call_details: Final = getattr(_logging_obj, "model_call_details", {})
|
||||
_first_handoff: Final = _model_call_details.get("first_api_call_start_time")
|
||||
if _first_handoff is not None:
|
||||
request_data["first_api_call_start_time"] = _first_handoff
|
||||
|
||||
# Lift recovered partial-stream usage, or an estimated input-side
|
||||
# usage for a dispatched failure, onto request_data so the
|
||||
# failure-path spend callbacks (which run after the logging object
|
||||
# is popped) record real token counts instead of zero.
|
||||
_usage_to_lift: Final = _failure_usage_to_lift(
|
||||
model_call_details=_model_call_details,
|
||||
request_body=request_data,
|
||||
dispatched=_first_handoff is not None,
|
||||
)
|
||||
if _usage_to_lift is not None:
|
||||
_lifted_usage, _lifted_cost = _usage_to_lift
|
||||
request_data["combined_usage_object"] = _lifted_usage
|
||||
request_data["response_cost"] = _lifted_cost
|
||||
request_data.update(_failure_fields_to_lift(request_data))
|
||||
|
||||
# Remove before callbacks iterate — not serialisable
|
||||
request_data.pop("litellm_logging_obj", None)
|
||||
|
|
@ -5960,15 +5970,14 @@ class ProxyUpdateSpend:
|
|||
)
|
||||
|
||||
break
|
||||
except DB_RETRY_SAFE_ERROR_TYPES as e:
|
||||
if i >= n_retry_times: # If we've reached the maximum number of retries
|
||||
_raise_failed_update_spend_exception(
|
||||
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
|
||||
)
|
||||
# Optionally, sleep for a bit before retrying
|
||||
await asyncio.sleep(2**i) # Exponential backoff
|
||||
except Exception as e:
|
||||
_raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj)
|
||||
await DBSpendUpdateWriter._handle_spend_update_failure(
|
||||
e=e,
|
||||
attempt=i,
|
||||
n_retry_times=n_retry_times,
|
||||
start_time=start_time,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def update_spend_logs(
|
||||
|
|
|
|||
|
|
@ -1996,6 +1996,12 @@ class LiteLLMCompletionResponsesConfig:
|
|||
output_items.append(item)
|
||||
return output_items
|
||||
|
||||
@staticmethod
|
||||
def _encode_thinking_blocks(message: Message) -> str | None:
|
||||
thinking_blocks: Final[Sequence[Mapping[str, object]]] = getattr(message, "thinking_blocks", None) or ()
|
||||
preserved: Final = tuple(block for block in thinking_blocks if block.get("signature") or block.get("data"))
|
||||
return json.dumps(preserved, separators=(",", ":")) if preserved else None
|
||||
|
||||
@staticmethod
|
||||
def _extract_reasoning_output_items(
|
||||
chat_completion_response: ModelResponse,
|
||||
|
|
@ -2004,12 +2010,14 @@ class LiteLLMCompletionResponsesConfig:
|
|||
for choice in choices:
|
||||
if hasattr(choice, "message") and choice.message:
|
||||
message = choice.message
|
||||
if hasattr(message, "reasoning_content") and message.reasoning_content:
|
||||
reasoning_content = getattr(message, "reasoning_content", None) or ""
|
||||
encrypted_content = LiteLLMCompletionResponsesConfig._encode_thinking_blocks(message)
|
||||
if reasoning_content or encrypted_content:
|
||||
# Only check the first choice for reasoning content
|
||||
return [
|
||||
GenericResponseOutputItem(
|
||||
type="reasoning",
|
||||
id=f"rs_{hash(str(message.reasoning_content))}",
|
||||
id=f"rs_{hash(reasoning_content or encrypted_content)}",
|
||||
status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status(
|
||||
choice.finish_reason
|
||||
),
|
||||
|
|
@ -2017,10 +2025,13 @@ class LiteLLMCompletionResponsesConfig:
|
|||
content=[
|
||||
OutputText(
|
||||
type="output_text",
|
||||
text=message.reasoning_content,
|
||||
text=text,
|
||||
annotations=[],
|
||||
)
|
||||
for text in (reasoning_content,)
|
||||
if text
|
||||
],
|
||||
encrypted_content=encrypted_content,
|
||||
)
|
||||
]
|
||||
return []
|
||||
|
|
@ -2292,18 +2303,19 @@ class LiteLLMCompletionResponsesConfig:
|
|||
# Translate completion_tokens_details to output_tokens_details
|
||||
if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details is not None:
|
||||
completion_details: Final = usage.completion_tokens_details
|
||||
output_details_dict: Final[dict[str, int]] = {}
|
||||
if hasattr(completion_details, "reasoning_tokens") and completion_details.reasoning_tokens is not None:
|
||||
output_details_dict["reasoning_tokens"] = completion_details.reasoning_tokens
|
||||
|
||||
if hasattr(completion_details, "text_tokens") and completion_details.text_tokens is not None:
|
||||
output_details_dict["text_tokens"] = completion_details.text_tokens
|
||||
|
||||
if hasattr(completion_details, "image_tokens") and completion_details.image_tokens is not None:
|
||||
output_details_dict["image_tokens"] = completion_details.image_tokens
|
||||
|
||||
if output_details_dict:
|
||||
response_usage.output_tokens_details = OutputTokensDetails(**output_details_dict)
|
||||
reasoning_token_count: Final = getattr(completion_details, "reasoning_tokens", None)
|
||||
optional_output_details: Final[dict[str, int]] = {
|
||||
field: value
|
||||
for field, value in (
|
||||
("text_tokens", getattr(completion_details, "text_tokens", None)),
|
||||
("image_tokens", getattr(completion_details, "image_tokens", None)),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
response_usage.output_tokens_details = OutputTokensDetails(
|
||||
reasoning_tokens=reasoning_token_count if reasoning_token_count is not None else 0,
|
||||
**optional_output_details,
|
||||
)
|
||||
|
||||
return response_usage
|
||||
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
|
|||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.ptu_pricing import zeroed_ptu_pricing
|
||||
from litellm.litellm_core_utils.request_timeout_resolver import (
|
||||
get_configured_request_timeout,
|
||||
)
|
||||
|
|
@ -7695,7 +7696,16 @@ class Router:
|
|||
- None: If the deployment is not active for the current environment (if 'supported_environments' is set in litellm_params)
|
||||
"""
|
||||
try:
|
||||
litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(**_litellm_params)
|
||||
zeroed_pricing: Final = (
|
||||
zeroed_ptu_pricing(_model_info, _litellm_params) if _model_info.get("db_model") is not True else None
|
||||
)
|
||||
litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(
|
||||
**(
|
||||
_litellm_params
|
||||
if zeroed_pricing is None
|
||||
else MappingProxyType({**_litellm_params, **zeroed_pricing})
|
||||
)
|
||||
)
|
||||
warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=_litellm_params)
|
||||
deployment = Deployment(
|
||||
**deployment_info,
|
||||
|
|
@ -10253,11 +10263,13 @@ class Router:
|
|||
returned_models.extend(self.get_model_list_from_routing_groups(model_name=model_name))
|
||||
|
||||
if len(returned_models) == 0: # check if wildcard route
|
||||
potential_wildcard_models: Final = self.pattern_router.route(model_name) or []
|
||||
potential_wildcard_models: Final = self.pattern_router.get_deployments_by_pattern(model=model_name or "")
|
||||
|
||||
## check for team-specific wildcard models
|
||||
if team_id is not None and team_id in self.team_pattern_routers:
|
||||
potential_team_only_wildcard_models: Final = self.team_pattern_routers[team_id].route(model_name) or []
|
||||
potential_team_only_wildcard_models: Final = self.team_pattern_routers[
|
||||
team_id
|
||||
].get_deployments_by_pattern(model=model_name or "")
|
||||
potential_wildcard_models.extend(potential_team_only_wildcard_models)
|
||||
|
||||
if model_name is not None and potential_wildcard_models is not None:
|
||||
|
|
@ -11189,6 +11201,8 @@ class Router:
|
|||
if pre_routing_hook_response is not None:
|
||||
model = pre_routing_hook_response.model
|
||||
messages = pre_routing_hook_response.messages
|
||||
if pre_routing_hook_response.litellm_params:
|
||||
request_kwargs.update(pre_routing_hook_response.litellm_params)
|
||||
#########################################################
|
||||
|
||||
# Resolve the strategy and logger AFTER the pre-routing hook, since
|
||||
|
|
@ -11298,6 +11312,8 @@ class Router:
|
|||
if pre_routing_hook_response is not None:
|
||||
model = pre_routing_hook_response.model
|
||||
messages = pre_routing_hook_response.messages
|
||||
if pre_routing_hook_response.litellm_params:
|
||||
request_kwargs.update(pre_routing_hook_response.litellm_params)
|
||||
|
||||
# 2. Get healthy deployments
|
||||
healthy_deployments: Final = await self.async_get_healthy_deployments(
|
||||
|
|
|
|||
|
|
@ -53,6 +53,21 @@ model_list:
|
|||
REASONING: o1-preview
|
||||
```
|
||||
|
||||
Each tier can also use a model entry with request parameter overrides. A tier value may be
|
||||
a model string, a single object, or a list mixing strings and objects. Object entries must
|
||||
contain a model name and may contain any LiteLLM request parameters. The model name must
|
||||
still resolve to a deployment in `model_list`; this configuration does not create one
|
||||
|
||||
```yaml
|
||||
tiers:
|
||||
COMPLEX: opus
|
||||
REASONING:
|
||||
- model_name: opus
|
||||
litellm_params:
|
||||
reasoning_effort: xhigh
|
||||
- abc
|
||||
```
|
||||
|
||||
### Renaming the tiers
|
||||
|
||||
`tier_labels` puts your own vocabulary on the four tiers:
|
||||
|
|
@ -165,7 +180,7 @@ response = litellm.completion(
|
|||
|
||||
### Reasoning Override
|
||||
|
||||
If 2+ reasoning markers are detected in the user message, the request is automatically routed to the REASONING tier regardless of the weighted score. This ensures complex reasoning tasks get the appropriate model.
|
||||
If 2+ reasoning markers are detected in the user message, the request is promoted to the REASONING tier even when the weighted score maps lower, so complex reasoning tasks get the appropriate model. The promotion requires the score to reach `reasoning_override_min_score`, which tracks `tier_boundaries.simple_medium` unless set, so stock phrases on an otherwise trivial prompt cannot buy the top tier. Set it to `0` to promote on the markers alone.
|
||||
|
||||
### System Prompt Handling
|
||||
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ from litellm.constants import EMPTY_MAPPING, RETURN_RAW_MODEL_NAME_METADATA_KEY
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
|
||||
from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
|
||||
from litellm.llms.base_llm.base_utils import type_to_response_format_param
|
||||
from litellm.types.utils import (
|
||||
AUTOROUTER_CLASSIFIER_CALL_ORIGIN,
|
||||
|
|
@ -663,6 +664,35 @@ class ClassificationOutcome(NamedTuple):
|
|||
classifier_cost: float | None = None
|
||||
|
||||
|
||||
class _SessionAffinityPin(NamedTuple):
|
||||
model: str
|
||||
tier: ComplexityTier | None
|
||||
|
||||
|
||||
def _parse_session_affinity_pin(value: object) -> _SessionAffinityPin | None:
|
||||
if isinstance(value, str):
|
||||
return _SessionAffinityPin(model=value, tier=None)
|
||||
parts: Final[tuple[object, object] | None] = (
|
||||
(value.get("model"), value.get("tier"))
|
||||
if isinstance(value, Mapping)
|
||||
else (value[0], value[1])
|
||||
if isinstance(value, (list, tuple)) and len(value) == 2
|
||||
else None
|
||||
)
|
||||
if parts is None:
|
||||
return None
|
||||
model, tier_value = parts
|
||||
if not isinstance(model, str):
|
||||
return None
|
||||
tier: Final = ComplexityTier(tier_value) if isinstance(tier_value, str) else None
|
||||
return _SessionAffinityPin(model=model, tier=tier)
|
||||
|
||||
|
||||
def _session_affinity_cache_value(model: str, tier: ComplexityTier | str | None) -> Mapping[str, str | None]:
|
||||
tier_value: Final = _tier_name(tier) if tier is not None else None
|
||||
return {"model": model, "tier": tier_value} # mutable-ok: cache requires JSON mapping
|
||||
|
||||
|
||||
class ComplexityRouter(CustomLogger):
|
||||
"""
|
||||
Complexity router that classifies requests and routes to appropriate models.
|
||||
|
|
@ -1020,13 +1050,14 @@ class ComplexityRouter(CustomLogger):
|
|||
weights: Final = self.config.dimension_weights
|
||||
weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions)
|
||||
|
||||
# Check for reasoning override (2+ reasoning markers)
|
||||
boundaries: Final = self._effective_tier_boundaries()
|
||||
clears_override_floor: Final = weighted_score >= self._effective_reasoning_override_min_score()
|
||||
|
||||
# Reuse match count from _score_keyword_match to avoid scanning twice
|
||||
if reasoning_match_count >= 2:
|
||||
if reasoning_match_count >= 2 and clears_override_floor:
|
||||
return ComplexityTier.REASONING, weighted_score, tuple(signals), "reasoning_override"
|
||||
|
||||
# Map score to tier
|
||||
boundaries: Final = self._effective_tier_boundaries()
|
||||
if weighted_score < boundaries["simple_medium"]:
|
||||
tier = ComplexityTier.SIMPLE
|
||||
elif weighted_score < boundaries["medium_complex"]:
|
||||
|
|
@ -1038,6 +1069,18 @@ class ComplexityRouter(CustomLogger):
|
|||
|
||||
return tier, weighted_score, tuple(signals), "heuristic_scorer"
|
||||
|
||||
def _effective_reasoning_override_min_score(self) -> float:
|
||||
"""The score a request must reach before the reasoning-marker override may promote it.
|
||||
|
||||
Unset tracks the SIMPLE/MEDIUM boundary, so moving that boundary moves this floor with it
|
||||
and the override still cannot rescue a request the mapping would call SIMPLE. An explicit
|
||||
0 is a real floor, not an absent one, so the comparison is against None.
|
||||
"""
|
||||
configured: Final = self.config.reasoning_override_min_score
|
||||
if configured is None:
|
||||
return self._effective_tier_boundaries()["simple_medium"]
|
||||
return configured
|
||||
|
||||
def _effective_tier_boundaries(self) -> StandardLoggingRoutingDecisionTierBoundaries:
|
||||
"""The tier boundaries in effect, with the documented defaults filled in.
|
||||
|
||||
|
|
@ -1065,6 +1108,7 @@ class ComplexityRouter(CustomLogger):
|
|||
classifier_model: str | None = None,
|
||||
classifier_cost: float | None = None,
|
||||
conversation_continuing: bool = True,
|
||||
tier_litellm_params: Mapping[str, object] | None = None,
|
||||
) -> StandardLoggingRoutingDecision:
|
||||
"""Assemble the per-request provenance record for this router's decision.
|
||||
|
||||
|
|
@ -1094,6 +1138,7 @@ class ComplexityRouter(CustomLogger):
|
|||
if score is not None:
|
||||
decision["score"] = score
|
||||
decision["tier_boundaries"] = self._effective_tier_boundaries()
|
||||
decision["reasoning_override_min_score"] = self._effective_reasoning_override_min_score()
|
||||
if signals:
|
||||
# Stored as a list because this record is serialized to JSON for the spend
|
||||
# log and read back as an array by the dashboard; a sequence type that only
|
||||
|
|
@ -1113,6 +1158,10 @@ class ComplexityRouter(CustomLogger):
|
|||
decision["classifier_model"] = classifier_model
|
||||
if classifier_cost is not None:
|
||||
decision["classifier_cost"] = classifier_cost
|
||||
if tier_litellm_params:
|
||||
masked_tier_litellm_params: Final = mask_credentials_in_payload(tier_litellm_params)
|
||||
if isinstance(masked_tier_litellm_params, Mapping):
|
||||
decision["tier_litellm_params"] = masked_tier_litellm_params
|
||||
return decision
|
||||
|
||||
async def aclassify(
|
||||
|
|
@ -1443,6 +1492,13 @@ class ComplexityRouter(CustomLogger):
|
|||
|
||||
raise ValueError(f"No model configured for tier {tier_key} and no default_model set")
|
||||
|
||||
def _litellm_params_for_model(self, tier: ComplexityTier | str | None, model: str) -> Mapping[str, object]:
|
||||
if tier is None:
|
||||
return MappingProxyType({})
|
||||
entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ())
|
||||
entry: Final = next((candidate for candidate in entries if candidate.model_name == model), None)
|
||||
return entry.litellm_params if entry is not None else MappingProxyType({})
|
||||
|
||||
@staticmethod
|
||||
def _pick_from_tier_value(model: str | list[str], tier_key: str) -> str:
|
||||
if isinstance(model, str):
|
||||
|
|
@ -2054,9 +2110,10 @@ class ComplexityRouter(CustomLogger):
|
|||
cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None
|
||||
|
||||
if cache_key is not None:
|
||||
pinned_model: Final = await self.litellm_router_instance.cache.async_get_cache(key=cache_key)
|
||||
if isinstance(pinned_model, str):
|
||||
routed_model: str | None = pinned_model
|
||||
pinned_value: Final = await self.litellm_router_instance.cache.async_get_cache(key=cache_key)
|
||||
pinned_pin: Final = _parse_session_affinity_pin(pinned_value)
|
||||
if pinned_pin is not None:
|
||||
routed_model: str | None = pinned_pin.model
|
||||
pin_escalation_keyword: str | None = None
|
||||
if self.escalation_keywords:
|
||||
user_message: Final = (
|
||||
|
|
@ -2065,16 +2122,21 @@ class ComplexityRouter(CustomLogger):
|
|||
if user_message is not None:
|
||||
pin_escalation_keyword = self._matched_escalation_keyword(user_message)
|
||||
if pin_escalation_keyword is not None:
|
||||
routed_model = self._escalated_pin(pinned_model)
|
||||
routed_model = self._escalated_pin(pinned_pin.model)
|
||||
if routed_model is not None:
|
||||
escalated: Final = routed_model != pinned_model
|
||||
escalated: Final = routed_model != pinned_pin.model
|
||||
resolved_pin_tier: Final = (
|
||||
pinned_pin.tier
|
||||
if not escalated and pinned_pin.tier is not None
|
||||
else self._tier_for_model(routed_model)
|
||||
)
|
||||
# The floor outranks the pin because plan mode is a transient state of the
|
||||
# session, not a request to move it: the turns carrying the sentinel route at
|
||||
# the floor, and the stored pin deliberately keeps the session's own model so
|
||||
# the first turn after plan mode exits auto-routes exactly as it would have.
|
||||
# Escalation is the opposite on purpose -- an explicit ask to re-pin higher.
|
||||
pin_plan_sentinel: Final = self._matched_plan_mode_signal(request_kwargs, resolved_messages)
|
||||
pinned_tier: Final = self._tier_for_model(routed_model) if pin_plan_sentinel is not None else None
|
||||
pinned_tier: Final = resolved_pin_tier if pin_plan_sentinel is not None else None
|
||||
plan_floored: Final = (
|
||||
pinned_tier is not None and self._apply_plan_mode_floor(pinned_tier) != pinned_tier
|
||||
)
|
||||
|
|
@ -2085,7 +2147,7 @@ class ComplexityRouter(CustomLogger):
|
|||
# pin mid-conversation just because it outlives the original write.
|
||||
await self.litellm_router_instance.cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=session_model,
|
||||
value=_session_affinity_cache_value(session_model, resolved_pin_tier),
|
||||
ttl=self.config.session_affinity_ttl_seconds,
|
||||
)
|
||||
if self.config.adaptive:
|
||||
|
|
@ -2104,19 +2166,23 @@ class ComplexityRouter(CustomLogger):
|
|||
verbose_router_logger.info(
|
||||
"ComplexityRouter: routing decision cause=%s, routed_model=%s", cause, routed_model
|
||||
)
|
||||
routed_pin_tier: Final = self._tier_for_model(routed_model) if plan_floored else resolved_pin_tier
|
||||
session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model)
|
||||
has_original_messages: Final = messages is not None and len(messages) > 0
|
||||
return self._with_session_deployment_affinity(
|
||||
PreRoutingHookResponse(
|
||||
model=routed_model,
|
||||
messages=messages if has_original_messages else None,
|
||||
litellm_params=session_tier_litellm_params,
|
||||
routing_decision=self._build_routing_decision(
|
||||
routed_model=routed_model,
|
||||
cause=cause,
|
||||
tier=self._tier_for_model(routed_model),
|
||||
tier=routed_pin_tier,
|
||||
matched_keyword=pin_plan_sentinel if plan_floored else None,
|
||||
escalation_keyword=pin_escalation_keyword,
|
||||
escalated=escalated,
|
||||
conversation_continuing=conversation_continuing,
|
||||
tier_litellm_params=session_tier_litellm_params,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
|
@ -2143,7 +2209,10 @@ class ComplexityRouter(CustomLogger):
|
|||
if pinnable and cache_key is not None and response is not None:
|
||||
await self.litellm_router_instance.cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=response.model,
|
||||
value=_session_affinity_cache_value(
|
||||
response.model,
|
||||
response.routing_decision.get("tier") if response.routing_decision is not None else None,
|
||||
),
|
||||
ttl=self.config.session_affinity_ttl_seconds,
|
||||
)
|
||||
return self._with_session_deployment_affinity(response)
|
||||
|
|
@ -2257,6 +2326,7 @@ class ComplexityRouter(CustomLogger):
|
|||
)
|
||||
keyword_plan_floored: Final = routed_tier != escalated_tier
|
||||
routed_model = await self._pick_model_for_tier(routed_tier, messages, resolved_messages, request_kwargs)
|
||||
keyword_tier_litellm_params: Final = self._litellm_params_for_model(routed_tier, routed_model)
|
||||
keyword_cause: Final[RoutingDecisionCause] = (
|
||||
"plan_mode"
|
||||
if keyword_plan_floored
|
||||
|
|
@ -2272,6 +2342,7 @@ class ComplexityRouter(CustomLogger):
|
|||
return PreRoutingHookResponse(
|
||||
model=routed_model,
|
||||
messages=messages if has_original_messages else None,
|
||||
litellm_params=keyword_tier_litellm_params,
|
||||
routing_decision=self._build_routing_decision(
|
||||
routed_model=routed_model,
|
||||
conversation_continuing=conversation_continuing,
|
||||
|
|
@ -2280,6 +2351,7 @@ class ComplexityRouter(CustomLogger):
|
|||
matched_keyword=plan_mode_sentinel if keyword_plan_floored else override.matched_keyword,
|
||||
escalation_keyword=escalation_keyword,
|
||||
escalated=keyword_escalated,
|
||||
tier_litellm_params=keyword_tier_litellm_params,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -2366,6 +2438,7 @@ class ComplexityRouter(CustomLogger):
|
|||
routed_model,
|
||||
)
|
||||
|
||||
tier_litellm_params: Final = self._litellm_params_for_model(tier, routed_model)
|
||||
classifier_model: Final = (
|
||||
self.config.classifier_llm_config.model
|
||||
if outcome.cause == "llm_classifier" and self.config.classifier_llm_config is not None
|
||||
|
|
@ -2391,6 +2464,7 @@ class ComplexityRouter(CustomLogger):
|
|||
return PreRoutingHookResponse(
|
||||
model=routed_model,
|
||||
messages=messages if has_original_messages else None,
|
||||
litellm_params=tier_litellm_params,
|
||||
routing_decision=self._build_routing_decision(
|
||||
routed_model=routed_model,
|
||||
conversation_continuing=conversation_continuing,
|
||||
|
|
@ -2403,5 +2477,6 @@ class ComplexityRouter(CustomLogger):
|
|||
escalated=escalated,
|
||||
classifier_model=classifier_model,
|
||||
classifier_cost=outcome.classifier_cost,
|
||||
tier_litellm_params=tier_litellm_params,
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -5,10 +5,12 @@ Contains default keyword lists, weights, tier boundaries, and configuration clas
|
|||
All values are configurable via proxy config.yaml.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from enum import Enum
|
||||
from typing import Final, Literal
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator
|
||||
|
||||
from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin
|
||||
|
||||
|
|
@ -159,6 +161,44 @@ class ReminderMarkerPair(BaseModel):
|
|||
return self
|
||||
|
||||
|
||||
class ComplexityTierModel(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
model_name: str
|
||||
litellm_params: Annotated[Mapping[str, object], SkipValidation()] = Field(
|
||||
default_factory=lambda: MappingProxyType({})
|
||||
)
|
||||
|
||||
@field_validator("litellm_params", mode="before")
|
||||
@classmethod
|
||||
def _freeze_litellm_params(cls, value: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return MappingProxyType(dict(value))
|
||||
|
||||
@field_serializer("litellm_params")
|
||||
def _serialize_litellm_params(self, value: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return dict(value) # mutable-ok: Pydantic JSON serialization requires a concrete mapping
|
||||
|
||||
|
||||
def _normalize_tier_entries(
|
||||
raw_value: object,
|
||||
tier: str,
|
||||
) -> tuple[str | list[str], tuple[ComplexityTierModel, ...]]:
|
||||
raw_entries: Final = raw_value if isinstance(raw_value, (list, tuple)) else (raw_value,)
|
||||
entries: Final = tuple(
|
||||
ComplexityTierModel(model_name=entry) if isinstance(entry, str) else ComplexityTierModel.model_validate(entry)
|
||||
for entry in raw_entries
|
||||
)
|
||||
model_names: Final = tuple(entry.model_name for entry in entries)
|
||||
if len(model_names) != len(frozenset(model_names)):
|
||||
raise ValueError(f"tier {tier} contains duplicate model_name values; each pool entry needs distinct parameters")
|
||||
normalized: Final = (
|
||||
entries[0].model_name
|
||||
if not isinstance(raw_value, (list, tuple))
|
||||
else list(model_names) # mutable-ok: config.tiers must preserve its existing list contract
|
||||
)
|
||||
return normalized, entries
|
||||
|
||||
|
||||
# ─── Default Keyword Lists ───
|
||||
# Note: Keywords should be full words/phrases to avoid substring false positives.
|
||||
# The matching logic uses word boundary detection for single-word keywords.
|
||||
|
|
@ -425,6 +465,9 @@ class ComplexityRouterConfig(BaseModel):
|
|||
"A list is randomly picked from when adaptive=False, and used as a soft-floor home pool when adaptive=True"
|
||||
),
|
||||
)
|
||||
tier_model_configs: Mapping[str, tuple[ComplexityTierModel, ...]] = Field(
|
||||
default_factory=dict,
|
||||
)
|
||||
|
||||
tier_definitions: tuple[TierDefinition, ...] | None = Field(
|
||||
default=None,
|
||||
|
|
@ -481,6 +524,15 @@ class ComplexityRouterConfig(BaseModel):
|
|||
),
|
||||
)
|
||||
|
||||
reasoning_override_min_score: float | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Minimum weighted score a request must reach before 2+ reasoning markers may promote it to the "
|
||||
"reasoning tier. Unset tracks tier_boundaries.simple_medium, so the override never rescues a "
|
||||
"request the scorer placed in the cheapest tier; 0 restores the unconditional override"
|
||||
),
|
||||
)
|
||||
|
||||
# Token count thresholds
|
||||
token_thresholds: dict[str, int] = Field(
|
||||
default_factory=lambda: DEFAULT_TOKEN_THRESHOLDS.copy(),
|
||||
|
|
@ -768,6 +820,55 @@ class ComplexityRouterConfig(BaseModel):
|
|||
coerced[key] = item
|
||||
return coerced
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def _normalize_tier_model_configs(cls, value: object) -> object:
|
||||
if not isinstance(value, dict):
|
||||
return value
|
||||
raw_tiers: Final = value.get("tiers")
|
||||
if not isinstance(raw_tiers, dict):
|
||||
return value
|
||||
existing_configs: Final = value.get("tier_model_configs")
|
||||
normalized_entries: Final = MappingProxyType(
|
||||
{tier: _normalize_tier_entries(raw_value, tier) for tier, raw_value in raw_tiers.items()}
|
||||
)
|
||||
normalized_tiers: Final = MappingProxyType(
|
||||
{tier: normalized for tier, (normalized, _) in normalized_entries.items()}
|
||||
)
|
||||
incoming_params: Final = (
|
||||
MappingProxyType(
|
||||
{
|
||||
(tier, entry.model_name): entry.litellm_params
|
||||
for tier, entries in existing_configs.items()
|
||||
for entry in (ComplexityTierModel.model_validate(item) for item in entries)
|
||||
}
|
||||
)
|
||||
if isinstance(existing_configs, dict)
|
||||
else MappingProxyType({})
|
||||
)
|
||||
tier_model_configs: Final = MappingProxyType(
|
||||
{
|
||||
tier: tuple(
|
||||
entry.model_copy(
|
||||
update=MappingProxyType(
|
||||
{
|
||||
"litellm_params": incoming_params.get((tier, entry.model_name), entry.litellm_params),
|
||||
}
|
||||
)
|
||||
)
|
||||
for entry in entries
|
||||
)
|
||||
for tier, (_, entries) in normalized_entries.items()
|
||||
if any(entry.litellm_params for entry in entries)
|
||||
or (isinstance(existing_configs, dict) and tier in existing_configs)
|
||||
}
|
||||
)
|
||||
return { # mutable-ok: Pydantic before-validator requires a concrete mapping
|
||||
**value,
|
||||
"tiers": normalized_tiers,
|
||||
"tier_model_configs": tier_model_configs,
|
||||
}
|
||||
|
||||
@field_validator("escalation_keywords")
|
||||
@classmethod
|
||||
def _normalize_escalation_keywords(cls, value: list[str] | None) -> list[str] | None:
|
||||
|
|
|
|||
|
|
@ -365,6 +365,22 @@ def get_secret(
|
|||
raise e
|
||||
|
||||
|
||||
def secret_manager_would_be_consulted(secret_name: str) -> bool:
|
||||
"""
|
||||
Returns True if a `get_secret` read for `secret_name` would actually reach the hosted manager.
|
||||
|
||||
Mirrors the gating `get_secret` applies below: the manager has to be up and readable, and
|
||||
`hosted_keys`, when set, is an allowlist of the names it is consulted for. Callers use this to
|
||||
tell "the manager does not have this key" apart from "the manager was never asked".
|
||||
"""
|
||||
if not _should_read_secret_from_secret_manager():
|
||||
return False
|
||||
key_management_settings: Final = litellm._key_management_settings
|
||||
if key_management_settings is None or key_management_settings.hosted_keys is None:
|
||||
return True
|
||||
return secret_name.removeprefix("os.environ/") in key_management_settings.hosted_keys
|
||||
|
||||
|
||||
def _should_read_secret_from_secret_manager() -> bool:
|
||||
"""
|
||||
Returns True if the secret manager should be used to read the secret, False otherwise
|
||||
|
|
@ -373,11 +389,7 @@ def _should_read_secret_from_secret_manager() -> bool:
|
|||
- If the `_key_management_settings` access mode is "read_only" or "read_and_write", return True
|
||||
- Otherwise, return False
|
||||
"""
|
||||
if litellm.secret_manager_client is not None:
|
||||
if litellm._key_management_settings is not None:
|
||||
if (
|
||||
litellm._key_management_settings.access_mode == "read_only"
|
||||
or litellm._key_management_settings.access_mode == "read_and_write"
|
||||
):
|
||||
return True
|
||||
return False
|
||||
key_management_settings: Final = litellm._key_management_settings
|
||||
if litellm.secret_manager_client is None or key_management_settings is None:
|
||||
return False
|
||||
return key_management_settings.access_mode in ("read_only", "read_and_write")
|
||||
|
|
|
|||
|
|
@ -620,6 +620,12 @@ class AnthropicResponseUsageBlock(BaseModel):
|
|||
output_tokens: int
|
||||
|
||||
|
||||
class AnthropicOutputTokensDetails(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
thinking_tokens: int | None = None
|
||||
|
||||
|
||||
AnthropicFinishReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from collections.abc import Mapping
|
|||
from datetime import datetime, timezone
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import AliasChoices, BaseModel, ConfigDict, Field, computed_field, field_validator, model_validator
|
||||
from pydantic import BaseModel, Field, computed_field, field_validator, model_validator
|
||||
|
||||
from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig
|
||||
from litellm.types.utils import StandardLoggingRoutingDecision
|
||||
|
|
@ -155,13 +155,17 @@ DEFAULT_SHADOW_EVAL_JUDGE_MODEL: Final[str] = "anthropic/claude-sonnet-5"
|
|||
|
||||
|
||||
class StartShadowEvalRequest(BaseModel):
|
||||
"""Start duplicating a key's traffic for blind comparison against an auto-router."""
|
||||
"""Start duplicating one or more keys' traffic for blind comparison against an auto-router."""
|
||||
|
||||
api_key_id: str = Field(
|
||||
api_key_ids: tuple[str, ...] = Field(
|
||||
min_length=1,
|
||||
max_length=100,
|
||||
description=(
|
||||
"The hashed virtual key whose traffic will be shadowed. Shadow evaluation runs ONLY on this "
|
||||
"key's traffic; requests made with any other key are not sampled."
|
||||
)
|
||||
"The hashed virtual keys whose traffic will be shadowed. Shadow evaluation runs ONLY on these "
|
||||
"keys' traffic; requests made with any other key are not sampled. Each key carries its own "
|
||||
"max_turns budget, so one key exhausting its budget leaves the others sampling. At most 100 "
|
||||
"keys per job, which also bounds every read the job's endpoints make."
|
||||
),
|
||||
)
|
||||
router_name: str = Field(description="The auto-router under evaluation, in either direction")
|
||||
direction: ShadowEvalDirection = Field(
|
||||
|
|
@ -204,8 +208,9 @@ class StartShadowEvalRequest(BaseModel):
|
|||
ge=1,
|
||||
le=2000,
|
||||
description=(
|
||||
"Sample budget: the job judges at most this many turns, then completes. This is also the spend "
|
||||
"bound; expected judge cost is roughly max_turns times one judge call"
|
||||
"Per-key sample budget: the job judges at most this many turns of EACH scoped key's traffic, "
|
||||
"so a job over N keys judges at most N times max_turns turns. This is also the spend bound; "
|
||||
"expected judge cost is roughly that turn ceiling times one judge call"
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -214,6 +219,12 @@ class StartShadowEvalRequest(BaseModel):
|
|||
def _round_percentage(cls, value: float) -> float:
|
||||
return round(value, 2)
|
||||
|
||||
@field_validator("api_key_ids")
|
||||
@classmethod
|
||||
def _dedupe_keys(cls, value: tuple[str, ...]) -> tuple[str, ...]:
|
||||
"""A key named twice would collide with itself on the one-active-per-(key, direction) index."""
|
||||
return tuple(dict.fromkeys(value))
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _baseline_model_matches_direction(self) -> "StartShadowEvalRequest":
|
||||
if self.direction == "reverse" and self.baseline_model is None:
|
||||
|
|
@ -251,24 +262,46 @@ class ShadowEvalResult(BaseModel):
|
|||
by_tier: tuple[ShadowEvalSlice, ...]
|
||||
by_current_model: tuple[ShadowEvalSlice, ...] = Field(
|
||||
description=(
|
||||
"Sliced by the model that served the real arm: the key's incumbent models in forward mode, "
|
||||
"Sliced by the model that served the real arm: the keys' incumbent models in forward mode, "
|
||||
"and in reverse the models the router itself picked"
|
||||
)
|
||||
)
|
||||
by_key: tuple[ShadowEvalSlice, ...] = Field(
|
||||
description=(
|
||||
"One slice per scoped key that has judged verdicts, grouped on the raw key hash. Keys the job "
|
||||
"scopes but has not judged a turn for yet are absent rather than reported as zero"
|
||||
),
|
||||
)
|
||||
overall_shadow_win_rate_pct: float
|
||||
overall_tie_rate_pct: float
|
||||
|
||||
|
||||
class ShadowEvalJobResponse(BaseModel):
|
||||
"""A shadow-eval job. Validates directly from the prisma record (job_id reads the
|
||||
row's id); status is derived from stopped_at and ends_at, never stored, so no writer
|
||||
anywhere can produce an inconsistent one. Aggregate fields are populated by the
|
||||
detail endpoint only and stay None on list responses."""
|
||||
class ShadowEvalJobKeyResponse(BaseModel):
|
||||
"""One key a job shadows, with its own budget and stop state."""
|
||||
|
||||
model_config = ConfigDict(from_attributes=True, populate_by_name=True)
|
||||
api_key_id: str = Field(description="The hashed virtual key whose traffic this entry scopes")
|
||||
max_turns: int = Field(description="This key's own sample budget, independent of its siblings'")
|
||||
stopped_at: datetime | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"When this key's slot was stamped free, whether its own budget ran out, the window closed, "
|
||||
"or an operator stopped the job; status is derived, so a spent budget reads completed even "
|
||||
"while this is still unset"
|
||||
),
|
||||
)
|
||||
attempt_count: int | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"This key's sampled attempts so far, judged and errored alike, the same count the sampler "
|
||||
"budgets against max_turns; populated on list and detail responses. Frozen at stopped_at "
|
||||
"once the key is stamped, so in-flight attempts landing after a stop never reclassify it"
|
||||
),
|
||||
)
|
||||
|
||||
@property
|
||||
def budget_spent(self) -> bool:
|
||||
return self.attempt_count is not None and self.attempt_count >= self.max_turns
|
||||
|
||||
job_id: str = Field(validation_alias=AliasChoices("id", "job_id"))
|
||||
api_key_id: str = Field(description="The hashed virtual key whose traffic this job evaluates, and only that key's")
|
||||
key_alias: str | None = Field(
|
||||
default=None,
|
||||
description="Alias of the shadowed key, resolved from the key row at read time; None when unset or deleted",
|
||||
|
|
@ -277,15 +310,34 @@ class ShadowEvalJobResponse(BaseModel):
|
|||
default=None,
|
||||
description="Masked display name (sk-...) of the shadowed key, resolved at read time like key_alias",
|
||||
)
|
||||
|
||||
|
||||
class ShadowEvalJobResponse(BaseModel):
|
||||
"""A shadow-eval job over one or more keys, each with its own budget and stop state;
|
||||
status is derived from stopped_by, the keys' stop and budget state, and ends_at,
|
||||
never stored, so no writer anywhere can produce an inconsistent one. Aggregate
|
||||
fields are populated by the detail endpoint only and stay None on list responses."""
|
||||
|
||||
job_id: str
|
||||
keys: tuple[ShadowEvalJobKeyResponse, ...] = Field(
|
||||
min_length=1,
|
||||
description="The keys whose traffic this job evaluates, and only those keys', each with its own budget",
|
||||
)
|
||||
router_name: str
|
||||
direction: ShadowEvalDirection = "forward"
|
||||
baseline_model: str | None = None
|
||||
judge_model: str
|
||||
shadow_percentage: float
|
||||
max_turns: int
|
||||
created_at: datetime
|
||||
ends_at: datetime
|
||||
stopped_at: datetime | None = None
|
||||
stopped_by: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"The operator who stopped the job early, recorded by the stop endpoint; 'unknown' backfilled "
|
||||
"by migration for jobs that displayed stopped when the column arrived; None when the job "
|
||||
"ended on its own. Its presence is what makes a job read stopped rather than completed"
|
||||
),
|
||||
)
|
||||
|
||||
judged_count: int | None = Field(default=None, description="Verdicts recorded; detail endpoint only")
|
||||
error_count: int | None = Field(default=None, description="Sampled attempts that errored; detail endpoint only")
|
||||
|
|
@ -296,12 +348,19 @@ class ShadowEvalJobResponse(BaseModel):
|
|||
@computed_field
|
||||
@property
|
||||
def status(self) -> ShadowEvalStatus:
|
||||
"""A job whose window has passed reads completed even if a later sweep stamped
|
||||
stopped_at; stopped means sampling ended before the window did."""
|
||||
"""Three recorded facts, no history-guessing: a stop is stopped_by (the migration
|
||||
backfills it for every job that displayed stopped when the column arrived, so the
|
||||
pre-column population is closed), completion is the window passing or every key
|
||||
spending its budget, and anything else is running. The all-keys-stamped fallback
|
||||
covers only stops written by pre-column pods during a rolling deploy."""
|
||||
if self.stopped_by is not None:
|
||||
return "stopped"
|
||||
if datetime.now(timezone.utc) >= (
|
||||
self.ends_at if self.ends_at.tzinfo else self.ends_at.replace(tzinfo=timezone.utc)
|
||||
):
|
||||
return "completed"
|
||||
if self.stopped_at is not None:
|
||||
if all(key.budget_spent for key in self.keys):
|
||||
return "completed"
|
||||
if all(key.stopped_at is not None for key in self.keys):
|
||||
return "stopped"
|
||||
return "running"
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc
|
|||
|
||||
import datetime
|
||||
import enum
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints
|
||||
|
||||
|
|
@ -897,6 +898,7 @@ class PreRoutingHookResponse(BaseModel):
|
|||
messages: list[dict[str, Any]] | None
|
||||
routing_decision: StandardLoggingRoutingDecision | None = None
|
||||
session_affinity_ttl_seconds: int | None = None
|
||||
litellm_params: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
_PreRoutingStrategyT_co = TypeVar("_PreRoutingStrategyT_co", covariant=True)
|
||||
|
|
|
|||
|
|
@ -248,6 +248,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
regional_processing_uplift_multiplier_us: (
|
||||
float | None
|
||||
) # OpenAI US data-residency uplift multiplier applied to all token costs (e.g. 1.10 = +10%)
|
||||
regional_endpoint_uplift_multiplier: ReadOnly[
|
||||
float | None
|
||||
] # Vertex AI non-global (regional) endpoint uplift multiplier applied to all token costs (e.g. 1.10 = +10%)
|
||||
output_cost_per_character: float | None # only for vertex ai models
|
||||
output_cost_per_audio_token: float | None
|
||||
output_cost_per_token_above_128k_tokens: float | None # only for vertex ai models
|
||||
|
|
@ -2836,9 +2839,11 @@ class StandardLoggingRoutingDecision(TypedDict, total=False):
|
|||
classifier_cost: float
|
||||
escalated: bool
|
||||
tier_boundaries: StandardLoggingRoutingDecisionTierBoundaries
|
||||
reasoning_override_min_score: ReadOnly[float]
|
||||
conversation_continuing: bool
|
||||
savings_baseline_model: str
|
||||
savings_baseline_deployment_id: str
|
||||
tier_litellm_params: Mapping[str, object] # writable-ok: Pydantic warns on ReadOnly TypedDict fields
|
||||
|
||||
|
||||
# Fields whose values quote the caller's prompt. Dropped when an operator turns message
|
||||
|
|
@ -2860,9 +2865,11 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset(
|
|||
"classifier_cost",
|
||||
"escalated",
|
||||
"tier_boundaries",
|
||||
"reasoning_override_min_score",
|
||||
"conversation_continuing",
|
||||
"savings_baseline_model",
|
||||
"savings_baseline_deployment_id",
|
||||
"tier_litellm_params",
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -3113,16 +3120,17 @@ class CostBreakdown(TypedDict, total=False):
|
|||
"""
|
||||
Detailed cost breakdown for a request.
|
||||
|
||||
``service_tier`` and ``data_residency`` record the pricing basis the cost was
|
||||
computed on, not what the caller asked for. A consumer that has to price a
|
||||
counterfactual against this request (what another model would have charged for
|
||||
it) needs the same basis to compare like with like, and re-deriving it from the
|
||||
request is not possible after the fact: the tier the biller used comes from
|
||||
``optional_params``, which no log record carries.
|
||||
``service_tier``, ``data_residency``, and ``vertex_location`` record the pricing
|
||||
basis the cost was computed on, not what the caller asked for. A consumer that has
|
||||
to price a counterfactual against this request (what another model would have
|
||||
charged for it) needs the same basis to compare like with like, and re-deriving it
|
||||
from the request is not possible after the fact: the tier the biller used comes
|
||||
from ``optional_params``, which no log record carries.
|
||||
"""
|
||||
|
||||
service_tier: str | None
|
||||
data_residency: str | None
|
||||
vertex_location: ReadOnly[str | None]
|
||||
input_cost: float # Cost of raw (non-cached) input tokens only
|
||||
cache_read_cost: float # Cost of cache-read tokens (discounted rate)
|
||||
cache_creation_cost: float # Cost of cache-write tokens (premium rate)
|
||||
|
|
@ -3388,6 +3396,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
|
|||
annotation_cost_per_page: float | None = None
|
||||
regional_processing_uplift_multiplier_eu: float | None = None
|
||||
regional_processing_uplift_multiplier_us: float | None = None
|
||||
regional_endpoint_uplift_multiplier: float | None = None
|
||||
|
||||
@classmethod
|
||||
def strip_custom_pricing_fields(cls, model_info: dict[str, Any]) -> dict[str, Any]:
|
||||
|
|
@ -3818,6 +3827,7 @@ class SearchProviders(str, Enum):
|
|||
YOU_COM = "you_com"
|
||||
APISERPENT = "apiserpent"
|
||||
TINYFISH = "tinyfish"
|
||||
AGENTCORE = "agentcore"
|
||||
NIMBLE = "nimble"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -5662,6 +5662,7 @@ def _get_model_info_helper(
|
|||
regional_processing_uplift_multiplier_us=_model_info.get(
|
||||
"regional_processing_uplift_multiplier_us", None
|
||||
),
|
||||
regional_endpoint_uplift_multiplier=_model_info.get("regional_endpoint_uplift_multiplier", None),
|
||||
output_cost_per_audio_token=_model_info.get("output_cost_per_audio_token", None),
|
||||
output_cost_per_character=_model_info.get("output_cost_per_character", None),
|
||||
output_cost_per_reasoning_token=_model_info.get("output_cost_per_reasoning_token", None),
|
||||
|
|
@ -9077,6 +9078,7 @@ class ProviderConfigManager:
|
|||
from litellm.llms.apiserpent.search.transformation import (
|
||||
APISerpentSearchConfig,
|
||||
)
|
||||
from litellm.llms.bedrock.search.transformation import AgentCoreSearchConfig
|
||||
from litellm.llms.brave.search.transformation import BraveSearchConfig
|
||||
from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig
|
||||
from litellm.llms.duckduckgo.search.transformation import DuckDuckGoSearchConfig
|
||||
|
|
@ -9115,6 +9117,7 @@ class ProviderConfigManager:
|
|||
SearchProviders.YOU_COM: YouComSearchConfig,
|
||||
SearchProviders.APISERPENT: APISerpentSearchConfig,
|
||||
SearchProviders.TINYFISH: TinyfishSearchConfig,
|
||||
SearchProviders.AGENTCORE: AgentCoreSearchConfig,
|
||||
SearchProviders.NIMBLE: NimbleSearchConfig,
|
||||
}
|
||||
config_class: Final = PROVIDER_TO_CONFIG_MAP.get(provider, None)
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -514,6 +514,11 @@
|
|||
"type": "object",
|
||||
"description": "Provider-internal routing hints (e.g. bedrock_invocation_schema)."
|
||||
},
|
||||
"regional_endpoint_uplift_multiplier": {
|
||||
"type": "number",
|
||||
"minimum": 1,
|
||||
"description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%)."
|
||||
},
|
||||
"regional_processing_uplift_multiplier_eu": {
|
||||
"type": "number",
|
||||
"minimum": 1,
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@
|
|||
"limit": 711
|
||||
},
|
||||
"ANN205": {
|
||||
"limit": 113
|
||||
"limit": 112
|
||||
},
|
||||
"ANN206": {
|
||||
"limit": 133
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue