mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
Merge upstream litellm_internal_staging; resolve mcp_tool_search_top_k conflicts
Keep configurable default top_k while adopting upstream typing on object permission models.
This commit is contained in:
commit
632b1859c0
2647 changed files with 84023 additions and 46600 deletions
|
|
@ -281,6 +281,32 @@ jobs:
|
|||
uv build --wheel --out-dir dist
|
||||
uv run --no-sync python tests/windows_tests/check_windows_wheel_install.py
|
||||
|
||||
base_sdk_install:
|
||||
docker:
|
||||
- image: cimg/python:3.12@sha256:9c796c23c84e84a66a964acb508d39dc5433c81a47e07efd56dccbbc2427e07c
|
||||
auth:
|
||||
username: ${DOCKERHUB_USERNAME}
|
||||
password: ${DOCKERHUB_PASSWORD}
|
||||
working_directory: ~/project
|
||||
steps:
|
||||
- checkout
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- run:
|
||||
name: Build the wheel
|
||||
environment:
|
||||
UV_HTTP_TIMEOUT: "300"
|
||||
command: |
|
||||
uv build --wheel --out-dir dist
|
||||
- run:
|
||||
name: Install the wheel with no extras and smoke-check it
|
||||
environment:
|
||||
UV_HTTP_TIMEOUT: "300"
|
||||
command: |
|
||||
uv venv /tmp/base-sdk --python 3.12
|
||||
VIRTUAL_ENV=/tmp/base-sdk uv pip install dist/*.whl
|
||||
/tmp/base-sdk/bin/python tests/base_sdk_tests/check_base_sdk_install.py
|
||||
|
||||
local_testing_part1:
|
||||
docker:
|
||||
- &python312_image
|
||||
|
|
@ -3031,6 +3057,8 @@ workflows:
|
|||
only:
|
||||
- main
|
||||
- /litellm_.*/
|
||||
- base_sdk_install:
|
||||
filters: *main_branches
|
||||
- local_testing_part1:
|
||||
filters: *main_branches
|
||||
- local_testing_part2:
|
||||
|
|
|
|||
1
.github/pull_request_template.md
vendored
1
.github/pull_request_template.md
vendored
|
|
@ -40,6 +40,7 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
|
|||
The proof must be completely e2e with no mocks, using, for example, actual LLM calls costing real $. `pytest` commands are not enough
|
||||
For bug fixes: show reproduction before the fix and passing behavior after
|
||||
Include the commit hash each proof was captured at, for both the before and the after runs
|
||||
If the change applies to all three LLM endpoints (/v1/responses, /v1/chat/completions, /v1/messages), include proof for every single one of them, not just one
|
||||
For new features: show the feature working end-to-end
|
||||
For UI changes: include before/after screenshots -->
|
||||
|
||||
|
|
|
|||
|
|
@ -24,9 +24,12 @@ jobs:
|
|||
- name: Update JSON Data
|
||||
run: |
|
||||
uv run --frozen --with 'aiohttp==3.13.3' python ".github/workflows/auto_update_price_and_context_window_file.py"
|
||||
- name: Regenerate JSON Schema
|
||||
run: |
|
||||
uv run --frozen python ci_cd/generate_model_prices_schema.py
|
||||
- name: Create Pull Request
|
||||
run: |
|
||||
git add model_prices_and_context_window.json
|
||||
git add model_prices_and_context_window.json model_prices_and_context_window.schema.json
|
||||
git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')"
|
||||
gh pr create --title "Update model_prices_and_context_window.json file" \
|
||||
--body "Automated update for model_prices_and_context_window.json" \
|
||||
|
|
|
|||
33
.github/workflows/image-scan.yml
vendored
33
.github/workflows/image-scan.yml
vendored
|
|
@ -9,6 +9,8 @@ on:
|
|||
- "litellm_**"
|
||||
paths:
|
||||
- docker/Dockerfile.non_root
|
||||
- migrations/Dockerfile
|
||||
- migrations/run.py
|
||||
- tests/proxy_migration_tests/test_offline_image_migration.py
|
||||
- uv.lock
|
||||
- ui/litellm-dashboard/package-lock.json
|
||||
|
|
@ -83,3 +85,34 @@ jobs:
|
|||
--only-fixed \
|
||||
--fail-on high \
|
||||
--output table
|
||||
|
||||
migrations-image:
|
||||
name: migrations-image
|
||||
runs-on: ubuntu-latest
|
||||
if: >-
|
||||
github.event_name != 'pull_request' ||
|
||||
github.event.pull_request.head.repo.full_name == github.repository
|
||||
timeout-minutes: 30
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Build migrations image
|
||||
run: docker build -f migrations/Dockerfile -t litellm-migrations-scan:${{ github.sha }} .
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Verify offline migration as a non-root uid
|
||||
env:
|
||||
LITELLM_IMAGE: litellm-migrations-scan:${{ github.sha }}
|
||||
LITELLM_MIGRATION_INTERPRETER: python3
|
||||
LITELLM_MIGRATION_SCRIPT: /app/run.py
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
|
||||
|
|
|
|||
8
.github/workflows/test-code-quality.yml
vendored
8
.github/workflows/test-code-quality.yml
vendored
|
|
@ -7,13 +7,17 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
code-quality:
|
||||
|
|
|
|||
9
.github/workflows/test-model-map.yaml
vendored
9
.github/workflows/test-model-map.yaml
vendored
|
|
@ -22,3 +22,12 @@ jobs:
|
|||
- name: Validate model_prices_and_context_window.json
|
||||
run: |
|
||||
jq empty model_prices_and_context_window.json
|
||||
|
||||
- name: Set up uv
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Check model_prices_and_context_window.schema.json is in sync
|
||||
run: |
|
||||
uv run --frozen python ci_cd/generate_model_prices_schema.py --check
|
||||
|
|
|
|||
8
.github/workflows/test-unit-core-utils.yml
vendored
8
.github/workflows/test-unit-core-utils.yml
vendored
|
|
@ -7,6 +7,10 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
@ -14,8 +18,8 @@ permissions:
|
|||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
core-utils:
|
||||
|
|
|
|||
|
|
@ -7,13 +7,17 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
documentation:
|
||||
|
|
|
|||
|
|
@ -7,6 +7,10 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
@ -14,8 +18,8 @@ permissions:
|
|||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
enterprise-routing:
|
||||
|
|
|
|||
8
.github/workflows/test-unit-integrations.yml
vendored
8
.github/workflows/test-unit-integrations.yml
vendored
|
|
@ -7,6 +7,10 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
@ -14,8 +18,8 @@ permissions:
|
|||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
integrations:
|
||||
|
|
|
|||
|
|
@ -7,13 +7,17 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
vertex-ai:
|
||||
|
|
|
|||
9
.github/workflows/test-unit-misc.yml
vendored
9
.github/workflows/test-unit-misc.yml
vendored
|
|
@ -7,6 +7,10 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
@ -14,8 +18,8 @@ permissions:
|
|||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
misc:
|
||||
|
|
@ -27,6 +31,7 @@ jobs:
|
|||
tests/test_litellm/a2a_protocol
|
||||
tests/test_litellm/anthropic_interface
|
||||
tests/test_litellm/completion_extras
|
||||
tests/test_litellm/compression
|
||||
tests/test_litellm/containers
|
||||
tests/test_litellm/experimental_mcp_client
|
||||
tests/test_litellm/models
|
||||
|
|
|
|||
8
.github/workflows/test-unit-proxy-auth.yml
vendored
8
.github/workflows/test-unit-proxy-auth.yml
vendored
|
|
@ -7,6 +7,10 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
@ -14,8 +18,8 @@ permissions:
|
|||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
proxy-auth:
|
||||
|
|
|
|||
8
.github/workflows/test-unit-proxy-db.yml
vendored
8
.github/workflows/test-unit-proxy-db.yml
vendored
|
|
@ -7,13 +7,17 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
# Semantic matrix: each shard groups tests by concern (auth, server, logging, …)
|
||||
# rather than alphabetical letter ranges. Adding a new test file means adding it
|
||||
|
|
|
|||
|
|
@ -7,14 +7,18 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
proxy-endpoints:
|
||||
|
|
|
|||
8
.github/workflows/test-unit-proxy-infra.yml
vendored
8
.github/workflows/test-unit-proxy-infra.yml
vendored
|
|
@ -7,6 +7,10 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
@ -14,8 +18,8 @@ permissions:
|
|||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
proxy-infra:
|
||||
|
|
|
|||
8
.github/workflows/test-unit-proxy-legacy.yml
vendored
8
.github/workflows/test-unit-proxy-legacy.yml
vendored
|
|
@ -7,13 +7,17 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
test:
|
||||
|
|
|
|||
|
|
@ -7,6 +7,10 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
@ -14,8 +18,8 @@ permissions:
|
|||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
responses-caching-types:
|
||||
|
|
|
|||
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -141,3 +141,4 @@ crash.*.log
|
|||
.coverage
|
||||
|
||||
ui/litellm-dashboard/out/
|
||||
litellm.log
|
||||
|
|
|
|||
|
|
@ -61,6 +61,8 @@ Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages
|
|||
|
||||
When working on a PR, keep the PR description in sync with new commits being made
|
||||
|
||||
Replies/rebuttals to AI PR review bots must be 15-25 word human-readable replies
|
||||
|
||||
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
|
||||
|
||||
Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers.
|
||||
|
|
|
|||
6
Makefile
6
Makefile
|
|
@ -75,7 +75,7 @@ install-dev:
|
|||
bootstrap:
|
||||
$(UV) sync --inexact --frozen --extra proxy --group proxy-dev --group e2e-dev
|
||||
$(UV_RUN) python scripts/prisma_generate_if_needed.py
|
||||
cd ui/litellm-dashboard && npm ci --no-audit --no-fund
|
||||
cd ui/litellm-dashboard && npm install --no-audit --no-fund
|
||||
@main_root=$$(git worktree list --porcelain | head -1 | sed 's/^worktree //'); \
|
||||
if [ "$$main_root" != "$$(git rev-parse --show-toplevel)" ] && [ -f "$$main_root/.env" ] && [ ! -f .env ]; then \
|
||||
cp "$$main_root/.env" .env && echo "bootstrap: copied .env from $$main_root"; \
|
||||
|
|
@ -176,6 +176,8 @@ lint-ruff-FULL-dev: install-dev
|
|||
if [ -n "$$files" ]; then echo "$$files" | xargs $(UV_RUN) ruff check; \
|
||||
else echo "No changed .py files to check."; fi
|
||||
|
||||
lint-basedpyright lint-basedpyright-budget-update: export NODE_OPTIONS := --max-old-space-size=12288
|
||||
|
||||
lint-basedpyright: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
|
||||
($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging
|
||||
|
||||
|
|
@ -237,7 +239,7 @@ lint-dev: lint-format-changed check-circular-imports check-import-safety
|
|||
# test-linting.yml (Python), test-litellm-ui-build.yml's frontend-lint (dashboard), and
|
||||
# check-ui-api-types.yml (API-type drift), skipping any whose files you didn't stage.
|
||||
# Not auto-installed as a git hook so it never slows an unrelated human commit.
|
||||
pre-commit:
|
||||
pre-commit: bootstrap
|
||||
./scripts/pre_commit_lint.sh
|
||||
|
||||
# Testing targets
|
||||
|
|
|
|||
|
|
@ -63,6 +63,8 @@ RUN mkdir -p /home/nonroot && \
|
|||
HOME=/home/nonroot prisma generate --schema=./schema.prisma && \
|
||||
chown -R nonroot:nonroot /home/nonroot/.cache
|
||||
|
||||
RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh
|
||||
|
||||
# ---------- Runtime ----------
|
||||
FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
||||
|
||||
|
|
@ -93,5 +95,5 @@ USER nonroot
|
|||
|
||||
EXPOSE 4001/tcp
|
||||
|
||||
ENTRYPOINT ["uvicorn", "backend.main:app"]
|
||||
ENTRYPOINT ["/app/docker/component_entrypoint.sh", "uvicorn", "backend.main:app"]
|
||||
CMD ["--host", "0.0.0.0", "--port", "4001"]
|
||||
|
|
|
|||
|
|
@ -1,12 +1,12 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"limit": 34906
|
||||
"limit": 29813
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2701
|
||||
"limit": 2645
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 330
|
||||
"limit": 329
|
||||
},
|
||||
"reportAttributeAccessIssue": {
|
||||
"limit": 516
|
||||
|
|
@ -18,13 +18,13 @@
|
|||
"limit": 59
|
||||
},
|
||||
"reportDeprecated": {
|
||||
"limit": 326
|
||||
"limit": 325
|
||||
},
|
||||
"reportDuplicateImport": {
|
||||
"limit": 42
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 10230
|
||||
"limit": 9473
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 11
|
||||
|
|
@ -33,7 +33,7 @@
|
|||
"limit": 227
|
||||
},
|
||||
"reportIncompatibleMethodOverride": {
|
||||
"limit": 78
|
||||
"limit": 77
|
||||
},
|
||||
"reportIncompatibleVariableOverride": {
|
||||
"limit": 12
|
||||
|
|
@ -54,10 +54,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"limit": 5893
|
||||
"limit": 5855
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15886
|
||||
"limit": 15852
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 41
|
||||
|
|
@ -72,7 +72,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportOptionalMemberAccess": {
|
||||
"limit": 1085
|
||||
"limit": 1079
|
||||
},
|
||||
"reportOptionalOperand": {
|
||||
"limit": 0
|
||||
|
|
@ -84,13 +84,13 @@
|
|||
"limit": 77
|
||||
},
|
||||
"reportPrivateUsage": {
|
||||
"limit": 2438
|
||||
"limit": 2437
|
||||
},
|
||||
"reportRedeclaration": {
|
||||
"limit": 12
|
||||
},
|
||||
"reportReturnType": {
|
||||
"limit": 225
|
||||
"limit": 219
|
||||
},
|
||||
"reportTypedDictNotRequiredAccess": {
|
||||
"limit": 27
|
||||
|
|
@ -99,31 +99,31 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 45870
|
||||
"limit": 45324
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 113
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 40525
|
||||
"limit": 40452
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 20384
|
||||
"limit": 20309
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 32099
|
||||
"limit": 31978
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 177
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"limit": 1023
|
||||
"limit": 1021
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"limit": 7
|
||||
},
|
||||
"reportUnnecessaryIsInstance": {
|
||||
"limit": 1206
|
||||
"limit": 1204
|
||||
},
|
||||
"reportUntypedBaseClass": {
|
||||
"limit": 165
|
||||
|
|
@ -135,10 +135,10 @@
|
|||
"limit": 33
|
||||
},
|
||||
"reportUnusedFunction": {
|
||||
"limit": 206
|
||||
"limit": 204
|
||||
},
|
||||
"reportUnusedImport": {
|
||||
"limit": 1005
|
||||
"limit": 1003
|
||||
},
|
||||
"reportUnusedVariable": {
|
||||
"limit": 1297
|
||||
|
|
|
|||
325
ci_cd/generate_model_prices_schema.py
Normal file
325
ci_cd/generate_model_prices_schema.py
Normal file
|
|
@ -0,0 +1,325 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import jsonschema
|
||||
|
||||
REPO_ROOT = Path(__file__).parent.parent
|
||||
PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json"
|
||||
SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json"
|
||||
|
||||
SPECIAL_ROOT_KEYS = frozenset({"sample_spec", "fallback_generalizations"})
|
||||
|
||||
JsonSchema = dict
|
||||
|
||||
NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
|
||||
NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
|
||||
BOOLEAN: JsonSchema = {"type": "boolean"}
|
||||
STRING: JsonSchema = {"type": "string"}
|
||||
|
||||
EXTRA_BOOLEAN_KEYS = frozenset(
|
||||
{
|
||||
"gemini_native_audio",
|
||||
"gemini_audio_only_live",
|
||||
"uses_embed_content",
|
||||
"use_openai_responses_path",
|
||||
"bedrock_converse_supports_strict_tools",
|
||||
}
|
||||
)
|
||||
|
||||
OBJECT_KEYS: dict[str, JsonSchema] = {
|
||||
"search_context_cost_per_query": {
|
||||
"type": "object",
|
||||
"description": "USD cost per web search query, keyed by search context size.",
|
||||
"properties": {
|
||||
"search_context_size_low": NONNEG_NUMBER,
|
||||
"search_context_size_medium": NONNEG_NUMBER,
|
||||
"search_context_size_high": NONNEG_NUMBER,
|
||||
},
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"metadata": {
|
||||
"type": "object",
|
||||
"description": "Free-form notes about the entry (e.g. pricing derivation).",
|
||||
},
|
||||
"provider_specific_entry": {
|
||||
"type": "object",
|
||||
"description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).",
|
||||
},
|
||||
}
|
||||
|
||||
ARRAY_KEYS: dict[str, JsonSchema] = {
|
||||
"supported_endpoints": {
|
||||
"type": "array",
|
||||
"description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.",
|
||||
"items": STRING,
|
||||
},
|
||||
"supported_modalities": {
|
||||
"type": "array",
|
||||
"description": "Input modalities the model accepts.",
|
||||
"items": {"type": "string", "enum": ["text", "image", "audio", "video"]},
|
||||
},
|
||||
"supported_output_modalities": {
|
||||
"type": "array",
|
||||
"description": "Output modalities the model can produce.",
|
||||
"items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]},
|
||||
},
|
||||
"supported_regions": {
|
||||
"type": "array",
|
||||
"description": "Cloud regions the model is available in ('global' or region ids).",
|
||||
"items": STRING,
|
||||
},
|
||||
"tiered_pricing": {
|
||||
"type": "array",
|
||||
"description": "Context-length or result-count tiered rates; each tier's costs apply within its range.",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"range": {
|
||||
"type": "array",
|
||||
"description": "[min, max] prompt-token span this tier applies to.",
|
||||
"items": NONNEG_NUMBER,
|
||||
"minItems": 2,
|
||||
"maxItems": 2,
|
||||
},
|
||||
"max_results_range": {
|
||||
"type": "array",
|
||||
"description": "[min, max] result-count span this tier applies to (search models).",
|
||||
"items": NONNEG_NUMBER,
|
||||
"minItems": 2,
|
||||
"maxItems": 2,
|
||||
},
|
||||
"input_cost_per_token": NONNEG_NUMBER,
|
||||
"output_cost_per_token": NONNEG_NUMBER,
|
||||
"output_cost_per_reasoning_token": NONNEG_NUMBER,
|
||||
"cache_read_input_token_cost": NONNEG_NUMBER,
|
||||
"input_cost_per_query": NONNEG_NUMBER,
|
||||
},
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
INTEGER_KEYS: dict[str, JsonSchema] = {
|
||||
"max_tokens": {
|
||||
**NONNEG_INTEGER,
|
||||
"description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.",
|
||||
},
|
||||
"max_input_tokens": {
|
||||
**NONNEG_INTEGER,
|
||||
"description": "Maximum prompt/context tokens the model accepts.",
|
||||
},
|
||||
"max_output_tokens": {
|
||||
**NONNEG_INTEGER,
|
||||
"description": "Maximum tokens the model can generate in one response.",
|
||||
},
|
||||
"output_vector_size": {
|
||||
**NONNEG_INTEGER,
|
||||
"description": "Embedding dimension for embedding models.",
|
||||
},
|
||||
"prompt_cache_min_tokens": {
|
||||
**NONNEG_INTEGER,
|
||||
"description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
|
||||
},
|
||||
"tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
|
||||
"rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
|
||||
}
|
||||
|
||||
NUMBER_KEYS: dict[str, JsonSchema] = {
|
||||
"regional_processing_uplift_multiplier_eu": {
|
||||
"type": "number",
|
||||
"minimum": 1,
|
||||
"description": "Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%).",
|
||||
},
|
||||
"regional_processing_uplift_multiplier_us": {
|
||||
"type": "number",
|
||||
"minimum": 1,
|
||||
"description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).",
|
||||
},
|
||||
}
|
||||
|
||||
COST_DESCRIPTIONS: dict[str, str] = {
|
||||
"input_cost_per_token": "USD per prompt token.",
|
||||
"output_cost_per_token": "USD per generated token.",
|
||||
"output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.",
|
||||
"cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.",
|
||||
"cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.",
|
||||
"input_cost_per_token_batches": "USD per prompt token via the provider's batch API.",
|
||||
"output_cost_per_token_batches": "USD per generated token via the provider's batch API.",
|
||||
}
|
||||
|
||||
|
||||
def cost_description(key: str) -> Optional[str]:
|
||||
if key in COST_DESCRIPTIONS:
|
||||
return COST_DESCRIPTIONS[key]
|
||||
if key.endswith("_flex"):
|
||||
return "Flex service-tier rate for the same-named base field."
|
||||
if key.endswith("_priority"):
|
||||
return "Priority service-tier rate for the same-named base field."
|
||||
if "_above_" in key:
|
||||
return "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
return None
|
||||
|
||||
|
||||
def cost_schema(key: str) -> JsonSchema:
|
||||
description = cost_description(key)
|
||||
return {**NONNEG_NUMBER, "description": description} if description else dict(NONNEG_NUMBER)
|
||||
|
||||
|
||||
def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
|
||||
return {
|
||||
"litellm_provider": {
|
||||
"type": "string",
|
||||
"description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.",
|
||||
},
|
||||
"mode": {
|
||||
"type": "string",
|
||||
"description": "Primary API surface / task type of the model.",
|
||||
"enum": list(modes),
|
||||
},
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "URL of the provider pricing/model page this entry was taken from.",
|
||||
},
|
||||
"deprecation_date": {
|
||||
"type": "string",
|
||||
"description": "Date the provider deprecates the model, YYYY-MM-DD.",
|
||||
"format": "date",
|
||||
"pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$",
|
||||
},
|
||||
"web_search_billing_unit": {
|
||||
"type": "string",
|
||||
"description": "Whether web search is billed per query or per prompt.",
|
||||
"enum": ["per_query", "per_prompt"],
|
||||
},
|
||||
"bedrock_output_config_effort_ceiling": {
|
||||
"type": "string",
|
||||
"description": "Highest reasoning effort the Bedrock output_config accepts for this model.",
|
||||
"enum": ["low", "medium", "high", "max", "xhigh"],
|
||||
},
|
||||
"comment": STRING,
|
||||
"audio_transcription_config": STRING,
|
||||
}
|
||||
|
||||
|
||||
def classify(key: str, modes: tuple) -> Optional[JsonSchema]:
|
||||
curated = {**OBJECT_KEYS, **ARRAY_KEYS, **string_key_schemas(modes), **INTEGER_KEYS, **NUMBER_KEYS}
|
||||
if key in curated:
|
||||
return curated[key]
|
||||
if key.startswith("supports_") or key in EXTRA_BOOLEAN_KEYS:
|
||||
return BOOLEAN
|
||||
if "cost" in key:
|
||||
return cost_schema(key)
|
||||
return None
|
||||
|
||||
|
||||
def build_schema(prices: dict) -> JsonSchema:
|
||||
entries = {name: entry for name, entry in prices.items() if name not in SPECIAL_ROOT_KEYS}
|
||||
all_keys = tuple(sorted({key for entry in entries.values() for key in entry}))
|
||||
modes = tuple(sorted({entry["mode"] for entry in entries.values() if "mode" in entry}))
|
||||
unclassified = tuple(key for key in all_keys if classify(key, modes) is None)
|
||||
if unclassified:
|
||||
raise SystemExit(
|
||||
f"Unclassified keys in {PRICES_PATH.name}: {', '.join(unclassified)}. "
|
||||
f"Add them to the key tables in {Path(__file__).name} and rerun it."
|
||||
)
|
||||
entry_properties = {key: classify(key, modes) for key in all_keys}
|
||||
return {
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"title": "LiteLLM model_prices_and_context_window.json",
|
||||
"description": (
|
||||
"Schema for LiteLLM's model price and context window registry "
|
||||
"(https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). "
|
||||
"Every top-level key except 'sample_spec' and 'fallback_generalizations' is a model id, "
|
||||
"optionally prefixed with its provider (e.g. 'azure/gpt-5.4'), mapping to a model entry. "
|
||||
"All costs are USD per unit. New optional fields are added regularly, so consumers should "
|
||||
"ignore unknown fields rather than reject them."
|
||||
),
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"sample_spec": {
|
||||
"type": "object",
|
||||
"description": (
|
||||
"Documentation placeholder illustrating the entry shape; not a real model and not "
|
||||
"schema-conformant (several values are prose)."
|
||||
),
|
||||
},
|
||||
"fallback_generalizations": {
|
||||
"type": "object",
|
||||
"description": "Regex rules that generalize unknown model ids to known families; not a model entry.",
|
||||
"properties": {
|
||||
"rules": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": STRING,
|
||||
"pattern": STRING,
|
||||
"description": STRING,
|
||||
},
|
||||
"required": ["name", "pattern"],
|
||||
"additionalProperties": True,
|
||||
},
|
||||
}
|
||||
},
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
"additionalProperties": {"$ref": "#/$defs/modelEntry"},
|
||||
"$defs": {
|
||||
"modelEntry": {
|
||||
"type": "object",
|
||||
"description": (
|
||||
"Pricing, limits, and capability flags for one model. Fields other than litellm_provider "
|
||||
"are optional; boolean capability flags are simply omitted when unknown or false."
|
||||
),
|
||||
"required": ["litellm_provider"],
|
||||
"properties": entry_properties,
|
||||
"additionalProperties": True,
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def render(schema: JsonSchema) -> str:
|
||||
return json.dumps(schema, indent=2) + "\n"
|
||||
|
||||
|
||||
def validation_errors(prices: dict, schema: JsonSchema) -> tuple:
|
||||
validator = jsonschema.Draft202012Validator(
|
||||
schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER
|
||||
)
|
||||
return tuple(
|
||||
f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
|
||||
for error in validator.iter_errors(prices)
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
check = "--check" in sys.argv[1:]
|
||||
prices = json.loads(PRICES_PATH.read_text())
|
||||
rendered = render(build_schema(prices))
|
||||
errors = validation_errors(prices, json.loads(rendered))
|
||||
if errors:
|
||||
print(f"{PRICES_PATH.name} does not validate against the generated schema:")
|
||||
print("\n".join(errors[:20]))
|
||||
return 1
|
||||
if not check:
|
||||
SCHEMA_PATH.write_text(rendered)
|
||||
print(f"wrote {SCHEMA_PATH}")
|
||||
return 0
|
||||
if not SCHEMA_PATH.exists() or SCHEMA_PATH.read_text() != rendered:
|
||||
print(
|
||||
f"{SCHEMA_PATH.name} is out of sync with {PRICES_PATH.name}. "
|
||||
f"Run `python {Path(__file__).relative_to(REPO_ROOT)}` and commit the result."
|
||||
)
|
||||
return 1
|
||||
print(f"{SCHEMA_PATH.name} is in sync and {PRICES_PATH.name} validates against it")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
|
@ -0,0 +1,523 @@
|
|||
{
|
||||
"annotations": {
|
||||
"list": []
|
||||
},
|
||||
"editable": true,
|
||||
"fiscalYearStartMonth": 0,
|
||||
"graphTooltip": 0,
|
||||
"links": [],
|
||||
"panels": [
|
||||
{
|
||||
"type": "stat",
|
||||
"title": "Requests",
|
||||
"gridPos": {
|
||||
"h": 4,
|
||||
"w": 6,
|
||||
"x": 0,
|
||||
"y": 0
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "short",
|
||||
"decimals": 0,
|
||||
"color": {
|
||||
"mode": "fixed",
|
||||
"fixedColor": "blue"
|
||||
}
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"reduceOptions": {
|
||||
"calcs": [
|
||||
"lastNotNull"
|
||||
],
|
||||
"fields": "",
|
||||
"values": false
|
||||
},
|
||||
"colorMode": "background",
|
||||
"graphMode": "none"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"editorMode": "code",
|
||||
"instant": true,
|
||||
"expr": "sum(increase(gen_ai_client_operation_duration_seconds_count{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range]))"
|
||||
}
|
||||
],
|
||||
"id": 1
|
||||
},
|
||||
{
|
||||
"type": "stat",
|
||||
"title": "Spend",
|
||||
"description": "LiteLLM's computed cost for the selected window, from gen_ai.usage.cost",
|
||||
"gridPos": {
|
||||
"h": 4,
|
||||
"w": 6,
|
||||
"x": 6,
|
||||
"y": 0
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "currencyUSD",
|
||||
"decimals": 4,
|
||||
"color": {
|
||||
"mode": "fixed",
|
||||
"fixedColor": "green"
|
||||
}
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"reduceOptions": {
|
||||
"calcs": [
|
||||
"lastNotNull"
|
||||
],
|
||||
"fields": "",
|
||||
"values": false
|
||||
},
|
||||
"colorMode": "background",
|
||||
"graphMode": "none"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"editorMode": "code",
|
||||
"instant": true,
|
||||
"expr": "sum(increase(gen_ai_usage_cost_USD_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range]))"
|
||||
}
|
||||
],
|
||||
"id": 2
|
||||
},
|
||||
{
|
||||
"type": "stat",
|
||||
"title": "Tokens",
|
||||
"gridPos": {
|
||||
"h": 4,
|
||||
"w": 6,
|
||||
"x": 12,
|
||||
"y": 0
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "short",
|
||||
"decimals": 0,
|
||||
"color": {
|
||||
"mode": "fixed",
|
||||
"fixedColor": "purple"
|
||||
}
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"reduceOptions": {
|
||||
"calcs": [
|
||||
"lastNotNull"
|
||||
],
|
||||
"fields": "",
|
||||
"values": false
|
||||
},
|
||||
"colorMode": "background",
|
||||
"graphMode": "none"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"editorMode": "code",
|
||||
"instant": true,
|
||||
"expr": "sum(increase(gen_ai_client_token_usage_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range]))"
|
||||
}
|
||||
],
|
||||
"id": 3
|
||||
},
|
||||
{
|
||||
"type": "stat",
|
||||
"title": "p95 request duration",
|
||||
"gridPos": {
|
||||
"h": 4,
|
||||
"w": 6,
|
||||
"x": 18,
|
||||
"y": 0
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "s",
|
||||
"decimals": 2,
|
||||
"color": {
|
||||
"mode": "fixed",
|
||||
"fixedColor": "orange"
|
||||
}
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"reduceOptions": {
|
||||
"calcs": [
|
||||
"lastNotNull"
|
||||
],
|
||||
"fields": "",
|
||||
"values": false
|
||||
},
|
||||
"colorMode": "background",
|
||||
"graphMode": "none"
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"editorMode": "code",
|
||||
"instant": true,
|
||||
"expr": "histogram_quantile(0.95, sum by (le) (rate(gen_ai_client_operation_duration_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range])))"
|
||||
}
|
||||
],
|
||||
"id": 4
|
||||
},
|
||||
{
|
||||
"type": "timeseries",
|
||||
"title": "Request rate by model",
|
||||
"gridPos": {
|
||||
"h": 8,
|
||||
"w": 12,
|
||||
"x": 0,
|
||||
"y": 4
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "reqpm",
|
||||
"custom": {
|
||||
"lineWidth": 2,
|
||||
"fillOpacity": 8,
|
||||
"showPoints": "never"
|
||||
}
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"legend": {
|
||||
"displayMode": "list",
|
||||
"placement": "bottom"
|
||||
},
|
||||
"tooltip": {
|
||||
"mode": "multi",
|
||||
"sort": "desc"
|
||||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"editorMode": "code",
|
||||
"legendFormat": "{{gen_ai_request_model}}",
|
||||
"expr": "sum by (gen_ai_request_model) (rate(gen_ai_client_operation_duration_seconds_count{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])) * 60"
|
||||
}
|
||||
],
|
||||
"id": 5
|
||||
},
|
||||
{
|
||||
"type": "timeseries",
|
||||
"title": "Spend rate by model",
|
||||
"description": "USD per hour, derived from the gen_ai.usage.cost histogram",
|
||||
"gridPos": {
|
||||
"h": 8,
|
||||
"w": 12,
|
||||
"x": 12,
|
||||
"y": 4
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "currencyUSD",
|
||||
"custom": {
|
||||
"lineWidth": 2,
|
||||
"fillOpacity": 8,
|
||||
"showPoints": "never"
|
||||
}
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"legend": {
|
||||
"displayMode": "list",
|
||||
"placement": "bottom"
|
||||
},
|
||||
"tooltip": {
|
||||
"mode": "multi",
|
||||
"sort": "desc"
|
||||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"editorMode": "code",
|
||||
"legendFormat": "{{gen_ai_request_model}}",
|
||||
"expr": "sum by (gen_ai_request_model) (rate(gen_ai_usage_cost_USD_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])) * 3600"
|
||||
}
|
||||
],
|
||||
"id": 6
|
||||
},
|
||||
{
|
||||
"type": "timeseries",
|
||||
"title": "Tokens per minute by model and type",
|
||||
"description": "gen_ai.client.token.usage split by the gen_ai.token.type attribute",
|
||||
"gridPos": {
|
||||
"h": 8,
|
||||
"w": 12,
|
||||
"x": 0,
|
||||
"y": 12
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "short",
|
||||
"custom": {
|
||||
"lineWidth": 2,
|
||||
"fillOpacity": 8,
|
||||
"showPoints": "never"
|
||||
}
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"legend": {
|
||||
"displayMode": "list",
|
||||
"placement": "bottom"
|
||||
},
|
||||
"tooltip": {
|
||||
"mode": "multi",
|
||||
"sort": "desc"
|
||||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"editorMode": "code",
|
||||
"legendFormat": "{{gen_ai_request_model}} {{gen_ai_token_type}}",
|
||||
"expr": "sum by (gen_ai_request_model, gen_ai_token_type) (rate(gen_ai_client_token_usage_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])) * 60"
|
||||
}
|
||||
],
|
||||
"id": 7
|
||||
},
|
||||
{
|
||||
"type": "timeseries",
|
||||
"title": "p95 request duration by model",
|
||||
"gridPos": {
|
||||
"h": 8,
|
||||
"w": 12,
|
||||
"x": 12,
|
||||
"y": 12
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "s",
|
||||
"custom": {
|
||||
"lineWidth": 2,
|
||||
"fillOpacity": 0,
|
||||
"showPoints": "never"
|
||||
}
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"legend": {
|
||||
"displayMode": "list",
|
||||
"placement": "bottom"
|
||||
},
|
||||
"tooltip": {
|
||||
"mode": "multi",
|
||||
"sort": "desc"
|
||||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"editorMode": "code",
|
||||
"legendFormat": "{{gen_ai_request_model}}",
|
||||
"expr": "histogram_quantile(0.95, sum by (le, gen_ai_request_model) (rate(gen_ai_client_operation_duration_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])))"
|
||||
}
|
||||
],
|
||||
"id": 8
|
||||
},
|
||||
{
|
||||
"type": "timeseries",
|
||||
"title": "p95 time to first token (streaming)",
|
||||
"description": "gen_ai.server.time_to_first_token, recorded only for streaming requests",
|
||||
"gridPos": {
|
||||
"h": 8,
|
||||
"w": 12,
|
||||
"x": 0,
|
||||
"y": 20
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "s",
|
||||
"custom": {
|
||||
"lineWidth": 2,
|
||||
"fillOpacity": 0,
|
||||
"showPoints": "never"
|
||||
}
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"legend": {
|
||||
"displayMode": "list",
|
||||
"placement": "bottom"
|
||||
},
|
||||
"tooltip": {
|
||||
"mode": "multi",
|
||||
"sort": "desc"
|
||||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"editorMode": "code",
|
||||
"legendFormat": "{{gen_ai_request_model}}",
|
||||
"expr": "histogram_quantile(0.95, sum by (le, gen_ai_request_model) (rate(gen_ai_server_time_to_first_token_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])))"
|
||||
}
|
||||
],
|
||||
"id": 9
|
||||
},
|
||||
{
|
||||
"type": "timeseries",
|
||||
"title": "p95 provider generation time",
|
||||
"description": "gen_ai.client.response.duration, upstream generation time excluding LiteLLM overhead",
|
||||
"gridPos": {
|
||||
"h": 8,
|
||||
"w": 12,
|
||||
"x": 12,
|
||||
"y": 20
|
||||
},
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"unit": "s",
|
||||
"custom": {
|
||||
"lineWidth": 2,
|
||||
"fillOpacity": 0,
|
||||
"showPoints": "never"
|
||||
}
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"options": {
|
||||
"legend": {
|
||||
"displayMode": "list",
|
||||
"placement": "bottom"
|
||||
},
|
||||
"tooltip": {
|
||||
"mode": "multi",
|
||||
"sort": "desc"
|
||||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"refId": "A",
|
||||
"editorMode": "code",
|
||||
"legendFormat": "{{gen_ai_request_model}}",
|
||||
"expr": "histogram_quantile(0.95, sum by (le, gen_ai_request_model) (rate(gen_ai_client_response_duration_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])))"
|
||||
}
|
||||
],
|
||||
"id": 10
|
||||
}
|
||||
],
|
||||
"preload": false,
|
||||
"refresh": "30s",
|
||||
"schemaVersion": 42,
|
||||
"tags": [
|
||||
"litellm",
|
||||
"genai",
|
||||
"opentelemetry"
|
||||
],
|
||||
"templating": {
|
||||
"list": [
|
||||
{
|
||||
"name": "datasource",
|
||||
"label": "Prometheus",
|
||||
"type": "datasource",
|
||||
"query": "prometheus",
|
||||
"current": {},
|
||||
"hide": 0
|
||||
},
|
||||
{
|
||||
"name": "service",
|
||||
"label": "Service",
|
||||
"type": "query",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"query": "label_values(gen_ai_client_operation_duration_seconds_count, service_name)",
|
||||
"refresh": 2,
|
||||
"includeAll": true,
|
||||
"multi": true,
|
||||
"current": {
|
||||
"text": "All",
|
||||
"value": "$__all"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "model",
|
||||
"label": "Model",
|
||||
"type": "query",
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${datasource}"
|
||||
},
|
||||
"query": "label_values(gen_ai_client_operation_duration_seconds_count{service_name=~\"$service\"}, gen_ai_request_model)",
|
||||
"refresh": 2,
|
||||
"includeAll": true,
|
||||
"multi": true,
|
||||
"current": {
|
||||
"text": "All",
|
||||
"value": "$__all"
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"time": {
|
||||
"from": "now-1h",
|
||||
"to": "now"
|
||||
},
|
||||
"timepicker": {},
|
||||
"timezone": "browser",
|
||||
"title": "LiteLLM GenAI (OpenTelemetry)",
|
||||
"uid": "litellm-genai-otel",
|
||||
"weekStart": ""
|
||||
}
|
||||
|
|
@ -0,0 +1,35 @@
|
|||
# LiteLLM GenAI dashboard (OpenTelemetry metrics)
|
||||
|
||||
Dashboard for the `gen_ai.*` metrics the OpenTelemetry v2 integration emits, as opposed to the `litellm_*` Prometheus metrics the other dashboards in this folder chart.
|
||||
|
||||
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source. Panels: request count, spend, token count, p95 duration, request rate by model, spend rate per hour by model, tokens per minute split by input and output, p95 duration by model, p95 time to first token, and p95 provider generation time. Template variables for data source, service, and model.
|
||||
|
||||
## Pre-requisites
|
||||
|
||||
Metrics are off by default. In the proxy environment:
|
||||
|
||||
```shell
|
||||
LITELLM_OTEL_V2=true
|
||||
LITELLM_OTEL_INTEGRATION_ENABLE_METRICS=true
|
||||
OTEL_EXPORTER="otlp_http"
|
||||
OTEL_ENDPOINT="<your OTLP endpoint>"
|
||||
```
|
||||
|
||||
You also need the metric attribute filter, or the panels will plot flat lines at zero. LiteLLM's default attribute set includes per-request fields, so nearly every request lands in its own time series with a single sample, and `rate()` has nothing to compute over:
|
||||
|
||||
```yaml title="config.yaml"
|
||||
callback_settings:
|
||||
otel:
|
||||
attributes:
|
||||
include_list:
|
||||
- gen_ai.operation.name
|
||||
- gen_ai.system
|
||||
- gen_ai.request.model
|
||||
- gen_ai.framework
|
||||
```
|
||||
|
||||
See [Grafana Cloud](https://docs.litellm.ai/docs/observability/grafana_cloud) for the full setup, and [OpenTelemetry v2](https://docs.litellm.ai/docs/observability/opentelemetry_v2#metrics) for the metric reference.
|
||||
|
||||
## Note on Grafana's AI Observability integration
|
||||
|
||||
Grafana Cloud ships prebuilt GenAI dashboards that query these same metric names, so they look like a drop-in alternative to this one. They are not: twenty of their twenty-two panels filter on `telemetry_sdk_name="openlit"`, a label LiteLLM does not carry and cannot be configured to add, so those panels stay empty.
|
||||
|
|
@ -2,6 +2,10 @@
|
|||
|
||||
This folder contains the `json` for creating Grafana Dashboards
|
||||
|
||||
## [LiteLLM GenAI Dashboard (OpenTelemetry)](./dashboard_genai_otel)
|
||||
|
||||
Charts the `gen_ai.*` metrics from the OpenTelemetry v2 integration: spend, tokens, request rate, and latency percentiles by model. Separate from the dashboards below, which chart the `litellm_*` Prometheus metrics.
|
||||
|
||||
## [LiteLLM v2 Dashboard](./dashboard_v2)
|
||||
|
||||
<img width="1316" alt="grafana_1" src="https://github.com/user-attachments/assets/d0df802d-0cb9-4906-a679-941c547789ab">
|
||||
|
|
|
|||
8
docker/component_entrypoint.sh
Executable file
8
docker/component_entrypoint.sh
Executable file
|
|
@ -0,0 +1,8 @@
|
|||
#!/bin/sh
|
||||
|
||||
if [ "$USE_DDTRACE" = "true" ]; then
|
||||
export DD_TRACE_OPENAI_ENABLED="False"
|
||||
exec ddtrace-run "$@"
|
||||
fi
|
||||
|
||||
exec "$@"
|
||||
|
|
@ -17,6 +17,7 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy._types import LiteLLM_ManagedObjectTable
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.router import Router
|
||||
from litellm.types.router import Deployment
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
|
||||
|
|
@ -281,6 +282,32 @@ class CheckBatchCost:
|
|||
return deployment_id
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _get_managed_file_model_name(
|
||||
cls,
|
||||
job: "LiteLLM_ManagedObjectTable",
|
||||
deployment_info: "Deployment",
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Public model group name to encode as ``target_model_names`` on unified output file ids.
|
||||
|
||||
Key model-access checks resolve a managed file id back to a model via its
|
||||
``target_model_names``, so this must be the model group the caller requested, never the
|
||||
underlying provider model (e.g. ``gpt-5.5``), which no key is allowed to call.
|
||||
"""
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
convert_b64_uid_to_unified_uid,
|
||||
get_models_from_unified_file_id,
|
||||
)
|
||||
|
||||
input_file_id = cls._get_input_file_id(job)
|
||||
target_model_names = (
|
||||
get_models_from_unified_file_id(convert_b64_uid_to_unified_uid(input_file_id)) if input_file_id else []
|
||||
)
|
||||
if target_model_names:
|
||||
return ",".join(target_model_names)
|
||||
return deployment_info.model_name or None
|
||||
|
||||
@staticmethod
|
||||
def _get_input_file_id(job: "LiteLLM_ManagedObjectTable") -> Optional[str]:
|
||||
import json
|
||||
|
|
@ -406,6 +433,10 @@ class CheckBatchCost:
|
|||
managed_files_hook = self.proxy_logging_obj.get_proxy_hook("managed_files")
|
||||
if managed_files_hook is not None:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
managed_file_model_name = self._get_managed_file_model_name(
|
||||
job=job, deployment_info=deployment_info
|
||||
)
|
||||
_minimal_auth = UserAPIKeyAuth(
|
||||
user_id=job.created_by or "default-user-id",
|
||||
team_id=getattr(job, "team_id", None),
|
||||
|
|
@ -417,7 +448,7 @@ class CheckBatchCost:
|
|||
_unified_file_id = managed_files_hook.get_unified_output_file_id(
|
||||
output_file_id=_raw_file_id,
|
||||
model_id=model_id,
|
||||
model_name=str(model_name) if model_name else deployment_info.model_name or None,
|
||||
model_name=managed_file_model_name,
|
||||
)
|
||||
await managed_files_hook.store_unified_file_id(
|
||||
file_id=_unified_file_id,
|
||||
|
|
|
|||
|
|
@ -215,7 +215,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
|
||||
if result:
|
||||
return LiteLLM_ManagedFileTable(**result)
|
||||
return LiteLLM_ManagedFileTable.model_validate(result)
|
||||
|
||||
## CHECK DB
|
||||
db_object = await self.prisma_client.db.litellm_managedfiletable.find_first(
|
||||
|
|
@ -223,7 +223,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
|
||||
if db_object:
|
||||
return LiteLLM_ManagedFileTable(**db_object.model_dump())
|
||||
return LiteLLM_ManagedFileTable.model_validate(db_object.model_dump())
|
||||
return None
|
||||
|
||||
async def delete_unified_file_id(
|
||||
|
|
@ -349,7 +349,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if isinstance(batch.file_object, str)
|
||||
else batch.file_object
|
||||
)
|
||||
batch_obj = LiteLLMBatch(**batch_data)
|
||||
batch_obj = LiteLLMBatch.model_validate(batch_data)
|
||||
batch_obj.id = batch.unified_object_id
|
||||
batch_objects.append(batch_obj)
|
||||
|
||||
|
|
@ -382,7 +382,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
"flat_model_file_ids": {"hasSome": model_object_ids},
|
||||
}
|
||||
)
|
||||
return [OpenAIFileObject(**file_object.file_object) for file_object in file_ids]
|
||||
return [OpenAIFileObject.model_validate(file_object.file_object) for file_object in file_ids]
|
||||
|
||||
async def check_managed_file_id_access(
|
||||
self, data: Dict, user_api_key_dict: UserAPIKeyAuth
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.52"
|
||||
version = "0.1.53"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.52"
|
||||
version = "0.1.53"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -65,6 +65,8 @@ RUN mkdir -p /home/nonroot && \
|
|||
HOME=/home/nonroot prisma generate --schema=./schema.prisma && \
|
||||
chown -R nonroot:nonroot /home/nonroot/.cache
|
||||
|
||||
RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh
|
||||
|
||||
# ---------- Runtime ----------
|
||||
FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
||||
|
||||
|
|
@ -95,5 +97,5 @@ USER nonroot
|
|||
|
||||
EXPOSE 4000/tcp
|
||||
|
||||
ENTRYPOINT ["sh", "-c", "exec uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
|
||||
ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
|
||||
CMD ["--host", "0.0.0.0", "--port", "4000"]
|
||||
|
|
|
|||
|
|
@ -5,5 +5,5 @@ dependencies:
|
|||
- name: redis
|
||||
repository: oci://registry-1.docker.io/bitnamicharts
|
||||
version: 18.19.1
|
||||
digest: sha256:8660fe6287f9941d08c0902f3f13731079b8cecd2a5da2fbc54e5b7aae4a6f62
|
||||
generated: "2024-03-10T02:28:52.275022+05:30"
|
||||
digest: sha256:38962e231f6596b93f82a8412bbe4cf5de696caecf5775dfbbd163383eb1c009
|
||||
generated: "2026-07-28T10:21:22.511401-07:00"
|
||||
|
|
|
|||
|
|
@ -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.0
|
||||
version: 1.1.1
|
||||
|
||||
# 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
|
||||
|
|
@ -32,10 +32,10 @@ annotations:
|
|||
|
||||
dependencies:
|
||||
- name: "postgresql"
|
||||
version: ">=13.3.0"
|
||||
version: "14.3.1"
|
||||
repository: oci://registry-1.docker.io/bitnamicharts
|
||||
condition: db.deployStandalone
|
||||
- name: redis
|
||||
version: ">=18.0.0"
|
||||
version: "18.19.1"
|
||||
repository: oci://registry-1.docker.io/bitnamicharts
|
||||
condition: redis.enabled
|
||||
|
|
|
|||
|
|
@ -130,6 +130,16 @@ Set `billingMetrics.caSecretName` only when the collector is a private or test o
|
|||
| `db.deployStandalone` | Deploy a standalone, single instance deployment of Postgres, using the Bitnami postgresql chart. This is useful for getting started but doesn't provide HA or (by default) data backups. | `true` |
|
||||
| `postgresql.*` | If `db.deployStandalone` is `true`, configuration passed to the Bitnami postgresql chart. See the [Bitnami Documentation](https://github.com/bitnami/charts/tree/main/bitnami/postgresql) for full configuration details. See [values.yaml](./values.yaml) for the default configuration. | See [values.yaml](./values.yaml) |
|
||||
| `postgresql.auth.*` | If `db.deployStandalone` is `true`, care should be taken to ensure the default `password` and `postgres-password` values are **NOT** used. | `NoTaGrEaTpAsSwOrD` |
|
||||
| `postgresql.image.*` | If `db.deployStandalone` is `true`, the image for the bundled Postgres. Pinned to a `docker.io/bitnamilegacy` build because Bitnami retired the versioned tags under `docker.io/bitnami`. | `bitnamilegacy/postgresql:16.2.0-debian-12-r6` |
|
||||
| `redis.image.*` | If `redis.enabled` is `true`, the image for the bundled Redis. Pinned to a `docker.io/bitnamilegacy` build for the same reason. | `bitnamilegacy/redis:7.2.4-debian-12-r9` |
|
||||
|
||||
#### Bundled Postgres image
|
||||
|
||||
Bitnami removed the versioned tags from `docker.io/bitnami` and republished the archived builds under `docker.io/bitnamilegacy`, so the image defaults that ship inside the `postgresql` and `redis` subcharts no longer pull. The chart pins both to the `bitnamilegacy` copies of the exact builds those subchart versions were released with, which keeps the on-disk data directory layout unchanged for existing installs.
|
||||
|
||||
Keep `postgresql.image.tag` pinned. `docker.io/bitnami/postgresql` still publishes a floating `latest`, and pointing the bundled Postgres at a different major version starts the server against a data directory it cannot read (`database files are incompatible with server`). There is no in-place way back, so crossing a major version means dumping the database with the old image and restoring it into the new one. The chart refuses to render when the tag is empty or `latest`.
|
||||
|
||||
Those images no longer receive updates. For anything beyond getting started, run Postgres outside the chart and point at it with `db.useExisting`.
|
||||
|
||||
#### Example Postgres `db.useExisting` Secret
|
||||
|
||||
|
|
|
|||
|
|
@ -146,3 +146,18 @@ Get redis service port
|
|||
{{ .Values.redis.master.service.ports.redis }}
|
||||
{{- end -}}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Reject an unpinned image tag for the bundled PostgreSQL.
|
||||
A floating tag lets a chart upgrade start a newer PostgreSQL major against the
|
||||
existing PersistentVolumeClaim. The server then refuses to start on a data
|
||||
directory written by another major version, and the only way back is a dump
|
||||
taken before the change, which by that point no longer exists.
|
||||
*/}}
|
||||
{{- define "litellm.validateBundledPostgresImageTag" -}}
|
||||
{{- $tag := .Values.postgresql.image.tag | default "" | toString -}}
|
||||
{{- $digest := .Values.postgresql.image.digest | default "" | toString -}}
|
||||
{{- if and (eq $digest "") (or (eq $tag "") (eq $tag "latest")) -}}
|
||||
{{- fail (printf "postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got %q). An unpinned tag can start a different PostgreSQL major against the existing data directory, which makes the database unreadable and is not recoverable in place. Crossing a major version requires a dump and restore." $tag) -}}
|
||||
{{- end -}}
|
||||
{{- end -}}
|
||||
|
|
|
|||
|
|
@ -35,6 +35,8 @@ spec:
|
|||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
serviceAccountName: {{ include "litellm.migrationServiceAccountName" . }}
|
||||
securityContext:
|
||||
{{- toYaml .Values.podSecurityContext | nindent 8 }}
|
||||
{{- with .Values.migrationJob.extraInitContainers }}
|
||||
initContainers:
|
||||
{{- tpl (toYaml .) $ | nindent 8 }}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
{{- if .Values.db.deployStandalone -}}
|
||||
{{- include "litellm.validateBundledPostgresImageTag" . -}}
|
||||
apiVersion: v1
|
||||
kind: Secret
|
||||
metadata:
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ metadata:
|
|||
spec:
|
||||
containers:
|
||||
- name: test
|
||||
image: bitnami/kubectl:latest
|
||||
image: docker.io/bitnamilegacy/kubectl:1.29.2-debian-12-r3
|
||||
command: ['sh', '-c']
|
||||
args:
|
||||
- |
|
||||
|
|
|
|||
94
helm/litellm-helm/tests/bundled_db_images_tests.yaml
Normal file
94
helm/litellm-helm/tests/bundled_db_images_tests.yaml
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
suite: test bundled database images
|
||||
templates:
|
||||
- charts/postgresql/templates/primary/statefulset.yaml
|
||||
- charts/redis/templates/master/application.yaml
|
||||
- charts/redis/templates/configmap.yaml
|
||||
- charts/redis/templates/health-configmap.yaml
|
||||
- charts/redis/templates/scripts-configmap.yaml
|
||||
- charts/redis/templates/secret.yaml
|
||||
- secret-dbcredentials.yaml
|
||||
- templates/tests/test-servicemonitor.yaml
|
||||
tests:
|
||||
- it: should pull the bundled postgres from a repository that still publishes the pinned tag
|
||||
template: charts/postgresql/templates/primary/statefulset.yaml
|
||||
set:
|
||||
db.deployStandalone: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: docker.io/bitnamilegacy/postgresql:16.2.0-debian-12-r6
|
||||
|
||||
- it: should pull the bundled postgres metrics exporter from the same repository
|
||||
template: charts/postgresql/templates/primary/statefulset.yaml
|
||||
set:
|
||||
db.deployStandalone: true
|
||||
postgresql.metrics.enabled: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[1].image
|
||||
value: docker.io/bitnamilegacy/postgres-exporter:0.15.0-debian-12-r14
|
||||
|
||||
- it: should run the bundled postgres init container from the same repository
|
||||
template: charts/postgresql/templates/primary/statefulset.yaml
|
||||
set:
|
||||
db.deployStandalone: true
|
||||
postgresql.volumePermissions.enabled: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.initContainers[0].image
|
||||
value: docker.io/bitnamilegacy/os-shell:12-debian-12-r16
|
||||
|
||||
- it: should pull the bundled redis from a repository that still publishes the pinned tag
|
||||
template: charts/redis/templates/master/application.yaml
|
||||
set:
|
||||
redis.enabled: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: docker.io/bitnamilegacy/redis:7.2.4-debian-12-r9
|
||||
|
||||
- it: should reject a floating postgres tag that could cross a major version on an existing volume
|
||||
template: secret-dbcredentials.yaml
|
||||
set:
|
||||
db.deployStandalone: true
|
||||
postgresql.image.tag: latest
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: 'postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got "latest"). An unpinned tag can start a different PostgreSQL major against the existing data directory, which makes the database unreadable and is not recoverable in place. Crossing a major version requires a dump and restore.'
|
||||
|
||||
- it: should reject an empty postgres tag
|
||||
template: secret-dbcredentials.yaml
|
||||
set:
|
||||
db.deployStandalone: true
|
||||
postgresql.image.tag: ""
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: 'postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got ""). An unpinned tag can start a different PostgreSQL major against the existing data directory, which makes the database unreadable and is not recoverable in place. Crossing a major version requires a dump and restore.'
|
||||
|
||||
- it: should accept an empty postgres tag when the image is pinned by digest
|
||||
template: secret-dbcredentials.yaml
|
||||
set:
|
||||
db.deployStandalone: true
|
||||
postgresql.image.tag: ""
|
||||
postgresql.image.digest: sha256:0d0e2f1a5b3c4d6e7f8091a2b3c4d5e6f708192a3b4c5d6e7f8091a2b3c4d5e6
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 1
|
||||
|
||||
- it: should run the servicemonitor test pod from a pinned image
|
||||
template: templates/tests/test-servicemonitor.yaml
|
||||
set:
|
||||
serviceMonitor.enabled: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.containers[0].image
|
||||
value: docker.io/bitnamilegacy/kubectl:1.29.2-debian-12-r3
|
||||
|
||||
- it: should not constrain the postgres tag when the bundled database is not deployed
|
||||
template: secret-dbcredentials.yaml
|
||||
set:
|
||||
db.deployStandalone: false
|
||||
postgresql.image.tag: latest
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 0
|
||||
|
|
@ -254,3 +254,39 @@ tests:
|
|||
content:
|
||||
name: sidecar-tpl
|
||||
image: "ghcr.io/berriai/litellm-database:test"
|
||||
- it: should render the pod-level securityContext from podSecurityContext
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
migrationJob:
|
||||
enabled: true
|
||||
podSecurityContext:
|
||||
fsGroup: 10000
|
||||
runAsUser: 10000
|
||||
runAsNonRoot: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.securityContext
|
||||
value:
|
||||
fsGroup: 10000
|
||||
runAsUser: 10000
|
||||
runAsNonRoot: true
|
||||
- it: should keep the pod-level and container-level securityContext separate
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
migrationJob:
|
||||
enabled: true
|
||||
podSecurityContext:
|
||||
fsGroup: 10000
|
||||
securityContext:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.securityContext
|
||||
value:
|
||||
fsGroup: 10000
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].securityContext
|
||||
value:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
|
|
|
|||
|
|
@ -328,8 +328,32 @@ lifecycle: {}
|
|||
|
||||
# Settings for Bitnami postgresql chart (if db.deployStandalone is true, ignored
|
||||
# otherwise)
|
||||
#
|
||||
# Bitnami retired the versioned tags under docker.io/bitnami and republished the
|
||||
# archived builds under docker.io/bitnamilegacy, so the subchart's own image
|
||||
# defaults no longer resolve. The repository below points at the same build the
|
||||
# subchart was released with, which keeps the on-disk data directory layout
|
||||
# identical for existing installs.
|
||||
#
|
||||
# Keep the tag pinned. docker.io/bitnami still publishes a floating `latest`,
|
||||
# and starting a newer PostgreSQL major against an existing data directory
|
||||
# leaves the server refusing to boot ("database files are incompatible with
|
||||
# server") with no way back other than a dump taken beforehand. Crossing a major
|
||||
# version is a dump-and-restore, not an image bump. The chart refuses to render
|
||||
# an unpinned tag for this reason
|
||||
postgresql:
|
||||
architecture: standalone
|
||||
image:
|
||||
repository: bitnamilegacy/postgresql
|
||||
tag: 16.2.0-debian-12-r6
|
||||
volumePermissions:
|
||||
image:
|
||||
repository: bitnamilegacy/os-shell
|
||||
tag: 12-debian-12-r16
|
||||
metrics:
|
||||
image:
|
||||
repository: bitnamilegacy/postgres-exporter
|
||||
tag: 0.15.0-debian-12-r14
|
||||
auth:
|
||||
username: litellm
|
||||
database: litellm
|
||||
|
|
@ -359,9 +383,36 @@ postgresql:
|
|||
# When `redis.sentinel.enabled` is set, the coordination block is rendered with
|
||||
# `sentinel_nodes` and `service_name` (from `redis.sentinel.masterSet`) instead
|
||||
# of host/port, because a plain Redis client cannot talk to the sentinel port
|
||||
#
|
||||
# The image repositories carry the same bitnamilegacy repoint as postgresql
|
||||
# above; the versioned tags the subchart ships with are gone from
|
||||
# docker.io/bitnami
|
||||
redis:
|
||||
enabled: false
|
||||
architecture: standalone
|
||||
image:
|
||||
repository: bitnamilegacy/redis
|
||||
tag: 7.2.4-debian-12-r9
|
||||
sentinel:
|
||||
image:
|
||||
repository: bitnamilegacy/redis-sentinel
|
||||
tag: 7.2.4-debian-12-r7
|
||||
metrics:
|
||||
image:
|
||||
repository: bitnamilegacy/redis-exporter
|
||||
tag: 1.58.0-debian-12-r4
|
||||
volumePermissions:
|
||||
image:
|
||||
repository: bitnamilegacy/os-shell
|
||||
tag: 12-debian-12-r16
|
||||
sysctl:
|
||||
image:
|
||||
repository: bitnamilegacy/os-shell
|
||||
tag: 12-debian-12-r16
|
||||
kubectl:
|
||||
image:
|
||||
repository: bitnamilegacy/kubectl
|
||||
tag: 1.29.2-debian-12-r3
|
||||
coordination:
|
||||
# Set to false to keep the bundled Redis for response caching only and leave
|
||||
# `general_settings.coordination_redis` out of the rendered config. A
|
||||
|
|
|
|||
|
|
@ -138,6 +138,59 @@ is false the chart uses the provided name, or the namespace `default` SA.
|
|||
{{- end -}}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
ServiceAccount name for the migrations Job.
|
||||
|
||||
The Job is a pre-install / pre-upgrade hook, so it is created before the
|
||||
chart's ordinary resources. A ServiceAccount the chart creates is one of
|
||||
those ordinary resources, which makes borrowing the backend name a cycle:
|
||||
the hook pod is rejected because the account does not exist yet. So when
|
||||
`serviceAccounts.backend.create` is true the Job falls back to the namespace
|
||||
`default` account unless the operator names one that already exists. With
|
||||
`create` false the backend name is either an operator-supplied existing
|
||||
account or `default`, both of which are safe for the hook, so the Job keeps
|
||||
sharing it.
|
||||
|
||||
`migrationJob.serviceAccountName` always wins when set, which is how a Job
|
||||
that needs credentials of its own (IRSA / Workload Identity for IAM database
|
||||
auth) gets them.
|
||||
*/}}
|
||||
{{- define "litellm.migrations.serviceAccountName" -}}
|
||||
{{- if .Values.migrationJob.serviceAccountName -}}
|
||||
{{ .Values.migrationJob.serviceAccountName }}
|
||||
{{- else if .Values.serviceAccounts.backend.create -}}
|
||||
default
|
||||
{{- else -}}
|
||||
{{ include "litellm.backend.serviceAccountName" . }}
|
||||
{{- end -}}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Extra pod labels for a component's Deployment, validated against its selector.
|
||||
|
||||
Invoke with a dict:
|
||||
(dict "podLabels" .Values.gateway.podLabels "componentName" "gateway")
|
||||
|
||||
The three selector keys are also emitted on the pod template, so a podLabels
|
||||
entry reusing one renders a duplicate YAML key whose later value wins. That
|
||||
leaves the pod template no longer matching the (immutable) selector and the
|
||||
apiserver rejects the Deployment. Fail at template time naming the key
|
||||
instead, so the operator gets the reason here rather than an opaque
|
||||
`selector does not match template labels` from the apiserver.
|
||||
|
||||
The migrations Job takes podLabels unvalidated: a Job's selector is generated
|
||||
by the controller rather than declared, so nothing there can collide.
|
||||
*/}}
|
||||
{{- define "litellm.podLabels" -}}
|
||||
{{- $componentName := .componentName -}}
|
||||
{{- range $key, $value := .podLabels }}
|
||||
{{- if has $key (list "app.kubernetes.io/name" "app.kubernetes.io/instance" "app.kubernetes.io/component") }}
|
||||
{{- fail (printf "%s.podLabels cannot set %s: it is part of the Deployment's immutable selector" $componentName $key) }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- toYaml .podLabels }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Master-key + database + redis env block — shared by gateway, backend, and the
|
||||
migrations Job.
|
||||
|
|
|
|||
|
|
@ -23,9 +23,16 @@ spec:
|
|||
{{- end }}
|
||||
labels:
|
||||
{{- include "litellm.backend.selectorLabels" . | nindent 8 }}
|
||||
{{- with .Values.backend.podLabels }}
|
||||
{{- include "litellm.podLabels" (dict "podLabels" . "componentName" "backend") | nindent 8 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
serviceAccountName: {{ include "litellm.backend.serviceAccountName" . }}
|
||||
automountServiceAccountToken: {{ .Values.serviceAccounts.backend.automount }}
|
||||
{{- with .Values.backend.podSecurityContext }}
|
||||
securityContext:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.imagePullSecrets }}
|
||||
imagePullSecrets:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
|
|
@ -34,6 +41,10 @@ spec:
|
|||
- name: backend
|
||||
image: "{{ .Values.backend.image.repository }}:{{ .Values.backend.image.tag | default .Chart.AppVersion }}"
|
||||
imagePullPolicy: {{ .Values.backend.image.pullPolicy }}
|
||||
{{- with .Values.backend.securityContext }}
|
||||
securityContext:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
ports:
|
||||
- name: http
|
||||
containerPort: 4001
|
||||
|
|
@ -70,8 +81,15 @@ spec:
|
|||
readinessProbe:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.backend.lifecycle }}
|
||||
lifecycle:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
resources:
|
||||
{{- toYaml .Values.backend.resources | nindent 12 }}
|
||||
{{- with .Values.backend.extraContainers }}
|
||||
{{- tpl (toYaml .) $ | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }}
|
||||
volumes:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
|
|
@ -102,4 +120,8 @@ spec:
|
|||
topologySpreadConstraints:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- $gracePeriod := .Values.backend.terminationGracePeriodSeconds }}
|
||||
{{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }}
|
||||
terminationGracePeriodSeconds: {{ $gracePeriod }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
|
|
|||
|
|
@ -21,9 +21,16 @@ spec:
|
|||
{{- end }}
|
||||
labels:
|
||||
{{- include "litellm.gateway.selectorLabels" . | nindent 8 }}
|
||||
{{- with .Values.gateway.podLabels }}
|
||||
{{- include "litellm.podLabels" (dict "podLabels" . "componentName" "gateway") | nindent 8 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
serviceAccountName: {{ include "litellm.gateway.serviceAccountName" . }}
|
||||
automountServiceAccountToken: {{ .Values.serviceAccounts.gateway.automount }}
|
||||
{{- with .Values.gateway.podSecurityContext }}
|
||||
securityContext:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.imagePullSecrets }}
|
||||
imagePullSecrets:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
|
|
@ -32,6 +39,10 @@ spec:
|
|||
- name: gateway
|
||||
image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
|
||||
imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
|
||||
{{- with .Values.gateway.securityContext }}
|
||||
securityContext:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
ports:
|
||||
- name: http
|
||||
containerPort: 4000
|
||||
|
|
@ -72,8 +83,15 @@ spec:
|
|||
readinessProbe:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.gateway.lifecycle }}
|
||||
lifecycle:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
resources:
|
||||
{{- toYaml .Values.gateway.resources | nindent 12 }}
|
||||
{{- with .Values.gateway.extraContainers }}
|
||||
{{- tpl (toYaml .) $ | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
|
||||
volumes:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
|
|
@ -104,4 +122,8 @@ spec:
|
|||
topologySpreadConstraints:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- $gracePeriod := .Values.gateway.terminationGracePeriodSeconds }}
|
||||
{{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }}
|
||||
terminationGracePeriodSeconds: {{ $gracePeriod }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
|
|
|||
|
|
@ -23,12 +23,21 @@ spec:
|
|||
ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }}
|
||||
template:
|
||||
metadata:
|
||||
{{- /* The Job's selector is generated by the controller rather than
|
||||
declared, so podLabels may override a chart label here. Merge
|
||||
instead of appending so an override replaces the key rather than
|
||||
rendering it twice. */}}
|
||||
{{- $chartLabels := merge (dict "app.kubernetes.io/component" "migrations") (fromYaml (include "litellm.commonLabels" .)) }}
|
||||
labels:
|
||||
{{- include "litellm.commonLabels" . | nindent 8 }}
|
||||
app.kubernetes.io/component: migrations
|
||||
{{- toYaml (merge (deepCopy .Values.migrationJob.podLabels) $chartLabels) | nindent 8 }}
|
||||
spec:
|
||||
restartPolicy: Never
|
||||
serviceAccountName: {{ include "litellm.backend.serviceAccountName" . }}
|
||||
serviceAccountName: {{ include "litellm.migrations.serviceAccountName" . }}
|
||||
automountServiceAccountToken: {{ .Values.migrationJob.automountServiceAccountToken }}
|
||||
{{- with .Values.migrationJob.podSecurityContext }}
|
||||
securityContext:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.imagePullSecrets }}
|
||||
imagePullSecrets:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
|
|
@ -37,10 +46,22 @@ spec:
|
|||
- name: prisma-migrations
|
||||
image: "{{ .Values.migrationJob.image.repository }}:{{ .Values.migrationJob.image.tag | default .Chart.AppVersion }}"
|
||||
imagePullPolicy: {{ .Values.migrationJob.image.pullPolicy }}
|
||||
{{- with .Values.migrationJob.securityContext }}
|
||||
securityContext:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
env:
|
||||
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.migrationJob) | nindent 12 }}
|
||||
{{- with .Values.migrationJob.volumeMounts }}
|
||||
volumeMounts:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.migrationJob.resources }}
|
||||
resources:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.migrationJob.volumes }}
|
||||
volumes:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
|
|
|||
|
|
@ -18,9 +18,16 @@ spec:
|
|||
{{- end }}
|
||||
labels:
|
||||
{{- include "litellm.ui.selectorLabels" . | nindent 8 }}
|
||||
{{- with .Values.ui.podLabels }}
|
||||
{{- include "litellm.podLabels" (dict "podLabels" . "componentName" "ui") | nindent 8 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
serviceAccountName: {{ include "litellm.ui.serviceAccountName" . }}
|
||||
automountServiceAccountToken: {{ .Values.serviceAccounts.ui.automount }}
|
||||
{{- with .Values.ui.podSecurityContext }}
|
||||
securityContext:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.imagePullSecrets }}
|
||||
imagePullSecrets:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
|
|
@ -29,6 +36,10 @@ spec:
|
|||
- name: ui
|
||||
image: "{{ .Values.ui.image.repository }}:{{ .Values.ui.image.tag | default .Chart.AppVersion }}"
|
||||
imagePullPolicy: {{ .Values.ui.image.pullPolicy }}
|
||||
{{- with .Values.ui.securityContext }}
|
||||
securityContext:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
ports:
|
||||
- name: http
|
||||
containerPort: 3000
|
||||
|
|
@ -58,8 +69,15 @@ spec:
|
|||
readinessProbe:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.ui.lifecycle }}
|
||||
lifecycle:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
resources:
|
||||
{{- toYaml .Values.ui.resources | nindent 12 }}
|
||||
{{- with .Values.ui.extraContainers }}
|
||||
{{- tpl (toYaml .) $ | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.ui.volumes }}
|
||||
volumes:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
|
|
@ -80,4 +98,8 @@ spec:
|
|||
topologySpreadConstraints:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- $gracePeriod := .Values.ui.terminationGracePeriodSeconds }}
|
||||
{{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }}
|
||||
terminationGracePeriodSeconds: {{ $gracePeriod }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
|
|
|||
169
helm/litellm/tests/migration_job_tests.yaml
Normal file
169
helm/litellm/tests/migration_job_tests.yaml
Normal file
|
|
@ -0,0 +1,169 @@
|
|||
suite: test migrations Job ServiceAccount resolution and pod hardening
|
||||
templates:
|
||||
- migrations-job.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: borrows the namespace default account when no ServiceAccount is configured
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.serviceAccountName
|
||||
value: default
|
||||
|
||||
- it: falls back to the namespace default account when the chart creates the backend ServiceAccount
|
||||
set:
|
||||
serviceAccounts.backend.create: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.serviceAccountName
|
||||
value: default
|
||||
- notEqual:
|
||||
path: spec.template.spec.serviceAccountName
|
||||
value: RELEASE-NAME-litellm-backend
|
||||
|
||||
- it: keeps sharing an existing backend ServiceAccount the chart does not create
|
||||
set:
|
||||
serviceAccounts.backend.create: false
|
||||
serviceAccounts.backend.name: existing-backend-sa
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.serviceAccountName
|
||||
value: existing-backend-sa
|
||||
|
||||
- it: prefers an explicit migration ServiceAccount over the created backend one
|
||||
set:
|
||||
serviceAccounts.backend.create: true
|
||||
migrationJob.serviceAccountName: migrations-sa
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.serviceAccountName
|
||||
value: migrations-sa
|
||||
|
||||
- it: prefers an explicit migration ServiceAccount over an existing backend one
|
||||
set:
|
||||
serviceAccounts.backend.create: false
|
||||
serviceAccounts.backend.name: existing-backend-sa
|
||||
migrationJob.serviceAccountName: migrations-sa
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.serviceAccountName
|
||||
value: migrations-sa
|
||||
|
||||
- it: mounts no ServiceAccount token by default
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.automountServiceAccountToken
|
||||
value: false
|
||||
|
||||
- it: mounts a ServiceAccount token when the operator asks for one
|
||||
set:
|
||||
migrationJob.automountServiceAccountToken: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.automountServiceAccountToken
|
||||
value: true
|
||||
|
||||
- it: keeps the token off the Job when the backend disables automounting
|
||||
set:
|
||||
serviceAccounts.backend.create: true
|
||||
serviceAccounts.backend.automount: false
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.serviceAccountName
|
||||
value: default
|
||||
- equal:
|
||||
path: spec.template.spec.automountServiceAccountToken
|
||||
value: false
|
||||
|
||||
- it: renders no hardening fields by default
|
||||
asserts:
|
||||
- isNull:
|
||||
path: spec.template.spec.securityContext
|
||||
- isNull:
|
||||
path: spec.template.spec.containers[0].securityContext
|
||||
- isNull:
|
||||
path: spec.template.spec.volumes
|
||||
- isNull:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
- equal:
|
||||
path: spec.template.metadata.labels
|
||||
value:
|
||||
app.kubernetes.io/name: litellm
|
||||
app.kubernetes.io/instance: RELEASE-NAME
|
||||
app.kubernetes.io/managed-by: Helm
|
||||
helm.sh/chart: litellm-0.1.0
|
||||
app.kubernetes.io/component: migrations
|
||||
|
||||
- it: renders pod-level and container-level securityContext in their own scopes
|
||||
set:
|
||||
migrationJob.podSecurityContext:
|
||||
runAsNonRoot: true
|
||||
runAsUser: 65532
|
||||
seccompProfile:
|
||||
type: RuntimeDefault
|
||||
migrationJob.securityContext:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
capabilities:
|
||||
drop:
|
||||
- ALL
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.securityContext
|
||||
value:
|
||||
runAsNonRoot: true
|
||||
runAsUser: 65532
|
||||
seccompProfile:
|
||||
type: RuntimeDefault
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].securityContext
|
||||
value:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
capabilities:
|
||||
drop:
|
||||
- ALL
|
||||
|
||||
- it: renders volumes on the pod and volumeMounts on the migration container
|
||||
set:
|
||||
migrationJob.volumes:
|
||||
- name: tmp
|
||||
emptyDir:
|
||||
sizeLimit: 64Mi
|
||||
migrationJob.volumeMounts:
|
||||
- name: tmp
|
||||
mountPath: /tmp
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.volumes
|
||||
value:
|
||||
- name: tmp
|
||||
emptyDir:
|
||||
sizeLimit: 64Mi
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
value:
|
||||
- name: tmp
|
||||
mountPath: /tmp
|
||||
|
||||
- it: merges podLabels with the chart labels on the Job pod
|
||||
set:
|
||||
migrationJob.podLabels:
|
||||
egress-policy: restricted
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.metadata.labels['egress-policy']
|
||||
value: restricted
|
||||
- equal:
|
||||
path: spec.template.metadata.labels['app.kubernetes.io/component']
|
||||
value: migrations
|
||||
|
||||
- it: accepts a podLabel that reuses a chart label, since the Job selector is controller-generated
|
||||
set:
|
||||
migrationJob.podLabels:
|
||||
app.kubernetes.io/component: batch-migrations
|
||||
asserts:
|
||||
- notFailedTemplate: {}
|
||||
- equal:
|
||||
path: spec.template.metadata.labels['app.kubernetes.io/component']
|
||||
value: batch-migrations
|
||||
298
helm/litellm/tests/pod_hardening_tests.yaml
Normal file
298
helm/litellm/tests/pod_hardening_tests.yaml
Normal file
|
|
@ -0,0 +1,298 @@
|
|||
suite: test pod hardening knobs on the component deployments
|
||||
templates:
|
||||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- backend/deployment.yaml
|
||||
- ui/deployment.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: gateway renders no hardening fields by default
|
||||
template: gateway/deployment.yaml
|
||||
asserts:
|
||||
- isNull:
|
||||
path: spec.template.spec.securityContext
|
||||
- isNull:
|
||||
path: spec.template.spec.containers[0].securityContext
|
||||
- isNull:
|
||||
path: spec.template.spec.containers[0].lifecycle
|
||||
- isNull:
|
||||
path: spec.template.spec.terminationGracePeriodSeconds
|
||||
- lengthEqual:
|
||||
path: spec.template.spec.containers
|
||||
count: 1
|
||||
- equal:
|
||||
path: spec.template.metadata.labels
|
||||
value:
|
||||
app.kubernetes.io/name: litellm
|
||||
app.kubernetes.io/instance: RELEASE-NAME
|
||||
app.kubernetes.io/component: gateway
|
||||
|
||||
- it: gateway renders pod-level and container-level securityContext in their own scopes
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.podSecurityContext:
|
||||
runAsNonRoot: true
|
||||
runAsUser: 65532
|
||||
fsGroup: 65532
|
||||
seccompProfile:
|
||||
type: RuntimeDefault
|
||||
gateway.securityContext:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
capabilities:
|
||||
drop:
|
||||
- ALL
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.securityContext
|
||||
value:
|
||||
runAsNonRoot: true
|
||||
runAsUser: 65532
|
||||
fsGroup: 65532
|
||||
seccompProfile:
|
||||
type: RuntimeDefault
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].securityContext
|
||||
value:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
capabilities:
|
||||
drop:
|
||||
- ALL
|
||||
|
||||
- it: gateway merges podLabels with the selector labels
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.podLabels:
|
||||
egress-policy: restricted
|
||||
team: platform
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.metadata.labels
|
||||
value:
|
||||
app.kubernetes.io/name: litellm
|
||||
app.kubernetes.io/instance: RELEASE-NAME
|
||||
app.kubernetes.io/component: gateway
|
||||
egress-policy: restricted
|
||||
team: platform
|
||||
- equal:
|
||||
path: spec.selector.matchLabels
|
||||
value:
|
||||
app.kubernetes.io/name: litellm
|
||||
app.kubernetes.io/instance: RELEASE-NAME
|
||||
app.kubernetes.io/component: gateway
|
||||
|
||||
- it: gateway rejects a podLabel that collides with the selector
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.podLabels:
|
||||
app.kubernetes.io/component: not-gateway
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "gateway.podLabels cannot set app.kubernetes.io/component: it is part of the Deployment's immutable selector"
|
||||
|
||||
- it: backend rejects a podLabel that collides with the selector
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
backend.podLabels:
|
||||
app.kubernetes.io/name: not-litellm
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "backend.podLabels cannot set app.kubernetes.io/name: it is part of the Deployment's immutable selector"
|
||||
|
||||
- it: ui rejects a podLabel that collides with the selector
|
||||
template: ui/deployment.yaml
|
||||
set:
|
||||
ui.podLabels:
|
||||
app.kubernetes.io/instance: not-the-release
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "ui.podLabels cannot set app.kubernetes.io/instance: it is part of the Deployment's immutable selector"
|
||||
|
||||
- it: gateway renders lifecycle hooks on the container
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.lifecycle:
|
||||
preStop:
|
||||
httpGet:
|
||||
path: /health/drain
|
||||
port: 4000
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].lifecycle
|
||||
value:
|
||||
preStop:
|
||||
httpGet:
|
||||
path: /health/drain
|
||||
port: 4000
|
||||
|
||||
- it: gateway renders terminationGracePeriodSeconds on the pod spec
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.terminationGracePeriodSeconds: 90
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.terminationGracePeriodSeconds
|
||||
value: 90
|
||||
|
||||
- it: gateway honors an explicit terminationGracePeriodSeconds of zero
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.terminationGracePeriodSeconds: 0
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.terminationGracePeriodSeconds
|
||||
value: 0
|
||||
|
||||
- it: gateway appends extraContainers after the gateway container
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.extraContainers:
|
||||
- name: auth-sidecar
|
||||
image: registry.example.com/auth-proxy:1.2.3
|
||||
args:
|
||||
- --upstream
|
||||
- http://127.0.0.1:4000
|
||||
asserts:
|
||||
- lengthEqual:
|
||||
path: spec.template.spec.containers
|
||||
count: 2
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].name
|
||||
value: gateway
|
||||
- equal:
|
||||
path: spec.template.spec.containers[1]
|
||||
value:
|
||||
name: auth-sidecar
|
||||
image: registry.example.com/auth-proxy:1.2.3
|
||||
args:
|
||||
- --upstream
|
||||
- http://127.0.0.1:4000
|
||||
|
||||
- it: gateway templates chart context inside extraContainers
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.extraContainers:
|
||||
- name: auth-sidecar
|
||||
image: registry.example.com/auth-proxy:1.2.3
|
||||
env:
|
||||
- name: RELEASE
|
||||
value: "{{ .Release.Name }}"
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[1].env[0].value
|
||||
value: RELEASE-NAME
|
||||
|
||||
- it: backend renders every hardening knob in the right scope
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
backend.podLabels:
|
||||
egress-policy: restricted
|
||||
backend.podSecurityContext:
|
||||
runAsNonRoot: true
|
||||
backend.securityContext:
|
||||
readOnlyRootFilesystem: true
|
||||
backend.lifecycle:
|
||||
preStop:
|
||||
exec:
|
||||
command:
|
||||
- sleep
|
||||
- "5"
|
||||
backend.terminationGracePeriodSeconds: 60
|
||||
backend.extraContainers:
|
||||
- name: auth-sidecar
|
||||
image: registry.example.com/auth-proxy:1.2.3
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.metadata.labels['egress-policy']
|
||||
value: restricted
|
||||
- equal:
|
||||
path: spec.template.spec.securityContext
|
||||
value:
|
||||
runAsNonRoot: true
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].securityContext
|
||||
value:
|
||||
readOnlyRootFilesystem: true
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].lifecycle
|
||||
value:
|
||||
preStop:
|
||||
exec:
|
||||
command:
|
||||
- sleep
|
||||
- "5"
|
||||
- equal:
|
||||
path: spec.template.spec.terminationGracePeriodSeconds
|
||||
value: 60
|
||||
- equal:
|
||||
path: spec.template.spec.containers[1].name
|
||||
value: auth-sidecar
|
||||
|
||||
- it: ui renders every hardening knob in the right scope
|
||||
template: ui/deployment.yaml
|
||||
set:
|
||||
ui.podLabels:
|
||||
egress-policy: restricted
|
||||
ui.podSecurityContext:
|
||||
runAsNonRoot: true
|
||||
fsGroup: 101
|
||||
ui.securityContext:
|
||||
readOnlyRootFilesystem: true
|
||||
ui.lifecycle:
|
||||
preStop:
|
||||
exec:
|
||||
command:
|
||||
- /bin/sh
|
||||
- -c
|
||||
- nginx -s quit
|
||||
ui.terminationGracePeriodSeconds: 30
|
||||
ui.extraContainers:
|
||||
- name: auth-sidecar
|
||||
image: registry.example.com/auth-proxy:1.2.3
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.metadata.labels['egress-policy']
|
||||
value: restricted
|
||||
- equal:
|
||||
path: spec.template.spec.securityContext
|
||||
value:
|
||||
runAsNonRoot: true
|
||||
fsGroup: 101
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].securityContext
|
||||
value:
|
||||
readOnlyRootFilesystem: true
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].lifecycle
|
||||
value:
|
||||
preStop:
|
||||
exec:
|
||||
command:
|
||||
- /bin/sh
|
||||
- -c
|
||||
- nginx -s quit
|
||||
- equal:
|
||||
path: spec.template.spec.terminationGracePeriodSeconds
|
||||
value: 30
|
||||
- equal:
|
||||
path: spec.template.spec.containers[1].name
|
||||
value: auth-sidecar
|
||||
|
||||
- it: backend and ui render no hardening fields by default
|
||||
templates:
|
||||
- backend/deployment.yaml
|
||||
- ui/deployment.yaml
|
||||
asserts:
|
||||
- isNull:
|
||||
path: spec.template.spec.securityContext
|
||||
- isNull:
|
||||
path: spec.template.spec.containers[0].securityContext
|
||||
- isNull:
|
||||
path: spec.template.spec.containers[0].lifecycle
|
||||
- isNull:
|
||||
path: spec.template.spec.terminationGracePeriodSeconds
|
||||
- lengthEqual:
|
||||
path: spec.template.spec.containers
|
||||
count: 1
|
||||
106
helm/litellm/tests/probe_tests.yaml
Normal file
106
helm/litellm/tests/probe_tests.yaml
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
suite: test liveness and readiness probe timeouts
|
||||
templates:
|
||||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- backend/deployment.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: gateway probes set an explicit timeout that outlasts a saturated event loop
|
||||
template: gateway/deployment.yaml
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].livenessProbe
|
||||
value:
|
||||
httpGet:
|
||||
path: /health/liveliness
|
||||
port: http
|
||||
initialDelaySeconds: 10
|
||||
periodSeconds: 15
|
||||
timeoutSeconds: 10
|
||||
failureThreshold: 6
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].readinessProbe
|
||||
value:
|
||||
httpGet:
|
||||
path: /health/readiness
|
||||
port: http
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
timeoutSeconds: 10
|
||||
|
||||
- it: backend probes set an explicit timeout that outlasts a saturated event loop
|
||||
template: backend/deployment.yaml
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].livenessProbe
|
||||
value:
|
||||
httpGet:
|
||||
path: /health/liveliness
|
||||
port: http
|
||||
initialDelaySeconds: 10
|
||||
periodSeconds: 15
|
||||
timeoutSeconds: 10
|
||||
failureThreshold: 6
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].readinessProbe
|
||||
value:
|
||||
httpGet:
|
||||
path: /health/readiness
|
||||
port: http
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
timeoutSeconds: 10
|
||||
|
||||
- it: no single-event-loop component is left on the kubernetes default 1s probe timeout
|
||||
templates:
|
||||
- gateway/deployment.yaml
|
||||
- backend/deployment.yaml
|
||||
asserts:
|
||||
- isNotNullOrEmpty:
|
||||
path: spec.template.spec.containers[0].livenessProbe.timeoutSeconds
|
||||
- isNotNullOrEmpty:
|
||||
path: spec.template.spec.containers[0].readinessProbe.timeoutSeconds
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].livenessProbe.timeoutSeconds
|
||||
value: 10
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].readinessProbe.timeoutSeconds
|
||||
value: 10
|
||||
|
||||
- it: gateway liveness tolerates a longer outage than readiness before acting
|
||||
template: gateway/deployment.yaml
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].livenessProbe.failureThreshold
|
||||
value: 6
|
||||
- notExists:
|
||||
path: spec.template.spec.containers[0].readinessProbe.failureThreshold
|
||||
|
||||
- it: probe timeouts and thresholds stay overridable per component
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.readinessProbe.timeoutSeconds: 3
|
||||
gateway.readinessProbe.periodSeconds: 20
|
||||
gateway.livenessProbe.timeoutSeconds: 4
|
||||
gateway.livenessProbe.failureThreshold: 3
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].readinessProbe
|
||||
value:
|
||||
httpGet:
|
||||
path: /health/readiness
|
||||
port: http
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 20
|
||||
timeoutSeconds: 3
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].livenessProbe
|
||||
value:
|
||||
httpGet:
|
||||
path: /health/liveliness
|
||||
port: http
|
||||
initialDelaySeconds: 10
|
||||
periodSeconds: 15
|
||||
timeoutSeconds: 4
|
||||
failureThreshold: 3
|
||||
|
|
@ -57,6 +57,42 @@ migrationJob:
|
|||
backoffLimit: 4
|
||||
ttlSecondsAfterFinished: 120
|
||||
resources: {}
|
||||
# ServiceAccount for the Job pod only.
|
||||
#
|
||||
# The Job is a pre-install / pre-upgrade hook, so it runs before the chart's
|
||||
# ordinary resources exist. With `serviceAccounts.backend.create: true` the
|
||||
# backend ServiceAccount is one of those ordinary resources, so a Job that
|
||||
# borrowed its name would reference an account that does not exist yet and
|
||||
# the first install would fail with a forbidden pod creation. The name set
|
||||
# here always wins; when it is empty the Job falls back to `default` if the
|
||||
# chart creates the backend ServiceAccount, and to the backend
|
||||
# ServiceAccount name otherwise (that name is either an existing account you
|
||||
# supplied or `default`).
|
||||
#
|
||||
# Point this at a pre-existing ServiceAccount when the Job needs credentials
|
||||
# of its own, e.g. the IRSA / Workload Identity annotations that
|
||||
# `database.writer.useIAMAuth` relies on. That is also the upgrade path to
|
||||
# watch: a release already running with `serviceAccounts.backend.create:
|
||||
# true` used to hand the Job the created backend account on every upgrade,
|
||||
# and now hands it `default` unless you name an account here.
|
||||
serviceAccountName: ""
|
||||
# The Job runs `prisma migrate deploy` against Postgres and never calls the
|
||||
# K8s API, so it defaults to no projected ServiceAccount token, the same
|
||||
# reasoning the ui SA above uses. Flip to true if your Job genuinely needs
|
||||
# one; IAM database auth does not, since EKS Pod Identity injects its own
|
||||
# projected token volume and GKE Workload Identity goes through the
|
||||
# metadata server, neither of which is the default token mount.
|
||||
automountServiceAccountToken: false
|
||||
# Standard k8s pod-level and container-level securityContext for the Job
|
||||
# pod. Same shape as gateway.podSecurityContext / gateway.securityContext.
|
||||
podSecurityContext: {}
|
||||
securityContext: {}
|
||||
# Extra pod labels on the Job pod, merged into the chart's common labels.
|
||||
podLabels: {}
|
||||
# Additional volumes on the Job pod and volumeMounts on its container, e.g.
|
||||
# the writable scratch space a read-only root filesystem needs.
|
||||
volumes: []
|
||||
volumeMounts: []
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-migrations
|
||||
tag: "" # defaults to .Chart.AppVersion
|
||||
|
|
@ -180,10 +216,13 @@ gateway:
|
|||
httpGet: { path: /health/liveliness, port: http }
|
||||
initialDelaySeconds: 10
|
||||
periodSeconds: 15
|
||||
timeoutSeconds: 10
|
||||
failureThreshold: 6
|
||||
readinessProbe:
|
||||
httpGet: { path: /health/readiness, port: http }
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
timeoutSeconds: 10
|
||||
hpa:
|
||||
enabled: true
|
||||
minReplicas: 1
|
||||
|
|
@ -200,6 +239,37 @@ gateway:
|
|||
minAvailable: ""
|
||||
maxUnavailable: ""
|
||||
podAnnotations: {}
|
||||
# Extra pod labels, merged into the chart's selector labels. Do not
|
||||
# re-declare `app.kubernetes.io/name` / `instance` / `component` here: they
|
||||
# form the Deployment's immutable selector.
|
||||
podLabels: {}
|
||||
# Pod-level securityContext, applied to every container in the pod
|
||||
# (runAsNonRoot, runAsUser, fsGroup, seccompProfile, ...). Empty by default
|
||||
# so the cluster's own defaults keep applying to existing installs; clusters
|
||||
# enforcing a restricted Pod Security Standard usually want at least
|
||||
# `runAsNonRoot: true` and `seccompProfile.type: RuntimeDefault`.
|
||||
podSecurityContext: {}
|
||||
# Container-level securityContext for the gateway container. Empty by
|
||||
# default for the same reason. Example:
|
||||
# allowPrivilegeEscalation: false
|
||||
# readOnlyRootFilesystem: true
|
||||
# capabilities:
|
||||
# drop:
|
||||
# - ALL
|
||||
# `readOnlyRootFilesystem: true` needs writable scratch space; supply it
|
||||
# through `volumes` / `volumeMounts` above rather than expecting the chart
|
||||
# to guess the paths your workload writes to.
|
||||
securityContext: {}
|
||||
# Extra sidecar containers appended to the gateway pod, e.g. an auth or
|
||||
# egress proxy. Rendered through `tpl`, so entries may reference chart
|
||||
# values and release metadata.
|
||||
extraContainers: []
|
||||
# Container lifecycle hooks (postStart / preStop) for the gateway container.
|
||||
lifecycle: {}
|
||||
# Grace period the kubelet allows between SIGTERM and SIGKILL. Leave empty
|
||||
# to inherit the Kubernetes default of 30s. Set it a few seconds above the
|
||||
# proxy's GRACEFUL_SHUTDOWN_TIMEOUT when you use a draining preStop hook.
|
||||
terminationGracePeriodSeconds: ""
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
|
|
@ -242,10 +312,13 @@ backend:
|
|||
httpGet: { path: /health/liveliness, port: http }
|
||||
initialDelaySeconds: 10
|
||||
periodSeconds: 15
|
||||
timeoutSeconds: 10
|
||||
failureThreshold: 6
|
||||
readinessProbe:
|
||||
httpGet: { path: /health/readiness, port: http }
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
timeoutSeconds: 10
|
||||
hpa:
|
||||
enabled: true
|
||||
minReplicas: 1
|
||||
|
|
@ -257,6 +330,13 @@ backend:
|
|||
minAvailable: ""
|
||||
maxUnavailable: ""
|
||||
podAnnotations: {}
|
||||
# Same shape as the gateway blocks of the same name.
|
||||
podLabels: {}
|
||||
podSecurityContext: {}
|
||||
securityContext: {}
|
||||
extraContainers: []
|
||||
lifecycle: {}
|
||||
terminationGracePeriodSeconds: ""
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
|
|
@ -310,6 +390,16 @@ ui:
|
|||
minAvailable: ""
|
||||
maxUnavailable: ""
|
||||
podAnnotations: {}
|
||||
# Same shape as the gateway blocks of the same name. The nginx runtime
|
||||
# writes its pid, cache, and proxy temp files under the image's root
|
||||
# filesystem, so `securityContext.readOnlyRootFilesystem: true` here needs
|
||||
# emptyDir volumes mounted over those paths.
|
||||
podLabels: {}
|
||||
podSecurityContext: {}
|
||||
securityContext: {}
|
||||
extraContainers: []
|
||||
lifecycle: {}
|
||||
terminationGracePeriodSeconds: ""
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
|
|
|
|||
106
litellm-proxy-extras/litellm_proxy_extras/replica_identity.py
Normal file
106
litellm-proxy-extras/litellm_proxy_extras/replica_identity.py
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
"""Optional post-migration step that raises Postgres REPLICA IDENTITY to FULL.
|
||||
|
||||
Logical-replication consumers (Neon / lakehouse sync and similar) need FULL
|
||||
replica identity to reconstruct the old row of an UPDATE or DELETE. Prisma
|
||||
leaves every table it creates at the Postgres default, so the setting has to be
|
||||
re-applied by hand after each migration run. Setting
|
||||
``LITELLM_SET_REPLICA_IDENTITY_FULL`` makes every migration run re-assert it.
|
||||
|
||||
The statement goes through the Prisma CLI rather than a Postgres driver because
|
||||
``litellm-proxy-extras`` has no runtime dependencies, while the CLI is already
|
||||
required for the migrations themselves.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
|
||||
REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL"
|
||||
|
||||
REPLICA_IDENTITY_FULL_SQL = r"""
|
||||
DO $$
|
||||
DECLARE
|
||||
target regclass;
|
||||
BEGIN
|
||||
SET LOCAL lock_timeout = '5s';
|
||||
FOR target IN
|
||||
SELECT c.oid::regclass
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relkind = 'r'
|
||||
AND c.relreplident <> 'f'
|
||||
AND n.nspname = ANY (current_schemas(false))
|
||||
AND c.relname LIKE 'LiteLLM\_%'
|
||||
LOOP
|
||||
BEGIN
|
||||
EXECUTE format('ALTER TABLE %s REPLICA IDENTITY FULL', target);
|
||||
EXCEPTION WHEN lock_not_available THEN
|
||||
RAISE WARNING 'REPLICA IDENTITY FULL skipped for %: table busy, retrying next run', target;
|
||||
END;
|
||||
END LOOP;
|
||||
END
|
||||
$$;
|
||||
"""
|
||||
|
||||
|
||||
def apply_replica_identity_full(
|
||||
schema_path: str,
|
||||
prisma_command: str,
|
||||
prisma_env: dict[str, str],
|
||||
) -> bool:
|
||||
"""Set REPLICA IDENTITY FULL on every LiteLLM table that is not already FULL.
|
||||
|
||||
Never raises. Replication metadata is not needed to serve requests, so
|
||||
every failure mode is reported and stepped over rather than taking down a
|
||||
migration run that already succeeded: a database that refuses the ALTER
|
||||
(most often because the runtime user does not own the tables), a missing
|
||||
or unrunnable Prisma CLI, a read-only temp directory, or a timeout.
|
||||
|
||||
Returns True when the statement was applied, False when it failed.
|
||||
"""
|
||||
logger.info("Applying REPLICA IDENTITY FULL to LiteLLM tables")
|
||||
try:
|
||||
with tempfile.TemporaryDirectory(prefix="litellm_replica_identity_") as tmp_dir:
|
||||
sql_path = Path(tmp_dir) / "replica_identity_full.sql"
|
||||
sql_path.write_text(REPLICA_IDENTITY_FULL_SQL)
|
||||
subprocess.run(
|
||||
[
|
||||
prisma_command,
|
||||
"db",
|
||||
"execute",
|
||||
"--file",
|
||||
str(sql_path),
|
||||
"--schema",
|
||||
schema_path,
|
||||
],
|
||||
timeout=60,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=prisma_env,
|
||||
)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.error(
|
||||
"Failed to set REPLICA IDENTITY FULL. Logical replication "
|
||||
"consumers may reject updates to these tables. Grant table "
|
||||
"ownership to the migration user, or apply "
|
||||
"`ALTER TABLE ... REPLICA IDENTITY FULL` by hand. Error: %s",
|
||||
e.stderr,
|
||||
)
|
||||
return False
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.error("Timed out setting REPLICA IDENTITY FULL on LiteLLM tables")
|
||||
return False
|
||||
except OSError as e:
|
||||
logger.error(
|
||||
"Could not run the REPLICA IDENTITY FULL statement. Logical "
|
||||
"replication consumers may reject updates to these tables. "
|
||||
"Error: %s",
|
||||
e,
|
||||
)
|
||||
return False
|
||||
|
||||
logger.info("REPLICA IDENTITY FULL applied to LiteLLM tables")
|
||||
return True
|
||||
|
|
@ -10,6 +10,10 @@ from pathlib import Path
|
|||
from typing import Optional
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.replica_identity import (
|
||||
REPLICA_IDENTITY_FULL_ENV_VAR,
|
||||
apply_replica_identity_full,
|
||||
)
|
||||
|
||||
|
||||
def str_to_bool(value: Optional[str]) -> bool:
|
||||
|
|
@ -676,6 +680,39 @@ class ProxyExtrasDBManager:
|
|||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
||||
@staticmethod
|
||||
def apply_replica_identity_full_if_requested() -> bool:
|
||||
"""
|
||||
Re-assert REPLICA IDENTITY FULL on LiteLLM's tables when the operator
|
||||
opted in via LITELLM_SET_REPLICA_IDENTITY_FULL.
|
||||
|
||||
Prisma leaves new tables at the Postgres default, which logical
|
||||
replication consumers reject, so the setting has to be re-applied after
|
||||
every migration run rather than once by hand.
|
||||
|
||||
Returns:
|
||||
bool: True if the setting was applied, False if it was not
|
||||
requested or could not be applied.
|
||||
"""
|
||||
if not str_to_bool(os.getenv(REPLICA_IDENTITY_FULL_ENV_VAR)):
|
||||
return False
|
||||
try:
|
||||
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
|
||||
prisma_command = _get_prisma_command()
|
||||
prisma_env = _get_prisma_env()
|
||||
except OSError as e:
|
||||
logger.error(
|
||||
"Could not resolve the migrations directory for the REPLICA "
|
||||
"IDENTITY FULL step, skipping it. Error: %s",
|
||||
e,
|
||||
)
|
||||
return False
|
||||
return apply_replica_identity_full(
|
||||
schema_path=schema_path,
|
||||
prisma_command=prisma_command,
|
||||
prisma_env=prisma_env,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def setup_database(
|
||||
use_migrate: bool = False, use_v2_resolver: bool = False
|
||||
|
|
@ -694,6 +731,15 @@ class ProxyExtrasDBManager:
|
|||
Returns:
|
||||
bool: True if setup was successful, False otherwise
|
||||
"""
|
||||
migrated = ProxyExtrasDBManager._run_migrations(
|
||||
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
|
||||
)
|
||||
if migrated:
|
||||
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
|
||||
return migrated
|
||||
|
||||
@staticmethod
|
||||
def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
|
||||
if use_v2_resolver:
|
||||
logger.info("Using v2 migration resolver (--use_v2_migration_resolver)")
|
||||
return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate)
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
# Adding a provider / route to litellm-rust
|
||||
|
||||
Three layers, same for every route (see `ocr` and `realtime` as references):
|
||||
Everything for a route lives in `crates/core/src/<route>/`; `crates/core/src/messages` is the reference. A host (the axum gateway, the Python bridge) only calls the route's entrypoint.
|
||||
|
||||
1. **Transform contract (pure)** — `crates/core/src/<route>/transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) + types in `types.rs`. No network, env, or auth.
|
||||
2. **Provider config (pure)** — `crates/providers/src/<provider>/<route>/transformation.rs`: implement that trait as a `const <PROVIDER>_<ROUTE>_CONFIG`, mirroring the Python provider tree. Add parity unit tests.
|
||||
3. **HTTP / transport (the host)** — `crates/providers/src/<route>.rs` (e.g. `ocr.rs`, `realtime.rs`): the callable fn (`run_ocr`, `realtime`). It resolves the key, builds the auth header, builds URL + transforms via the config, then does the network call. This is the only layer allowed to do I/O.
|
||||
1. **Entrypoint** — `mod.rs`: `pub async fn <route>(request) -> CoreResult<Response>`, the Rust equivalent of `litellm.<route>()`, plus a `<route>_stream` variant when the route streams. It is the only thing a host touches.
|
||||
2. **Transform contract** — `transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) with types in `types.rs`.
|
||||
3. **Provider config** — `crates/core/src/providers/<provider>/<route>/transformation.rs`: implement that trait as a `const <PROVIDER>_<ROUTE>_CONFIG`, mirroring the Python provider tree. Add parity unit tests.
|
||||
4. **Prepare + handler** — `prepare.rs` resolves provider/model, credentials, auth headers, and URL, then transforms the request; `handler.rs` performs the provider call through the shared client in `client.rs` and transforms the response.
|
||||
|
||||
## Coding standards
|
||||
|
||||
|
|
@ -25,4 +26,4 @@ variants of it. The test for a good abstraction is that adding the next provider
|
|||
is a few declarative lines, not a new file of duplicated flow. Only diverge from
|
||||
the base when behavior is genuinely different, and say so explicitly in the PR.
|
||||
|
||||
**Calling:** the host invokes the route fn — the Python bridge calls `run_ocr`; the `ai-gateway` server calls `realtime`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`.
|
||||
**Calling:** hosts invoke the core entrypoint — the Python bridge and the `ai-gateway` route service both call `litellm_core::messages::messages`. Never add a provider handler to `ai-gateway`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`.
|
||||
|
|
|
|||
|
|
@ -4,14 +4,30 @@ litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes (
|
|||
|
||||
## Crates
|
||||
|
||||
| Crate | Role | Pure / I/O |
|
||||
|-------|------|------------|
|
||||
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure |
|
||||
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
|
||||
| Crate | Role |
|
||||
|-------|------|
|
||||
| litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. |
|
||||
| litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
|
||||
## Where a route lives
|
||||
|
||||
A top-level LiteLLM call is a module under `crates/core/src/<route>/`, shaped like `messages`:
|
||||
|
||||
```
|
||||
core/src/messages/
|
||||
mod.rs # pub async fn messages(..) -> CoreResult<..> (+ messages_stream for SSE)
|
||||
types.rs # request/response types, MessagesRequest
|
||||
transformation.rs # the provider template trait
|
||||
prepare.rs # provider resolution, auth headers, URL
|
||||
handler.rs # the provider call
|
||||
client.rs # the shared reqwest client
|
||||
```
|
||||
|
||||
Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched.
|
||||
|
||||
Adding a crate: default to a MODULE. New crate ONLY on a real trigger — separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these.
|
||||
|
||||
Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional.
|
||||
|
|
|
|||
|
|
@ -23,21 +23,34 @@ the base when behavior is genuinely different, and say so explicitly in the PR.
|
|||
|
||||
## Crates (exactly three — see AGENTS.md)
|
||||
|
||||
`litellm-core` describes work; `litellm-ai-gateway` executes it; `litellm-python-bridge`
|
||||
exposes it to the Python SDK. A crate is a **layer**, not a route — add modules, not crates.
|
||||
`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call.
|
||||
`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and
|
||||
`litellm-python-bridge` exposes it to the Python SDK. A crate is a **layer**, not
|
||||
a route — add modules, not crates.
|
||||
|
||||
## Core Boundary
|
||||
|
||||
`litellm-core` is the pure translation layer; the `litellm-ai-gateway` host executes work.
|
||||
`litellm-core` owns the whole call. The Rust equivalent of `litellm.messages()`
|
||||
is `litellm_core::messages::messages(request).await`: you call it, it does the
|
||||
provider call, and you get a typed non-streaming response back.
|
||||
|
||||
Route-level Rust structure mirrors LiteLLM's Python responsibilities:
|
||||
- `core/src/<route>/` owns the route contract, shared types, and provider
|
||||
template traits. For OCR, this means `core/src/ocr`.
|
||||
- `core/src/<route>/` owns the route end to end: the public entrypoint fn named
|
||||
after the route in `mod.rs`, the request/response types (`types.rs`), the
|
||||
provider template trait (`transformation.rs`), the provider/auth/URL
|
||||
resolution (`prepare.rs`), the HTTP client (`client.rs`), and the handler that
|
||||
performs the call (`handler.rs`). `core/src/messages` is the reference.
|
||||
- `core/src/providers/<provider>/<route>/transformation.rs` owns the
|
||||
provider-specific transform. For Mistral OCR, this means
|
||||
`core/src/providers/mistral/ocr/transformation.rs`.
|
||||
- Network execution lives in the host crate `ai-gateway` (`ai-gateway/src/io/`),
|
||||
never inside `core`.
|
||||
provider-specific transform. For Anthropic Messages, this means
|
||||
`core/src/providers/anthropic/messages/transformation.rs`.
|
||||
- Handlers live in `core`, never in a host. `ai-gateway` must not contain a
|
||||
route handler that talks to a provider; its axum route reads the HTTP request,
|
||||
picks a deployment, and calls the `core` entrypoint. `python-bridge` marshals
|
||||
Python objects and calls the same entrypoint.
|
||||
|
||||
Streaming keeps the same shape: the route entrypoint has a `<route>_stream`
|
||||
variant in `core` that returns the upstream response so a host can splice it to
|
||||
its own caller; the host still owns no provider logic.
|
||||
|
||||
Call-hook and lifecycle instrumentation, including phase timing, usage
|
||||
accumulation, and callback payload construction, always lives in `core`.
|
||||
|
|
@ -45,21 +58,31 @@ Hosts feed observed events into core and dispatch the completed payloads through
|
|||
their I/O logger; hosts must not own callback orchestration.
|
||||
|
||||
Allowed in `core`:
|
||||
- Pure request transforms
|
||||
- Pure response transforms
|
||||
- Pure stream chunk normalization
|
||||
- The public entrypoint for a top-level LiteLLM call
|
||||
- Request/response transforms and stream chunk normalization
|
||||
- Provider resolution, auth header construction, and URL building
|
||||
- The provider HTTP call itself, through a shared reused client with connect and
|
||||
request timeouts
|
||||
- Shared data types and validation errors
|
||||
- Deterministic token/cost helper logic
|
||||
|
||||
Not allowed in `core`:
|
||||
- Network calls
|
||||
- Environment variable or secret reads
|
||||
- Serving HTTP: axum routes, extractors, and transport concerns stay in the host
|
||||
- Filesystem access
|
||||
- Database or cache access
|
||||
- Provider SDK signing or auth flows
|
||||
- Database access
|
||||
- Config file reading and rollout state
|
||||
- Logging callbacks, spend writes, or custom callbacks
|
||||
- Global mutable runtime state
|
||||
|
||||
Env reads in `core` are limited to credential fallback inside a route's
|
||||
`prepare.rs` (the `env_lookup` closure), mirroring what the Python SDK does when
|
||||
no key is passed. Everything else config-shaped is resolved by the host and
|
||||
passed in.
|
||||
|
||||
Routes still hosted in `ai-gateway` (`ocr`, `audio_transcription`, `realtime`)
|
||||
predate this rule and are being moved into `core` route modules; do not add new
|
||||
ones there, and prefer moving one when you touch it.
|
||||
|
||||
Python owns rollout state and fallback while Rust is being introduced. Rust
|
||||
paths must be off by default until parity tests prove equivalence with Python.
|
||||
A new provider/route may instead be implemented rust-only with no Python
|
||||
|
|
@ -93,10 +116,10 @@ the first PR:
|
|||
- Preserve Python output shape intentionally. If a field is always serialized as
|
||||
`null` for Python parity, leave a short comment explaining that parity choice.
|
||||
|
||||
## Host I/O Rules
|
||||
## Network I/O Rules
|
||||
|
||||
These rules apply when adding future crates or modules that execute network I/O,
|
||||
such as `ai-gateway`, router hosts, or standalone servers:
|
||||
These rules apply to every module that executes network I/O, whether it is a
|
||||
`core` route handler or a host such as `ai-gateway`:
|
||||
|
||||
- Set connect and full-request timeouts. No unbounded waits.
|
||||
- Reuse HTTP clients; do not construct clients per request.
|
||||
|
|
|
|||
|
|
@ -2,18 +2,31 @@
|
|||
|
||||
This workspace contains the staged Rust implementation for LiteLLM.
|
||||
|
||||
Rust starts as a pure transform core used by the existing Python host. Python
|
||||
continues to own auth, configuration, network I/O, retries, routing, logging,
|
||||
`litellm-core` is the LiteLLM SDK in Rust: one entrypoint per top-level call
|
||||
that makes the LLM call and hands back a typed response, the same shape as
|
||||
`litellm.messages()` in Python.
|
||||
|
||||
```rust
|
||||
let response = litellm_core::messages::messages(MessagesRequest {
|
||||
model: "claude-sonnet-4-5",
|
||||
body,
|
||||
api_key: Some(key),
|
||||
..
|
||||
})
|
||||
.await?;
|
||||
```
|
||||
|
||||
Python continues to own configuration, retries, routing policy, logging,
|
||||
callbacks, spend tracking, and customer plugins until each Rust path has parity
|
||||
coverage and production evidence.
|
||||
|
||||
## Crates
|
||||
|
||||
| Crate | Role | Pure / I/O |
|
||||
|-------|------|------------|
|
||||
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure |
|
||||
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
|
||||
| Crate | Role |
|
||||
|-------|------|
|
||||
| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. |
|
||||
| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
|
||||
|
|
@ -21,16 +34,16 @@ Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-
|
|||
|
||||
```text
|
||||
crates/
|
||||
core/ Route contracts, shared pure types, errors, and templates.
|
||||
src/ocr/
|
||||
providers/ Provider-specific pure transforms.
|
||||
src/mistral/ocr/transformation.rs
|
||||
core/ The SDK: route modules + provider transforms.
|
||||
src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client
|
||||
src/providers/anthropic/messages/transformation.rs
|
||||
ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints.
|
||||
python-bridge/ PyO3 bridge for Python LiteLLM.
|
||||
```
|
||||
|
||||
The folder shape should follow the Python provider tree:
|
||||
`providers/src/<provider>/<route>/transformation.rs`. The bridge should expose
|
||||
one function per top-level route, starting with `ocr(payload)`.
|
||||
The folder shape follows the Python provider tree:
|
||||
`core/src/providers/<provider>/<route>/transformation.rs`. The bridge exposes one
|
||||
function per top-level route, mirroring the core entrypoints.
|
||||
|
||||
## Checks
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
# Provider coding standards (litellm-rust)
|
||||
|
||||
Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MISTRAL_OCR_CONFIG`) is the reference; `messages` (`ANTHROPIC_MESSAGES_CONFIG`) is the next port.
|
||||
Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` (`core/src/messages`, `ANTHROPIC_MESSAGES_CONFIG`) is the reference: a route is a `core` module with a public entrypoint that makes the call and returns a typed response.
|
||||
|
||||
## Provider resolution
|
||||
|
||||
|
|
@ -16,10 +16,10 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MIST
|
|||
|
||||
## Boundaries
|
||||
|
||||
7. Layers never cross: `core` = pure transforms/types (no network, env, secrets, auth, logging, global mutable state); `ai-gateway` = all I/O, auth headers, HTTP/SSE, lifecycle hooks; `python-bridge` = thin PyO3 adapter.
|
||||
7. Layers never cross: `core` = the call itself (entrypoint, types, transforms, provider resolution, auth headers, provider HTTP, lifecycle hooks); `ai-gateway` = serving HTTP/WS (routing, extractors, auth of *our* callers, streaming to the client); `python-bridge` = thin PyO3 adapter. Hosts call the core entrypoint; they never build a provider request.
|
||||
8. Generic/route files contain zero provider-specific branches. A provider is one module under `core/src/providers/<provider>/<route>/`; a route is a module, never a new crate.
|
||||
9. Route entry point stays thin: `<route>()` -> `prepare_*` -> `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing. Handlers validate and delegate; no business logic in them.
|
||||
10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Env reads happen only at the host/config layer, with the `DEFAULT_*` fallback defined in `constants.rs`.
|
||||
9. Route entry point stays thin: `core::<route>::<route>()` -> `prepare_*` -> handler (or `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing). Axum handlers validate and delegate to a service that calls the entrypoint; no business logic in them.
|
||||
10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Config-shaped env reads happen at the host/config layer with the `DEFAULT_*` fallback defined in `constants.rs`; the only env read in `core` is the credential fallback in a route's `prepare.rs`.
|
||||
|
||||
## Types and errors
|
||||
|
||||
|
|
@ -33,7 +33,7 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MIST
|
|||
|
||||
16. Never log request/response bodies, base64 payloads, document contents, or secrets. Truncate and bound any upstream body before it crosses a host boundary.
|
||||
17. Treat empty/whitespace credentials, URLs, and config values as absent at the host resolution layer.
|
||||
18. Host I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS.
|
||||
18. Network I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS.
|
||||
|
||||
## Tests and rollout
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
# ai-gateway — folder architecture
|
||||
|
||||
The Axum server that fronts the Rust gateway. It owns transport + config + auth
|
||||
only; deployment selection lives in `core::router`, transforms in `core`/`providers`.
|
||||
only; deployment selection lives in `core::router`, and the LLM call itself
|
||||
(transforms, auth headers, provider HTTP) lives behind a `core` route entrypoint
|
||||
such as `litellm_core::messages::messages`. No provider handler lives here.
|
||||
|
||||
```
|
||||
src/
|
||||
|
|
@ -32,6 +34,11 @@ src/
|
|||
args; it runs during extraction. Never re-implement the check per route.
|
||||
- **Handlers are thin.** A handler validates and delegates to its `service`. No
|
||||
business logic, no provider calls, no transforms in handlers.
|
||||
- **Services call `core`, they don't reimplement it.** A `service` picks the
|
||||
deployment and calls the `core` route entrypoint. Provider resolution, auth
|
||||
headers, URL building, and the HTTP call are `core`'s job; a service that
|
||||
builds a provider request itself is a bug (`routes/messages/service.rs` is
|
||||
the reference).
|
||||
- **State is shared and cheap to clone.** Long-lived handles live behind `Arc` in
|
||||
`state.rs`; read env/config only in `main.rs` when building state.
|
||||
|
||||
|
|
|
|||
|
|
@ -8,11 +8,11 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame.
|
|||
|
||||
`litellm-rust` is exactly three crates (a crate is a **layer**, not a route):
|
||||
|
||||
| Crate | Role | Pure / I/O |
|
||||
|-------|------|------------|
|
||||
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under `providers/`), and the router. Builds requests/responses; no network. | Pure |
|
||||
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under `io/`) plus the Axum server binary (behind the `server` feature). | I/O |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
|
||||
| Crate | Role |
|
||||
|-------|------|
|
||||
| litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. |
|
||||
| litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
|
||||
|
|
|
|||
|
|
@ -29,18 +29,6 @@ pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500;
|
|||
#[cfg(feature = "server")]
|
||||
pub(crate) const DEFAULT_PROVIDER: &str = "openai";
|
||||
|
||||
/// Full-request timeout ceiling for Anthropic Messages provider calls, in
|
||||
/// seconds. Mirrors the Python Anthropic Messages default. The per-request
|
||||
/// timeout from `litellm_params` still overrides this on the request builder.
|
||||
pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
/// Connect timeout for Anthropic Messages provider calls, in seconds.
|
||||
pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10;
|
||||
|
||||
/// Max characters of an upstream error body echoed across the host boundary
|
||||
/// before truncation, so provider bodies are bounded and data-minimized.
|
||||
pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
|
||||
pub(crate) const DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS: u64 = 10;
|
||||
pub(crate) const DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS: u64 = 300;
|
||||
|
||||
|
|
@ -48,10 +36,6 @@ pub(crate) const DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS: u64 = 300;
|
|||
#[cfg(feature = "server")]
|
||||
pub(crate) const MESSAGES_ROUTE_PATH: &str = "/v1/messages";
|
||||
|
||||
/// Provider name used by the Anthropic Messages route when a deployment's
|
||||
/// provider model does not carry an explicit provider prefix.
|
||||
pub(crate) const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic";
|
||||
|
||||
/// Request headers owned by the gateway and never forwarded upstream.
|
||||
#[cfg(feature = "server")]
|
||||
pub(crate) const MESSAGES_HEADERS_NOT_FORWARDED: &[&str] =
|
||||
|
|
|
|||
|
|
@ -1 +0,0 @@
|
|||
pub use crate::messages::{MessagesRequest, messages};
|
||||
|
|
@ -1,5 +1,4 @@
|
|||
pub mod audio_transcription;
|
||||
pub mod messages;
|
||||
pub mod ocr;
|
||||
pub mod realtime;
|
||||
pub mod realtime_pool;
|
||||
|
|
|
|||
|
|
@ -4,7 +4,9 @@
|
|||
//! without pulling in the HTTP server:
|
||||
//!
|
||||
//! - Call-type modules such as [`ocr`]: provider transforms, lifecycle hooks,
|
||||
//! and provider I/O. Always available — no feature required.
|
||||
//! and provider I/O. Always available — no feature required. These predate the
|
||||
//! rule that a route's entrypoint and handler live in `litellm-core` (see
|
||||
//! `litellm_core::messages`) and move there as they are touched.
|
||||
//! - [`io`]: compatibility exports and realtime WebSocket splice helpers.
|
||||
//! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling
|
||||
//! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway`
|
||||
|
|
@ -14,7 +16,6 @@
|
|||
pub mod audio_transcription;
|
||||
mod client;
|
||||
pub mod io;
|
||||
pub mod messages;
|
||||
pub mod ocr;
|
||||
|
||||
/// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and
|
||||
|
|
|
|||
|
|
@ -1,49 +0,0 @@
|
|||
use litellm_core::CoreResult;
|
||||
use serde_json::Value;
|
||||
|
||||
mod client;
|
||||
mod common_utils;
|
||||
mod handler;
|
||||
mod prepare;
|
||||
mod types;
|
||||
|
||||
pub use types::MessagesRequest;
|
||||
|
||||
use handler::{execute_messages_provider_call, execute_messages_provider_stream};
|
||||
use prepare::prepare_messages_call;
|
||||
|
||||
pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<Value> {
|
||||
match execute_messages(request, false).await? {
|
||||
MessagesResponse::Json(body) => Ok(body),
|
||||
MessagesResponse::Stream(response) => {
|
||||
drop(response);
|
||||
Err(litellm_core::CoreError::InvalidResponse(
|
||||
"non-streaming messages execution returned a stream".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) enum MessagesResponse {
|
||||
Json(Value),
|
||||
Stream(reqwest::Response),
|
||||
}
|
||||
|
||||
pub(crate) async fn execute_messages(
|
||||
request: MessagesRequest<'_>,
|
||||
stream: bool,
|
||||
) -> CoreResult<MessagesResponse> {
|
||||
let prepared = prepare_messages_call(request)?;
|
||||
if stream {
|
||||
execute_messages_provider_stream(prepared)
|
||||
.await
|
||||
.map(MessagesResponse::Stream)
|
||||
} else {
|
||||
execute_messages_provider_call(prepared)
|
||||
.await
|
||||
.map(MessagesResponse::Json)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
|
@ -1,24 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub struct MessagesRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub body: Value,
|
||||
pub api_key: Option<&'a str>,
|
||||
pub api_base: Option<&'a str>,
|
||||
pub custom_llm_provider: Option<&'a str>,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
pub(crate) struct ProviderMessagesRequest {
|
||||
pub(crate) provider: String,
|
||||
pub(crate) model: String,
|
||||
pub(crate) config: &'static dyn AnthropicMessagesProviderConfig,
|
||||
pub(crate) url: String,
|
||||
pub(crate) body: Value,
|
||||
pub(crate) upstream_headers: Vec<(String, String)>,
|
||||
pub(crate) timeout: Option<Duration>,
|
||||
}
|
||||
|
|
@ -19,7 +19,10 @@ async fn handle(...) -> impl IntoResponse { ... }
|
|||
When a route has business logic worth testing without axum, put it in a sibling
|
||||
`service` (a file, or a folder if the route grows). The route file stays the
|
||||
**axum surface** (router + handler + any socket/SSE adapter); `service` is plain
|
||||
Rust with **no axum types**. `realtime/` is the example:
|
||||
Rust with **no axum types**, and its job is to pick the deployment and call the
|
||||
`core` route entrypoint (see `messages/service.rs` calling
|
||||
`litellm_core::messages::messages`). Never build a provider request, resolve a
|
||||
key, or perform the provider call here. `realtime/` is the older example:
|
||||
```
|
||||
realtime/
|
||||
mod.rs # axum surface: router() + handler + the WS<->events adapter
|
||||
|
|
@ -33,6 +36,8 @@ genuinely gets hard to read.
|
|||
`crate::auth::RequireMasterKey` to its arguments; it runs during extraction.
|
||||
Never re-implement the check per route.
|
||||
- **Handlers contain no business logic; `service` contains no axum types.**
|
||||
- **No provider handlers in this crate.** Transforms, auth headers, and the
|
||||
provider HTTP call live in `core/src/<route>/`.
|
||||
- A route owns its paths in its own `router()`; `mod.rs` only merges.
|
||||
- Cross-cutting concerns (logging, CORS, timeouts) → Tower layers in `mod.rs`,
|
||||
not duplicated in handlers.
|
||||
|
|
|
|||
|
|
@ -1,12 +1,12 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_core::constants::ANTHROPIC_MESSAGES_PROVIDER;
|
||||
use litellm_core::messages::types::MessagesRequest;
|
||||
use litellm_core::messages::{messages, messages_stream};
|
||||
use litellm_core::router::Router;
|
||||
use litellm_core::{CoreError, CoreResult};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
|
||||
use crate::messages::{MessagesRequest, execute_messages};
|
||||
|
||||
pub(crate) enum MessagesResponse {
|
||||
Json(Value),
|
||||
Stream(reqwest::Response),
|
||||
|
|
@ -52,13 +52,14 @@ pub async fn run(
|
|||
extra_headers,
|
||||
timeout: None,
|
||||
};
|
||||
let stream = request.body.get("stream").and_then(Value::as_bool) == Some(true);
|
||||
execute_messages(request, stream)
|
||||
.await
|
||||
.map(|response| match response {
|
||||
crate::messages::MessagesResponse::Json(body) => MessagesResponse::Json(body),
|
||||
crate::messages::MessagesResponse::Stream(upstream) => {
|
||||
MessagesResponse::Stream(upstream)
|
||||
}
|
||||
if request.body.get("stream").and_then(Value::as_bool) == Some(true) {
|
||||
return messages_stream(request).await.map(MessagesResponse::Stream);
|
||||
}
|
||||
|
||||
let response = messages(request).await?;
|
||||
serde_json::to_value(response)
|
||||
.map(MessagesResponse::Json)
|
||||
.map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("failed to serialize messages response: {err}"))
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,3 +1,7 @@
|
|||
litellm-core is the PURE translation layer — types, route contracts (traits), provider transforms (modules under `providers/`), and the router. No network, no I/O, no env reads.
|
||||
litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-level call is a module under `src/<route>/` exposing a public entrypoint named after the route (`messages::messages()`, the Rust equivalent of `litellm.messages()`): you call it and get a typed non-streaming response back.
|
||||
|
||||
Routes (ocr, realtime) and providers (mistral, openai) are modules, not crates.
|
||||
A route module owns everything the call needs: types, the provider template trait, provider transforms (under `providers/`), provider/auth/URL resolution, and the handler that performs the HTTP call. Handlers belong here, never in a host crate.
|
||||
|
||||
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback dispatch. Env reads are limited to credential fallback in a route's `prepare.rs`.
|
||||
|
||||
Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates.
|
||||
|
|
|
|||
|
|
@ -4,20 +4,28 @@ Rules for `litellm-rust/crates/core`.
|
|||
|
||||
## Responsibility
|
||||
|
||||
`core` owns shared data types, typed errors, and deterministic helper contracts.
|
||||
It must stay pure and host-independent.
|
||||
`core` is the LiteLLM SDK in Rust: it makes the LLM call. Every top-level
|
||||
LiteLLM call has a public entrypoint here, named after the route
|
||||
(`messages::messages()` is the Rust equivalent of `litellm.messages()`), and
|
||||
calling it returns a typed non-streaming response.
|
||||
|
||||
Allowed:
|
||||
- The public entrypoint for a route, plus its `<route>_stream` variant when the
|
||||
route supports streaming.
|
||||
- Provider resolution, auth header construction, URL building, and the provider
|
||||
HTTP call (shared reused client, connect + request timeouts).
|
||||
- Shared request/response structs.
|
||||
- Typed errors with stable, non-sensitive messages.
|
||||
- Deterministic validation helpers.
|
||||
- Serialization helpers that intentionally mirror Python output shape.
|
||||
- Route templates that match Python base config responsibilities, such as
|
||||
`ocr::transformation::OcrProviderConfig`.
|
||||
`messages::transformation::AnthropicMessagesProviderConfig`.
|
||||
|
||||
Not allowed:
|
||||
- Network, filesystem, database, cache, or environment access.
|
||||
- Secret reads or auth/header construction.
|
||||
- Serving HTTP: axum routers, extractors, and other transport concerns.
|
||||
- Filesystem, database, or cache access.
|
||||
- Config file reading or rollout state; the host resolves those and passes them
|
||||
in. Env reads are limited to credential fallback in a route's `prepare.rs`.
|
||||
- Logging callbacks, tracing spans, spend writes, or customer callbacks.
|
||||
- Provider-specific branching that belongs in `providers`.
|
||||
- Panics for user/provider-controlled input.
|
||||
|
|
@ -33,10 +41,21 @@ typed field on a struct, not a raw string threaded through the API.
|
|||
|
||||
## Structure
|
||||
|
||||
Use route names directly under `src/`: `ocr`, future `messages`,
|
||||
Use route names directly under `src/`: `messages`, `ocr`, future
|
||||
`chat_completions`, `embeddings`, and similar top-level LiteLLM calls. Do not
|
||||
invent broad names like `engine` for route contracts.
|
||||
|
||||
`src/messages` is the reference shape for a route module:
|
||||
|
||||
```
|
||||
mod.rs pub async fn messages(..) (+ messages_stream)
|
||||
types.rs request/response types
|
||||
transformation.rs the provider template trait
|
||||
prepare.rs provider resolution, auth headers, URL
|
||||
handler.rs the provider call
|
||||
client.rs the shared reqwest client
|
||||
```
|
||||
|
||||
## Parity Rules
|
||||
|
||||
- Every shared type used by a provider transform needs unit tests for
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ repository.workspace = true
|
|||
|
||||
[dependencies]
|
||||
rand.workspace = true
|
||||
reqwest.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
|
@ -30,5 +31,4 @@ bedrock-auth = [
|
|||
]
|
||||
|
||||
[dev-dependencies]
|
||||
reqwest.workspace = true
|
||||
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
|
||||
|
|
|
|||
|
|
@ -1,3 +1,19 @@
|
|||
pub const OPENAI_DEFAULT_API_BASE: &str = "https://api.openai.com";
|
||||
pub const OPENAI_RESPONSES_DEFAULT_API_BASE: &str = "https://api.openai.com/v1";
|
||||
pub const OPENAI_RESPONSES_PATH: &str = "/responses";
|
||||
|
||||
/// Full-request timeout ceiling for Anthropic Messages provider calls, in
|
||||
/// seconds. Mirrors the Python Anthropic Messages default. The per-request
|
||||
/// timeout from the caller still overrides this on the request builder.
|
||||
pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
/// Connect timeout for Anthropic Messages provider calls, in seconds.
|
||||
pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10;
|
||||
|
||||
/// Max characters of an upstream error body echoed across the call boundary
|
||||
/// before truncation, so provider bodies are bounded and data-minimized.
|
||||
pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
|
||||
/// Provider name used for Anthropic Messages when a deployment's provider model
|
||||
/// does not carry an explicit provider prefix.
|
||||
pub const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic";
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::{CoreError, json_type_name};
|
||||
use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
|
||||
use litellm_core::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
|
||||
use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS;
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
|
||||
use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
|
||||
|
||||
use super::transformation::AnthropicMessagesProviderConfig;
|
||||
|
||||
pub(super) fn truncate_error_body(body: &str) -> String {
|
||||
if body.chars().count() <= MESSAGES_ERROR_BODY_MAX_CHARS {
|
||||
|
|
@ -1,15 +1,13 @@
|
|||
use litellm_core::CoreResult;
|
||||
use litellm_core::error::CoreError;
|
||||
use serde_json::Value;
|
||||
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
|
||||
use super::client::http_client;
|
||||
use super::common_utils::truncate_error_body;
|
||||
use super::types::ProviderMessagesRequest;
|
||||
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
|
||||
use super::types::{AnthropicMessagesResponse, ProviderMessagesRequest};
|
||||
|
||||
pub(super) async fn execute_messages_provider_call(
|
||||
request: ProviderMessagesRequest,
|
||||
) -> CoreResult<Value> {
|
||||
) -> CoreResult<AnthropicMessagesResponse> {
|
||||
let mut request_builder = http_client().post(&request.url).json(&request.body);
|
||||
for (key, value) in &request.upstream_headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
|
|
@ -39,12 +37,7 @@ pub(super) async fn execute_messages_provider_call(
|
|||
let response = serde_json::from_str(&text).map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("invalid messages response JSON: {err}"))
|
||||
})?;
|
||||
let transformed = request
|
||||
.config
|
||||
.transform_response(&request.model, response)?;
|
||||
serde_json::to_value(transformed).map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("failed to serialize messages response: {err}"))
|
||||
})
|
||||
request.config.transform_response(&request.model, response)
|
||||
}
|
||||
|
||||
pub(super) async fn execute_messages_provider_stream(
|
||||
|
|
@ -1,2 +1,32 @@
|
|||
//! The Anthropic Messages call, the Rust equivalent of Python's
|
||||
//! `litellm.messages()`.
|
||||
//!
|
||||
//! [`messages`] is the top-level entrypoint: give it a model, a body, and
|
||||
//! credentials, and it resolves the provider, transforms the request, calls the
|
||||
//! provider, and returns a typed non-streaming response. [`messages_stream`]
|
||||
//! is the streaming variant; it hands the raw upstream response back so a host
|
||||
//! can splice the event stream to its own caller.
|
||||
|
||||
mod client;
|
||||
mod common_utils;
|
||||
mod handler;
|
||||
mod prepare;
|
||||
pub mod transformation;
|
||||
pub mod types;
|
||||
|
||||
use crate::error::CoreResult;
|
||||
|
||||
use handler::{execute_messages_provider_call, execute_messages_provider_stream};
|
||||
use prepare::prepare_messages_call;
|
||||
use types::{AnthropicMessagesResponse, MessagesRequest};
|
||||
|
||||
pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<AnthropicMessagesResponse> {
|
||||
execute_messages_provider_call(prepare_messages_call(request)?).await
|
||||
}
|
||||
|
||||
pub async fn messages_stream(request: MessagesRequest<'_>) -> CoreResult<reqwest::Response> {
|
||||
execute_messages_provider_stream(prepare_messages_call(request)?).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
|
|
|||
|
|
@ -1,9 +1,8 @@
|
|||
use litellm_core::CoreError;
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_core::messages::transformation::MessagesAuthStrategy;
|
||||
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
|
||||
|
||||
use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};
|
||||
use super::transformation::MessagesAuthStrategy;
|
||||
use super::types::{MessagesRequest, ProviderMessagesRequest};
|
||||
|
||||
pub(super) fn prepare_messages_call(
|
||||
|
|
@ -1,14 +1,16 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_core::error::CoreError;
|
||||
use serde_json::{Map, Value, json};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
|
||||
use crate::error::CoreError;
|
||||
|
||||
use super::common_utils::{
|
||||
has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body,
|
||||
};
|
||||
use super::{MessagesRequest, messages};
|
||||
use super::messages;
|
||||
use super::types::MessagesRequest;
|
||||
|
||||
async fn read_http_request(socket: &mut TcpStream) -> String {
|
||||
let mut request = Vec::new();
|
||||
|
|
@ -152,8 +154,8 @@ async fn messages_round_trip_builds_azure_request_and_passes_response_through()
|
|||
.await
|
||||
.expect("messages request succeeds");
|
||||
|
||||
assert_eq!(response["content"][0]["text"], "hi");
|
||||
assert_eq!(response["stop_reason"], "end_turn");
|
||||
assert_eq!(response.content[0]["text"], "hi");
|
||||
assert_eq!(response.stop_reason.as_deref(), Some("end_turn"));
|
||||
|
||||
let request = server.await.expect("server task completes");
|
||||
let (head, body) = request.split_once("\r\n\r\n").expect("has body");
|
||||
|
|
@ -208,8 +210,8 @@ async fn messages_round_trip_builds_native_anthropic_request() {
|
|||
.await
|
||||
.expect("messages request succeeds");
|
||||
|
||||
assert_eq!(response["content"][0]["text"], "hi");
|
||||
assert_eq!(response["stop_reason"], "end_turn");
|
||||
assert_eq!(response.content[0]["text"], "hi");
|
||||
assert_eq!(response.stop_reason.as_deref(), Some("end_turn"));
|
||||
|
||||
let request = server.await.expect("server task completes");
|
||||
let (head, _) = request.split_once("\r\n\r\n").expect("has body");
|
||||
|
|
@ -1,6 +1,30 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::transformation::AnthropicMessagesProviderConfig;
|
||||
|
||||
pub struct MessagesRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub body: Value,
|
||||
pub api_key: Option<&'a str>,
|
||||
pub api_base: Option<&'a str>,
|
||||
pub custom_llm_provider: Option<&'a str>,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
pub(super) struct ProviderMessagesRequest {
|
||||
pub(super) provider: String,
|
||||
pub(super) model: String,
|
||||
pub(super) config: &'static dyn AnthropicMessagesProviderConfig,
|
||||
pub(super) url: String,
|
||||
pub(super) body: Value,
|
||||
pub(super) upstream_headers: Vec<(String, String)>,
|
||||
pub(super) timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum SystemPrompt {
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
litellm-python-bridge is the PyO3 cdylib that exposes Rust to the litellm Python SDK — a thin adapter (Python objects → Rust calls → Python results) over litellm-ai-gateway.
|
||||
litellm-python-bridge is the PyO3 cdylib that exposes Rust to the litellm Python SDK — a thin adapter (Python objects → Rust calls → Python results) over the litellm-core route entrypoints (e.g. `litellm_core::messages::messages`).
|
||||
|
||||
Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call into litellm-ai-gateway.
|
||||
Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call the core entrypoint.
|
||||
|
|
|
|||
|
|
@ -11,11 +11,11 @@ Python-compatible dictionaries.
|
|||
## Bridge Shape
|
||||
|
||||
- Prefer one stable method per top-level LiteLLM route, for example
|
||||
`ocr(payload)`.
|
||||
`messages(...)`, calling the matching `litellm-core` entrypoint.
|
||||
- Do not add one exported PyO3 function per provider helper unless there is a
|
||||
measured reason.
|
||||
- Provider dispatch belongs in Rust route modules such as
|
||||
`litellm_providers::ocr`, not in this PyO3 crate.
|
||||
- Provider dispatch belongs in the `litellm-core` route module (e.g.
|
||||
`litellm_core::messages`), not in this PyO3 crate.
|
||||
- Python owns rollout state and fallback. Rust should return errors; Python
|
||||
decides whether to raise or fall back. For a rust-only provider/route (no
|
||||
Python reference), the Python side is a thin dispatch that calls Rust and
|
||||
|
|
|
|||
|
|
@ -4,10 +4,11 @@ use std::time::Duration;
|
|||
use litellm_ai_gateway::io::audio_transcription::{
|
||||
AudioTranscriptionRequest, audio_transcription as run_audio_transcription,
|
||||
};
|
||||
use litellm_ai_gateway::io::messages::{MessagesRequest, messages as run_messages};
|
||||
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
|
||||
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::messages::messages as run_messages;
|
||||
use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest};
|
||||
use pyo3::exceptions::{PyRuntimeError, PyValueError};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyAny, PyDict};
|
||||
|
|
@ -35,6 +36,15 @@ fn json_to_py(py: Python<'_>, value: Value) -> PyResult<Py<PyAny>> {
|
|||
Ok(json.call_method1("loads", (encoded,))?.unbind())
|
||||
}
|
||||
|
||||
fn messages_response_to_py(
|
||||
py: Python<'_>,
|
||||
response: AnthropicMessagesResponse,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let value =
|
||||
serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?;
|
||||
json_to_py(py, value)
|
||||
}
|
||||
|
||||
fn core_error_to_pyerr(err: CoreError) -> PyErr {
|
||||
match err {
|
||||
CoreError::Auth(message) => PyValueError::new_err(message),
|
||||
|
|
@ -382,7 +392,7 @@ fn messages(
|
|||
});
|
||||
|
||||
match result {
|
||||
Ok(value) => json_to_py(py, value),
|
||||
Ok(response) => messages_response_to_py(py, response),
|
||||
Err(err) => Err(core_error_to_pyerr(err)),
|
||||
}
|
||||
}
|
||||
|
|
@ -404,7 +414,7 @@ fn amessages(
|
|||
marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?;
|
||||
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
let value = run_messages(MessagesRequest {
|
||||
let response = run_messages(MessagesRequest {
|
||||
model: &model,
|
||||
body,
|
||||
api_key: api_key.as_deref(),
|
||||
|
|
@ -416,7 +426,7 @@ fn amessages(
|
|||
.await
|
||||
.map_err(core_error_to_pyerr)?;
|
||||
|
||||
Python::attach(|py| json_to_py(py, value))
|
||||
Python::attach(|py| messages_response_to_py(py, response))
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -449,6 +449,8 @@ enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None
|
|||
custom_prometheus_metadata_labels: List[str] = []
|
||||
custom_prometheus_tags: List[str] = []
|
||||
prometheus_metrics_config: Optional[List] = None
|
||||
prometheus_exclude_metrics: Optional[List[str]] = None
|
||||
prometheus_exclude_labels: Optional[List[str]] = None
|
||||
prometheus_emit_stream_label: bool = False
|
||||
# Opt-in: emit `rate_limit_category` and `rate_limit_type` labels on
|
||||
# `litellm_proxy_failed_requests_metric`. Off by default to preserve the
|
||||
|
|
|
|||
|
|
@ -17,39 +17,40 @@ until they're actually needed.
|
|||
|
||||
import importlib
|
||||
import sys
|
||||
from typing import Any, Optional, cast, Callable
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast
|
||||
|
||||
# Import all the data structures that define what can be lazy-loaded
|
||||
# These are just lists of names and maps of where to find them
|
||||
from ._lazy_imports_registry import (
|
||||
# Name tuples
|
||||
COST_CALCULATOR_NAMES,
|
||||
LITELLM_LOGGING_NAMES,
|
||||
UTILS_NAMES,
|
||||
TOKEN_COUNTER_NAMES,
|
||||
LLM_CLIENT_CACHE_NAMES,
|
||||
BEDROCK_TYPES_NAMES,
|
||||
TYPES_UTILS_NAMES,
|
||||
CACHING_NAMES,
|
||||
HTTP_HANDLER_NAMES,
|
||||
DOTPROMPT_NAMES,
|
||||
LLM_CONFIG_NAMES,
|
||||
TYPES_NAMES,
|
||||
LLM_PROVIDER_LOGIC_NAMES,
|
||||
UTILS_MODULE_NAMES,
|
||||
# Import maps
|
||||
_UTILS_IMPORT_MAP,
|
||||
_COST_CALCULATOR_IMPORT_MAP,
|
||||
_TYPES_UTILS_IMPORT_MAP,
|
||||
_TOKEN_COUNTER_IMPORT_MAP,
|
||||
_BEDROCK_TYPES_IMPORT_MAP,
|
||||
_CACHING_IMPORT_MAP,
|
||||
_LITELLM_LOGGING_IMPORT_MAP,
|
||||
_COST_CALCULATOR_IMPORT_MAP,
|
||||
_DOTPROMPT_IMPORT_MAP,
|
||||
_TYPES_IMPORT_MAP,
|
||||
_LITELLM_LOGGING_IMPORT_MAP,
|
||||
_LLM_CONFIGS_IMPORT_MAP,
|
||||
_LLM_PROVIDER_LOGIC_IMPORT_MAP,
|
||||
_TOKEN_COUNTER_IMPORT_MAP,
|
||||
_TYPES_IMPORT_MAP,
|
||||
_TYPES_UTILS_IMPORT_MAP,
|
||||
# Import maps
|
||||
_UTILS_IMPORT_MAP,
|
||||
_UTILS_MODULE_IMPORT_MAP,
|
||||
BEDROCK_TYPES_NAMES,
|
||||
CACHING_NAMES,
|
||||
# Name tuples
|
||||
COST_CALCULATOR_NAMES,
|
||||
DOTPROMPT_NAMES,
|
||||
HTTP_HANDLER_NAMES,
|
||||
LITELLM_LOGGING_NAMES,
|
||||
LLM_CLIENT_CACHE_NAMES,
|
||||
LLM_CONFIG_NAMES,
|
||||
LLM_PROVIDER_LOGIC_NAMES,
|
||||
TOKEN_COUNTER_NAMES,
|
||||
TYPES_NAMES,
|
||||
TYPES_UTILS_NAMES,
|
||||
UTILS_MODULE_NAMES,
|
||||
UTILS_NAMES,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -77,7 +78,7 @@ def _get_utils_globals() -> dict:
|
|||
# They're separate from the main lazy import system because they have specific use cases
|
||||
|
||||
# Lazy loader for default encoding - avoids importing heavy tiktoken library at startup
|
||||
_default_encoding: Optional[Any] = None
|
||||
_default_encoding: Any | None = None
|
||||
|
||||
|
||||
def _get_default_encoding() -> Any:
|
||||
|
|
@ -99,7 +100,7 @@ def _get_default_encoding() -> Any:
|
|||
|
||||
|
||||
# Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time
|
||||
_get_modified_max_tokens_func: Optional[Any] = None
|
||||
_get_modified_max_tokens_func: Any | None = None
|
||||
|
||||
|
||||
def _get_modified_max_tokens() -> Any:
|
||||
|
|
@ -123,7 +124,7 @@ def _get_modified_max_tokens() -> Any:
|
|||
|
||||
|
||||
# Lazy loader for token_counter to avoid importing token_counter module at module import time
|
||||
_token_counter_new_func: Optional[Any] = None
|
||||
_token_counter_new_func: Any | None = None
|
||||
|
||||
|
||||
def _get_token_counter_new() -> Any:
|
||||
|
|
@ -153,7 +154,7 @@ def _get_token_counter_new() -> Any:
|
|||
# This registry maps attribute names (like "ModelResponse") to handler functions
|
||||
# It's built once the first time someone accesses a lazy-loaded attribute
|
||||
# Example: {"ModelResponse": _lazy_import_utils, "Cache": _lazy_import_caching, ...}
|
||||
_LAZY_IMPORT_REGISTRY: Optional[dict[str, Callable[[str], Any]]] = None
|
||||
_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], Any]] | None = None
|
||||
|
||||
|
||||
def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
|
||||
|
|
|
|||
|
|
@ -4,11 +4,11 @@ import os
|
|||
import sys
|
||||
from datetime import datetime
|
||||
from logging import Formatter
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
|
||||
set_verbose = False
|
||||
|
||||
|
|
@ -86,7 +86,7 @@ handler.setLevel(numeric_level)
|
|||
handler.addFilter(_secret_filter)
|
||||
|
||||
|
||||
def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]:
|
||||
def _try_parse_json_message(message: str) -> dict[str, Any] | None:
|
||||
"""
|
||||
Try to parse a log message as JSON. Returns parsed dict if valid, else None.
|
||||
Handles messages that are entirely valid JSON (e.g. json.dumps output).
|
||||
|
|
@ -103,7 +103,7 @@ def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]:
|
|||
return parsed
|
||||
|
||||
|
||||
def _try_parse_embedded_python_dict(message: str) -> Optional[Dict[str, Any]]:
|
||||
def _try_parse_embedded_python_dict(message: str) -> dict[str, Any] | None:
|
||||
"""
|
||||
Try to find and parse a Python dict repr (e.g. str(d) or repr(d)) embedded in
|
||||
the message. Handles patterns like:
|
||||
|
|
@ -149,7 +149,7 @@ _STANDARD_RECORD_ATTRS = _get_standard_record_attrs()
|
|||
|
||||
class JsonFormatter(Formatter):
|
||||
def __init__(self):
|
||||
super(JsonFormatter, self).__init__()
|
||||
super().__init__()
|
||||
|
||||
def formatTime(self, record, datefmt=None):
|
||||
# Use datetime to format the timestamp in ISO 8601 format
|
||||
|
|
@ -158,7 +158,7 @@ class JsonFormatter(Formatter):
|
|||
|
||||
def format(self, record):
|
||||
message_str = record.getMessage()
|
||||
json_record: Dict[str, Any] = {
|
||||
json_record: dict[str, Any] = {
|
||||
"message": message_str,
|
||||
"level": record.levelname,
|
||||
"timestamp": self.formatTime(record),
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import json
|
|||
|
||||
# s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation
|
||||
import os
|
||||
from typing import Callable, List, Optional, Union
|
||||
from collections.abc import Callable
|
||||
|
||||
import redis # type: ignore
|
||||
import redis.asyncio as async_redis # type: ignore
|
||||
|
|
@ -61,23 +61,51 @@ def _get_redis_kwargs():
|
|||
return available_args
|
||||
|
||||
|
||||
def _get_redis_url_kwargs(client=None):
|
||||
def _init_arg_names(cls: type) -> frozenset[str]:
|
||||
"""Every ``__init__`` parameter accepted anywhere in a class's MRO.
|
||||
|
||||
Keyword-only parameters are included, and the MRO is walked because redis-py splits a
|
||||
connection's parameters between ``AbstractConnection`` and its concrete subclasses.
|
||||
"""
|
||||
return frozenset(
|
||||
name
|
||||
for klass in inspect.getmro(cls)
|
||||
if klass is not object
|
||||
for spec in (inspect.getfullargspec(klass.__init__),)
|
||||
for name in spec.args + spec.kwonlyargs
|
||||
)
|
||||
|
||||
|
||||
def _get_redis_url_kwargs(client: type | None = None) -> tuple[str, ...]:
|
||||
"""Connection kwargs that redis-py forwards from ``from_url`` down to the connection.
|
||||
|
||||
``from_url`` is declared as ``(cls, url, **kwargs)``, so introspecting it yields no
|
||||
connection kwargs at all. What it really does is hand its kwargs to the connection
|
||||
class, so that class's signature is the allowlist.
|
||||
|
||||
Taking the client's signature instead would be wrong in both directions: it omits
|
||||
nothing useful, but it admits client-only parameters such as
|
||||
``single_connection_client`` and ``auto_close_connection_pool``, plus the ``ssl_*``
|
||||
family that only ``SSLConnection`` accepts. Those reach ``AbstractConnection`` and
|
||||
raise ``TypeError`` the first time a connection is created. TLS on a url config is
|
||||
selected by the ``rediss://`` scheme, which picks ``SSLConnection`` on its own.
|
||||
"""
|
||||
if client is None:
|
||||
client = redis.Redis.from_url
|
||||
arg_spec = inspect.getfullargspec(redis.Redis.from_url)
|
||||
client = redis.Redis
|
||||
connection_cls = async_redis.Connection if client is async_redis.Redis else redis.Connection
|
||||
|
||||
exclude_args = frozenset(
|
||||
{
|
||||
"self",
|
||||
"connection_pool",
|
||||
"retry",
|
||||
}
|
||||
)
|
||||
|
||||
# Only allow primitive arguments
|
||||
exclude_args = {
|
||||
"self",
|
||||
"connection_pool",
|
||||
"retry",
|
||||
}
|
||||
include_args = ("url", "max_connections")
|
||||
|
||||
include_args = ["url"]
|
||||
|
||||
available_args = [x for x in arg_spec.args if x not in exclude_args] + include_args
|
||||
|
||||
return available_args
|
||||
return tuple(x for x in _init_arg_names(connection_cls) if x not in exclude_args) + include_args
|
||||
|
||||
|
||||
def _get_redis_cluster_kwargs(client=None):
|
||||
|
|
@ -132,7 +160,7 @@ def _redis_kwargs_from_environment():
|
|||
|
||||
def create_gcp_iam_redis_connect_func(
|
||||
service_account: str,
|
||||
ssl_ca_certs: Optional[str] = None,
|
||||
ssl_ca_certs: str | None = None,
|
||||
) -> Callable:
|
||||
"""
|
||||
Creates a custom Redis connection function for GCP IAM authentication.
|
||||
|
|
@ -175,9 +203,9 @@ def create_gcp_iam_redis_connect_func(
|
|||
|
||||
|
||||
def _build_azure_credential(
|
||||
azure_client_id: Optional[str] = None,
|
||||
azure_tenant_id: Optional[str] = None,
|
||||
azure_client_secret: Optional[str] = None,
|
||||
azure_client_id: str | None = None,
|
||||
azure_tenant_id: str | None = None,
|
||||
azure_client_secret: str | None = None,
|
||||
):
|
||||
"""
|
||||
Build a long-lived Azure credential object.
|
||||
|
|
@ -213,9 +241,9 @@ def _build_azure_credential(
|
|||
|
||||
|
||||
def _generate_azure_ad_redis_token(
|
||||
azure_client_id: Optional[str] = None,
|
||||
azure_tenant_id: Optional[str] = None,
|
||||
azure_client_secret: Optional[str] = None,
|
||||
azure_client_id: str | None = None,
|
||||
azure_tenant_id: str | None = None,
|
||||
azure_client_secret: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
One-shot helper that builds a credential and fetches a single Azure AD
|
||||
|
|
@ -235,9 +263,9 @@ def _generate_azure_ad_redis_token(
|
|||
|
||||
|
||||
def create_azure_ad_redis_connect_func(
|
||||
azure_client_id: Optional[str] = None,
|
||||
azure_tenant_id: Optional[str] = None,
|
||||
azure_client_secret: Optional[str] = None,
|
||||
azure_client_id: str | None = None,
|
||||
azure_tenant_id: str | None = None,
|
||||
azure_client_secret: str | None = None,
|
||||
) -> Callable:
|
||||
"""
|
||||
Creates a custom Redis connection function for Azure AD authentication.
|
||||
|
|
@ -341,7 +369,7 @@ def _get_redis_client_logic(**env_overrides):
|
|||
**env_overrides,
|
||||
}
|
||||
|
||||
_startup_nodes: Optional[Union[str, list]] = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore
|
||||
_startup_nodes: str | list | None = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore
|
||||
"REDIS_CLUSTER_NODES"
|
||||
)
|
||||
|
||||
|
|
@ -352,21 +380,21 @@ def _get_redis_client_logic(**env_overrides):
|
|||
elif _startup_nodes is None:
|
||||
redis_kwargs.pop("startup_nodes", None)
|
||||
|
||||
_sentinel_nodes: Optional[Union[str, list]] = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore
|
||||
_sentinel_nodes: str | list | None = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore
|
||||
"REDIS_SENTINEL_NODES"
|
||||
)
|
||||
|
||||
if _sentinel_nodes is not None and isinstance(_sentinel_nodes, str):
|
||||
redis_kwargs["sentinel_nodes"] = json.loads(_sentinel_nodes)
|
||||
|
||||
_sentinel_password: Optional[str] = redis_kwargs.get("sentinel_password", None) or get_secret_str(
|
||||
_sentinel_password: str | None = redis_kwargs.get("sentinel_password", None) or get_secret_str(
|
||||
"REDIS_SENTINEL_PASSWORD"
|
||||
)
|
||||
|
||||
if _sentinel_password is not None:
|
||||
redis_kwargs["sentinel_password"] = _sentinel_password
|
||||
|
||||
_service_name: Optional[str] = redis_kwargs.get("service_name", None) or get_secret( # type: ignore
|
||||
_service_name: str | None = redis_kwargs.get("service_name", None) or get_secret( # type: ignore
|
||||
"REDIS_SERVICE_NAME"
|
||||
)
|
||||
|
||||
|
|
@ -437,9 +465,12 @@ def _get_redis_client_logic(**env_overrides):
|
|||
redis_kwargs.pop("port", None)
|
||||
redis_kwargs.pop("db", None)
|
||||
redis_kwargs.pop("password", None)
|
||||
elif "startup_nodes" in redis_kwargs and redis_kwargs["startup_nodes"] is not None:
|
||||
pass
|
||||
elif "sentinel_nodes" in redis_kwargs and redis_kwargs["sentinel_nodes"] is not None:
|
||||
elif (
|
||||
"startup_nodes" in redis_kwargs
|
||||
and redis_kwargs["startup_nodes"] is not None
|
||||
or "sentinel_nodes" in redis_kwargs
|
||||
and redis_kwargs["sentinel_nodes"] is not None
|
||||
):
|
||||
pass
|
||||
elif "host" not in redis_kwargs or redis_kwargs["host"] is None:
|
||||
raise ValueError("Either 'host' or 'url' must be specified for redis.")
|
||||
|
|
@ -449,7 +480,7 @@ def _get_redis_client_logic(**env_overrides):
|
|||
|
||||
|
||||
def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
|
||||
_redis_cluster_nodes_in_env: Optional[str] = get_secret("REDIS_CLUSTER_NODES") # type: ignore
|
||||
_redis_cluster_nodes_in_env: str | None = get_secret("REDIS_CLUSTER_NODES") # type: ignore
|
||||
if _redis_cluster_nodes_in_env is not None:
|
||||
try:
|
||||
redis_kwargs["startup_nodes"] = json.loads(_redis_cluster_nodes_in_env)
|
||||
|
|
@ -467,7 +498,7 @@ def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
|
|||
if arg in args:
|
||||
cluster_kwargs[arg] = redis_kwargs[arg]
|
||||
|
||||
new_startup_nodes: List[ClusterNode] = []
|
||||
new_startup_nodes: list[ClusterNode] = []
|
||||
|
||||
for item in redis_kwargs["startup_nodes"]:
|
||||
new_startup_nodes.append(ClusterNode(**item))
|
||||
|
|
@ -559,9 +590,9 @@ def get_redis_client(**env_overrides):
|
|||
|
||||
|
||||
def get_redis_async_client(
|
||||
connection_pool: Optional[async_redis.BlockingConnectionPool] = None,
|
||||
connection_pool: async_redis.BlockingConnectionPool | None = None,
|
||||
**env_overrides,
|
||||
) -> Union[async_redis.Redis, async_redis.RedisCluster]:
|
||||
) -> async_redis.Redis | async_redis.RedisCluster:
|
||||
redis_kwargs = _get_redis_client_logic(**env_overrides)
|
||||
|
||||
if "startup_nodes" in redis_kwargs:
|
||||
|
|
@ -590,7 +621,7 @@ def get_redis_async_client(
|
|||
username=os.environ.get("REDIS_USERNAME") or None,
|
||||
)
|
||||
|
||||
new_startup_nodes: List[ClusterNode] = []
|
||||
new_startup_nodes: list[ClusterNode] = []
|
||||
|
||||
for item in redis_kwargs["startup_nodes"]:
|
||||
new_startup_nodes.append(ClusterNode(**item))
|
||||
|
|
@ -614,15 +645,13 @@ def get_redis_async_client(
|
|||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
if connection_pool is not None:
|
||||
return async_redis.Redis(connection_pool=connection_pool)
|
||||
args = _get_redis_url_kwargs(client=async_redis.Redis.from_url)
|
||||
args = _get_redis_url_kwargs(client=async_redis.Redis)
|
||||
url_kwargs = {}
|
||||
for arg in redis_kwargs:
|
||||
if arg in args:
|
||||
url_kwargs[arg] = redis_kwargs[arg]
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
"REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format(arg)
|
||||
)
|
||||
verbose_logger.debug(f"REDIS: ignoring argument: {arg}. Not an allowed async_redis.Redis.from_url arg.")
|
||||
return async_redis.Redis.from_url(**url_kwargs)
|
||||
|
||||
# Check for Redis Sentinel
|
||||
|
|
@ -654,7 +683,7 @@ def get_redis_async_client(
|
|||
|
||||
def get_redis_connection_pool(
|
||||
**env_overrides,
|
||||
) -> Optional[async_redis.BlockingConnectionPool]:
|
||||
) -> async_redis.BlockingConnectionPool | None:
|
||||
redis_kwargs = _get_redis_client_logic(**env_overrides)
|
||||
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
|
||||
|
||||
|
|
@ -662,10 +691,10 @@ def get_redis_connection_pool(
|
|||
return None
|
||||
|
||||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
pool_kwargs = {
|
||||
"timeout": REDIS_CONNECTION_POOL_TIMEOUT,
|
||||
"url": redis_kwargs["url"],
|
||||
}
|
||||
allowed_args = _get_redis_url_kwargs(client=async_redis.Redis)
|
||||
pool_kwargs = {k: v for k, v in redis_kwargs.items() if k in allowed_args and k != "max_connections"}
|
||||
pool_kwargs["timeout"] = REDIS_CONNECTION_POOL_TIMEOUT
|
||||
pool_kwargs["url"] = redis_kwargs["url"]
|
||||
if "max_connections" in redis_kwargs:
|
||||
try:
|
||||
pool_kwargs["max_connections"] = int(redis_kwargs["max_connections"])
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Dict, Optional, Tuple, Union
|
||||
from typing import Any
|
||||
|
||||
from redis.credentials import CredentialProvider # type: ignore[attr-defined]
|
||||
|
||||
|
|
@ -14,7 +14,7 @@ _GCP_IAM_TOKEN_TTL_SECONDS = 3300
|
|||
# Module-level cache shared across all GCPIAMCredentialProvider instances for the
|
||||
# same service account, so multiple Redis connections on the same pod share one token.
|
||||
# Keyed by service_account → (token, expiry_monotonic_timestamp).
|
||||
_token_cache: Dict[str, Tuple[str, float]] = {}
|
||||
_token_cache: dict[str, tuple[str, float]] = {}
|
||||
_token_cache_lock = threading.Lock()
|
||||
|
||||
|
||||
|
|
@ -95,11 +95,11 @@ class GCPIAMCredentialProvider(CredentialProvider):
|
|||
def __init__(self, gcp_service_account: str) -> None:
|
||||
self._gcp_service_account = gcp_service_account
|
||||
|
||||
def get_credentials(self) -> Tuple[str]:
|
||||
def get_credentials(self) -> tuple[str]:
|
||||
token = _get_cached_gcp_iam_token(self._gcp_service_account)
|
||||
return (token,)
|
||||
|
||||
async def get_credentials_async(self) -> Tuple[str]:
|
||||
async def get_credentials_async(self) -> tuple[str]:
|
||||
token = await asyncio.to_thread(_get_cached_gcp_iam_token, self._gcp_service_account)
|
||||
return (token,)
|
||||
|
||||
|
|
@ -115,17 +115,17 @@ class AzureADCredentialProvider(CredentialProvider):
|
|||
fail authentication after the initial token expired (~1 hour TTL).
|
||||
"""
|
||||
|
||||
def __init__(self, credential: Any, username: Optional[str] = None) -> None:
|
||||
def __init__(self, credential: Any, username: str | None = None) -> None:
|
||||
self._credential = credential
|
||||
self._username = username
|
||||
|
||||
def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:
|
||||
def get_credentials(self) -> tuple[str] | tuple[str, str]:
|
||||
token = self._credential.get_token(AZURE_REDIS_SCOPE).token
|
||||
if self._username:
|
||||
return (self._username, token)
|
||||
return (token,)
|
||||
|
||||
async def get_credentials_async(self) -> Union[Tuple[str], Tuple[str, str]]:
|
||||
async def get_credentials_async(self) -> tuple[str] | tuple[str, str]:
|
||||
token_obj = await asyncio.to_thread(self._credential.get_token, AZURE_REDIS_SCOPE)
|
||||
if self._username:
|
||||
return (self._username, token_obj.token)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
from typing import TYPE_CHECKING, Any, Optional, Union
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -24,7 +24,7 @@ else:
|
|||
UserAPIKeyAuth = Any
|
||||
|
||||
|
||||
def _get_otel_v2_class() -> Optional[type]:
|
||||
def _get_otel_v2_class() -> type | None:
|
||||
"""Return the ``OpenTelemetryV2`` class, or ``None`` if the OTel SDK is absent.
|
||||
|
||||
Imported lazily: ``litellm.integrations.otel.logger`` imports the OpenTelemetry
|
||||
|
|
@ -54,7 +54,7 @@ class ServiceLogging(CustomLogger):
|
|||
if "prometheus_system" in litellm.service_callback:
|
||||
self.prometheusServicesLogger = PrometheusServicesLogger()
|
||||
|
||||
def _resolve_otel_service_logger(self, callback: Any) -> Optional[Any]:
|
||||
def _resolve_otel_service_logger(self, callback: Any) -> Any | None:
|
||||
"""Resolve the OTel logger (legacy or V2) to emit a service span on.
|
||||
|
||||
Returns the logger instance whose ``async_service_*_hook`` should fire for
|
||||
|
|
@ -88,9 +88,9 @@ class ServiceLogging(CustomLogger):
|
|||
service: ServiceTypes,
|
||||
duration: float,
|
||||
call_type: str,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
start_time: Optional[Union[datetime, float]] = None,
|
||||
end_time: Optional[Union[float, datetime]] = None,
|
||||
parent_otel_span: Span | None = None,
|
||||
start_time: datetime | float | None = None,
|
||||
end_time: float | datetime | None = None,
|
||||
):
|
||||
"""
|
||||
Handles both sync and async monitoring by checking for existing event loop.
|
||||
|
|
@ -152,10 +152,10 @@ class ServiceLogging(CustomLogger):
|
|||
service: ServiceTypes,
|
||||
call_type: str,
|
||||
duration: float,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
start_time: Optional[Union[datetime, float]] = None,
|
||||
end_time: Optional[Union[datetime, float]] = None,
|
||||
event_metadata: Optional[dict] = None,
|
||||
parent_otel_span: Span | None = None,
|
||||
start_time: datetime | float | None = None,
|
||||
end_time: datetime | float | None = None,
|
||||
event_metadata: dict | None = None,
|
||||
):
|
||||
"""
|
||||
- For counting if the redis, postgres call is successful
|
||||
|
|
@ -218,7 +218,6 @@ class ServiceLogging(CustomLogger):
|
|||
self.prometheusServicesLogger = PrometheusServicesLogger()
|
||||
elif self.prometheusServicesLogger is None:
|
||||
self.prometheusServicesLogger = self.prometheusServicesLogger()
|
||||
return
|
||||
|
||||
async def init_datadog_logger_if_none(self):
|
||||
"""
|
||||
|
|
@ -230,8 +229,6 @@ class ServiceLogging(CustomLogger):
|
|||
if not hasattr(self, "dd_logger"):
|
||||
self.dd_logger: DataDogLogger = DataDogLogger()
|
||||
|
||||
return
|
||||
|
||||
async def init_otel_logger_if_none(self):
|
||||
"""
|
||||
initializes otel_logger if it is None or no attribute exists on ServiceLogging Object
|
||||
|
|
@ -246,18 +243,17 @@ class ServiceLogging(CustomLogger):
|
|||
verbose_logger.warning(
|
||||
"ServiceLogger: open_telemetry_logger is None or not an instance of OpenTelemetry"
|
||||
)
|
||||
return
|
||||
|
||||
async def async_service_failure_hook(
|
||||
self,
|
||||
service: ServiceTypes,
|
||||
duration: float,
|
||||
error: Union[str, Exception],
|
||||
error: str | Exception,
|
||||
call_type: str,
|
||||
parent_otel_span: Optional[Span] = None,
|
||||
start_time: Optional[Union[datetime, float]] = None,
|
||||
end_time: Optional[Union[float, datetime]] = None,
|
||||
event_metadata: Optional[dict] = None,
|
||||
parent_otel_span: Span | None = None,
|
||||
start_time: datetime | float | None = None,
|
||||
end_time: float | datetime | None = None,
|
||||
event_metadata: dict | None = None,
|
||||
):
|
||||
"""
|
||||
- For counting if the redis, postgres call is unsuccessful
|
||||
|
|
@ -324,7 +320,7 @@ class ServiceLogging(CustomLogger):
|
|||
request_data: dict,
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: Optional[str] = None,
|
||||
traceback_str: str | None = None,
|
||||
):
|
||||
"""
|
||||
Hook to track failed litellm-service calls
|
||||
|
|
@ -347,7 +343,7 @@ class ServiceLogging(CustomLogger):
|
|||
pass
|
||||
else:
|
||||
raise Exception(
|
||||
"Duration={} is not a float or timedelta object. type={}".format(_duration, type(_duration))
|
||||
f"Duration={_duration} is not a float or timedelta object. type={type(_duration)}"
|
||||
) # invalid _duration value
|
||||
# Batch polling callbacks (check_batch_cost) don't include call_type in kwargs.
|
||||
# Use .get() to avoid KeyError.
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ Custom A2A Card Resolver for LiteLLM.
|
|||
Extends the A2A SDK's card resolver to support multiple well-known paths.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Dict
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import LOCALHOST_URL_PATTERNS
|
||||
|
|
@ -114,7 +114,7 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
|
|||
async def get_agent_card(
|
||||
self,
|
||||
relative_card_path: str | None = None,
|
||||
http_kwargs: Dict[str, Any] | None = None,
|
||||
http_kwargs: dict[str, Any] | None = None,
|
||||
) -> "AgentCard":
|
||||
"""
|
||||
Fetch the agent card, trying multiple well-known paths.
|
||||
|
|
|
|||
|
|
@ -4,7 +4,8 @@ LiteLLM A2A Client class.
|
|||
Provides a class-based interface for A2A agent invocation.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, AsyncIterator, Dict, Optional
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from litellm.types.agents import LiteLLMSendMessageResponse
|
||||
|
||||
|
|
@ -50,7 +51,7 @@ class A2AClient:
|
|||
self,
|
||||
base_url: str,
|
||||
timeout: float = 60.0,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize the A2A client wrapper.
|
||||
|
|
@ -63,7 +64,7 @@ class A2AClient:
|
|||
self.base_url = base_url
|
||||
self.timeout = timeout
|
||||
self.extra_headers = extra_headers
|
||||
self._a2a_client: Optional["A2AClientType"] = None
|
||||
self._a2a_client: A2AClientType | None = None
|
||||
|
||||
async def _get_client(self) -> "A2AClientType":
|
||||
"""Get or create the underlying A2A client."""
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ Supports dynamic cost parameters that allow platform owners
|
|||
to define custom costs per agent query or per token.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
|
|
@ -18,7 +18,7 @@ else:
|
|||
class A2ACostCalculator:
|
||||
@staticmethod
|
||||
def calculate_a2a_cost(
|
||||
litellm_logging_obj: Optional[LitellmLoggingObject],
|
||||
litellm_logging_obj: LitellmLoggingObject | None,
|
||||
) -> float:
|
||||
"""
|
||||
Calculate the cost of an A2A send_message call.
|
||||
|
|
@ -73,8 +73,8 @@ class A2ACostCalculator:
|
|||
@staticmethod
|
||||
def _calculate_token_based_cost(
|
||||
model_call_details: dict,
|
||||
input_cost_per_token: Optional[float],
|
||||
output_cost_per_token: Optional[float],
|
||||
input_cost_per_token: float | None,
|
||||
output_cost_per_token: float | None,
|
||||
) -> float:
|
||||
"""
|
||||
Calculate cost based on token usage and per-token pricing.
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ A2A Protocol Exception Mapping Utils.
|
|||
Maps A2A SDK exceptions to LiteLLM A2A exception types.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.card_resolver import (
|
||||
|
|
@ -57,7 +57,7 @@ class A2AExceptionCheckers:
|
|||
return any(pattern in error_str_lower for pattern in CONNECTION_ERROR_PATTERNS)
|
||||
|
||||
@staticmethod
|
||||
def is_localhost_url(url: Optional[str]) -> bool:
|
||||
def is_localhost_url(url: str | None) -> bool:
|
||||
"""
|
||||
Check if a URL is a localhost/internal URL.
|
||||
|
||||
|
|
@ -96,9 +96,9 @@ class A2AExceptionCheckers:
|
|||
|
||||
def map_a2a_exception(
|
||||
original_exception: Exception,
|
||||
card_url: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
card_url: str | None = None,
|
||||
api_base: str | None = None,
|
||||
model: str | None = None,
|
||||
) -> Exception:
|
||||
"""
|
||||
Map an A2A SDK exception to a LiteLLM A2A exception type.
|
||||
|
|
|
|||
|
|
@ -4,8 +4,6 @@ A2A Protocol Exceptions.
|
|||
Custom exception types for A2A protocol operations, following LiteLLM's exception pattern.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
|
|
@ -21,11 +19,11 @@ class A2AError(Exception):
|
|||
message: str,
|
||||
status_code: int = 500,
|
||||
llm_provider: str = "a2a_agent",
|
||||
model: Optional[str] = None,
|
||||
response: Optional[httpx.Response] = None,
|
||||
litellm_debug_info: Optional[str] = None,
|
||||
max_retries: Optional[int] = None,
|
||||
num_retries: Optional[int] = None,
|
||||
model: str | None = None,
|
||||
response: httpx.Response | None = None,
|
||||
litellm_debug_info: str | None = None,
|
||||
max_retries: int | None = None,
|
||||
num_retries: int | None = None,
|
||||
):
|
||||
self.status_code = status_code
|
||||
self.message = f"litellm.A2AError: {message}"
|
||||
|
|
@ -65,12 +63,12 @@ class A2AConnectionError(A2AError):
|
|||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
url: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
response: Optional[httpx.Response] = None,
|
||||
litellm_debug_info: Optional[str] = None,
|
||||
max_retries: Optional[int] = None,
|
||||
num_retries: Optional[int] = None,
|
||||
url: str | None = None,
|
||||
model: str | None = None,
|
||||
response: httpx.Response | None = None,
|
||||
litellm_debug_info: str | None = None,
|
||||
max_retries: int | None = None,
|
||||
num_retries: int | None = None,
|
||||
):
|
||||
self.url = url
|
||||
super().__init__(
|
||||
|
|
@ -98,10 +96,10 @@ class A2AAgentCardError(A2AError):
|
|||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
url: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
response: Optional[httpx.Response] = None,
|
||||
litellm_debug_info: Optional[str] = None,
|
||||
url: str | None = None,
|
||||
model: str | None = None,
|
||||
response: httpx.Response | None = None,
|
||||
litellm_debug_info: str | None = None,
|
||||
):
|
||||
self.url = url
|
||||
super().__init__(
|
||||
|
|
@ -132,8 +130,8 @@ class A2ALocalhostURLError(A2AConnectionError):
|
|||
self,
|
||||
localhost_url: str,
|
||||
base_url: str,
|
||||
original_error: Optional[Exception] = None,
|
||||
model: Optional[str] = None,
|
||||
original_error: Exception | None = None,
|
||||
model: str | None = None,
|
||||
):
|
||||
self.localhost_url = localhost_url
|
||||
self.base_url = base_url
|
||||
|
|
|
|||
|
|
@ -16,8 +16,8 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
|
|||
)
|
||||
|
||||
__all__ = [
|
||||
"A2ACompletionBridgeTransformation",
|
||||
"A2ACompletionBridgeHandler",
|
||||
"A2ACompletionBridgeTransformation",
|
||||
"handle_a2a_completion",
|
||||
"handle_a2a_completion_streaming",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -10,7 +10,8 @@ A2A Streaming Events (in order):
|
|||
4. Status update (kind: "status-update") - Final status "completed" with final=true
|
||||
"""
|
||||
|
||||
from typing import Any, AsyncIterator, Dict, Optional
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -46,13 +47,13 @@ class A2ACompletionBridgeHandler:
|
|||
@staticmethod
|
||||
async def handle_non_streaming(
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
params: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
*,
|
||||
_skip_a2a_provider_routing: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Handle non-streaming A2A request via litellm.acompletion.
|
||||
|
||||
|
|
@ -105,7 +106,7 @@ class A2ACompletionBridgeHandler:
|
|||
verbose_logger.info(f"A2A completion bridge: model={full_model}, api_base={api_base}")
|
||||
|
||||
# Build completion params dict
|
||||
completion_params: Dict[str, Any] = {
|
||||
completion_params: dict[str, Any] = {
|
||||
"model": full_model,
|
||||
"messages": openai_messages,
|
||||
"api_base": api_base,
|
||||
|
|
@ -149,13 +150,13 @@ class A2ACompletionBridgeHandler:
|
|||
@staticmethod
|
||||
async def handle_streaming(
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
params: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
*,
|
||||
_skip_a2a_provider_routing: bool = False,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
"""
|
||||
Handle streaming A2A request via litellm.acompletion with stream=True.
|
||||
|
||||
|
|
@ -223,7 +224,7 @@ class A2ACompletionBridgeHandler:
|
|||
verbose_logger.info(f"A2A completion bridge streaming: model={full_model}, api_base={api_base}")
|
||||
|
||||
# Build completion params dict
|
||||
completion_params: Dict[str, Any] = {
|
||||
completion_params: dict[str, Any] = {
|
||||
"model": full_model,
|
||||
"messages": openai_messages,
|
||||
"api_base": api_base,
|
||||
|
|
@ -305,11 +306,11 @@ class A2ACompletionBridgeHandler:
|
|||
# Convenience functions that delegate to the class methods
|
||||
async def handle_a2a_completion(
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
params: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Convenience function for non-streaming A2A completion."""
|
||||
return await A2ACompletionBridgeHandler.handle_non_streaming(
|
||||
request_id=request_id,
|
||||
|
|
@ -322,11 +323,11 @@ async def handle_a2a_completion(
|
|||
|
||||
async def handle_a2a_completion_streaming(
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
params: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
"""Convenience function for streaming A2A completion."""
|
||||
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
|
||||
request_id=request_id,
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ A2A Streaming Events:
|
|||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -30,7 +30,7 @@ class A2AStreamingContext:
|
|||
Tracks task_id, context_id, and message accumulation.
|
||||
"""
|
||||
|
||||
def __init__(self, request_id: str, input_message: Dict[str, Any]):
|
||||
def __init__(self, request_id: str, input_message: dict[str, Any]):
|
||||
self.request_id = request_id
|
||||
self.task_id = str(uuid4())
|
||||
self.context_id = str(uuid4())
|
||||
|
|
@ -46,9 +46,9 @@ class A2ACompletionBridgeTransformation:
|
|||
"""
|
||||
|
||||
@staticmethod
|
||||
def _extract_text_from_a2a_parts(parts: List[Dict[str, Any]]) -> str:
|
||||
def _extract_text_from_a2a_parts(parts: list[dict[str, Any]]) -> str:
|
||||
"""Extract text from A2A parts (with or without explicit ``kind``)."""
|
||||
content_parts: List[str] = []
|
||||
content_parts: list[str] = []
|
||||
for part in parts:
|
||||
if not isinstance(part, dict):
|
||||
continue
|
||||
|
|
@ -62,16 +62,16 @@ class A2ACompletionBridgeTransformation:
|
|||
|
||||
@staticmethod
|
||||
def get_forward_metadata(
|
||||
a2a_message: Dict[str, Any],
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
a2a_message: dict[str, Any],
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Merge A2A metadata from MessageSendParams and the message for downstream providers.
|
||||
|
||||
Forwarded once on the LangGraph run payload (``metadata``), not duplicated on
|
||||
each input message — see ``apply_forward_metadata_to_completion_params``.
|
||||
"""
|
||||
merged: Dict[str, Any] = {}
|
||||
merged: dict[str, Any] = {}
|
||||
if params and isinstance(params.get("metadata"), dict):
|
||||
merged.update(params["metadata"])
|
||||
message_metadata = a2a_message.get("metadata")
|
||||
|
|
@ -81,9 +81,9 @@ class A2ACompletionBridgeTransformation:
|
|||
|
||||
@staticmethod
|
||||
def apply_forward_metadata_to_completion_params(
|
||||
completion_params: Dict[str, Any],
|
||||
a2a_message: Dict[str, Any],
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
completion_params: dict[str, Any],
|
||||
a2a_message: dict[str, Any],
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Attach A2A metadata to completion kwargs for provider bridges (e.g. LangGraph).
|
||||
|
|
@ -104,8 +104,8 @@ class A2ACompletionBridgeTransformation:
|
|||
# ``extra_body.metadata`` so the configured keys remain authoritative
|
||||
# and an A2A caller cannot overwrite server-set run metadata.
|
||||
existing_metadata = extra_body.get("metadata")
|
||||
existing_dict: Dict[str, Any] = existing_metadata if isinstance(existing_metadata, dict) else {}
|
||||
merged_metadata: Dict[str, Any] = {**forward_metadata, **existing_dict}
|
||||
existing_dict: dict[str, Any] = existing_metadata if isinstance(existing_metadata, dict) else {}
|
||||
merged_metadata: dict[str, Any] = {**forward_metadata, **existing_dict}
|
||||
extra_body = {**extra_body, "metadata": merged_metadata}
|
||||
completion_params["extra_body"] = extra_body
|
||||
|
||||
|
|
@ -113,8 +113,8 @@ class A2ACompletionBridgeTransformation:
|
|||
|
||||
@staticmethod
|
||||
def a2a_message_to_openai_messages(
|
||||
a2a_message: Dict[str, Any],
|
||||
) -> List[Dict[str, Any]]:
|
||||
a2a_message: dict[str, Any],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Transform an A2A message to OpenAI message format.
|
||||
|
||||
|
|
@ -143,7 +143,7 @@ class A2ACompletionBridgeTransformation:
|
|||
|
||||
# Do not attach A2A message.metadata here — the completion bridge forwards it
|
||||
# once at run level via extra_body.metadata (LangGraph POST /runs/wait shape).
|
||||
openai_message: Dict[str, Any] = {"role": openai_role, "content": content}
|
||||
openai_message: dict[str, Any] = {"role": openai_role, "content": content}
|
||||
|
||||
verbose_logger.debug(f"A2A -> OpenAI transform: role={role} -> {openai_role}, content_length={len(content)}")
|
||||
|
||||
|
|
@ -152,8 +152,8 @@ class A2ACompletionBridgeTransformation:
|
|||
@staticmethod
|
||||
def openai_response_to_a2a_response(
|
||||
response: Any,
|
||||
request_id: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
request_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Transform a LiteLLM ModelResponse to A2A SendMessageResponse format.
|
||||
|
||||
|
|
@ -198,7 +198,7 @@ class A2ACompletionBridgeTransformation:
|
|||
@staticmethod
|
||||
def create_task_event(
|
||||
ctx: A2AStreamingContext,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Create the initial task event with status 'submitted'.
|
||||
|
||||
|
|
@ -232,8 +232,8 @@ class A2ACompletionBridgeTransformation:
|
|||
ctx: A2AStreamingContext,
|
||||
state: str,
|
||||
final: bool = False,
|
||||
message_text: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
message_text: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Create a status update event.
|
||||
|
||||
|
|
@ -243,7 +243,7 @@ class A2ACompletionBridgeTransformation:
|
|||
final: Whether this is the final event
|
||||
message_text: Optional message text for 'working' status
|
||||
"""
|
||||
status: Dict[str, Any] = {
|
||||
status: dict[str, Any] = {
|
||||
"state": state,
|
||||
"timestamp": A2ACompletionBridgeTransformation._get_timestamp(),
|
||||
}
|
||||
|
|
@ -275,7 +275,7 @@ class A2ACompletionBridgeTransformation:
|
|||
def create_artifact_update_event(
|
||||
ctx: A2AStreamingContext,
|
||||
text: str,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Create an artifact update event with content.
|
||||
|
||||
|
|
|
|||
|
|
@ -12,14 +12,11 @@ Provides standalone functions with @client decorator for LiteLLM logging integra
|
|||
import asyncio
|
||||
import datetime
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator, Coroutine
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
AsyncIterator,
|
||||
Coroutine,
|
||||
Dict,
|
||||
Optional,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
|
|
@ -87,7 +84,7 @@ A2ACardResolver = LiteLLMA2ACardResolver
|
|||
|
||||
|
||||
def _set_usage_on_logging_obj(
|
||||
kwargs: Dict[str, Any],
|
||||
kwargs: dict[str, Any],
|
||||
prompt_tokens: int,
|
||||
completion_tokens: int,
|
||||
) -> None:
|
||||
|
|
@ -110,7 +107,7 @@ def _set_usage_on_logging_obj(
|
|||
|
||||
|
||||
def _set_agent_id_on_logging_obj(
|
||||
kwargs: Dict[str, Any],
|
||||
kwargs: dict[str, Any],
|
||||
agent_id: str | None,
|
||||
) -> None:
|
||||
"""
|
||||
|
|
@ -156,7 +153,7 @@ def _set_litellm_params_on_logging_obj(
|
|||
logging_obj.model_call_details["litellm_params"] = {**existing, **cost_params}
|
||||
|
||||
|
||||
def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str:
|
||||
def _get_a2a_model_info(a2a_client: Any, kwargs: dict[str, Any]) -> str:
|
||||
"""
|
||||
Extract agent info and set model/custom_llm_provider for cost tracking.
|
||||
|
||||
|
|
@ -199,8 +196,8 @@ async def _send_message_via_completion_bridge(
|
|||
request: "SendMessageRequest",
|
||||
custom_llm_provider: str,
|
||||
api_base: str | None,
|
||||
litellm_params: Dict[str, Any],
|
||||
agent_extra_headers: Dict[str, str] | None = None,
|
||||
litellm_params: dict[str, Any],
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> LiteLLMSendMessageResponse:
|
||||
"""
|
||||
Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore).
|
||||
|
|
@ -370,9 +367,9 @@ async def asend_message(
|
|||
a2a_client: Optional["A2AClientType"] = None,
|
||||
request: Optional["SendMessageRequest"] = None,
|
||||
api_base: str | None = None,
|
||||
litellm_params: Dict[str, Any] | None = None,
|
||||
litellm_params: dict[str, Any] | None = None,
|
||||
agent_id: str | None = None,
|
||||
agent_extra_headers: Dict[str, str] | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> LiteLLMSendMessageResponse:
|
||||
"""
|
||||
|
|
@ -453,7 +450,7 @@ async def asend_message(
|
|||
if api_base is None:
|
||||
raise ValueError("Either a2a_client or api_base is required for standard A2A flow")
|
||||
trace_id = trace_id or str(uuid.uuid4())
|
||||
extra_headers: Dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id}
|
||||
extra_headers: dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id}
|
||||
if agent_id:
|
||||
extra_headers["X-LiteLLM-Agent-Id"] = agent_id
|
||||
# Overlay agent-level headers (agent headers take precedence over LiteLLM internal ones)
|
||||
|
|
@ -518,7 +515,7 @@ def send_message(
|
|||
a2a_client: "A2AClientType",
|
||||
request: "SendMessageRequest",
|
||||
**kwargs: Any,
|
||||
) -> Union[LiteLLMSendMessageResponse, Coroutine[Any, Any, LiteLLMSendMessageResponse]]:
|
||||
) -> LiteLLMSendMessageResponse | Coroutine[Any, Any, LiteLLMSendMessageResponse]:
|
||||
"""
|
||||
Sync: Send a message to an A2A agent.
|
||||
|
||||
|
|
@ -547,9 +544,9 @@ def _build_streaming_logging_obj(
|
|||
request: "SendStreamingMessageRequest",
|
||||
agent_name: str,
|
||||
agent_id: str | None,
|
||||
litellm_params: Dict[str, Any] | None,
|
||||
metadata: Dict[str, Any] | None,
|
||||
proxy_server_request: Dict[str, Any] | None,
|
||||
litellm_params: dict[str, Any] | None,
|
||||
metadata: dict[str, Any] | None,
|
||||
proxy_server_request: dict[str, Any] | None,
|
||||
) -> Logging:
|
||||
"""Build logging object for streaming A2A requests."""
|
||||
start_time = datetime.datetime.now()
|
||||
|
|
@ -568,6 +565,7 @@ def _build_streaming_logging_obj(
|
|||
logging_obj.custom_llm_provider = "a2a_agent"
|
||||
logging_obj.model_call_details["model"] = model
|
||||
logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent"
|
||||
logging_obj.model_call_details["call_type"] = logging_obj.call_type
|
||||
if agent_id:
|
||||
logging_obj.model_call_details["agent_id"] = agent_id
|
||||
|
||||
|
|
@ -589,11 +587,11 @@ async def asend_message_streaming(
|
|||
a2a_client: Optional["A2AClientType"] = None,
|
||||
request: Optional["SendStreamingMessageRequest"] = None,
|
||||
api_base: str | None = None,
|
||||
litellm_params: Dict[str, Any] | None = None,
|
||||
litellm_params: dict[str, Any] | None = None,
|
||||
agent_id: str | None = None,
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
proxy_server_request: Dict[str, Any] | None = None,
|
||||
agent_extra_headers: Dict[str, str] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
proxy_server_request: dict[str, Any] | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
**kwargs: object,
|
||||
) -> AsyncIterator[Any]:
|
||||
"""
|
||||
|
|
@ -727,7 +725,7 @@ async def asend_message_streaming(
|
|||
async def create_a2a_client(
|
||||
base_url: str,
|
||||
timeout: float = DEFAULT_A2A_AGENT_TIMEOUT,
|
||||
extra_headers: Dict[str, str] | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
streaming: bool = False,
|
||||
) -> "A2AClientType":
|
||||
"""
|
||||
|
|
@ -808,7 +806,7 @@ async def create_a2a_client(
|
|||
async def aget_agent_card(
|
||||
base_url: str,
|
||||
timeout: float = DEFAULT_A2A_AGENT_TIMEOUT,
|
||||
extra_headers: Dict[str, str] | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
) -> "AgentCard":
|
||||
"""
|
||||
Fetch the agent card from an A2A agent.
|
||||
|
|
|
|||
|
|
@ -7,4 +7,4 @@ This module contains provider-specific implementations for the A2A protocol.
|
|||
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
|
||||
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
|
||||
|
||||
__all__ = ["BaseA2AProviderConfig", "A2AProviderConfigManager"]
|
||||
__all__ = ["A2AProviderConfigManager", "BaseA2AProviderConfig"]
|
||||
|
|
|
|||
|
|
@ -3,7 +3,8 @@ Base configuration for A2A protocol providers.
|
|||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, AsyncIterator, Dict, Optional
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
|
||||
class BaseA2AProviderConfig(ABC):
|
||||
|
|
@ -18,10 +19,10 @@ class BaseA2AProviderConfig(ABC):
|
|||
async def handle_non_streaming(
|
||||
self,
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
**kwargs,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Handle non-streaming A2A request.
|
||||
|
||||
|
|
@ -34,16 +35,15 @@ class BaseA2AProviderConfig(ABC):
|
|||
Returns:
|
||||
A2A SendMessageResponse dict
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def handle_streaming(
|
||||
self,
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
**kwargs,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
"""
|
||||
Handle streaming A2A request.
|
||||
|
||||
|
|
|
|||
|
|
@ -2,7 +2,8 @@
|
|||
Bedrock AgentCore A2A provider configuration.
|
||||
"""
|
||||
|
||||
from typing import Any, AsyncIterator, Dict, Optional
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
|
||||
from litellm.a2a_protocol.providers.bedrock_agentcore.handler import (
|
||||
|
|
@ -22,10 +23,10 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
|
|||
async def handle_non_streaming(
|
||||
self,
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
**kwargs,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""Handle non-streaming request to AgentCore A2A agent."""
|
||||
litellm_params = kwargs.get("litellm_params")
|
||||
if not litellm_params:
|
||||
|
|
@ -42,10 +43,10 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
|
|||
async def handle_streaming(
|
||||
self,
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
**kwargs,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
"""Handle streaming request to AgentCore A2A agent."""
|
||||
litellm_params = kwargs.get("litellm_params")
|
||||
if not litellm_params:
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue