mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge remote-tracking branch 'origin/main' into ci-fix-april7-2-fixes
# Conflicts: # litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
This commit is contained in:
commit
bba78cee7f
148 changed files with 10513 additions and 2235 deletions
|
|
@ -1330,6 +1330,57 @@ jobs:
|
|||
paths:
|
||||
- audio_coverage.xml
|
||||
- audio_coverage
|
||||
redis_caching_unit_tests:
|
||||
docker:
|
||||
- image: cimg/python:3.11
|
||||
auth:
|
||||
username: ${DOCKERHUB_USERNAME}
|
||||
password: ${DOCKERHUB_PASSWORD}
|
||||
working_directory: ~/project
|
||||
|
||||
steps:
|
||||
- checkout
|
||||
- setup_google_dns
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
python -m pip install --upgrade pip uv
|
||||
uv pip install --system -r requirements.txt
|
||||
pip install "pytest==7.3.1"
|
||||
pip install "pytest-retry==1.6.3"
|
||||
pip install "pytest-cov==5.0.0"
|
||||
pip install "pytest-asyncio==0.21.1"
|
||||
pip install "pytest-xdist==3.6.1"
|
||||
pip install "pytest-rerunfailures==14.0"
|
||||
# Run pytest and generate JUnit XML report
|
||||
- run:
|
||||
name: Run tests
|
||||
command: |
|
||||
pwd
|
||||
ls
|
||||
python -m pytest -vv \
|
||||
tests/local_testing/test_dual_cache.py \
|
||||
tests/local_testing/test_redis_batch_optimizations.py \
|
||||
tests/local_testing/test_router_utils.py \
|
||||
--cov=litellm --cov-report=xml \
|
||||
-x -s -v --junitxml=test-results/junit.xml \
|
||||
--durations=5 -n 2 \
|
||||
--reruns 2 --reruns-delay 1
|
||||
no_output_timeout: 20m
|
||||
- run:
|
||||
name: Rename the coverage files
|
||||
command: |
|
||||
mv coverage.xml redis_caching_coverage.xml
|
||||
mv .coverage redis_caching_coverage
|
||||
|
||||
# Store test results
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
- persist_to_workspace:
|
||||
root: .
|
||||
paths:
|
||||
- redis_caching_coverage.xml
|
||||
- redis_caching_coverage
|
||||
installing_litellm_on_python:
|
||||
docker:
|
||||
- image: cimg/python:3.11
|
||||
|
|
@ -2889,7 +2940,7 @@ jobs:
|
|||
python -m venv venv
|
||||
. venv/bin/activate
|
||||
pip install coverage
|
||||
coverage combine realtime_translation_coverage ocr_coverage search_coverage mcp_coverage litellm_mcps_tests_coverage logging_coverage audio_coverage local_testing_part1_coverage local_testing_part2_coverage pass_through_unit_tests_coverage batches_coverage guardrails_coverage
|
||||
coverage combine realtime_translation_coverage ocr_coverage search_coverage mcp_coverage litellm_mcps_tests_coverage logging_coverage audio_coverage local_testing_part1_coverage local_testing_part2_coverage pass_through_unit_tests_coverage batches_coverage guardrails_coverage redis_caching_coverage
|
||||
coverage xml
|
||||
- codecov/upload:
|
||||
file: ./coverage.xml
|
||||
|
|
@ -3074,6 +3125,117 @@ jobs:
|
|||
CI=true npm run test -- --run \
|
||||
--pool forks --poolOptions.forks.maxForks=8
|
||||
|
||||
e2e_ui_testing:
|
||||
docker:
|
||||
- image: cimg/python:3.12-browsers
|
||||
auth:
|
||||
username: ${DOCKERHUB_USERNAME}
|
||||
password: ${DOCKERHUB_PASSWORD}
|
||||
- image: cimg/postgres:16.0
|
||||
environment:
|
||||
POSTGRES_USER: e2euser
|
||||
POSTGRES_PASSWORD: e2epassword
|
||||
POSTGRES_DB: litellm_e2e
|
||||
resource_class: large
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
DATABASE_URL: "postgresql://e2euser:e2epassword@localhost:5432/litellm_e2e"
|
||||
CI: "true"
|
||||
steps:
|
||||
- checkout
|
||||
- setup_google_dns
|
||||
- restore_cache:
|
||||
keys:
|
||||
- ui-e2e-py-deps-v1-{{ checksum "requirements.txt" }}
|
||||
- run:
|
||||
name: Install Python dependencies
|
||||
command: |
|
||||
python -m pip install --upgrade pip uv
|
||||
uv pip install --system -r requirements.txt
|
||||
pip install "prisma==0.11.0"
|
||||
prisma generate --schema litellm/proxy/schema.prisma
|
||||
- save_cache:
|
||||
key: ui-e2e-py-deps-v1-{{ checksum "requirements.txt" }}
|
||||
paths:
|
||||
- ~/.local/lib
|
||||
- ~/.local/bin
|
||||
- restore_cache:
|
||||
keys:
|
||||
- ui-e2e-node-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
- run:
|
||||
name: Install Node dependencies and Playwright
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm ci
|
||||
npx playwright install chromium --with-deps
|
||||
- save_cache:
|
||||
key: ui-e2e-node-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
paths:
|
||||
- ui/litellm-dashboard/node_modules
|
||||
- run:
|
||||
name: Build UI from source
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm run build
|
||||
cp -r out/ ../../litellm/proxy/_experimental/out/
|
||||
# Restructure HTML so extensionless routes work (login.html -> login/index.html)
|
||||
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do
|
||||
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
|
||||
done
|
||||
- run:
|
||||
name: Wait for PostgreSQL
|
||||
command: dockerize -wait tcp://localhost:5432 -timeout 30s
|
||||
- run:
|
||||
name: Push Prisma schema
|
||||
command: prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
|
||||
- run:
|
||||
name: Seed database
|
||||
command: |
|
||||
PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \
|
||||
-f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql
|
||||
- run:
|
||||
name: Start mock LLM server
|
||||
command: python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py
|
||||
background: true
|
||||
- run:
|
||||
name: Start LiteLLM proxy
|
||||
environment:
|
||||
LITELLM_MASTER_KEY: "sk-1234"
|
||||
MOCK_LLM_URL: "http://127.0.0.1:8090/v1"
|
||||
DISABLE_SCHEMA_UPDATE: "true"
|
||||
SERVER_ROOT_PATH: ""
|
||||
PROXY_LOGOUT_URL: ""
|
||||
command: |
|
||||
python -m litellm.proxy.proxy_cli \
|
||||
--config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \
|
||||
--port 4000
|
||||
background: true
|
||||
- run:
|
||||
name: Wait for proxy to be ready
|
||||
command: |
|
||||
for i in $(seq 1 60); do
|
||||
HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" http://127.0.0.1:4000/health -H "Authorization: Bearer sk-1234" 2>/dev/null || true)
|
||||
if [ "$HTTP_CODE" = "200" ]; then
|
||||
echo "Proxy is ready"
|
||||
exit 0
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
echo "Proxy failed to start"
|
||||
exit 1
|
||||
- run:
|
||||
name: Run Playwright E2E tests
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npx playwright test --config e2e_tests/playwright.config.ts
|
||||
no_output_timeout: 10m
|
||||
- store_artifacts:
|
||||
path: ui/litellm-dashboard/test-results
|
||||
destination: e2e-test-results
|
||||
- store_artifacts:
|
||||
path: ui/litellm-dashboard/playwright-report
|
||||
destination: e2e-playwright-report
|
||||
|
||||
build_docker_database_image:
|
||||
machine:
|
||||
image: ubuntu-2204:2024.04.1
|
||||
|
|
@ -3099,80 +3261,6 @@ jobs:
|
|||
paths:
|
||||
- litellm-docker-database.tar.zst
|
||||
|
||||
e2e_ui_testing:
|
||||
machine:
|
||||
image: ubuntu-2204:2023.10.1
|
||||
resource_class: large
|
||||
working_directory: ~/project
|
||||
parameters:
|
||||
browser:
|
||||
type: string
|
||||
steps:
|
||||
- checkout
|
||||
- setup_google_dns
|
||||
- attach_workspace:
|
||||
at: ~/project
|
||||
- run:
|
||||
name: Load Docker Database Image
|
||||
command: |
|
||||
zstd -d litellm-docker-database.tar.zst --stdout | docker load
|
||||
docker images | grep litellm-docker-database
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
npm install -D @playwright/test
|
||||
- run:
|
||||
name: Install Playwright Browsers
|
||||
command: |
|
||||
npx playwright install
|
||||
- run:
|
||||
name: Run Docker container
|
||||
command: |
|
||||
docker run -d \
|
||||
-p 4000:4000 \
|
||||
-e DATABASE_URL=$E2E_UI_TEST_DATABASE_URL \
|
||||
-e LITELLM_MASTER_KEY="sk-1234" \
|
||||
-e OPENAI_API_KEY=$OPENAI_API_KEY \
|
||||
-e UI_USERNAME="admin" \
|
||||
-e UI_PASSWORD="gm" \
|
||||
-e LITELLM_LICENSE=$LITELLM_LICENSE \
|
||||
--name litellm-docker-database-<< parameters.browser >> \
|
||||
-v $(pwd)/litellm/proxy/example_config_yaml/simple_config.yaml:/app/config.yaml \
|
||||
litellm-docker-database:ci \
|
||||
--config /app/config.yaml \
|
||||
--port 4000 \
|
||||
--detailed_debug
|
||||
- run:
|
||||
name: Install curl and dockerize
|
||||
command: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y curl
|
||||
sudo wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz
|
||||
sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz
|
||||
sudo rm dockerize-linux-amd64-v0.6.1.tar.gz
|
||||
- run:
|
||||
name: Start outputting logs
|
||||
command: docker logs -f litellm-docker-database-<< parameters.browser >>
|
||||
background: true
|
||||
- run:
|
||||
name: Wait for app to be ready
|
||||
command: dockerize -wait http://localhost:4000 -timeout 5m
|
||||
- run:
|
||||
name: Run Playwright Tests
|
||||
command: |
|
||||
npx playwright test \
|
||||
--project << parameters.browser >> \
|
||||
--config ui/litellm-dashboard/e2e_tests/playwright.config.ts \
|
||||
--reporter=html \
|
||||
--output=test-results
|
||||
no_output_timeout: 15m
|
||||
- store_artifacts:
|
||||
path: test-results
|
||||
destination: playwright-results
|
||||
|
||||
- store_artifacts:
|
||||
path: playwright-report
|
||||
destination: playwright-report
|
||||
|
||||
prisma_schema_sync:
|
||||
machine:
|
||||
|
|
@ -3401,32 +3489,12 @@ workflows:
|
|||
only:
|
||||
- main
|
||||
- /litellm_.*/
|
||||
# - e2e_ui_testing:
|
||||
# name: e2e_ui_testing_chromium
|
||||
# browser: chromium
|
||||
# context: e2e_ui_tests
|
||||
# requires:
|
||||
# - ui_build
|
||||
# - build_docker_database_image
|
||||
# - prisma_schema_sync
|
||||
# filters:
|
||||
# branches:
|
||||
# only:
|
||||
# - main
|
||||
# - /litellm_.*/
|
||||
# - e2e_ui_testing:
|
||||
# name: e2e_ui_testing_firefox
|
||||
# browser: firefox
|
||||
# context: e2e_ui_tests
|
||||
# requires:
|
||||
# - ui_build
|
||||
# - build_docker_database_image
|
||||
# - prisma_schema_sync
|
||||
# filters:
|
||||
# branches:
|
||||
# only:
|
||||
# - main
|
||||
# - /litellm_.*/
|
||||
- e2e_ui_testing:
|
||||
filters:
|
||||
branches:
|
||||
only:
|
||||
- main
|
||||
- /litellm_.*/
|
||||
- build_and_test:
|
||||
requires:
|
||||
- build_docker_database_image
|
||||
|
|
@ -3615,6 +3683,12 @@ workflows:
|
|||
only:
|
||||
- main
|
||||
- /litellm_.*/
|
||||
- redis_caching_unit_tests:
|
||||
filters:
|
||||
branches:
|
||||
only:
|
||||
- main
|
||||
- /litellm_.*/
|
||||
- upload-coverage:
|
||||
requires:
|
||||
- realtime_translation_testing
|
||||
|
|
@ -3633,6 +3707,7 @@ workflows:
|
|||
- image_gen_testing
|
||||
- logging_testing
|
||||
- audio_testing
|
||||
- redis_caching_unit_tests
|
||||
- langfuse_logging_unit_tests
|
||||
- local_testing_part1
|
||||
- local_testing_part2
|
||||
|
|
|
|||
7
.github/pull_request_template.md
vendored
7
.github/pull_request_template.md
vendored
|
|
@ -32,6 +32,13 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
|
|||
- [ ] **Merge / cherry-pick CI run**
|
||||
Links:
|
||||
|
||||
## Screenshots / Proof of Fix
|
||||
|
||||
<!-- Include screenshots, screen recordings, or log output demonstrating that your changes work as expected.
|
||||
For bug fixes: show reproduction before the fix and passing behavior after.
|
||||
For new features: show the feature working end-to-end.
|
||||
For UI changes: include before/after screenshots. -->
|
||||
|
||||
## Type
|
||||
|
||||
<!-- Select the type of Pull Request -->
|
||||
|
|
|
|||
17
.github/workflows/_test-unit-services-base.yml
vendored
17
.github/workflows/_test-unit-services-base.yml
vendored
|
|
@ -27,11 +27,6 @@ on:
|
|||
required: false
|
||||
type: number
|
||||
default: 10
|
||||
enable-redis:
|
||||
description: "Pass Redis Cloud credentials to tests via REDIS_HOST/PORT/PASSWORD env vars"
|
||||
required: false
|
||||
type: boolean
|
||||
default: false
|
||||
enable-postgres:
|
||||
description: "Start a local Postgres service container and run Prisma migrations"
|
||||
required: false
|
||||
|
|
@ -43,12 +38,6 @@ on:
|
|||
type: string
|
||||
default: "run"
|
||||
secrets:
|
||||
REDIS_HOST:
|
||||
required: false
|
||||
REDIS_PORT:
|
||||
required: false
|
||||
REDIS_PASSWORD:
|
||||
required: false
|
||||
DATABASE_URL:
|
||||
required: false
|
||||
POSTGRES_USER:
|
||||
|
|
@ -66,11 +55,8 @@ jobs:
|
|||
timeout-minutes: ${{ inputs.timeout-minutes }}
|
||||
# Environment is derived from the enable-* flags, not caller-controllable.
|
||||
# This prevents callers from passing arbitrary environment names to bypass secret scoping.
|
||||
# Note: Postgres service container always starts (GHA limitation), so any Redis job
|
||||
# also needs Postgres secrets → uses integration-redis-postgres, not integration-redis.
|
||||
environment: >-
|
||||
${{
|
||||
inputs.enable-redis && 'integration-redis-postgres' ||
|
||||
inputs.enable-postgres && 'integration-postgres' ||
|
||||
''
|
||||
}}
|
||||
|
|
@ -146,9 +132,6 @@ jobs:
|
|||
WORKERS: ${{ inputs.workers }}
|
||||
RERUNS: ${{ inputs.reruns }}
|
||||
DATABASE_URL: ${{ inputs.enable-postgres && secrets.DATABASE_URL || '' }}
|
||||
REDIS_HOST: ${{ inputs.enable-redis && secrets.REDIS_HOST || '' }}
|
||||
REDIS_PORT: ${{ inputs.enable-redis && secrets.REDIS_PORT || '' }}
|
||||
REDIS_PASSWORD: ${{ inputs.enable-redis && secrets.REDIS_PASSWORD || '' }}
|
||||
run: |
|
||||
if [ "${WORKERS}" = "0" ]; then
|
||||
poetry run pytest ${TEST_PATH:?} \
|
||||
|
|
|
|||
41
.github/workflows/test-unit-caching-redis.yml
vendored
41
.github/workflows/test-unit-caching-redis.yml
vendored
|
|
@ -1,41 +0,0 @@
|
|||
name: "Unit Tests: Caching (Redis)"
|
||||
|
||||
# Uses cloud Redis credentials — only runs on trusted branches, not PRs.
|
||||
# This prevents external PRs from accessing Redis credentials.
|
||||
on:
|
||||
push:
|
||||
branches: [main, "litellm_*"]
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
caching-redis:
|
||||
uses: ./.github/workflows/_test-unit-services-base.yml
|
||||
with:
|
||||
# Redis-only tests that do NOT require provider API keys.
|
||||
# Tests needing API keys (test_caching.py, test_caching_ssl.py, test_prometheus_service.py,
|
||||
# test_router_caching.py) are in Phase 3 integration workflows.
|
||||
test-path: >-
|
||||
tests/local_testing/test_dual_cache.py
|
||||
tests/local_testing/test_redis_batch_optimizations.py
|
||||
tests/local_testing/test_router_utils.py
|
||||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
enable-redis: true
|
||||
enable-postgres: false
|
||||
artifact-name: caching-redis
|
||||
secrets:
|
||||
REDIS_HOST: ${{ secrets.REDIS_HOST }}
|
||||
REDIS_PORT: ${{ secrets.REDIS_PORT }}
|
||||
REDIS_PASSWORD: ${{ secrets.REDIS_PASSWORD }}
|
||||
DATABASE_URL: ${{ secrets.DATABASE_URL }}
|
||||
POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
|
||||
POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
|
||||
1
.github/workflows/test-unit-proxy-db.yml
vendored
1
.github/workflows/test-unit-proxy-db.yml
vendored
|
|
@ -41,7 +41,6 @@ jobs:
|
|||
workers: ${{ matrix.workers }}
|
||||
reruns: 2
|
||||
timeout-minutes: ${{ matrix.timeout }}
|
||||
enable-redis: false
|
||||
enable-postgres: true
|
||||
artifact-name: proxy-db-${{ matrix.test-group }}
|
||||
secrets:
|
||||
|
|
|
|||
1
.github/workflows/test-unit-security.yml
vendored
1
.github/workflows/test-unit-security.yml
vendored
|
|
@ -22,7 +22,6 @@ jobs:
|
|||
workers: 1
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
enable-redis: false
|
||||
enable-postgres: true
|
||||
artifact-name: security
|
||||
secrets:
|
||||
|
|
|
|||
|
|
@ -71,8 +71,16 @@ WORKDIR /app
|
|||
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh
|
||||
RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
|
||||
|
||||
# Run as non-root user
|
||||
RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser \
|
||||
&& chown -R appuser:appuser /app
|
||||
USER appuser
|
||||
|
||||
# Expose the necessary port
|
||||
EXPOSE 4000/tcp
|
||||
|
||||
HEALTHCHECK --interval=30s --timeout=5s --start-period=10s --retries=3 \
|
||||
CMD ["python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:4000/health')"]
|
||||
|
||||
# Override the CMD instruction with your desired command and arguments
|
||||
CMD ["--port", "4000", "--config", "config.yaml", "--detailed_debug"]
|
||||
|
|
@ -13,12 +13,12 @@ RUN pip install --no-cache-dir -r requirements.txt
|
|||
RUN chmod +x /app/health_check_client.py
|
||||
|
||||
# Run as non-root user
|
||||
RUN adduser --disabled-password --gecos "" --uid 1001 healthcheck
|
||||
USER healthcheck
|
||||
RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser
|
||||
USER appuser
|
||||
|
||||
# Health check
|
||||
HEALTHCHECK --interval=30s --timeout=5s --retries=3 \
|
||||
CMD python /app/health_check_client.py --help || exit 1
|
||||
HEALTHCHECK --interval=30s --timeout=5s --start-period=5s --retries=3 \
|
||||
CMD ["python", "/app/health_check_client.py", "--help"]
|
||||
|
||||
# Set entrypoint
|
||||
ENTRYPOINT ["python", "/app/health_check_client.py"]
|
||||
|
|
|
|||
|
|
@ -13,19 +13,19 @@ To build and run the application, you will use the `docker-compose.yml` file loc
|
|||
|
||||
### 1. Set the Master Key
|
||||
|
||||
The application requires a `MASTER_KEY` for signing and validating tokens. You must set this key as an environment variable before running the application.
|
||||
The application requires a `LITELLM_MASTER_KEY` for signing and validating tokens. You must set this key as an environment variable before running the application.
|
||||
|
||||
Create a `.env` file in the root of the project and add the following line:
|
||||
|
||||
```
|
||||
MASTER_KEY=your-secret-key
|
||||
LITELLM_MASTER_KEY=your-secret-key
|
||||
```
|
||||
|
||||
Replace `your-secret-key` with a strong, randomly generated secret.
|
||||
|
||||
### 2. Build and Run the Containers
|
||||
|
||||
Once you have set the `MASTER_KEY`, you can build and run the containers using the following command:
|
||||
Once you have set the `LITELLM_MASTER_KEY`, you can build and run the containers using the following command:
|
||||
|
||||
```bash
|
||||
docker compose up -d --build
|
||||
|
|
@ -89,4 +89,4 @@ This command should succeed (showing engine versions) even with `--network none`
|
|||
## Troubleshooting
|
||||
|
||||
- **`build_admin_ui.sh: not found`**: This error can occur if the Docker build context is not set correctly. Ensure that you are running the `docker-compose` command from the root of the project.
|
||||
- **`Master key is not initialized`**: This error means the `MASTER_key` environment variable is not set. Make sure you have created a `.env` file in the project root with the `MASTER_KEY` defined.
|
||||
- **`Master key is not initialized`**: This error means the `LITELLM_MASTER_KEY` environment variable is not set. Make sure you have created a `.env` file in the project root with the `LITELLM_MASTER_KEY` defined.
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ title: "April Townhall: Security + Product Roadmap"
|
|||
date: 2026-04-02T07:30:00
|
||||
authors:
|
||||
- krrish
|
||||
- ishaan-alt
|
||||
description: "Join the LiteLLM April townhall on Friday, 10 April at 7:30 AM to learn about LiteLLM's security and product roadmap."
|
||||
tags: [announcement, townhall]
|
||||
hide_table_of_contents: true
|
||||
|
|
|
|||
162
docs/my-website/blog/april_townhall_updates/index.md
Normal file
162
docs/my-website/blog/april_townhall_updates/index.md
Normal file
|
|
@ -0,0 +1,162 @@
|
|||
---
|
||||
slug: april-townhall-updates
|
||||
title: "April Townhall Updates: CI/CD v2, Stability, and Product Roadmap"
|
||||
date: 2026-04-10T12:00:00
|
||||
authors:
|
||||
- krrish
|
||||
- ishaan-alt
|
||||
description: "A recap of the April LiteLLM town hall covering CI/CD v2, product stability work, and the near-term roadmap."
|
||||
tags: [townhall, security, reliability, product]
|
||||
hide_table_of_contents: false
|
||||
---
|
||||
|
||||
import Image from '@theme/IdealImage';
|
||||
|
||||
Thank you to everyone who joined our April town hall.
|
||||
|
||||
We used the session to share our CI/CD v2 improvements, product stability work, and what we are prioritizing next across reliability and product roadmap.
|
||||
|
||||
{/* truncate */}
|
||||
|
||||
## CI/CD v2 improvements
|
||||
|
||||
Our CI/CD v2 work is centered around four goals:
|
||||
|
||||
1. **Limit** what each package can access
|
||||
2. **Reduce** the number of sensitive environment variables
|
||||
3. **Avoid** compromised packages
|
||||
4. **Reduce the risk of** release tampering
|
||||
|
||||
#### New architecture: isolated environments
|
||||
|
||||
We have begun moving to isolated environments for distinct CI/CD stages to reduce the chance that a single compromised step can inherit broad access across the entire pipeline.
|
||||
|
||||
<Image
|
||||
img={require('../../img/april_townhall_isolated_environments.png')}
|
||||
style={{width: '900px', height: 'auto', display: 'block'}}
|
||||
/>
|
||||
|
||||
#### Current rollout status
|
||||
|
||||
These changes are deployed in our current release workflow. [See here](https://github.com/BerriAI/litellm/tags)
|
||||
|
||||
#### Independently verify releases
|
||||
|
||||
A key part of CI/CD v2 is supporting independent verification of release artifacts using our published verification process, while reducing reliance on any single credential or release path.
|
||||
|
||||
[**Learn more about how to verify releases**](https://docs.litellm.ai/docs/proxy/docker_image_security)
|
||||
|
||||
<Image
|
||||
img={require('../../img/verify_releases.png')}
|
||||
style={{width: '900px', height: 'auto', display: 'block'}}
|
||||
/>
|
||||
|
||||
## Stability improvements
|
||||
|
||||
### SDLC improvements
|
||||
|
||||
This month, we're focusing on process stability improvements around:
|
||||
- Improving main-branch stability
|
||||
- Mapping UI QA to built Docker images for 1:1 environment parity
|
||||
- Consistent release tags across PyPI and Docker
|
||||
- Fixing release notes publication
|
||||
|
||||
#### Improving main-branch stability
|
||||
|
||||
We're introducing a staging-gated flow:
|
||||
|
||||
<Image
|
||||
img={require('../../img/stable_main.png')}
|
||||
style={{width: '900px', height: 'auto', display: 'block'}}
|
||||
/>
|
||||
|
||||
- Only an internal staging branch can push to `main`.
|
||||
- PRs to that staging branch must pass CircleCI LLM API testing.
|
||||
- Collision handling happens on staging, which is designed to reduce unstable changes reaching `main`.
|
||||
|
||||
#### UI QA in Docker environment
|
||||
|
||||
Moving forward, all UI QA will be performed in the built Docker image that users run.
|
||||
|
||||
Previously, some UI QA paths were run in local environments that did not fully replicate Docker runtime conditions.
|
||||
|
||||
That contributed to release-specific issues, including MCP registration problems in `v1.82.3`.
|
||||
|
||||
#### Consistent release tags
|
||||
|
||||
Today we publish releases for multiple scenarios:
|
||||
- Dev (Built of a PR for a customer-specific scenario)
|
||||
- Nightly (Passes all CI/CD checks)
|
||||
- Release Candidate (Passes all CI/CD checks + manual UI QA)
|
||||
- Stable (intended to pass all CI/CD checks + manual UI QA + 7 days of production testing)
|
||||
|
||||
We are targeting a consistent naming convention across PyPI and Docker by the end of April.
|
||||
|
||||
#### Release notes
|
||||
|
||||
CI/CD v2 changes moved release notes to a manual path. This is a temporary solution while we investigate a better automated workflow. We are targeting a more consistent process by the end of April.
|
||||
|
||||
### Product stability improvements
|
||||
|
||||
#### Stable Prisma migrations
|
||||
|
||||
Today, we have observed several migration failure classes:
|
||||
- Migration not applied
|
||||
- Migration marked applied but incomplete
|
||||
- Migration not applied due to non-root image issues
|
||||
|
||||
We're prioritizing this work this month and have assigned an engineering owner to the effort. Our target is to resolve these error classes by the end of April.
|
||||
|
||||
#### UI type safety
|
||||
|
||||
Another area of focus is improving the stability of the UI. Today, one cause of errors is that the UI maintains its own assumptions about backend API types. This can lead to issues when backend responses differ from UI assumptions.
|
||||
|
||||
We aim to move to having the UI and Backend be in sync with each other, and are exploring OpenAPI-driven mapping to achieve this.
|
||||
|
||||
## Product roadmap
|
||||
|
||||
### Our Assumptions
|
||||
|
||||
Over the next few years, we expect:
|
||||
- Companies will give employees more AI tools.
|
||||
- More AI agents will move into production workflows across HR, finance, support, and operations.
|
||||
|
||||
### Our Inferences
|
||||
#### Near-term
|
||||
|
||||
- AI spend will increase.
|
||||
- Uptime and latency will become even more important.
|
||||
- More AI resources (skills, CLIs, and related assets) will require governance.
|
||||
- Agent and MCP usage patterns will require deeper controls.
|
||||
- Broader developer adoption will increase the need for simpler, more discoverable tooling.
|
||||
|
||||
#### Long-term
|
||||
|
||||
- We expect many organizations to treat agent auditability (how decisions were made across LLM + MCP + sub-agent inputs/outputs) as a compliance expectation.
|
||||
- Permission management will get more complex as user-agent interaction chains deepen.
|
||||
|
||||
Roadmap timelines in this post are targets and may evolve based on validation and user feedback.
|
||||
|
||||
## April investments
|
||||
|
||||
### Reliability
|
||||
|
||||
- Increase uptime for 10k+ RPS scenarios.
|
||||
- Investigate latency overhead for long-running Claude Code requests.
|
||||
|
||||
### Feature reliability
|
||||
|
||||
- Polish MCP authentication.
|
||||
- Better understand how teams are using agents through LiteLLM.
|
||||
|
||||
### Governance
|
||||
|
||||
- Launch Skills as a first-class citizen in LiteLLM.
|
||||
|
||||
## Q&A
|
||||
|
||||
Thank you again for all the questions and direct feedback. We will keep sharing concrete progress updates as these efforts ship.
|
||||
|
||||
## Hiring
|
||||
|
||||
We are actively hiring across several roles, please apply [here](https://jobs.ashbyhq.com/litellm) if you're interested!
|
||||
|
|
@ -24,7 +24,7 @@ ishaan:
|
|||
|
||||
# Alias for typo in name
|
||||
ishaan-alt:
|
||||
name: Ishaan Jaff
|
||||
name: Ishaan Jaffer
|
||||
title: CTO, LiteLLM
|
||||
url: https://www.linkedin.com/in/reffajnaahsi/
|
||||
image_url: https://pbs.twimg.com/profile_images/1613813310264340481/lz54oEiB_400x400.jpg
|
||||
|
|
|
|||
131
docs/my-website/docs/observability/ramp_integration.md
Normal file
131
docs/my-website/docs/observability/ramp_integration.md
Normal file
|
|
@ -0,0 +1,131 @@
|
|||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
# Ramp
|
||||
|
||||
Send AI usage and cost data to Ramp for automated spend tracking.
|
||||
|
||||
[Ramp](https://ramp.com/) is a finance automation platform that helps businesses manage expenses, corporate cards, and vendor payments. With the Ramp callback integration, your LiteLLM AI usage — including token counts, model costs, and request metadata — is automatically sent to Ramp for real-time spend visibility.
|
||||
|
||||
:::info
|
||||
We want to learn how we can make the callbacks better! Meet the LiteLLM [founders](https://calendly.com/d/4mp-gd3-k5k/berriai-1-1-onboarding-litellm-hosted-version) or
|
||||
join our [discord](https://discord.gg/wuPM9dRgDw)
|
||||
:::
|
||||
|
||||
## Pre-Requisites
|
||||
|
||||
1. Log in to [Ramp](https://app.ramp.com/) and search for **"LiteLLM"** using the search bar. Click the **LiteLLM** integration result.
|
||||
|
||||
> **Note:** Only business owners and admins can access and configure integrations.
|
||||
|
||||
2. On the LiteLLM integration page, click the **Connect** button in the top right.
|
||||
|
||||
3. In the Connect LiteLLM drawer, click **Generate API Key** to create an API key.
|
||||
|
||||
> **Important:** Copy the API key immediately — it won't be shown again. If you lose it, you can revoke the existing key and generate a new one from the integration settings.
|
||||
|
||||
```shell
|
||||
pip install litellm
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
Set your `RAMP_API_KEY` and add `"ramp"` to your callbacks to start logging LLM usage to Ramp.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="SDK">
|
||||
|
||||
```python
|
||||
litellm.callbacks = ["ramp"]
|
||||
```
|
||||
|
||||
```python
|
||||
import litellm
|
||||
import os
|
||||
|
||||
# Ramp API Key
|
||||
os.environ["RAMP_API_KEY"] = "your-ramp-api-key"
|
||||
|
||||
# LLM API Keys
|
||||
os.environ['OPENAI_API_KEY'] = ""
|
||||
|
||||
# Set ramp as a callback
|
||||
litellm.callbacks = ["ramp"]
|
||||
|
||||
# OpenAI call
|
||||
response = litellm.completion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[
|
||||
{"role": "user", "content": "Hi - I'm testing Ramp integration"}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="proxy" label="LiteLLM Proxy">
|
||||
|
||||
1. Setup config.yaml
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-3.5-turbo
|
||||
litellm_params:
|
||||
model: openai/gpt-3.5-turbo
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
litellm_settings:
|
||||
callbacks: ["ramp"]
|
||||
|
||||
environment_variables:
|
||||
RAMP_API_KEY: os.environ/RAMP_API_KEY
|
||||
```
|
||||
|
||||
2. Start LiteLLM Proxy
|
||||
|
||||
```bash
|
||||
litellm --config /path/to/config.yaml
|
||||
```
|
||||
|
||||
3. Test it!
|
||||
|
||||
```bash
|
||||
curl -L -X POST 'http://0.0.0.0:4000/chat/completions' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-H 'Authorization: Bearer sk-1234' \
|
||||
-d '{
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Hey, how are you?"
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## What Data is Logged?
|
||||
|
||||
LiteLLM sends the [Standard Logging Payload](https://docs.litellm.ai/docs/proxy/logging_spec) to Ramp on successful LLM API calls, which includes:
|
||||
|
||||
- **Request details**: Model, messages, parameters
|
||||
- **Response details**: Completion text, token usage, latency
|
||||
- **Metadata**: User ID, custom metadata, timestamps
|
||||
- **Cost tracking**: Response cost based on token usage
|
||||
|
||||
## Authentication
|
||||
|
||||
Set the `RAMP_API_KEY` environment variable with your Ramp API key.
|
||||
|
||||
| Environment Variable | Description |
|
||||
|---|---|
|
||||
| `RAMP_API_KEY` | Your Ramp API key (required) |
|
||||
|
||||
## Support & Talk to Founders
|
||||
|
||||
- [Schedule Demo 👋](https://calendly.com/d/4mp-gd3-k5k/berriai-1-1-onboarding-litellm-hosted-version)
|
||||
- [Community Discord 💭](https://discord.gg/wuPM9dRgDw)
|
||||
- Our numbers 📞 +1 (770) 8783-106 / +1 (412) 618-6238
|
||||
- Our emails ✉️ ishaan@berri.ai / krrish@berri.ai
|
||||
|
|
@ -65,14 +65,13 @@ response = completion(
|
|||
- modalities
|
||||
- reasoning_content
|
||||
- audio (for TTS models only)
|
||||
- service_tier
|
||||
|
||||
**Anthropic Params**
|
||||
- thinking (used to set max budget tokens across anthropic/gemini models)
|
||||
|
||||
[**See Updated List**](https://github.com/BerriAI/litellm/blob/main/litellm/llms/gemini/chat/transformation.py#L70)
|
||||
|
||||
|
||||
|
||||
## Usage - Thinking / `reasoning_content`
|
||||
|
||||
LiteLLM translates OpenAI's `reasoning_effort` to Gemini's `thinking` parameter. [Code](https://github.com/BerriAI/litellm/blob/620664921902d7a9bfb29897a7b27c1a7ef4ddfb/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py#L362)
|
||||
|
|
@ -298,6 +297,19 @@ curl http://0.0.0.0:4000/v1/chat/completions \
|
|||
|
||||
|
||||
|
||||
## Usage - `service_tier`
|
||||
|
||||
LiteLLM propagates OpenAI's `service_tier` parameter to Gemini, and also extracts it from the response headers (`x-gemini-service-tier`) into `model_response.service_tier`.
|
||||
|
||||
| OpenAI `service_tier` | Gemini `service_tier` | Notes |
|
||||
| --------------------- | --------------------- | ----- |
|
||||
| `"auto"` | `"priority"` | LiteLLM maps OpenAI's `"auto"` to Gemini's `"priority"` tier, as `priority` will fall back on Gemini. |
|
||||
| `"flex"` | `"flex"` | Direct mapping. |
|
||||
| `"priority"` | `"priority"` | Direct mapping. |
|
||||
| `"default"` | `"standard"` | LiteLLM maps `"default"` to `"standard"`. |
|
||||
| Any other value | Passed as-is (lowercased) | Values are case-insensitive and normalized to lowercase. |
|
||||
|
||||
On the response, LiteLLM maps `"standard"` back to `"default"` for the Gemini API.
|
||||
|
||||
|
||||
## Text-to-Speech (TTS) Audio Output
|
||||
|
|
|
|||
|
|
@ -55,24 +55,33 @@ pip install litellm
|
|||
```
|
||||
|
||||
### Step 2: Set Your Credentials
|
||||
|
||||
Choose **one** of these authentication methods:
|
||||
|
||||
> **Breaking change**: credential resolution is "first-source-wins"
|
||||
>
|
||||
> Credential resolution no longer merges individual fields across sources.
|
||||
>
|
||||
> Resolution order is:
|
||||
`kwargs` → `service key` → `env (AICORE_*)` → `config` → `VCAP service`
|
||||
>
|
||||
> **Important behavior:** once LiteLLM finds *any* credential value in a source, it takes **all** credentials from that source exclusively (except `resource_group`, which may still be resolved separately).
|
||||
|
||||
Choose **one** of these authentication methods:
|
||||
<Tabs>
|
||||
<TabItem value="service-key" label="Service Key JSON (Recommended)">
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="service-key" label="Service Key JSON (Recommended)">
|
||||
The simplest approach - paste your entire service key as a single environment variable.
|
||||
|
||||
The simplest approach - paste your entire service key as a single environment variable. The service key must be wrapped in a `credentials` object:
|
||||
> **Note:** the service key no more needs to be wrapped in a "credentials" key.
|
||||
|
||||
```bash
|
||||
export AICORE_SERVICE_KEY='{
|
||||
"credentials": {
|
||||
"clientid": "your-client-id",
|
||||
"clientsecret": "your-client-secret",
|
||||
"url": "https://<your-instance>.authentication.sap.hana.ondemand.com",
|
||||
"serviceurls": {
|
||||
"AI_API_URL": "https://api.ai.<your-region>.aws.ml.hana.ondemand.com"
|
||||
}
|
||||
}
|
||||
}'
|
||||
export AICORE_RESOURCE_GROUP="default"
|
||||
```
|
||||
|
|
@ -220,6 +229,17 @@ model="sap/gemini-2.5-pro"
|
|||
# Incorrect - missing prefix
|
||||
model="gpt-4o" # ❌ Won't work
|
||||
```
|
||||
3. **Environment variables** - Set the following list of credentials in .env file
|
||||
<pre>
|
||||
AICORE_AUTH_URL = "https://* * * .authentication.sap.hana.ondemand.com/oauth/token",
|
||||
AICORE_CLIENT_ID = " *** ",
|
||||
AICORE_CLIENT_SECRET = " *** ",
|
||||
AICORE_RESOURCE_GROUP = " *** ",
|
||||
AICORE_BASE_URL = "https://api.ai.***.cfapps.sap.hana.ondemand.com/v2"
|
||||
</pre>
|
||||
|
||||
Other credential configuration options are also available. For more information, see the [SAP AI Core Documentation](https://help.sap.com/doc/generative-ai-hub-sdk/CLOUD/en-US/_reference/README_sphynx.html#configuration).
|
||||
## Usage - LiteLLM Python SDK
|
||||
|
||||
### Proxy Usage
|
||||
|
||||
|
|
@ -506,6 +526,241 @@ response = embedding(
|
|||
print(response.data[0]["embedding"]) # Vector representation
|
||||
```
|
||||
|
||||
### Additional Modules
|
||||
The SAP Gen AI Hub includes additional modules for advanced use cases:
|
||||
- [Grounding](https://help.sap.com/docs/sap-ai-core/generative-ai/grounding-035c455a5a424697b60f4a24b6d791fe?locale=en-US)
|
||||
- [Translation](https://help.sap.com/docs/sap-ai-core/generative-ai/translation?locale=en-US)
|
||||
- [Data Masking](https://help.sap.com/docs/sap-ai-core/generative-ai/data-masking-d9a54d9ca54b40beacbd24e1663ec3b4?locale=en-US)
|
||||
- [Content Filtering](https://help.sap.com/docs/sap-ai-core/generative-ai/content-filtering?locale=en-US)
|
||||
|
||||
#### Grounding
|
||||
Grounding is a service designed to handle data-related tasks, such as grounding and retrieval, using vector databases. It provides specialized data retrieval through these databases, grounding the retrieval process with your own external and context-relevant data. Grounding combines generative AI capabilities with the ability to use real-time, precise data to improve decision-making and business operations for specific AI-driven business solutions.
|
||||
##### Prerequisites
|
||||
To use the Grounding module in the orchestration pipeline, you need to prepare the knowledge base in advance.
|
||||
|
||||
Generative AI hub offers multiple options for users to provide data (prepare a knowledge base):
|
||||
- For Option 1: Upload the documents to a supported data repository and run the data pipeline to vectorize the documents.
|
||||
- For Option 2: Provide the chunks of document via Vector API directly.
|
||||
|
||||
To use grounding, choose from one of the following options.
|
||||
|
||||
Usage example:
|
||||
```python showLineNumbers title="Grounding Example"
|
||||
from litellm import completion
|
||||
|
||||
grounding_config = {
|
||||
'type': 'document_grounding_service',
|
||||
'config': {
|
||||
'filters': [
|
||||
{'id': 's3-docs',
|
||||
'data_repository_type': 'vector',
|
||||
'search_config': {'max_chunk_count': 2},
|
||||
'data_repositories': ['012345-6789-0123-4567-890123456789']
|
||||
}
|
||||
],
|
||||
'placeholders': {'input': ['user_query'], 'output': 'grounding_response'},
|
||||
'metadata_params': ['source', 'webUrl', 'title', 'mimeType', 'fileSuffix']
|
||||
}
|
||||
}
|
||||
|
||||
response = completion(model="sap/gpt-4o",
|
||||
messages=[
|
||||
{"content":"""Facility Solutions Company provides services to luxury residential complexes,
|
||||
apartments, individual homes, and commercial properties such as office buildings, retail
|
||||
spaces, industrial facilities, and educational institutions. Customers are encouraged to
|
||||
reach out with maintenance requests, service deficiencies, follow-ups, or any issues they
|
||||
need by email.""", "role": "system"},
|
||||
{"content":"""You are a helpful assistant for any queries for answering questions.
|
||||
Answer the request by providing relevant answers that fit to the request.
|
||||
Request: {{ ?user_query }}
|
||||
Context:{{ ?grounding_response }}""", "role": "user"}
|
||||
],
|
||||
placeholder_values={"user_query": "Is there a complaint?"},
|
||||
grounding=grounding_config
|
||||
)
|
||||
print(response.choices[0].message.content)
|
||||
```
|
||||
For more information about all available grounding configurations, see the [documentation](https://help.sap.com/docs/sap-ai-core/generative-ai/using-grounding-module-e1c4dd100dfb42ab890e1d95f3516187?locale=en-US).
|
||||
|
||||
#### Translation
|
||||
The translation module allows you to translate LLM text prompts into a chosen target language.
|
||||
|
||||
```python showLineNumbers title="Translation Example"
|
||||
from litellm import completion
|
||||
|
||||
translation_config = {
|
||||
'input':
|
||||
{'type': 'sap_document_translation',
|
||||
'config':
|
||||
{'source_language': 'en-US',
|
||||
'target_language': 'de-DE'}
|
||||
},
|
||||
'output':
|
||||
{'type': 'sap_document_translation',
|
||||
'config':
|
||||
{'source_language': 'de-DE',
|
||||
'target_language': 'fr-FR'}
|
||||
}
|
||||
}
|
||||
|
||||
response = completion(model="sap/gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello world!"}],
|
||||
translation=translation_config)
|
||||
|
||||
print(response.choices[0].message.content)
|
||||
```
|
||||
For more information about all available translation configurations, see the [documentation](https://help.sap.com/docs/sap-ai-core/generative-ai/translation?locale=en-US)
|
||||
|
||||
#### Data Masking
|
||||
The data masking module serves to anonymize or pseudonymize personally identifiable information from the input for selected entities.
|
||||
|
||||
```python showLineNumbers title="Data Masking Example"
|
||||
from litellm import completion, embedding
|
||||
masking_config = {
|
||||
'providers':
|
||||
[
|
||||
{
|
||||
'type': 'sap_data_privacy_integration',
|
||||
'method': 'anonymization',
|
||||
'entities': [
|
||||
{'type': 'profile-address'},
|
||||
{'type': 'profile-email'},
|
||||
{'type': 'profile-phone'},
|
||||
{'type': 'profile-person'},
|
||||
{'type': 'profile-location'}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
mock_cv = "some text with personal information"
|
||||
|
||||
response = completion(model="sap/gpt-4o",
|
||||
messages=[{"role": "user", "content": "Give a one sentence summary of the CV. CV: {{?cv}}?"}],
|
||||
placeholder_values={"cv": mock_cv},
|
||||
masking=masking_config)
|
||||
print(response.choices[0].message.content)
|
||||
|
||||
# Data masking module also available for embedding
|
||||
response = embedding(model="sap/text-embedding-3-small",
|
||||
input=mock_cv,
|
||||
masking=masking_config)
|
||||
print(response.data[0])
|
||||
```
|
||||
For more information about all available data masking configurations, see the [documentation](https://help.sap.com/docs/sap-ai-core/generative-ai/enhancing-model-consumption-with-data-masking-66ad6f469afc4c2cbaa91a27a33f7b21?locale=en-US)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
#### Content Filtering
|
||||
The content filtering module allows you to filter input and output based on content safety criteria.
|
||||
|
||||
The module supports two services:
|
||||
* Azure Content Safety
|
||||
* Llama Guard 3
|
||||
|
||||
```python showLineNumbers title="Content Filtering Example"
|
||||
from litellm import completion
|
||||
|
||||
filtering_config_azure = {
|
||||
'input':
|
||||
{
|
||||
'filters':
|
||||
[
|
||||
{'type': 'azure_content_safety',
|
||||
'config':
|
||||
{'hate': 0,
|
||||
'sexual': 0,
|
||||
'violence': 0,
|
||||
'self_harm': 0
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
'output':
|
||||
{
|
||||
'filters':
|
||||
[
|
||||
{'type': 'azure_content_safety',
|
||||
'config': {'hate': 0,
|
||||
'sexual': 0,
|
||||
'violence': 0,
|
||||
'self_harm': 0
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
response = completion(model="sap/gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello world!"}],
|
||||
filtering=filtering_config_azure)
|
||||
print(response.choices[0].message.content)
|
||||
# The model responds normally because the content does not violate any safety rules.
|
||||
|
||||
try:
|
||||
response = completion(model="sap/gpt-4o",
|
||||
messages=[{"role": "user", "content": "I hate you"}],
|
||||
filtering=filtering_config_azure)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
# The service raises an error:
|
||||
# "Input Filter: Content filtered due to safety violations. Please modify the prompt and try again."
|
||||
```
|
||||
For more information about all available content filtering configurations, see the [documentation](https://help.sap.com/docs/sap-ai-core/generative-ai/content-filtering?locale=en-US)
|
||||
|
||||
#### List of modules configuration for fallback
|
||||
SAP GEN AI Hub supports a fallback mechanism for handling errors. This mechanism allows you to specify a list of fallback modules to use in case of errors. The fallback modules should contain all parameters that are required for configuring the request.
|
||||
|
||||
Required parameters:
|
||||
- `model`
|
||||
- `messages`
|
||||
|
||||
Optional parameters:
|
||||
- `filtering`
|
||||
- `grounding`
|
||||
- `translation`
|
||||
- `masking`
|
||||
- `tools`
|
||||
|
||||
- and any of model's specific parameters.
|
||||
|
||||
|
||||
```python showLineNumbers title="Fallback Example"
|
||||
from litellm import completion
|
||||
|
||||
translation_config = {
|
||||
'input':
|
||||
{'type': 'sap_document_translation',
|
||||
'config':
|
||||
{'source_language': 'en-US',
|
||||
'target_language': 'de-DE'}
|
||||
},
|
||||
'output':
|
||||
{'type': 'sap_document_translation',
|
||||
'config':
|
||||
{'source_language': 'de-DE',
|
||||
'target_language': 'fr-FR'}
|
||||
}
|
||||
}
|
||||
|
||||
response = completion(model="sap/gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello world!"}],
|
||||
translation=translation_config,
|
||||
fallback_sap_modules=[{
|
||||
"model":"sap/gemini-2.5-flash",
|
||||
"messages":[{"role": "user", "content": "Hello world!"}],
|
||||
"translation":translation_config
|
||||
}])
|
||||
|
||||
# In case of error with the first configuration (model gpt-4o), the fallback module is used.
|
||||
|
||||
print(response.choices[0].message.content)
|
||||
|
||||
```
|
||||
|
||||
|
||||
## Reference
|
||||
|
||||
### Supported Parameters
|
||||
|
|
|
|||
|
|
@ -597,10 +597,13 @@ router_settings:
|
|||
| LITELLM_MCP_TOOL_LISTING_TIMEOUT | Timeout in seconds for listing tools from an MCP server. Default is 30
|
||||
| LITELLM_MCP_METADATA_TIMEOUT | HTTP client timeout in seconds for OAuth metadata fetching. Default is 10
|
||||
| LITELLM_MCP_HEALTH_CHECK_TIMEOUT | Health check timeout in seconds for MCP servers. Default is 10
|
||||
| LITELLM_MCP_STDIO_EXTRA_COMMANDS | Comma-separated extra command basenames allowed for MCP stdio transport beyond the built-in allowlist. Example: `my-mcp-bin`. Empty by default
|
||||
| MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL | Default TTL in seconds for MCP OAuth2 token cache. Default is 3600
|
||||
| MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE | Maximum number of entries in MCP OAuth2 token cache. Default is 200
|
||||
| MCP_OAUTH2_TOKEN_CACHE_MIN_TTL | Minimum TTL in seconds for MCP OAuth2 token cache. Default is 10
|
||||
| MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS | Seconds to subtract from token expiry when computing cache TTL. Default is 60
|
||||
| MCP_PER_USER_TOKEN_DEFAULT_TTL | Default TTL in seconds for per-user MCP OAuth tokens stored in Redis. Default is 43200 (12 hours)
|
||||
| MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS | Seconds to subtract from per-user MCP OAuth token expiry when computing Redis TTL. Default is 60
|
||||
| DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT | Default token count for mock response completions. Default is 20
|
||||
| DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT | Default token count for mock response prompts. Default is 10
|
||||
| DEFAULT_MODEL_CREATED_AT_TIME | Default creation timestamp for models. Default is 1677610602
|
||||
|
|
|
|||
274
docs/my-website/docs/proxy/credential_routing.md
Normal file
274
docs/my-website/docs/proxy/credential_routing.md
Normal file
|
|
@ -0,0 +1,274 @@
|
|||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
# Per-Team/Project Credential Routing
|
||||
|
||||
Route the same model to different LLM provider endpoints (e.g. different Azure instances) based on which team or project makes the request.
|
||||
|
||||
## Overview
|
||||
|
||||
In multi-tenant deployments, different teams often need the same model name (e.g., `gpt-4`) to hit different provider endpoints — for example, separate Azure OpenAI instances per business unit for cost isolation, data residency, or rate limit separation.
|
||||
|
||||
**Credential routing** lets you configure this in team/project metadata using the existing [credentials table](./ui_credentials.md), without duplicating model definitions or creating separate model groups per team.
|
||||
|
||||
```
|
||||
Hotel Team → gpt-4 → https://hotel-eastus.openai.azure.com/
|
||||
Flight Team → gpt-4 → https://flight-centralus.openai.azure.com/
|
||||
```
|
||||
|
||||
### Precedence Chain
|
||||
|
||||
When a request comes in, the system walks this precedence chain (first match wins):
|
||||
|
||||
1. **Clientside credentials** — `api_base`/`api_key` passed in the request body ([docs](./clientside_auth.md))
|
||||
2. **Project model-specific** — override for this exact model in the project's `model_config`
|
||||
3. **Project default** — `defaultconfig` in the project's `model_config`
|
||||
4. **Team model-specific** — override for this exact model in the team's `model_config`
|
||||
5. **Team default** — `defaultconfig` in the team's `model_config`
|
||||
6. **Deployment default** — the model's `litellm_params` as configured in `config.yaml`
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Step 1: Create Credentials
|
||||
|
||||
Store your Azure endpoint credentials in the credentials table. You can do this via the [UI](./ui_credentials.md) or API:
|
||||
|
||||
```bash showLineNumbers
|
||||
# Create credential for Hotel team's Azure endpoint
|
||||
curl -X POST 'http://0.0.0.0:4000/credentials' \
|
||||
-H 'Authorization: Bearer sk-1234' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"credential_name": "hotel-azure-eastus",
|
||||
"credential_values": {
|
||||
"api_base": "https://hotel-eastus.openai.azure.com/",
|
||||
"api_key": "sk-azure-hotel-key-xxx"
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
```bash showLineNumbers
|
||||
# Create credential for Flight team's Azure endpoint
|
||||
curl -X POST 'http://0.0.0.0:4000/credentials' \
|
||||
-H 'Authorization: Bearer sk-1234' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"credential_name": "flight-azure-centralus",
|
||||
"credential_values": {
|
||||
"api_base": "https://flight-centralus.openai.azure.com/",
|
||||
"api_key": "sk-azure-flight-key-xxx"
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
### Step 2: Set `model_config` on Teams
|
||||
|
||||
Add a `model_config` key to the team's metadata referencing the credential by name:
|
||||
|
||||
```bash showLineNumbers
|
||||
# Hotel team — default Azure endpoint for all models
|
||||
curl -X PATCH 'http://0.0.0.0:4000/team/update' \
|
||||
-H 'Authorization: Bearer sk-1234' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"team_id": "hotel-team-id",
|
||||
"metadata": {
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {
|
||||
"litellm_credentials": "hotel-azure-eastus"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
```bash showLineNumbers
|
||||
# Flight team — default Azure endpoint for all models
|
||||
curl -X PATCH 'http://0.0.0.0:4000/team/update' \
|
||||
-H 'Authorization: Bearer sk-1234' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"team_id": "flight-team-id",
|
||||
"metadata": {
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {
|
||||
"litellm_credentials": "flight-azure-centralus"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
### Step 3: Make Requests
|
||||
|
||||
Requests are automatically routed to the correct Azure endpoint based on the API key's team:
|
||||
|
||||
```bash showLineNumbers
|
||||
# Request using Hotel team's API key → routes to hotel-eastus.openai.azure.com
|
||||
curl http://localhost:4000/v1/chat/completions \
|
||||
-H 'Content-Type: application/json' \
|
||||
-H 'Authorization: Bearer sk-hotel-team-key' \
|
||||
-d '{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}'
|
||||
|
||||
# Request using Flight team's API key → routes to flight-centralus.openai.azure.com
|
||||
curl http://localhost:4000/v1/chat/completions \
|
||||
-H 'Content-Type: application/json' \
|
||||
-H 'Authorization: Bearer sk-flight-team-key' \
|
||||
-d '{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}'
|
||||
```
|
||||
|
||||
## Per-Model Overrides
|
||||
|
||||
You can set different credentials for specific models while keeping a default for everything else:
|
||||
|
||||
```bash showLineNumbers
|
||||
curl -X PATCH 'http://0.0.0.0:4000/team/update' \
|
||||
-H 'Authorization: Bearer sk-1234' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"team_id": "hotel-team-id",
|
||||
"metadata": {
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {
|
||||
"litellm_credentials": "hotel-azure-eastus"
|
||||
}
|
||||
},
|
||||
"gpt-4": {
|
||||
"azure": {
|
||||
"litellm_credentials": "hotel-azure-westus"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
With this config:
|
||||
- `gpt-4` requests → `hotel-azure-westus` credential (model-specific)
|
||||
- All other models → `hotel-azure-eastus` credential (default)
|
||||
|
||||
## Project-Level Overrides
|
||||
|
||||
Projects inherit their team's `model_config` but can override at the project level. Project overrides take precedence over team overrides.
|
||||
|
||||
```bash showLineNumbers
|
||||
# Project overrides the team default for all models
|
||||
curl -X PATCH 'http://0.0.0.0:4000/project/update' \
|
||||
-H 'Authorization: Bearer sk-1234' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"project_id": "hotel-rec-app-id",
|
||||
"metadata": {
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {
|
||||
"litellm_credentials": "hotel-rec-azure"
|
||||
}
|
||||
},
|
||||
"gpt-4-vision": {
|
||||
"azure": {
|
||||
"litellm_credentials": "hotel-rec-vision"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
### Full Example: Hotel Team with Two Projects
|
||||
|
||||
**Setup:**
|
||||
- **Hotel Team**: default `hotel-azure-eastus`, GPT-4 override to `hotel-azure-westus`
|
||||
- **Hotel Rec App** (project): default `hotel-rec-azure`, GPT-4-Vision override to `hotel-rec-vision`
|
||||
- **Hotel Review App** (project): no overrides — inherits team config
|
||||
|
||||
**Resolution:**
|
||||
|
||||
| Request | Resolved Credential | Why |
|
||||
|---|---|---|
|
||||
| Hotel Rec App → `gpt-4` | `hotel-rec-azure` | Project default (no project model-specific match for gpt-4) |
|
||||
| Hotel Rec App → `gpt-4-vision` | `hotel-rec-vision` | Project model-specific |
|
||||
| Hotel Review App → `gpt-3.5` | `hotel-azure-eastus` | Team default (no project config) |
|
||||
| Hotel Review App → `gpt-4` | `hotel-azure-westus` | Team model-specific |
|
||||
|
||||
## `model_config` Schema
|
||||
|
||||
The `model_config` key is a JSON object in team/project `metadata`:
|
||||
|
||||
```json
|
||||
{
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"<provider>": {
|
||||
"litellm_credentials": "<credential-name>"
|
||||
}
|
||||
},
|
||||
"<model-name>": {
|
||||
"<provider>": {
|
||||
"litellm_credentials": "<credential-name>"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Description |
|
||||
|---|---|
|
||||
| `defaultconfig` | Fallback credential for any model not explicitly listed |
|
||||
| `<model-name>` | Model-specific override — must match the LiteLLM model group name |
|
||||
| `<provider>` | Provider key (e.g. `azure`, `openai`, `bedrock`). When the model name includes a provider prefix (e.g. `azure/gpt-4`), the system prefers the matching provider key |
|
||||
| `litellm_credentials` | Name of a credential in the [credentials table](./ui_credentials.md) |
|
||||
|
||||
### Credential Values
|
||||
|
||||
The referenced credential can contain any combination of:
|
||||
|
||||
| Key | Description |
|
||||
|---|---|
|
||||
| `api_base` | Provider endpoint URL |
|
||||
| `api_key` | API key for the provider |
|
||||
| `api_version` | API version (e.g. for Azure) |
|
||||
|
||||
Only keys present in the credential are applied. Keys already in the request (e.g. clientside `api_version`) are never overwritten.
|
||||
|
||||
## Enabling the Feature
|
||||
|
||||
This feature is **disabled by default** and must be explicitly enabled. To enable it:
|
||||
|
||||
<Tabs>
|
||||
|
||||
<TabItem value="config" label="config.yaml">
|
||||
|
||||
```yaml
|
||||
litellm_settings:
|
||||
enable_model_config_credential_overrides: true
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="env" label="Environment Variable">
|
||||
|
||||
```bash
|
||||
export LITELLM_ENABLE_MODEL_CONFIG_CREDENTIAL_OVERRIDES=true
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
</Tabs>
|
||||
|
||||
:::info
|
||||
The feature flag must be enabled before `model_config` entries in team/project metadata take effect. Without it, credential routing is completely inert — no metadata is read, no credentials are resolved.
|
||||
:::
|
||||
|
||||
## Related Documentation
|
||||
|
||||
- [Adding LLM Credentials](./ui_credentials.md) — Create and manage reusable credentials
|
||||
- [Project Management](./project_management.md) — Project hierarchy and API
|
||||
- [Team Budgets](./team_budgets.md) — Team-level budget management
|
||||
- [Clientside LLM Credentials](./clientside_auth.md) — Passing credentials in the request body
|
||||
- [Credential Usage Tracking](./credential_usage_tracking.md) — Track spend by credential
|
||||
|
|
@ -99,7 +99,7 @@ The following checks were performed on each of these signatures:
|
|||
- The signatures were verified against the specified public key
|
||||
```
|
||||
|
||||
Learn more about LiteLLM's release signing in the [CI/CD v2 announcement](https://docs.litellm.ai/blog/ci-cd-v2-improvements#verify-docker-image-signatures).
|
||||
Learn more about LiteLLM's release signing in the [CI/CD v2 announcement](https://docs.litellm.ai/blog/ci-cd-v2-improvements#verify-docker-image-signatures). For a complete guide covering all image variants, CI/CD enforcement, and deployment best practices, see the [Docker Image Security Guide](./docker_image_security.md).
|
||||
|
||||
### Docker Run
|
||||
|
||||
|
|
|
|||
189
docs/my-website/docs/proxy/docker_image_security.md
Normal file
189
docs/my-website/docs/proxy/docker_image_security.md
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
# Docker Image Security Guide
|
||||
|
||||
LiteLLM signs every Docker image published to GHCR with [cosign](https://docs.sigstore.dev/cosign/overview/) starting from **v1.83.0**. This page covers how to verify signatures, enforce verification in CI/CD, and follow recommended deployment patterns.
|
||||
|
||||
## Signed images
|
||||
|
||||
All image variants published to `ghcr.io/berriai/` are signed with the same cosign key:
|
||||
|
||||
| Image | Description |
|
||||
|---|---|
|
||||
| `ghcr.io/berriai/litellm` | Core proxy |
|
||||
| `ghcr.io/berriai/litellm-database` | Proxy with Postgres dependencies |
|
||||
| `ghcr.io/berriai/litellm-non_root` | Non-root variant |
|
||||
| `ghcr.io/berriai/litellm-spend_logs` | Spend-logs sidecar |
|
||||
|
||||
The signing key was introduced in [commit `0112e53`](https://github.com/BerriAI/litellm/commit/0112e53046018d726492c814b3644b7d376029d0) and the public key is checked into the repository at [`cosign.pub`](https://github.com/BerriAI/litellm/blob/main/cosign.pub).
|
||||
|
||||
:::info Enterprise images
|
||||
Enterprise images (`litellm-ee`) follow the same signing process. Contact [support@berri.ai](mailto:support@berri.ai) to confirm coverage for your specific enterprise image tag.
|
||||
:::
|
||||
|
||||
## Verify image signatures
|
||||
|
||||
Install cosign following the [official instructions](https://docs.sigstore.dev/cosign/system_config/installation/).
|
||||
|
||||
### Verify with the pinned commit hash (recommended)
|
||||
|
||||
A commit hash is cryptographically immutable, making this the strongest verification method:
|
||||
|
||||
```bash
|
||||
cosign verify \
|
||||
--key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \
|
||||
ghcr.io/berriai/litellm:v1.83.0-stable
|
||||
```
|
||||
|
||||
Replace the image reference with any signed variant:
|
||||
|
||||
```bash
|
||||
# litellm-database
|
||||
cosign verify \
|
||||
--key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \
|
||||
ghcr.io/berriai/litellm-database:v1.83.0-stable
|
||||
|
||||
# litellm-non_root
|
||||
cosign verify \
|
||||
--key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \
|
||||
ghcr.io/berriai/litellm-non_root:v1.83.0-stable
|
||||
```
|
||||
|
||||
### Verify with a release tag (convenience)
|
||||
|
||||
Tags are protected in this repository and resolve to the same key:
|
||||
|
||||
```bash
|
||||
cosign verify \
|
||||
--key https://raw.githubusercontent.com/BerriAI/litellm/v1.83.0-stable/cosign.pub \
|
||||
ghcr.io/berriai/litellm-database:v1.83.0-stable
|
||||
```
|
||||
|
||||
### Expected output
|
||||
|
||||
```
|
||||
The following checks were performed on each of these signatures:
|
||||
- The cosign claims were validated
|
||||
- The signatures were verified against the specified public key
|
||||
```
|
||||
|
||||
## Enforce verification in CI/CD
|
||||
|
||||
### Kubernetes — Sigstore Policy Controller
|
||||
|
||||
The [Sigstore Policy Controller](https://docs.sigstore.dev/policy-controller/overview/) rejects pods whose images fail cosign verification.
|
||||
|
||||
1. Install the controller:
|
||||
|
||||
```bash
|
||||
helm repo add sigstore https://sigstore.github.io/helm-charts
|
||||
helm install policy-controller sigstore/policy-controller \
|
||||
-n cosign-system --create-namespace
|
||||
```
|
||||
|
||||
2. Create a `ClusterImagePolicy` with the LiteLLM public key:
|
||||
|
||||
```yaml
|
||||
apiVersion: policy.sigstore.dev/v1beta1
|
||||
kind: ClusterImagePolicy
|
||||
metadata:
|
||||
name: litellm-signed-images
|
||||
spec:
|
||||
images:
|
||||
- glob: "ghcr.io/berriai/litellm*"
|
||||
authorities:
|
||||
- key:
|
||||
data: |
|
||||
-----BEGIN PUBLIC KEY-----
|
||||
MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEKi4ivqGpE231OGH50PKbqy1Y1Kkb
|
||||
POJC8+i2Wko82gBOUCe3M0Vw86H/4rhUhfoYEti4gdJ9wZbYmK0I2EE96g==
|
||||
-----END PUBLIC KEY-----
|
||||
```
|
||||
|
||||
3. Label the namespace to enable enforcement:
|
||||
|
||||
```bash
|
||||
kubectl label namespace litellm policy.sigstore.dev/include=true
|
||||
```
|
||||
|
||||
Any pod in that namespace using an unsigned `ghcr.io/berriai/litellm*` image will be rejected at admission.
|
||||
|
||||
### GCP — Binary Authorization
|
||||
|
||||
[Binary Authorization](https://cloud.google.com/binary-authorization/docs) can enforce cosign signatures on Cloud Run and GKE.
|
||||
|
||||
1. Create a cosign-based attestor using the LiteLLM public key:
|
||||
|
||||
```bash
|
||||
# Import the public key into a Cloud KMS keyring or use a PGP/PKIX attestor.
|
||||
# See: https://cloud.google.com/binary-authorization/docs/creating-attestors-console
|
||||
```
|
||||
|
||||
2. Configure a Binary Authorization policy that requires the attestor for `ghcr.io/berriai/litellm*` images.
|
||||
|
||||
3. Enable the policy on your Cloud Run service or GKE cluster.
|
||||
|
||||
Refer to the [GCP Binary Authorization docs](https://cloud.google.com/binary-authorization/docs/setting-up) for full setup steps.
|
||||
|
||||
### AWS — ECS / ECR
|
||||
|
||||
AWS does not natively verify cosign signatures at deploy time. Common approaches:
|
||||
|
||||
- **CI/CD gate**: Run `cosign verify` in your deployment pipeline before pushing to ECR or updating the ECS task definition. Fail the pipeline if verification fails.
|
||||
- **OPA/Gatekeeper on EKS**: If running on EKS, use the Sigstore Policy Controller (same as the Kubernetes approach above).
|
||||
|
||||
### GitHub Actions gate
|
||||
|
||||
Add a verification step before any deployment job:
|
||||
|
||||
```yaml
|
||||
- name: Verify LiteLLM image signature
|
||||
run: |
|
||||
cosign verify \
|
||||
--key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \
|
||||
ghcr.io/berriai/litellm-database:${{ env.LITELLM_VERSION }}
|
||||
```
|
||||
|
||||
## Recommended deployment patterns
|
||||
|
||||
### Pin by digest
|
||||
|
||||
Digest pinning guarantees the exact image content regardless of tag mutations:
|
||||
|
||||
```yaml
|
||||
image: ghcr.io/berriai/litellm-database@sha256:<digest>
|
||||
```
|
||||
|
||||
Get the digest after pulling:
|
||||
|
||||
```bash
|
||||
docker inspect --format='{{index .RepoDigests 0}}' \
|
||||
ghcr.io/berriai/litellm-database:v1.83.0-stable
|
||||
```
|
||||
|
||||
Cosign verification works with digests too:
|
||||
|
||||
```bash
|
||||
cosign verify \
|
||||
--key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \
|
||||
ghcr.io/berriai/litellm-database@sha256:<digest>
|
||||
```
|
||||
|
||||
### Use stable release tags
|
||||
|
||||
If digest pinning is too rigid for your workflow, use `-stable` release tags (e.g. `v1.83.0-stable`). These are immutable release tags that will not be overwritten.
|
||||
|
||||
Avoid `main-latest` or `main-stable` in production — these rolling tags point to the most recent build and can change between deployments.
|
||||
|
||||
### Safe upgrade checklist
|
||||
|
||||
1. **Verify the new image** — run `cosign verify` against the new release tag or digest.
|
||||
2. **Test in staging** — deploy the verified image to a non-production environment.
|
||||
3. **Update your pinned reference** — change the digest or tag in your deployment manifest.
|
||||
4. **Deploy to production** — roll out using your standard deployment process.
|
||||
5. **Monitor `/health`** — confirm the proxy is healthy after the upgrade.
|
||||
|
||||
## Further reading
|
||||
|
||||
- [CI/CD v2 announcement](https://docs.litellm.ai/blog/ci-cd-v2-improvements) — background on LiteLLM's signing infrastructure
|
||||
- [Docker deployment guide](./deploy.md) — full Docker, Helm, and Terraform setup
|
||||
- [cosign documentation](https://docs.sigstore.dev/cosign/overview/) — cosign usage and key management
|
||||
- [Sigstore Policy Controller](https://docs.sigstore.dev/policy-controller/overview/) — Kubernetes admission control
|
||||
BIN
docs/my-website/img/april_townhall_isolated_environments.png
Normal file
BIN
docs/my-website/img/april_townhall_isolated_environments.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 312 KiB |
BIN
docs/my-website/img/stable_main.png
Normal file
BIN
docs/my-website/img/stable_main.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 236 KiB |
BIN
docs/my-website/img/verify_releases.png
Normal file
BIN
docs/my-website/img/verify_releases.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 1.3 MiB |
|
|
@ -349,6 +349,7 @@ const sidebars = {
|
|||
"proxy/debugging",
|
||||
"proxy/error_diagnosis",
|
||||
"proxy/deploy",
|
||||
"proxy/docker_image_security",
|
||||
"proxy/health",
|
||||
"proxy/master_key_rotations",
|
||||
"proxy/model_management",
|
||||
|
|
@ -563,7 +564,8 @@ const sidebars = {
|
|||
"proxy/model_access",
|
||||
"proxy/model_access_groups",
|
||||
"proxy/access_groups",
|
||||
"proxy/team_model_add"
|
||||
"proxy/team_model_add",
|
||||
"proxy/credential_routing"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
|
|
|||
|
|
@ -114,6 +114,7 @@ class PagerDutyAlerting(SlackAlerting):
|
|||
user_api_key_max_budget=_meta.get("user_api_key_max_budget"),
|
||||
user_api_key_budget_reset_at=_meta.get("user_api_key_budget_reset_at"),
|
||||
user_api_key_org_id=_meta.get("user_api_key_org_id"),
|
||||
user_api_key_org_alias=_meta.get("user_api_key_org_alias"),
|
||||
user_api_key_team_id=_meta.get("user_api_key_team_id"),
|
||||
user_api_key_project_id=_meta.get("user_api_key_project_id"),
|
||||
user_api_key_project_alias=_meta.get("user_api_key_project_alias"),
|
||||
|
|
@ -196,6 +197,7 @@ class PagerDutyAlerting(SlackAlerting):
|
|||
else None
|
||||
),
|
||||
user_api_key_org_id=user_api_key_dict.org_id,
|
||||
user_api_key_org_alias=user_api_key_dict.organization_alias,
|
||||
user_api_key_team_id=user_api_key_dict.team_id,
|
||||
user_api_key_project_id=user_api_key_dict.project_id,
|
||||
user_api_key_project_alias=user_api_key_dict.project_alias,
|
||||
|
|
|
|||
|
|
@ -318,6 +318,7 @@ return_response_headers: bool = (
|
|||
False # get response headers from LLM Api providers - example x-remaining-requests,
|
||||
)
|
||||
enable_json_schema_validation: bool = False
|
||||
enable_model_config_credential_overrides: bool = False
|
||||
enable_key_alias_format_validation: bool = (
|
||||
False # opt-in validation of key_alias format on /key/generate and /key/update
|
||||
)
|
||||
|
|
|
|||
|
|
@ -243,6 +243,12 @@ class JsonFormatter(Formatter):
|
|||
if key not in _STANDARD_RECORD_ATTRS and key not in json_record:
|
||||
json_record[key] = value
|
||||
|
||||
# Set component/logger only if not already supplied via extra={...}
|
||||
if "component" not in json_record:
|
||||
json_record["component"] = record.name
|
||||
if "logger" not in json_record:
|
||||
json_record["logger"] = f"{record.filename}:{record.lineno}"
|
||||
|
||||
if record.exc_info:
|
||||
json_record["stacktrace"] = record.exc_text or self.formatException(
|
||||
record.exc_info
|
||||
|
|
|
|||
|
|
@ -21,11 +21,11 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
|
|||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
**kwargs,
|
||||
**kwargs: Any,
|
||||
) -> Dict[str, Any]:
|
||||
"""Handle non-streaming request to Pydantic AI agent."""
|
||||
if not api_base:
|
||||
raise ValueError("api_base is required for Pydantic AI agents")
|
||||
if api_base is None:
|
||||
raise ValueError("api_base is required for PydanticAIProviderConfig")
|
||||
return await PydanticAIHandler.handle_non_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
|
|
|
|||
|
|
@ -135,12 +135,32 @@ MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL = int(
|
|||
MCP_NPM_CACHE_DIR = os.getenv("MCP_NPM_CACHE_DIR", "/tmp/.npm_mcp_cache")
|
||||
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", "10"))
|
||||
|
||||
# Per-user OAuth token Redis cache (for server-side token storage)
|
||||
MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX = "mcp:per_user_token"
|
||||
MCP_PER_USER_TOKEN_DEFAULT_TTL = int(
|
||||
os.getenv("MCP_PER_USER_TOKEN_DEFAULT_TTL", "43200") # 12 hours
|
||||
)
|
||||
MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS = int(
|
||||
os.getenv("MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", "60")
|
||||
)
|
||||
|
||||
# MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers.
|
||||
MCP_CLIENT_TIMEOUT = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"))
|
||||
MCP_TOOL_LISTING_TIMEOUT = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
|
||||
MCP_METADATA_TIMEOUT = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
|
||||
MCP_HEALTH_CHECK_TIMEOUT = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
|
||||
|
||||
# Allowlist of commands permitted for MCP stdio transport.
|
||||
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.
|
||||
# Note: allowlisted runtimes can still execute code via args (e.g. python -c "...").
|
||||
# This is an accepted residual risk since these endpoints require PROXY_ADMIN.
|
||||
# Extend via LITELLM_MCP_STDIO_EXTRA_COMMANDS env var (comma-separated).
|
||||
_MCP_STDIO_EXTRA_COMMANDS = os.getenv("LITELLM_MCP_STDIO_EXTRA_COMMANDS", "")
|
||||
MCP_STDIO_ALLOWED_COMMANDS: frozenset = frozenset(
|
||||
{"npx", "uvx", "python", "python3", "node", "docker", "deno"}
|
||||
| (set(_MCP_STDIO_EXTRA_COMMANDS.split(",")) - {""})
|
||||
)
|
||||
|
||||
LITELLM_UI_ALLOW_HEADERS = [
|
||||
"x-litellm-semantic-filter",
|
||||
"x-litellm-semantic-filter-tools",
|
||||
|
|
|
|||
|
|
@ -7,7 +7,8 @@ from pathlib import Path
|
|||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import yaml
|
||||
from jinja2 import DictLoader, Environment, select_autoescape
|
||||
from jinja2 import DictLoader, select_autoescape
|
||||
from jinja2.sandbox import ImmutableSandboxedEnvironment
|
||||
|
||||
|
||||
class PromptTemplate:
|
||||
|
|
@ -59,7 +60,10 @@ class PromptManager:
|
|||
self.prompt_directory = Path(prompt_directory) if prompt_directory else None
|
||||
self.prompts: Dict[str, PromptTemplate] = {}
|
||||
self.prompt_file = prompt_file
|
||||
self.jinja_env = Environment(
|
||||
# Sandboxed env: templates can come from user input via /prompts/test,
|
||||
# so we must block access to unsafe Python attributes and mutation of
|
||||
# caller-supplied mutables.
|
||||
self.jinja_env = ImmutableSandboxedEnvironment(
|
||||
loader=DictLoader({}),
|
||||
autoescape=select_autoescape(["html", "xml"]),
|
||||
# Use Handlebars-style delimiters to match Dotprompt spec
|
||||
|
|
|
|||
|
|
@ -33,5 +33,14 @@
|
|||
"X-Qualifire-API-Key": "{{environment_variables.QUALIFIRE_API_KEY}}"
|
||||
},
|
||||
"environment_variables": ["QUALIFIRE_API_KEY", "QUALIFIRE_WEBHOOK_URL"]
|
||||
},
|
||||
"ramp": {
|
||||
"event_types": ["llm_api_success"],
|
||||
"endpoint": "https://api.ramp.com/developer/v1/ai-usage/litellm",
|
||||
"headers": {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": "Bearer {{environment_variables.RAMP_API_KEY}}"
|
||||
},
|
||||
"environment_variables": ["RAMP_API_KEY"]
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1031,6 +1031,9 @@ class PrometheusLogger(CustomLogger):
|
|||
user_api_key_org_id = standard_logging_payload["metadata"].get(
|
||||
"user_api_key_org_id"
|
||||
)
|
||||
user_api_key_org_alias = standard_logging_payload["metadata"].get(
|
||||
"user_api_key_org_alias"
|
||||
)
|
||||
output_tokens = standard_logging_payload["completion_tokens"]
|
||||
tokens_used = standard_logging_payload["total_tokens"]
|
||||
response_cost = standard_logging_payload["response_cost"]
|
||||
|
|
@ -1068,6 +1071,8 @@ class PrometheusLogger(CustomLogger):
|
|||
model_group=standard_logging_payload["model_group"],
|
||||
team=user_api_team,
|
||||
team_alias=user_api_team_alias,
|
||||
org_id=user_api_key_org_id,
|
||||
org_alias=user_api_key_org_alias,
|
||||
user=user_id,
|
||||
user_email=standard_logging_payload["metadata"]["user_api_key_user_email"],
|
||||
status_code="200",
|
||||
|
|
@ -1746,6 +1751,8 @@ class PrometheusLogger(CustomLogger):
|
|||
api_key_alias=user_api_key_dict.key_alias,
|
||||
team=user_api_key_dict.team_id,
|
||||
team_alias=user_api_key_dict.team_alias,
|
||||
org_id=user_api_key_dict.org_id,
|
||||
org_alias=user_api_key_dict.organization_alias,
|
||||
requested_model=request_data.get("model", ""),
|
||||
status_code=str(status_code),
|
||||
exception_status=str(status_code),
|
||||
|
|
|
|||
|
|
@ -4754,6 +4754,7 @@ class StandardLoggingPayloadSetup:
|
|||
user_api_key_budget_reset_at=None,
|
||||
user_api_key_team_id=None,
|
||||
user_api_key_org_id=None,
|
||||
user_api_key_org_alias=None,
|
||||
user_api_key_project_id=None,
|
||||
user_api_key_project_alias=None,
|
||||
user_api_key_user_id=None,
|
||||
|
|
@ -5586,6 +5587,7 @@ def get_standard_logging_metadata(
|
|||
user_api_key_budget_reset_at=None,
|
||||
user_api_key_team_id=None,
|
||||
user_api_key_org_id=None,
|
||||
user_api_key_org_alias=None,
|
||||
user_api_key_project_id=None,
|
||||
user_api_key_project_alias=None,
|
||||
user_api_key_user_id=None,
|
||||
|
|
|
|||
|
|
@ -322,9 +322,8 @@ class StandardBuiltInToolCostTracking:
|
|||
)
|
||||
if has_url_citations:
|
||||
return True
|
||||
# Fallback: Check usage object for providers that use usage instead of annotations
|
||||
# (e.g., Vertex AI Gemini uses usage.prompt_tokens_details.web_search_requests)
|
||||
if usage is not None:
|
||||
# Vertex AI Gemini uses usage.prompt_tokens_details.web_search_requests
|
||||
if (
|
||||
hasattr(usage, "prompt_tokens_details")
|
||||
and usage.prompt_tokens_details is not None
|
||||
|
|
@ -335,6 +334,15 @@ class StandardBuiltInToolCostTracking:
|
|||
and usage.prompt_tokens_details.web_search_requests is not None
|
||||
):
|
||||
return True
|
||||
# Anthropic Claude (direct API and Vertex AI) uses server_tool_use.web_search_requests.
|
||||
# Without this check, Claude ModelResponse always falls through to return False
|
||||
# and _handle_web_search_cost() is never called.
|
||||
if (
|
||||
hasattr(usage, "server_tool_use")
|
||||
and usage.server_tool_use is not None
|
||||
and usage.server_tool_use.web_search_requests is not None
|
||||
):
|
||||
return True
|
||||
return False
|
||||
elif isinstance(response_object, ResponsesAPIResponse):
|
||||
# response api explicitly includes web_search_call in the output
|
||||
|
|
|
|||
|
|
@ -34,6 +34,18 @@ def get_cost_for_web_search_request(
|
|||
|
||||
return get_cost_for_anthropic_web_search(model_info=model_info, usage=usage)
|
||||
elif custom_llm_provider.startswith("vertex_ai"):
|
||||
# Anthropic Claude models on Vertex AI populate server_tool_use.web_search_requests
|
||||
# (same as the direct Anthropic API), not prompt_tokens_details.web_search_requests
|
||||
# (which is the Gemini field). Route claude-* models to the Anthropic calculator.
|
||||
model_key: str = model_info.get("key", "") if model_info else ""
|
||||
if "claude" in model_key.lower():
|
||||
from .anthropic.cost_calculation import get_cost_for_anthropic_web_search
|
||||
|
||||
verbose_logger.debug(
|
||||
"vertex_ai/claude model detected — routing web search cost to Anthropic calculator"
|
||||
)
|
||||
return get_cost_for_anthropic_web_search(model_info=model_info, usage=usage)
|
||||
|
||||
from .vertex_ai.gemini.cost_calculator import (
|
||||
cost_per_web_search_request as cost_per_web_search_request_vertex_ai,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -89,7 +89,12 @@ async def make_call(
|
|||
|
||||
try:
|
||||
response = await client.post(
|
||||
api_base, headers=headers, data=data, stream=True, timeout=timeout
|
||||
api_base,
|
||||
headers=headers,
|
||||
data=data,
|
||||
stream=True,
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_headers = getattr(e, "headers", None)
|
||||
|
|
@ -142,7 +147,12 @@ def make_sync_call(
|
|||
|
||||
try:
|
||||
response = client.post(
|
||||
api_base, headers=headers, data=data, stream=True, timeout=timeout
|
||||
api_base,
|
||||
headers=headers,
|
||||
data=data,
|
||||
stream=True,
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_headers = getattr(e, "headers", None)
|
||||
|
|
@ -266,7 +276,11 @@ class AnthropicChatCompletion(BaseLLM):
|
|||
|
||||
try:
|
||||
response = await async_handler.post(
|
||||
api_base, headers=headers, json=data, timeout=timeout
|
||||
api_base,
|
||||
headers=headers,
|
||||
json=data,
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
## LOGGING
|
||||
|
|
@ -469,6 +483,7 @@ class AnthropicChatCompletion(BaseLLM):
|
|||
headers=headers,
|
||||
data=json.dumps(data),
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
status_code = getattr(e, "status_code", 500)
|
||||
|
|
|
|||
|
|
@ -538,9 +538,9 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
merges usage from message_start and message_delta but ignores
|
||||
message_stop. This method buffers message_delta and, when
|
||||
message_stop arrives with cache usage, merges those fields into the
|
||||
message_delta usage and also updates the input_tokens on
|
||||
message_delta to include the full count (uncached + cache_creation +
|
||||
cache_read).
|
||||
message_delta usage. input_tokens is kept as the uncached-only
|
||||
count; downstream calculate_usage adds cache tokens to
|
||||
prompt_tokens.
|
||||
"""
|
||||
_CACHE_FIELDS = ("cache_creation_input_tokens", "cache_read_input_tokens")
|
||||
pending_delta = None
|
||||
|
|
@ -569,12 +569,7 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
|
||||
raw_input = stop_usage.get("input_tokens")
|
||||
if raw_input is not None:
|
||||
uncached = raw_input if isinstance(raw_input, int) else 0
|
||||
raw_cc = delta_usage.get("cache_creation_input_tokens", 0)
|
||||
cache_creation = raw_cc if isinstance(raw_cc, int) else 0
|
||||
raw_cr = delta_usage.get("cache_read_input_tokens", 0)
|
||||
cache_read = raw_cr if isinstance(raw_cr, int) else 0
|
||||
delta_usage["input_tokens"] = uncached + cache_creation + cache_read
|
||||
delta_usage["input_tokens"] = raw_input if isinstance(raw_input, int) else 0
|
||||
|
||||
if delta_usage:
|
||||
pending_delta["usage"] = delta_usage # type: ignore[arg-type]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import json
|
||||
import ssl
|
||||
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
|
|
@ -5027,6 +5028,16 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params={},
|
||||
)
|
||||
ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://")
|
||||
# OpenAI's WebSocket responses endpoint requires ?model= in the URL,
|
||||
# matching the Realtime API convention (wss://.../v1/realtime?model=...).
|
||||
# Use urllib.parse so existing query params (e.g. api-version) are preserved.
|
||||
_parsed = urlparse(ws_url)
|
||||
_qs = parse_qs(_parsed.query)
|
||||
if "model" not in _qs:
|
||||
_qs["model"] = [model]
|
||||
ws_url = urlunparse(
|
||||
_parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()}))
|
||||
)
|
||||
|
||||
try:
|
||||
ssl_context = get_shared_realtime_ssl_context()
|
||||
|
|
|
|||
|
|
@ -91,6 +91,7 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
|
|||
"modalities",
|
||||
"parallel_tool_calls",
|
||||
"web_search_options",
|
||||
"service_tier",
|
||||
]
|
||||
if supports_reasoning(model, custom_llm_provider="gemini"):
|
||||
supported_params.append("reasoning_effort")
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ Handles extraction of skill content (SKILL.md) from stored ZIP files
|
|||
and injection into the system prompt for non-Anthropic models.
|
||||
"""
|
||||
|
||||
import posixpath
|
||||
import zipfile
|
||||
from io import BytesIO
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
|
@ -103,8 +104,18 @@ class SkillPromptInjectionHandler:
|
|||
else:
|
||||
clean_path = name
|
||||
|
||||
if clean_path:
|
||||
files[clean_path] = zf.read(name)
|
||||
if not clean_path:
|
||||
continue
|
||||
|
||||
# Ensure the path stays within the intended directory
|
||||
normalized = posixpath.normpath(clean_path)
|
||||
if normalized.startswith("..") or posixpath.isabs(normalized):
|
||||
verbose_logger.warning(
|
||||
f"SkillPromptInjectionHandler: Skipping entry with invalid path in skill {skill.skill_id}: {name}"
|
||||
)
|
||||
continue
|
||||
|
||||
files[normalized] = zf.read(name)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"SkillPromptInjectionHandler: Error extracting files from skill {skill.skill_id}: {e}"
|
||||
|
|
|
|||
|
|
@ -94,9 +94,15 @@ class SkillsSandboxExecutor:
|
|||
|
||||
# Create a temp directory to stage files
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmpdir_abs = os.path.abspath(tmpdir)
|
||||
for path, content in skill_files.items():
|
||||
# Create the file in temp directory
|
||||
local_path = os.path.join(tmpdir, path)
|
||||
local_path = os.path.abspath(os.path.join(tmpdir, path))
|
||||
if not local_path.startswith(tmpdir_abs + os.sep):
|
||||
verbose_logger.warning(
|
||||
f"SkillsSandboxExecutor: Skipping file with invalid path: {path}"
|
||||
)
|
||||
continue
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
with open(local_path, "wb") as f:
|
||||
f.write(content)
|
||||
|
|
|
|||
0
litellm/llms/sap/__init__.py
Normal file
0
litellm/llms/sap/__init__.py
Normal file
|
|
@ -1,6 +1,8 @@
|
|||
from typing import Union, Literal
|
||||
from typing import Union, Literal, Optional
|
||||
from enum import Enum
|
||||
import warnings
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
|
||||
|
||||
def validate_different_content(v: Union[str, dict, list]) -> str:
|
||||
|
|
@ -20,7 +22,7 @@ def validate_different_content(v: Union[str, dict, list]) -> str:
|
|||
elif isinstance(v, str):
|
||||
return v
|
||||
raise ValueError("Content must be a string")
|
||||
return v
|
||||
|
||||
|
||||
|
||||
class TextContent(BaseModel):
|
||||
|
|
@ -49,6 +51,10 @@ class FunctionTool(BaseModel):
|
|||
parameters: dict = {"type": "object", "properties": {}}
|
||||
strict: bool = False
|
||||
|
||||
def model_dump(self, **kwargs) -> dict:
|
||||
kwargs["exclude_unset"] = False
|
||||
return super().model_dump(**kwargs)
|
||||
|
||||
@field_validator("parameters", mode="before")
|
||||
@classmethod
|
||||
def ensure_object_type(cls, v: dict) -> dict:
|
||||
|
|
@ -66,6 +72,10 @@ class ChatCompletionTool(BaseModel):
|
|||
type_: Literal["function"] = Field(default="function", alias="type")
|
||||
function: FunctionTool
|
||||
|
||||
def model_dump(self, **kwargs) -> dict:
|
||||
kwargs["exclude_unset"] = False
|
||||
return super().model_dump(**kwargs)
|
||||
|
||||
|
||||
class MessageToolCall(BaseModel):
|
||||
id: str
|
||||
|
|
@ -114,6 +124,9 @@ class SAPToolChatMessage(BaseModel):
|
|||
)
|
||||
|
||||
|
||||
ChatMessage = Union[SAPMessage, SAPUserMessage, SAPAssistantMessage, SAPToolChatMessage]
|
||||
|
||||
|
||||
class ResponseFormat(BaseModel):
|
||||
type_: Literal["text", "json_object"] = Field(default="text", alias="type")
|
||||
|
||||
|
|
@ -128,3 +141,607 @@ class JSONResponseSchema(BaseModel):
|
|||
class ResponseFormatJSONSchema(BaseModel):
|
||||
type_: Literal["json_schema"] = Field(default="json_schema", alias="type")
|
||||
json_schema: JSONResponseSchema
|
||||
|
||||
|
||||
class KeyValueListPair(BaseModel):
|
||||
key: str
|
||||
value: list[str]
|
||||
|
||||
|
||||
class DocumentMetadataKeyValueListPairs(KeyValueListPair):
|
||||
select_mode: Optional[list[Literal["ignoreIfKeyAbsent"]]] = None
|
||||
|
||||
|
||||
class GroundingSearchConfig(BaseModel):
|
||||
max_chunk_count: Optional[int] = Field(default=None, ge=0)
|
||||
max_document_count: Optional[int] = Field(default=None, ge=0)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_max_chunk_count_and_max_document_count(self):
|
||||
if self.max_chunk_count is not None and self.max_document_count is not None:
|
||||
raise ValueError("Cannot specify both maxChunkCount and maxDocumentCount.")
|
||||
return self
|
||||
|
||||
|
||||
class DocumentGroundingFilter(BaseModel):
|
||||
id_: Optional[str] = Field(default=None, alias="id")
|
||||
data_repository_type: Literal["vector", "help.sap.com"]
|
||||
search_config: Optional[GroundingSearchConfig] = None
|
||||
data_repositories: Optional[list[str]] = None
|
||||
data_repository_metadata: Optional[list[KeyValueListPair]] = None
|
||||
document_metadata: Optional[list[DocumentMetadataKeyValueListPairs]] = None
|
||||
chunk_metadata: Optional[list[KeyValueListPair]] = None
|
||||
|
||||
|
||||
class DocumentGroundingPlaceholders(BaseModel):
|
||||
input: list[str] = Field(min_length=1)
|
||||
output: str
|
||||
|
||||
|
||||
class DocumentGroundingConfig(BaseModel):
|
||||
filters: Optional[list[DocumentGroundingFilter]] = None
|
||||
placeholders: DocumentGroundingPlaceholders
|
||||
metadata_params: Optional[list[str]] = None
|
||||
|
||||
|
||||
class GroundingModuleConfig(BaseModel):
|
||||
type_: Literal["document_grounding_service"] = Field(
|
||||
default="document_grounding_service", alias="type"
|
||||
)
|
||||
config: DocumentGroundingConfig
|
||||
|
||||
|
||||
class Template(BaseModel):
|
||||
template: list[ChatMessage]
|
||||
defaults: Optional[dict[str, str]] = None
|
||||
response_format: Optional[Union[ResponseFormat, ResponseFormatJSONSchema]] = None
|
||||
tools: Optional[list[ChatCompletionTool]] = None
|
||||
|
||||
|
||||
class LLMModelDetails(BaseModel):
|
||||
name: str
|
||||
version: str = "latest"
|
||||
params: Optional[dict] = None
|
||||
|
||||
|
||||
class PromptTemplatingModuleConfig(BaseModel):
|
||||
prompt: Template
|
||||
model: LLMModelDetails
|
||||
|
||||
|
||||
class SAPMaskingProfileEntity(str, Enum):
|
||||
"""
|
||||
Enumerates the entity categories that can be masked by the SAP Data Privacy Integration service.
|
||||
|
||||
This enum lists different types of personal or sensitive information (PII) that can be detected and masked
|
||||
by the data masking module, such as personal details, organizational data, contact information, and identifiers.
|
||||
|
||||
Values:
|
||||
PERSON: Represents personal names.
|
||||
|
||||
ORG: Represents organizational names.
|
||||
|
||||
UNIVERSITY: Represents educational institutions.
|
||||
|
||||
LOCATION: Represents geographical locations.
|
||||
|
||||
EMAIL: Represents email addresses.
|
||||
|
||||
PHONE: Represents phone numbers.
|
||||
|
||||
ADDRESS: Represents physical addresses.
|
||||
|
||||
SAP_IDS_INTERNAL: Represents internal SAP identifiers.
|
||||
|
||||
SAP_IDS_PUBLIC: Represents public SAP identifiers.
|
||||
|
||||
URL: Represents URLs.
|
||||
|
||||
USERNAME_PASSWORD: Represents usernames and passwords.
|
||||
|
||||
NATIONAL_ID: Represents national identification numbers.
|
||||
|
||||
IBAN: Represents International Bank Account Numbers.
|
||||
|
||||
SSN: Represents Social Security Numbers.
|
||||
|
||||
CREDIT_CARD_NUMBER: Represents credit card numbers.
|
||||
|
||||
PASSPORT: Represents passport numbers.
|
||||
|
||||
DRIVING_LICENSE: Represents driving license numbers.
|
||||
|
||||
NATIONALITY: Represents nationality information.
|
||||
|
||||
RELIGIOUS_GROUP: Represents religious group affiliation.
|
||||
|
||||
POLITICAL_GROUP: Represents political group affiliation.
|
||||
|
||||
PRONOUNS_GENDER: Represents pronouns and gender identity.
|
||||
|
||||
GENDER: Represents gender information.
|
||||
|
||||
SEXUAL_ORIENTATION: Represents sexual orientation.
|
||||
|
||||
TRADE_UNION: Represents trade union membership.
|
||||
|
||||
SENSITIVE_DATA: Represents any other sensitive information.
|
||||
"""
|
||||
|
||||
PERSON = "profile-person"
|
||||
ORG = "profile-org"
|
||||
UNIVERSITY = "profile-university"
|
||||
LOCATION = "profile-location"
|
||||
EMAIL = "profile-email"
|
||||
PHONE = "profile-phone"
|
||||
ADDRESS = "profile-address"
|
||||
SAP_IDS_INTERNAL = "profile-sapids-internal"
|
||||
SAP_IDS_PUBLIC = "profile-sapids-public"
|
||||
URL = "profile-url"
|
||||
USERNAME_PASSWORD = "profile-username-password"
|
||||
NATIONAL_ID = "profile-nationalid"
|
||||
IBAN = "profile-iban"
|
||||
SSN = "profile-ssn"
|
||||
CREDIT_CARD_NUMBER = "profile-credit-card-number"
|
||||
PASSPORT = "profile-passport"
|
||||
DRIVING_LICENSE = "profile-driverlicense"
|
||||
NATIONALITY = "profile-nationality"
|
||||
RELIGIOUS_GROUP = "profile-religious-group"
|
||||
POLITICAL_GROUP = "profile-political-group"
|
||||
PRONOUNS_GENDER = "profile-pronouns-gender"
|
||||
GENDER = "profile-gender"
|
||||
SEXUAL_ORIENTATION = "profile-sexual-orientation"
|
||||
TRADE_UNION = "profile-trade-union"
|
||||
SENSITIVE_DATA = "profile-sensitive-data"
|
||||
ETHNICITY = "profile-ethnicity"
|
||||
|
||||
|
||||
class DPIMethodConstant(BaseModel):
|
||||
"""
|
||||
Replaces the entity with the specified value followed by an incrementing number
|
||||
"""
|
||||
|
||||
method: Literal["constant"] = "constant"
|
||||
value: str
|
||||
|
||||
|
||||
class DPIMethodFabricatedData(BaseModel):
|
||||
"""
|
||||
Replaces the entity with a randomly generated value appropriate to its type.
|
||||
"""
|
||||
|
||||
method: Literal["fabricated_data"] = "fabricated_data"
|
||||
|
||||
|
||||
class DPICustomEntity(BaseModel):
|
||||
"""
|
||||
regex: Regular expression to match the entity
|
||||
replacement_strategy: Replacement strategy to be used for the entity
|
||||
"""
|
||||
|
||||
regex: str
|
||||
replacement_strategy: DPIMethodConstant
|
||||
|
||||
|
||||
class DPIStandardEntity(BaseModel):
|
||||
"""
|
||||
type: Standard entity type to be masked
|
||||
replacement_strategy: Replacement strategy to be used for the entity
|
||||
"""
|
||||
|
||||
type_: SAPMaskingProfileEntity = Field(..., alias="type")
|
||||
replacement_strategy: Optional[
|
||||
Union[DPIMethodConstant, DPIMethodFabricatedData]
|
||||
] = None
|
||||
|
||||
|
||||
class MaskGroundingInput(BaseModel):
|
||||
"""
|
||||
Controls whether the input to the grounding module will be masked with the configuration
|
||||
supplied in the masking module
|
||||
"""
|
||||
|
||||
enabled: bool = False
|
||||
|
||||
|
||||
class MaskingProviderConfig(BaseModel):
|
||||
"""
|
||||
SAP Data Privacy Integration provider for data masking.
|
||||
|
||||
This class implements the SAP Data Privacy Integration service, which can anonymize or pseudonymize
|
||||
specified entity categories in the input data. It supports masking sensitive information like personal names,
|
||||
contact details, and identifiers.
|
||||
|
||||
Args:
|
||||
method: The method of masking to apply (anonymization or pseudonymization).
|
||||
|
||||
entities: A list of entity categories to be masked, such as names, locations, or emails.
|
||||
|
||||
allowlist: A list of strings that should not be masked.
|
||||
|
||||
mask_grounding_input: A flag indicating whether to mask input to the grounding module.
|
||||
"""
|
||||
|
||||
type_: Literal["sap_data_privacy_integration"] = Field(
|
||||
default="sap_data_privacy_integration", alias="type"
|
||||
)
|
||||
method: Literal["anonymization", "pseudonymization"]
|
||||
entities: list[Union[DPIStandardEntity, DPICustomEntity]]
|
||||
allowlist: Optional[list[str]] = None
|
||||
mask_grounding_input: Optional[MaskGroundingInput] = None
|
||||
|
||||
|
||||
class MaskingModuleConfig(BaseModel):
|
||||
"""
|
||||
Configuration for the data masking module.
|
||||
|
||||
Args:
|
||||
providers: list of masking service provider configurations
|
||||
masking_providers: list of masking provider configurations
|
||||
IMPORTANT: use exactly one of the parameters to set the list of masking provider configurations.
|
||||
DEPRECATED: parameter 'masking_providers' will be removed Sept 15, 2026. Use 'providers' instead.
|
||||
"""
|
||||
|
||||
providers: Optional[list[MaskingProviderConfig]] = Field(min_length=1, default=None)
|
||||
masking_providers: Optional[list[MaskingProviderConfig]] = Field(
|
||||
min_length=1, default=None
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def enforce_exactly_one_provider_list(self):
|
||||
has_providers = self.providers is not None
|
||||
has_masking_providers = self.masking_providers is not None
|
||||
|
||||
if not has_providers and not has_masking_providers:
|
||||
raise ValueError(
|
||||
"For SAP Masking Module Config you must provide 'providers'."
|
||||
)
|
||||
if has_providers and has_masking_providers:
|
||||
raise ValueError(
|
||||
"For SAP Masking Module Config you must set exactly one of: 'providers' or 'masking_providers', not both."
|
||||
)
|
||||
|
||||
if has_masking_providers:
|
||||
warnings.warn(
|
||||
"The 'masking_providers' parameter is deprecated and will be removed on Sept 15, 2026. "
|
||||
"Use 'providers' instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=5,
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class AzureThreshold(int, Enum):
|
||||
"""
|
||||
Enumerates the threshold levels for the Azure Content Safety service.
|
||||
|
||||
This enum defines the various threshold levels that can be used to filter
|
||||
content based on its safety score. Each threshold value represents a specific
|
||||
level of content moderation.
|
||||
|
||||
Values:
|
||||
ALLOW_SAFE: Allows only Safe content.
|
||||
|
||||
ALLOW_SAFE_LOW: Allows Safe and Low content.
|
||||
|
||||
ALLOW_SAFE_LOW_MEDIUM: Allows Safe, Low, and Medium content.
|
||||
|
||||
ALLOW_ALL: Allows all content (Safe, Low, Medium, and High).
|
||||
"""
|
||||
|
||||
ALLOW_SAFE = 0
|
||||
ALLOW_SAFE_LOW = 2
|
||||
ALLOW_SAFE_LOW_MEDIUM = 4
|
||||
ALLOW_ALL = 6
|
||||
|
||||
|
||||
class AzureContentFilter(BaseModel):
|
||||
"""
|
||||
Specific filter configuration for Azure Content Safety.
|
||||
|
||||
This class configures content filtering based on Azure's categories and
|
||||
severity levels. It allows setting thresholds for hate speech, sexual content,
|
||||
violence, and self-harm content.
|
||||
|
||||
Values:
|
||||
hate: Threshold for hate speech content.
|
||||
|
||||
sexual: Threshold for sexual content.
|
||||
|
||||
violence: Threshold for violent content.
|
||||
|
||||
self_harm: Threshold for self-harm content.
|
||||
"""
|
||||
|
||||
hate: Optional[Union[AzureThreshold, Literal[0, 2, 4, 6]]] = None
|
||||
sexual: Optional[Union[AzureThreshold, Literal[0, 2, 4, 6]]] = None
|
||||
violence: Optional[Union[AzureThreshold, Literal[0, 2, 4, 6]]] = None
|
||||
self_harm: Optional[Union[AzureThreshold, Literal[0, 2, 4, 6]]] = None
|
||||
|
||||
|
||||
class AzureContentSafetyInput(AzureContentFilter):
|
||||
"""
|
||||
Filter configuration for Azure Content Safety Input
|
||||
|
||||
Values:
|
||||
hate: Threshold for hate speech content.
|
||||
|
||||
sexual: Threshold for sexual content.
|
||||
|
||||
violence: Threshold for violent content.
|
||||
|
||||
self_harm: Threshold for self-harm content.
|
||||
|
||||
prompt_shield: A flag to use prompt shield
|
||||
"""
|
||||
|
||||
prompt_shield: Optional[bool] = False
|
||||
|
||||
|
||||
class AzureContentSafetyOutput(AzureContentFilter):
|
||||
"""
|
||||
Filter configuration for Azure Content Safety Output
|
||||
|
||||
Values:
|
||||
hate: Threshold for hate speech content.
|
||||
|
||||
sexual: Threshold for sexual content.
|
||||
|
||||
violence: Threshold for violent content.
|
||||
|
||||
self_harm: Threshold for self-harm content.
|
||||
|
||||
protected_material_code: Detect protected code content from known GitHub repositories.
|
||||
The scan includes software libraries, source code, algorithms,
|
||||
and other proprietary programming content.
|
||||
"""
|
||||
|
||||
protected_material_code: Optional[bool] = False
|
||||
|
||||
|
||||
class LlamaGuard38bFilter(BaseModel):
|
||||
"""
|
||||
Specific implementation of ContentFilter for Llama Guard 3. Llama Guard 3 is a
|
||||
Llama-3.1-8B pretrained model, fine-tuned for content safety classification.
|
||||
|
||||
Args:
|
||||
violent_crimes: Responses that enable, encourage, or endorse the commission of violent crimes.
|
||||
|
||||
non_violent_crimes: Responses that enable, encourage, or endorse the commission of non-violent crimes.
|
||||
|
||||
sex_crimes: Responses that enable, encourage, or endorse the commission of sex-related crimes.
|
||||
|
||||
child_exploitation: Responses that contain, describe, enable or endorse the sexual abuse of children.
|
||||
|
||||
defamation: Responses that are both verifiably false and likely to injure a living person's reputation.
|
||||
|
||||
specialized_advice: Responses that contain specialized financial, medical or legal advice.
|
||||
|
||||
privacy: Responses that contain sensitive or nonpublic personal information.
|
||||
|
||||
intellectual_property: Responses that may violate the intellectual property rights of any third party.
|
||||
|
||||
indiscriminate_weapons: Responses that enable, encourage, or endorse the creation of indiscriminate weapons.
|
||||
|
||||
hate: Responses that demean or dehumanize people on the basis of their sensitive, personal characteristics.
|
||||
|
||||
self_harm: Responses that enable, encourage, or endorse acts of intentional self-harm.
|
||||
|
||||
sexual_content: Responses that contain erotica.
|
||||
|
||||
elections: Responses that contain factually incorrect information about electoral systems and processes.
|
||||
|
||||
code_interpreter_abuse: Responses that seek to abuse code interpreters.
|
||||
"""
|
||||
|
||||
violent_crimes: bool = Field(default=False)
|
||||
non_violent_crimes: bool = Field(default=False)
|
||||
sex_crimes: bool = Field(default=False)
|
||||
child_exploitation: bool = Field(default=False)
|
||||
defamation: bool = Field(default=False)
|
||||
specialized_advice: bool = Field(default=False)
|
||||
privacy: bool = Field(default=False)
|
||||
intellectual_property: bool = Field(default=False)
|
||||
indiscriminate_weapons: bool = Field(default=False)
|
||||
hate: bool = Field(default=False)
|
||||
self_harm: bool = Field(default=False)
|
||||
sexual_content: bool = Field(default=False)
|
||||
elections: bool = Field(default=False)
|
||||
code_interpreter_abuse: bool = Field(default=False)
|
||||
|
||||
|
||||
class LlamaGuard38bFilterConfig(BaseModel):
|
||||
type_: Literal["llama_guard_3_8b"] = Field(default="llama_guard_3_8b", alias="type")
|
||||
config: LlamaGuard38bFilter
|
||||
|
||||
|
||||
class AzureContentSafetyInputFilterConfig(BaseModel):
|
||||
type_: Literal["azure_content_safety"] = Field(
|
||||
default="azure_content_safety", alias="type"
|
||||
)
|
||||
config: Optional[AzureContentSafetyInput] = None
|
||||
|
||||
|
||||
class AzureContentSafetyOutputFilterConfig(BaseModel):
|
||||
type_: Literal["azure_content_safety"] = Field(
|
||||
default="azure_content_safety", alias="type"
|
||||
)
|
||||
config: Optional[AzureContentSafetyOutput] = None
|
||||
|
||||
|
||||
class FilteringStreamOptions(BaseModel):
|
||||
"""
|
||||
overlap: Number of characters that should be additionally sent to content filtering services
|
||||
from previous chunks as additional context.
|
||||
"""
|
||||
|
||||
overlap: Optional[int] = Field(default=0, ge=0, le=10000)
|
||||
|
||||
|
||||
class InputFiltering(BaseModel):
|
||||
"""Module for managing and applying input content filters.
|
||||
|
||||
Args:
|
||||
filters: List of ContentFilter objects to be applied to input content.
|
||||
"""
|
||||
|
||||
filters: list[
|
||||
Union[AzureContentSafetyInputFilterConfig, LlamaGuard38bFilterConfig]
|
||||
] = Field(min_length=1)
|
||||
|
||||
|
||||
class OutputFiltering(BaseModel):
|
||||
"""Module for managing and applying output content filters.
|
||||
|
||||
Args:
|
||||
filters: List of ContentFilter objects to be applied to output content.
|
||||
|
||||
stream_options: Module-specific streaming options.
|
||||
"""
|
||||
|
||||
filters: list[
|
||||
Union[AzureContentSafetyOutputFilterConfig, LlamaGuard38bFilterConfig]
|
||||
] = Field(min_length=1)
|
||||
stream_options: Optional[FilteringStreamOptions] = None
|
||||
|
||||
|
||||
class FilteringModuleConfig(BaseModel):
|
||||
"""Module for managing and applying content filters.
|
||||
|
||||
Args:
|
||||
input: Module for filtering and validating input content before processing.
|
||||
|
||||
output: Module for filtering and validating output content after generation.
|
||||
"""
|
||||
|
||||
input: Optional[InputFiltering] = None
|
||||
output: Optional[OutputFiltering] = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def enforce_min_properties(self) -> "FilteringModuleConfig":
|
||||
"""
|
||||
Ensure at least one of input or output filtering is provided.
|
||||
"""
|
||||
if self.input is None and self.output is None:
|
||||
raise ValueError(
|
||||
"For using SAP Filtering Module you must provide at least one property: input or output filters."
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class SAPDocumentTranslationApplyToSelector(BaseModel):
|
||||
"""
|
||||
This selector allows you to define the scope of translation, such as specific placeholders or
|
||||
messages with specific roles.
|
||||
For example, {"category": "placeholders",
|
||||
"items": ["user_input"],
|
||||
"source_language": "de-DE"}
|
||||
targets the value of "user_input" in placeholder_values specified in the request payload;
|
||||
and considers the value to be in German.
|
||||
"""
|
||||
|
||||
category: Literal["placeholders", "template_roles"]
|
||||
items: list[str]
|
||||
source_language: str
|
||||
|
||||
|
||||
class InputTranslationConfig(BaseModel):
|
||||
"""
|
||||
Configuration for input translation.
|
||||
|
||||
Args:
|
||||
source_language: Language of the text to be translated. Example: de-DE
|
||||
target_language: Language to which the text should be translated. Example: en-US
|
||||
apply_to: List of selectors that define the scope of translation.
|
||||
"""
|
||||
|
||||
source_language: Optional[str] = None
|
||||
target_language: str
|
||||
apply_to: Optional[list[SAPDocumentTranslationApplyToSelector]] = None
|
||||
|
||||
|
||||
class OutputTranslationConfig(BaseModel):
|
||||
source_language: Optional[str] = None
|
||||
target_language: Union[str, SAPDocumentTranslationApplyToSelector]
|
||||
|
||||
|
||||
class SAPDocumentTranslationInput(BaseModel):
|
||||
"""
|
||||
Configuration for input translation
|
||||
|
||||
Args:
|
||||
type: The type of translation module (e.g., 'sap_document_translation').
|
||||
|
||||
translate_messages_history: If true, the messages history will be translated as well.
|
||||
|
||||
config: Configuration object for the translation module.
|
||||
"""
|
||||
|
||||
type_: Literal["sap_document_translation"] = Field(
|
||||
default="sap_document_translation", alias="type"
|
||||
)
|
||||
translate_messages_history: Optional[bool] = None
|
||||
config: InputTranslationConfig
|
||||
|
||||
|
||||
class SAPDocumentTranslationOutput(BaseModel):
|
||||
"""
|
||||
Configuration for output translation
|
||||
|
||||
Args:
|
||||
type: The type of translation module (e.g., 'sap_document_translation').
|
||||
|
||||
config: Configuration object for the translation module.
|
||||
"""
|
||||
|
||||
type_: Literal["sap_document_translation"] = Field(
|
||||
default="sap_document_translation", alias="type"
|
||||
)
|
||||
config: OutputTranslationConfig
|
||||
|
||||
|
||||
class TranslationModuleConfig(BaseModel):
|
||||
"""
|
||||
Configuration for translation module
|
||||
|
||||
Args:
|
||||
input: Configuration for input translation
|
||||
|
||||
output: Configuration for output translation
|
||||
"""
|
||||
|
||||
input: Optional[SAPDocumentTranslationInput] = None
|
||||
output: Optional[SAPDocumentTranslationOutput] = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def enforce_min_properties(self) -> "TranslationModuleConfig":
|
||||
if self.input is None and self.output is None:
|
||||
raise ValueError(
|
||||
"TranslationModuleConfig requires at least one of 'input' or 'output'."
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class ModuleConfig(BaseModel):
|
||||
prompt_templating: PromptTemplatingModuleConfig
|
||||
filtering: Optional[FilteringModuleConfig] = None
|
||||
masking: Optional[MaskingModuleConfig] = None
|
||||
grounding: Optional[GroundingModuleConfig] = None
|
||||
translation: Optional[TranslationModuleConfig] = None
|
||||
|
||||
|
||||
class GlobalStreamOptions(BaseModel):
|
||||
enabled: bool = False
|
||||
chunk_size: Optional[int] = Field(default=None, ge=1)
|
||||
delimiters: Optional[list[str]] = None
|
||||
|
||||
|
||||
class OrchestrationConfig(BaseModel):
|
||||
modules: Union[ModuleConfig, list[ModuleConfig]]
|
||||
stream: Optional[GlobalStreamOptions] = None
|
||||
|
||||
|
||||
class OrchestrationRequest(BaseModel):
|
||||
config: OrchestrationConfig
|
||||
placeholder_values: Optional[dict[str, str]] = None
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from typing import (
|
|||
TYPE_CHECKING,
|
||||
Iterator,
|
||||
AsyncIterator,
|
||||
FrozenSet,
|
||||
)
|
||||
from functools import cached_property
|
||||
import litellm
|
||||
|
|
@ -31,12 +32,13 @@ else:
|
|||
|
||||
from ..credentials import get_token_creator
|
||||
from .models import (
|
||||
SAPMessage,
|
||||
SAPAssistantMessage,
|
||||
SAPToolChatMessage,
|
||||
ChatCompletionTool,
|
||||
ResponseFormatJSONSchema,
|
||||
OrchestrationRequest,
|
||||
ResponseFormat,
|
||||
ResponseFormatJSONSchema,
|
||||
SAPAssistantMessage,
|
||||
SAPMessage,
|
||||
SAPToolChatMessage,
|
||||
SAPUserMessage,
|
||||
)
|
||||
from .handler import (
|
||||
|
|
@ -45,9 +47,65 @@ from .handler import (
|
|||
SAPStreamIterator,
|
||||
)
|
||||
|
||||
# Keys routed outside SAP orchestration `model.params` (prompt, stream, fallbacks, etc.)
|
||||
_SAP_MODEL_PARAMS_EXCLUDED_KEYS: FrozenSet[str] = frozenset(
|
||||
{
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"stream_options",
|
||||
"fallback_sap_modules",
|
||||
"placeholder_values",
|
||||
"model_version",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def validate_dict(data: dict, model) -> dict:
|
||||
return model(**data).model_dump(by_alias=True)
|
||||
return model(**data).model_dump(by_alias=True, exclude_unset=True)
|
||||
|
||||
|
||||
def _messages_to_sap_template(messages: List[Dict[str, str]]) -> list: # type: ignore[type-arg]
|
||||
template = []
|
||||
for message in messages:
|
||||
if message["role"] == "user":
|
||||
template.append(validate_dict(message, SAPUserMessage))
|
||||
elif message["role"] == "assistant":
|
||||
template.append(validate_dict(message, SAPAssistantMessage))
|
||||
elif message["role"] == "tool":
|
||||
template.append(validate_dict(message, SAPToolChatMessage))
|
||||
else:
|
||||
template.append(validate_dict(message, SAPMessage))
|
||||
return template
|
||||
|
||||
|
||||
def _tools_response_format_and_stream(
|
||||
optional_params: dict, model_params: dict
|
||||
) -> Tuple[dict, dict, dict]:
|
||||
tools_ = optional_params.pop("tools", [])
|
||||
tools_ = [validate_dict(tool, ChatCompletionTool) for tool in tools_]
|
||||
tools: dict = {"tools": tools_} if tools_ else {}
|
||||
|
||||
response_format = model_params.pop("response_format", {})
|
||||
resp_type = response_format.get("type", None)
|
||||
if resp_type:
|
||||
if resp_type == "json_schema":
|
||||
response_format = validate_dict(
|
||||
response_format, ResponseFormatJSONSchema
|
||||
)
|
||||
else:
|
||||
response_format = validate_dict(response_format, ResponseFormat)
|
||||
response_format = {"response_format": response_format}
|
||||
|
||||
model_params.pop("stream", False)
|
||||
stream_config: dict = {}
|
||||
if "stream_options" in optional_params:
|
||||
stream_options = optional_params.pop("stream_options", {})
|
||||
if "chunk_size" in stream_options:
|
||||
stream_config["chunk_size"] = stream_options.get("chunk_size")
|
||||
if "delimiters" in stream_options:
|
||||
stream_config["delimiters"] = stream_options.get("delimiters")
|
||||
|
||||
return tools, response_format, stream_config
|
||||
|
||||
|
||||
class GenAIHubOrchestrationConfig(OpenAIGPTConfig):
|
||||
|
|
@ -208,48 +266,25 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig):
|
|||
api_base_ = f"{self.deployment_url}/v2/completion"
|
||||
return api_base_
|
||||
|
||||
def transform_request(
|
||||
def _build_prompt_module(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[Dict[str, str]], # type: ignore
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
headers: dict,
|
||||
model_name: str,
|
||||
template_messages: List[Dict[str, str]],
|
||||
params: dict,
|
||||
) -> dict:
|
||||
# Filter out parameters that are not valid model params for SAP Orchestration API
|
||||
# - tools, model_version, deployment_url: handled separately
|
||||
excluded_params = {"tools", "model_version", "deployment_url"}
|
||||
|
||||
# Filter strict for GPT models only - SAP AI Core doesn't accept it as a model param
|
||||
# LangChain agents pass strict=true at top level, which fails for GPT models
|
||||
# Anthropic models accept strict, so preserve it for them
|
||||
if model.startswith("gpt"):
|
||||
excluded_params.add("strict")
|
||||
if model_name.startswith("gpt") and "strict" in params:
|
||||
params.pop("strict")
|
||||
|
||||
model_params = {
|
||||
k: v for k, v in optional_params.items() if k not in excluded_params
|
||||
}
|
||||
model_version = params.pop("model_version", "latest")
|
||||
|
||||
model_version = optional_params.pop("model_version", "latest")
|
||||
template = []
|
||||
for message in messages:
|
||||
if message["role"] == "user":
|
||||
template.append(validate_dict(message, SAPUserMessage))
|
||||
elif message["role"] == "assistant":
|
||||
template.append(validate_dict(message, SAPAssistantMessage))
|
||||
elif message["role"] == "tool":
|
||||
template.append(validate_dict(message, SAPToolChatMessage))
|
||||
else:
|
||||
template.append(validate_dict(message, SAPMessage))
|
||||
|
||||
tools_ = optional_params.pop("tools", [])
|
||||
tools_ = params.pop("tools", [])
|
||||
tools_ = [validate_dict(tool, ChatCompletionTool) for tool in tools_]
|
||||
if tools_ != []:
|
||||
tools = {"tools": tools_}
|
||||
else:
|
||||
tools = {}
|
||||
tools = {"tools": tools_} if tools_ else {}
|
||||
|
||||
response_format = model_params.pop("response_format", {})
|
||||
response_format = params.pop("response_format", {})
|
||||
resp_type = response_format.get("type", None)
|
||||
if resp_type:
|
||||
if resp_type == "json_schema":
|
||||
|
|
@ -259,33 +294,104 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig):
|
|||
else:
|
||||
response_format = validate_dict(response_format, ResponseFormat)
|
||||
response_format = {"response_format": response_format}
|
||||
model_params.pop("stream", False)
|
||||
stream_config = {}
|
||||
if "stream_options" in model_params:
|
||||
# stream_config["enabled"] = True
|
||||
stream_options = model_params.pop("stream_options", {})
|
||||
stream_config["chunk_size"] = stream_options.get("chunk_size", 100)
|
||||
if "delimiters" in stream_options:
|
||||
stream_config["delimiters"] = stream_options.get("delimiters")
|
||||
# else:
|
||||
# stream_config["enabled"] = False
|
||||
config = {
|
||||
"config": {
|
||||
"modules": {
|
||||
"prompt_templating": {
|
||||
"prompt": {"template": template, **tools, **response_format},
|
||||
"model": {
|
||||
"name": model,
|
||||
"params": model_params,
|
||||
"version": model_version,
|
||||
},
|
||||
},
|
||||
else:
|
||||
response_format = {}
|
||||
|
||||
placeholder_defaults = params.pop("placeholder_defaults", {})
|
||||
placeholder_defaults = (
|
||||
{"defaults": placeholder_defaults} if placeholder_defaults else {}
|
||||
)
|
||||
|
||||
optional_modules = {}
|
||||
optional_modules_lst = ["grounding", "masking", "filtering", "translation"]
|
||||
for module in optional_modules_lst:
|
||||
if params.get(module, None) is not None:
|
||||
optional_modules[module] = params.pop(module)
|
||||
|
||||
return {
|
||||
"prompt_templating": {
|
||||
"prompt": {
|
||||
"template": template_messages,
|
||||
**placeholder_defaults,
|
||||
**tools,
|
||||
**response_format,
|
||||
},
|
||||
"stream": stream_config,
|
||||
}
|
||||
"model": {
|
||||
"name": model_name,
|
||||
"params": params,
|
||||
"version": model_version,
|
||||
},
|
||||
},
|
||||
**optional_modules,
|
||||
}
|
||||
|
||||
return config
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[Dict[str, str]], # type: ignore
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
optional_params = dict(optional_params)
|
||||
optional_params.pop("deployment_url", None)
|
||||
|
||||
template = _messages_to_sap_template(messages)
|
||||
|
||||
placeholder_values = optional_params.pop("placeholder_values", None)
|
||||
fallback_modules = optional_params.pop("fallback_sap_modules", [])
|
||||
|
||||
optional_params.pop("stream", None)
|
||||
stream_config: dict = {}
|
||||
if "stream_options" in optional_params:
|
||||
stream_options = optional_params.pop("stream_options", {})
|
||||
if "chunk_size" in stream_options:
|
||||
stream_config["chunk_size"] = stream_options["chunk_size"]
|
||||
if "delimiters" in stream_options:
|
||||
stream_config["delimiters"] = stream_options["delimiters"]
|
||||
|
||||
optional_params.pop("tool_choice", None)
|
||||
|
||||
modules = [
|
||||
self._build_prompt_module(
|
||||
model_name=model,
|
||||
template_messages=template,
|
||||
params=dict(optional_params),
|
||||
)
|
||||
]
|
||||
|
||||
for modules_dict in fallback_modules:
|
||||
modules_dict = dict(modules_dict)
|
||||
fallback_model = modules_dict.pop("model", None)
|
||||
if fallback_model is None:
|
||||
raise ValueError(
|
||||
"Each entry in `fallback_sap_modules` must include a 'model' key."
|
||||
)
|
||||
if fallback_model.startswith("sap/"):
|
||||
fallback_model = fallback_model[4:]
|
||||
fallback_template = modules_dict.pop("messages", [])
|
||||
|
||||
modules.append(
|
||||
self._build_prompt_module(
|
||||
model_name=fallback_model,
|
||||
template_messages=fallback_template,
|
||||
params=modules_dict,
|
||||
)
|
||||
)
|
||||
|
||||
config_payload: Dict[str, Any] = {
|
||||
"modules": modules if len(modules) > 1 else modules[0],
|
||||
}
|
||||
if stream_config:
|
||||
config_payload["stream"] = stream_config
|
||||
|
||||
request_body: Dict[str, Any] = {"config": config_payload}
|
||||
if placeholder_values is not None:
|
||||
request_body["placeholder_values"] = placeholder_values
|
||||
|
||||
body = validate_dict(request_body, OrchestrationRequest)
|
||||
|
||||
return body
|
||||
|
||||
def transform_response(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from __future__ import annotations
|
||||
from typing import Any, Callable, Dict, Final, List, Optional, Sequence, Tuple
|
||||
from typing import Any, Callable, Dict, Final, List, Optional, Sequence, Tuple, Union
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from threading import Lock
|
||||
from pathlib import Path
|
||||
|
|
@ -7,9 +7,11 @@ from dataclasses import dataclass
|
|||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import httpx
|
||||
|
||||
from litellm import sap_service_key
|
||||
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
|
||||
from litellm.llms.custom_httpx.http_handler import _get_httpx_client, HTTPHandler
|
||||
from litellm._logging import verbose_logger
|
||||
import litellm
|
||||
|
||||
AUTH_ENDPOINT_SUFFIX = "/oauth/token"
|
||||
|
||||
|
|
@ -28,11 +30,25 @@ def _get_home() -> str:
|
|||
return os.getenv(HOME_PATH_ENV_VAR, DEFAULT_HOME_PATH)
|
||||
|
||||
|
||||
def _get_nested(d: Dict[str, Any], path: Sequence[str]) -> Any:
|
||||
def _get_nested(d: Union[Dict[str, Any], str], path: Sequence[str]) -> Any:
|
||||
cur: Any = d
|
||||
if isinstance(cur, str):
|
||||
# This shouldn't happen if service keys are pre-parsed correctly
|
||||
try:
|
||||
cur = json.loads(cur)
|
||||
except json.JSONDecodeError:
|
||||
verbose_logger.warning(
|
||||
"SAP service key or VCAP service is a string but not valid JSON."
|
||||
)
|
||||
return None
|
||||
for k in path:
|
||||
if not isinstance(cur, dict) or k not in cur:
|
||||
raise KeyError(".".join(path))
|
||||
if not isinstance(cur, dict):
|
||||
verbose_logger.warning(
|
||||
f"SAP service key or VCAP service traversal hit non-dict type '{type(cur).__name__}' at key '{k}'."
|
||||
)
|
||||
return None
|
||||
if k not in cur:
|
||||
return None
|
||||
cur = cur[k]
|
||||
return cur
|
||||
|
||||
|
|
@ -47,6 +63,13 @@ def _load_json_env(var_name: str) -> Optional[Dict[str, Any]]:
|
|||
return None
|
||||
|
||||
|
||||
def _str_or_none(value) -> Optional[str]:
|
||||
try:
|
||||
return str(value) if value is not None else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _load_vcap() -> Dict[str, Any]:
|
||||
return _load_json_env(VCAP_SERVICES_ENV_VAR) or {}
|
||||
|
||||
|
|
@ -59,6 +82,12 @@ def _get_vcap_service(label: str) -> Optional[Dict[str, Any]]:
|
|||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Source:
|
||||
name: str
|
||||
get: Callable[[CredentialsValue], Optional[str]]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CredentialsValue:
|
||||
name: str
|
||||
|
|
@ -82,7 +111,6 @@ CREDENTIAL_VALUES: Final[List[CredentialsValue]] = [
|
|||
transform_fn=lambda url: url.rstrip("/")
|
||||
+ ("" if url.endswith("/v2") else "/v2"),
|
||||
),
|
||||
CredentialsValue("resource_group", default="default"),
|
||||
CredentialsValue(
|
||||
"cert_url",
|
||||
("certurl",),
|
||||
|
|
@ -145,81 +173,239 @@ def _env_name(name: str) -> str:
|
|||
return f"AICORE_{name.upper()}"
|
||||
|
||||
|
||||
def _resolve_value(
|
||||
cred: CredentialsValue,
|
||||
*,
|
||||
kwargs: Dict[str, Any],
|
||||
env: Dict[str, str],
|
||||
config: Dict[str, Any],
|
||||
service_like: Optional[Dict[str, Any]],
|
||||
) -> Optional[str]:
|
||||
# 1) explicit kwargs
|
||||
if cred.name in kwargs and kwargs[cred.name] is not None:
|
||||
return kwargs[cred.name]
|
||||
def extract_credentials(source: Source) -> Dict[str, str]:
|
||||
"""Extract all credentials from a source."""
|
||||
credentials = {}
|
||||
for cv in CREDENTIAL_VALUES:
|
||||
value = source.get(cv)
|
||||
if value is not None:
|
||||
credentials[cv.name] = cv.transform_fn(value) if cv.transform_fn else value
|
||||
return credentials
|
||||
|
||||
# 2) environment variables (primary name)
|
||||
env_key = _env_name(cred.name)
|
||||
if env_key in env and env[env_key] is not None:
|
||||
return env[env_key]
|
||||
|
||||
# 3) config file (accept both prefixed and plain keys)
|
||||
for key in (env_key, cred.name):
|
||||
if key in config and config[key] is not None:
|
||||
return config[key]
|
||||
def resolve_credentials(sources: List[Source]) -> Dict[str, str]:
|
||||
"""Extract credentials from the first source that has any defined."""
|
||||
for source in sources:
|
||||
credentials = extract_credentials(source)
|
||||
if credentials:
|
||||
verbose_logger.debug(f"Resolved SAP credentials from source {source.name}")
|
||||
return credentials
|
||||
raise ValueError("No credentials found in any source")
|
||||
|
||||
# 4) service-like source (AICORE_SERVICE_KEY first, else VCAP)
|
||||
if service_like and cred.vcap_key:
|
||||
|
||||
def resolve_resource_group(sources: List[Source]) -> Optional[str]:
|
||||
"""Find resource_group from the first source that defines it."""
|
||||
rg_cred = CredentialsValue("resource_group", default="default")
|
||||
for source in sources:
|
||||
value = source.get(rg_cred)
|
||||
if value is not None:
|
||||
verbose_logger.debug(
|
||||
f"Resolved GEN AI Hub resource_group from source {source.name}"
|
||||
)
|
||||
return value
|
||||
return rg_cred.default
|
||||
|
||||
|
||||
def _parse_service_key_once(
|
||||
service_key: Optional[Union[str, dict]]
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Pre-parse service_key if it's a string to avoid repeated JSON parsing.
|
||||
|
||||
Returns None if parsing fails (other credential sources may still work).
|
||||
"""
|
||||
if service_key is None:
|
||||
return None
|
||||
if isinstance(service_key, dict):
|
||||
return service_key
|
||||
if isinstance(service_key, str):
|
||||
try:
|
||||
val = _get_nested(service_like, ("credentials",) + cred.vcap_key)
|
||||
if val is not None:
|
||||
return val
|
||||
except KeyError:
|
||||
pass
|
||||
return json.loads(service_key)
|
||||
except json.JSONDecodeError:
|
||||
verbose_logger.warning(
|
||||
"SAP service key is a string but not valid JSON. Skipping this source."
|
||||
)
|
||||
return None
|
||||
verbose_logger.warning(
|
||||
f"SAP service key has unexpected type '{type(service_key).__name__}'. Expected str or dict. Ignoring."
|
||||
)
|
||||
return None
|
||||
|
||||
# 5) default
|
||||
return cred.default
|
||||
|
||||
def _resolve_credential_from_service_key(
|
||||
service_key: Optional[Union[str, dict]], cv: CredentialsValue
|
||||
) -> Optional[str]:
|
||||
if service_key is None:
|
||||
return None
|
||||
val = _str_or_none(
|
||||
_get_nested(
|
||||
service_key, (("credentials",) + cv.vcap_key) if cv.vcap_key else (cv.name,)
|
||||
)
|
||||
)
|
||||
if val is None:
|
||||
return _str_or_none(
|
||||
_get_nested(service_key, cv.vcap_key if cv.vcap_key else (cv.name,))
|
||||
)
|
||||
return val
|
||||
|
||||
|
||||
def fetch_credentials(
|
||||
service_key: Optional[str] = None, profile: Optional[str] = None, **kwargs
|
||||
service_key: Optional[Union[str, dict]] = None,
|
||||
profile: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
Resolution order per key:
|
||||
Resolution order (first-source-wins):
|
||||
|
||||
Sources are checked in this order:
|
||||
kwargs
|
||||
> service key
|
||||
> env (AICORE_<NAME>)
|
||||
> config (AICORE_<NAME> or plain <name>)
|
||||
> service-like source from JSON in $AICORE_SERVICE_KEY (same structure as a VCAP service object)
|
||||
falling back to service entry in $VCAP_SERVICES with label 'aicore'
|
||||
> vcap service key
|
||||
> default
|
||||
|
||||
Important:
|
||||
- Credentials are extracted from the FIRST source that provides any credential value.
|
||||
- Values are NOT merged per key across sources. Except resource_group, which is merged.
|
||||
|
||||
Warning:
|
||||
- This function does NOT validate the returned credentials just parsed it from the sources.
|
||||
- Callers MUST explicitly call validate_credentials() on the returned dict
|
||||
"""
|
||||
config = init_conf(profile)
|
||||
env = os.environ # snapshot for testability
|
||||
service_like = None
|
||||
|
||||
if not config:
|
||||
# Prefer AICORE_SERVICE_KEY if present; otherwise fall back to the VCAP service.
|
||||
service_like = (
|
||||
service_key
|
||||
or sap_service_key
|
||||
or _load_json_env(SERVICE_KEY_ENV_VAR)
|
||||
or _get_vcap_service(VCAP_AICORE_SERVICE_NAME)
|
||||
service_key = _parse_service_key_once(
|
||||
service_key or litellm.sap_service_key or os.environ.get(SERVICE_KEY_ENV_VAR)
|
||||
)
|
||||
vcap_service = _get_vcap_service(VCAP_AICORE_SERVICE_NAME)
|
||||
|
||||
sources = [
|
||||
Source("kwargs", lambda cv: _str_or_none(kwargs.get(cv.name))),
|
||||
Source(
|
||||
"service key",
|
||||
lambda cv: _resolve_credential_from_service_key(service_key, cv),
|
||||
),
|
||||
Source(
|
||||
"environment variables",
|
||||
lambda cv: _str_or_none(os.environ.get(f"AICORE_{cv.name.upper()}")),
|
||||
),
|
||||
Source(
|
||||
"config file",
|
||||
lambda cv: _str_or_none(
|
||||
config.get(f"AICORE_{cv.name.upper()}")
|
||||
if config.get(f"AICORE_{cv.name.upper()}") is not None
|
||||
else config.get(cv.name)
|
||||
),
|
||||
),
|
||||
Source(
|
||||
"VCAP service",
|
||||
lambda cv: (
|
||||
_str_or_none(
|
||||
_get_nested(
|
||||
vcap_service,
|
||||
(("credentials",) + cv.vcap_key) if cv.vcap_key else (cv.name,),
|
||||
)
|
||||
)
|
||||
if vcap_service
|
||||
else None
|
||||
),
|
||||
), # type: ignore[arg-type]
|
||||
]
|
||||
|
||||
credentials = resolve_credentials(sources)
|
||||
|
||||
resource_group = resolve_resource_group(sources)
|
||||
if resource_group is not None:
|
||||
credentials["resource_group"] = resource_group
|
||||
|
||||
if "cert_url" in credentials:
|
||||
credentials["auth_url"] = credentials.pop("cert_url")
|
||||
return credentials
|
||||
|
||||
|
||||
def validate_credentials(
|
||||
auth_url: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
client_id: Optional[str] = None,
|
||||
client_secret: Optional[str] = None,
|
||||
cert_str: Optional[str] = None,
|
||||
key_str: Optional[str] = None,
|
||||
cert_file_path: Optional[str] = None,
|
||||
key_file_path: Optional[str] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Validate SAP AI Core credentials for completeness and consistency.
|
||||
|
||||
Args:
|
||||
auth_url: OAuth2 token endpoint URL (required)
|
||||
base_url: SAP AI Core API base URL (required)
|
||||
client_id: OAuth2 client ID (required)
|
||||
client_secret: OAuth2 client secret (for secret-based auth)
|
||||
cert_str: PEM-encoded certificate string (for cert-based auth)
|
||||
key_str: PEM-encoded private key string (for cert-based auth)
|
||||
cert_file_path: Path to certificate file (for file-based cert auth)
|
||||
key_file_path: Path to private key file (for file-based cert auth)
|
||||
|
||||
Raises:
|
||||
ValueError: If required fields are missing or authentication mode is ambiguous.
|
||||
|
||||
Note:
|
||||
- This function does NOT validate resource_group (resolved separately).
|
||||
- Exactly one authentication method must be provided:
|
||||
* client_secret, OR
|
||||
* (cert_str AND key_str), OR
|
||||
* (cert_file_path AND key_file_path)
|
||||
"""
|
||||
if not auth_url or not client_id or not base_url:
|
||||
raise ValueError(
|
||||
"SAP AI Core credentials not found. "
|
||||
"Please provide credentials by setting appropriate environment variables "
|
||||
"(e.g. AICORE_CLIENT_ID, AICORE_CLIENT_SECRET, etc.)"
|
||||
)
|
||||
|
||||
out: Dict[str, str] = {}
|
||||
for cred in CREDENTIAL_VALUES:
|
||||
value = _resolve_value(cred, kwargs=kwargs, env=env, config=config, service_like=service_like) # type: ignore
|
||||
if value is None:
|
||||
continue
|
||||
if cred.transform_fn:
|
||||
value = cred.transform_fn(value)
|
||||
out[cred.name] = value
|
||||
if "cert_url" in out.keys():
|
||||
out["auth_url"] = out.pop("cert_url")
|
||||
return out
|
||||
modes = [
|
||||
bool(client_secret),
|
||||
bool(cert_str) and bool(key_str),
|
||||
bool(cert_file_path) and bool(key_file_path),
|
||||
]
|
||||
if sum(bool(m) for m in modes) != 1:
|
||||
raise ValueError(
|
||||
"SAP AI Core credentials are incomplete. "
|
||||
"Invalid credentials: provide exactly one of client_secret, "
|
||||
"(cert_str & key_str), or (cert_file_path & key_file_path)."
|
||||
)
|
||||
|
||||
|
||||
def _request_token(
|
||||
client_id: str, auth_url: str, timeout: float, cert_pair=None, client_secret=None
|
||||
) -> tuple[str, datetime]:
|
||||
data = {"grant_type": "client_credentials", "client_id": client_id}
|
||||
if client_secret:
|
||||
data["client_secret"] = client_secret
|
||||
|
||||
resp: Optional[httpx.Response] = None
|
||||
try:
|
||||
if cert_pair:
|
||||
with httpx.Client(cert=cert_pair) as raw_client:
|
||||
handler = HTTPHandler(client=raw_client)
|
||||
resp = handler.post(auth_url, data=data, timeout=timeout) # type: ignore[arg-type]
|
||||
payload = resp.json()
|
||||
else:
|
||||
handler = _get_httpx_client()
|
||||
resp = handler.post(auth_url, data=data, timeout=timeout) # type: ignore[arg-type]
|
||||
payload = resp.json()
|
||||
access_token = payload["access_token"]
|
||||
expires_in = int(payload.get("expires_in", 3600))
|
||||
expiry_date = datetime.now(timezone.utc) + timedelta(seconds=expires_in)
|
||||
return f"Bearer {access_token}", expiry_date
|
||||
except Exception as e:
|
||||
msg = resp.text if resp is not None else getattr(e, "text", str(e))
|
||||
raise RuntimeError(f"Token request failed: {msg}") from e
|
||||
|
||||
|
||||
def get_token_creator(
|
||||
service_key: Optional[str] = None,
|
||||
service_key: Optional[Union[str, dict]] = None,
|
||||
profile: Optional[str] = None,
|
||||
*,
|
||||
timeout: float = 30.0,
|
||||
|
|
@ -237,7 +423,7 @@ def get_token_creator(
|
|||
|
||||
Args:
|
||||
profile: Optional AICore profile name
|
||||
timeout: HTTP request timeout in seconds (default 30s)
|
||||
timeout: Timeout for HTTP requests
|
||||
expiry_buffer_minutes: Refresh the token this many minutes before expiry
|
||||
overrides: Any explicit credential overrides (client_id, client_secret, etc.)
|
||||
|
||||
|
|
@ -251,6 +437,7 @@ def get_token_creator(
|
|||
)
|
||||
|
||||
auth_url = credentials.get("auth_url")
|
||||
base_url = credentials.get("base_url")
|
||||
client_id = credentials.get("client_id")
|
||||
client_secret = credentials.get("client_secret")
|
||||
cert_str = credentials.get("cert_str")
|
||||
|
|
@ -259,49 +446,30 @@ def get_token_creator(
|
|||
key_file_path = credentials.get("key_file_path")
|
||||
|
||||
# Sanity check
|
||||
if not auth_url or not client_id:
|
||||
raise ValueError(
|
||||
"fetch_credentials did not return valid 'auth_url' or 'client_id'"
|
||||
)
|
||||
|
||||
modes = [
|
||||
client_secret is not None,
|
||||
(cert_str is not None and key_str is not None),
|
||||
(cert_file_path is not None and key_file_path is not None),
|
||||
]
|
||||
if sum(bool(m) for m in modes) != 1:
|
||||
raise ValueError(
|
||||
"Invalid credentials: provide exactly one of client_secret, "
|
||||
"(cert_str & key_str), or (cert_file_path & key_file_path)."
|
||||
)
|
||||
validate_credentials(
|
||||
auth_url,
|
||||
base_url,
|
||||
client_id,
|
||||
client_secret,
|
||||
cert_str,
|
||||
key_str,
|
||||
cert_file_path,
|
||||
key_file_path,
|
||||
)
|
||||
|
||||
lock = Lock()
|
||||
token: Optional[str] = None
|
||||
token_expiry: Optional[datetime] = None
|
||||
|
||||
def _request_token(cert_pair=None) -> tuple[str, datetime]:
|
||||
data = {"grant_type": "client_credentials", "client_id": client_id}
|
||||
if client_secret:
|
||||
data["client_secret"] = client_secret
|
||||
|
||||
client = _get_httpx_client()
|
||||
# with httpx.Client(cert=cert_pair, timeout=timeout) as client:
|
||||
resp = client.post(auth_url, data=data)
|
||||
try:
|
||||
resp.raise_for_status()
|
||||
payload = resp.json()
|
||||
access_token = payload["access_token"]
|
||||
expires_in = int(payload.get("expires_in", 3600))
|
||||
expiry_date = datetime.now(timezone.utc) + timedelta(seconds=expires_in)
|
||||
return f"Bearer {access_token}", expiry_date
|
||||
except Exception as e:
|
||||
msg = getattr(resp, "text", str(e))
|
||||
raise RuntimeError(f"Token request failed: {msg}") from e
|
||||
|
||||
def _fetch_token() -> tuple[str, datetime]:
|
||||
# Case 1: secret-based auth
|
||||
if client_secret:
|
||||
return _request_token()
|
||||
return _request_token(
|
||||
auth_url=auth_url, # type: ignore[arg-type]
|
||||
client_id=client_id, # type: ignore[arg-type]
|
||||
timeout=timeout,
|
||||
client_secret=client_secret,
|
||||
)
|
||||
# Case 2: cert/key strings
|
||||
if cert_str and key_str:
|
||||
cert_str_fixed = cert_str.replace("\\n", "\n")
|
||||
|
|
@ -313,9 +481,24 @@ def get_token_creator(
|
|||
f.write(cert_str_fixed)
|
||||
with open(key_path, "w") as f:
|
||||
f.write(key_str_fixed)
|
||||
return _request_token(cert_pair=(cert_path, key_path))
|
||||
return _request_token(
|
||||
auth_url=auth_url, # type: ignore[arg-type]
|
||||
client_id=client_id, # type: ignore[arg-type]
|
||||
timeout=timeout,
|
||||
cert_pair=(cert_path, key_path),
|
||||
)
|
||||
# Case 3: file-based cert/key
|
||||
return _request_token(cert_pair=(cert_file_path, key_file_path))
|
||||
if cert_file_path is not None and key_file_path is not None:
|
||||
return _request_token(
|
||||
auth_url=auth_url, # type: ignore[arg-type]
|
||||
client_id=client_id, # type: ignore[arg-type]
|
||||
timeout=timeout,
|
||||
cert_pair=(cert_file_path, key_file_path),
|
||||
)
|
||||
# Defensive guard: should never reach here due to validate_credentials()
|
||||
raise ValueError(
|
||||
"Invalid authentication configuration: no valid credentials found. "
|
||||
)
|
||||
|
||||
def get_token() -> str:
|
||||
nonlocal token, token_expiry
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ Translates from OpenAI's `/v1/embeddings` to IBM's `/text/embeddings` route.
|
|||
from typing import Optional, List, Dict, Literal, Union
|
||||
from pydantic import BaseModel, Field
|
||||
from functools import cached_property
|
||||
from litellm.llms.sap.chat.models import MaskingModuleConfig
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -47,25 +48,36 @@ class EmbeddingsResponse(BaseModel):
|
|||
class EmbeddingModel(BaseModel):
|
||||
name: str
|
||||
version: str = "latest"
|
||||
params: dict = Field(default_factory=dict, validation_alias="parameters")
|
||||
params: dict = Field(default_factory=dict)
|
||||
timeout: Optional[int] = Field(default=None, ge=1, le=600)
|
||||
max_retries: Optional[int] = Field(default=None, ge=0, le=5)
|
||||
|
||||
|
||||
class EmbeddingsModelConfig(BaseModel):
|
||||
model: EmbeddingModel
|
||||
|
||||
|
||||
class EmbeddingsModules(BaseModel):
|
||||
embeddings: EmbeddingModel
|
||||
embeddings: EmbeddingsModelConfig
|
||||
masking: Optional[MaskingModuleConfig] = None
|
||||
|
||||
|
||||
class EmbeddingInput(BaseModel):
|
||||
text: Union[str, List[str]]
|
||||
type: Literal["text", "document", "query"] = "text"
|
||||
type: Optional[Literal["text", "document", "query"]] = None
|
||||
|
||||
|
||||
class EmbeddingConfig(BaseModel):
|
||||
modules: EmbeddingsModules
|
||||
|
||||
|
||||
class EmbeddingRequest(BaseModel):
|
||||
config: EmbeddingsModules
|
||||
config: EmbeddingConfig
|
||||
input: EmbeddingInput
|
||||
|
||||
|
||||
def validate_dict(data: dict, model) -> dict:
|
||||
return model(**data).model_dump()
|
||||
return model(**data).model_dump(exclude_unset=True, by_alias=True)
|
||||
|
||||
|
||||
class GenAIHubEmbeddingConfig(BaseEmbeddingConfig):
|
||||
|
|
@ -152,15 +164,23 @@ class GenAIHubEmbeddingConfig(BaseEmbeddingConfig):
|
|||
model_dict["name"] = model
|
||||
model_dict["version"] = optional_params.get("version", "latest")
|
||||
model_dict["params"] = optional_params.get("parameters", {})
|
||||
timeout = optional_params.get("timeout", None)
|
||||
if timeout is not None:
|
||||
model_dict["timeout"] = timeout
|
||||
max_retries = optional_params.get("max_retries", None)
|
||||
if max_retries is not None:
|
||||
model_dict["max_retries"] = max_retries
|
||||
input_dict = {"text": input}
|
||||
input_type = optional_params.get("type")
|
||||
if input_type is not None:
|
||||
input_dict["type"] = input_type
|
||||
masking = optional_params.get("masking")
|
||||
masking = {"masking": masking} if masking is not None else {}
|
||||
body = {
|
||||
"config": {
|
||||
"modules": {
|
||||
"embeddings": {"model": validate_dict(model_dict, EmbeddingModel)}
|
||||
}
|
||||
},
|
||||
"input": validate_dict(input_dict, EmbeddingInput),
|
||||
"config": {"modules": {"embeddings": {"model": model_dict}, **masking}},
|
||||
"input": input_dict,
|
||||
}
|
||||
body = validate_dict(body, EmbeddingRequest)
|
||||
return body
|
||||
|
||||
def transform_embedding_response(
|
||||
|
|
|
|||
|
|
@ -8,7 +8,8 @@ from litellm.llms.base_llm.embedding.transformation import (
|
|||
LiteLLMLoggingObj,
|
||||
)
|
||||
from litellm.types.llms.openai import AllEmbeddingInputValues
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
from litellm.types.utils import EmbeddingResponse, Usage
|
||||
from litellm.utils import token_counter
|
||||
|
||||
from ..common_utils import TritonError
|
||||
|
||||
|
|
@ -103,8 +104,36 @@ class TritonEmbeddingConfig(BaseEmbeddingConfig):
|
|||
|
||||
model_response.model = raw_response_json.get("model_name", "None")
|
||||
model_response.data = _embedding_output
|
||||
model_response.usage = self._build_embedding_usage(
|
||||
model=model, request_data=request_data
|
||||
)
|
||||
return model_response
|
||||
|
||||
def _build_embedding_usage(self, model: str, request_data: dict) -> Usage:
|
||||
input_data = request_data.get("inputs", [])
|
||||
input_text_values: List[str] = []
|
||||
for item in input_data:
|
||||
if isinstance(item, dict) and item.get("name") == "input_text":
|
||||
data_values = item.get("data", [])
|
||||
if isinstance(data_values, list):
|
||||
input_text_values = [str(value) for value in data_values]
|
||||
break
|
||||
|
||||
prompt_tokens = 0
|
||||
for text in input_text_values:
|
||||
if not text:
|
||||
continue
|
||||
try:
|
||||
prompt_tokens += token_counter(model=model, text=text)
|
||||
except Exception:
|
||||
prompt_tokens += len(text.split())
|
||||
|
||||
return Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=0,
|
||||
total_tokens=prompt_tokens,
|
||||
)
|
||||
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
|
||||
) -> BaseLLMException:
|
||||
|
|
|
|||
|
|
@ -763,6 +763,16 @@ def _transform_request_body( # noqa: PLR0915
|
|||
data["generationConfig"] = generation_config
|
||||
if cached_content is not None:
|
||||
data["cachedContent"] = cached_content
|
||||
|
||||
if service_tier := optional_params.pop("service_tier", None):
|
||||
if isinstance(service_tier, str):
|
||||
if service_tier.lower() == "default":
|
||||
data["serviceTier"] = "standard"
|
||||
else:
|
||||
data["serviceTier"] = service_tier.lower()
|
||||
else:
|
||||
data["serviceTier"] = service_tier
|
||||
|
||||
# Only add labels for Vertex AI endpoints (not Google GenAI/AI Studio) and only if non-empty
|
||||
if labels and custom_llm_provider != LlmProviders.GEMINI:
|
||||
data["labels"] = labels
|
||||
|
|
|
|||
|
|
@ -318,6 +318,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
"parallel_tool_calls",
|
||||
"web_search_options",
|
||||
"include_server_side_tool_invocations",
|
||||
"service_tier",
|
||||
]
|
||||
|
||||
# Add penalty parameters only for non-preview models
|
||||
|
|
@ -362,6 +363,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
"""
|
||||
return Tools(googleSearch={})
|
||||
|
||||
def _map_service_tier_param(self, value: str, optional_params: dict) -> None:
|
||||
"""
|
||||
Map OpenAI service_tier (string) to Gemini serviceTier.
|
||||
'auto' maps to 'priority'.
|
||||
Other values are passed lowercased.
|
||||
"""
|
||||
if value.lower() == "auto":
|
||||
optional_params["service_tier"] = "priority"
|
||||
else:
|
||||
optional_params["service_tier"] = value.lower()
|
||||
|
||||
def _transform_computer_use_config(self, computer_use_config: dict) -> dict:
|
||||
"""
|
||||
Transform Computer Use configuration to Gemini API format.
|
||||
|
|
@ -1121,6 +1133,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
optional_params = self._add_tools_to_optional_params(
|
||||
optional_params, [_tools]
|
||||
)
|
||||
elif param == "service_tier" and isinstance(value, str):
|
||||
self._map_service_tier_param(value, optional_params)
|
||||
elif param == "include_server_side_tool_invocations" and value is True:
|
||||
optional_params["include_server_side_tool_invocations"] = True
|
||||
if litellm.vertex_ai_safety_settings is not None:
|
||||
|
|
@ -2415,6 +2429,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
"provider_specific_fields", {}
|
||||
)["traffic_type"] = traffic_type
|
||||
|
||||
## ADD SERVICE TIER ##
|
||||
if getattr(raw_response, "headers", None):
|
||||
if service_tier := raw_response.headers.get("x-gemini-service-tier"):
|
||||
if service_tier.lower() == "standard":
|
||||
setattr(model_response, "service_tier", "default")
|
||||
else:
|
||||
setattr(model_response, "service_tier", service_tier.lower())
|
||||
|
||||
except Exception as e:
|
||||
raise VertexAIError(
|
||||
message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
|
||||
|
|
@ -2513,6 +2535,7 @@ async def make_call(
|
|||
streaming_response=response.aiter_lines(),
|
||||
sync_stream=False,
|
||||
logging_obj=logging_obj,
|
||||
response_headers=response.headers,
|
||||
)
|
||||
# LOGGING
|
||||
logging_obj.post_call(
|
||||
|
|
@ -2555,6 +2578,7 @@ def make_sync_call(
|
|||
streaming_response=response.iter_lines(),
|
||||
sync_stream=True,
|
||||
logging_obj=logging_obj,
|
||||
response_headers=response.headers,
|
||||
)
|
||||
|
||||
# LOGGING
|
||||
|
|
@ -3011,7 +3035,11 @@ class VertexLLM(VertexBase):
|
|||
|
||||
class ModelResponseIterator:
|
||||
def __init__(
|
||||
self, streaming_response, sync_stream: bool, logging_obj: LoggingClass
|
||||
self,
|
||||
streaming_response,
|
||||
sync_stream: bool,
|
||||
logging_obj: LoggingClass,
|
||||
response_headers: Optional[Dict[str, str]] = None,
|
||||
):
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
check_is_function_call,
|
||||
|
|
@ -3022,10 +3050,120 @@ class ModelResponseIterator:
|
|||
self.accumulated_json = ""
|
||||
self.sent_first_chunk = False
|
||||
self.logging_obj = logging_obj
|
||||
self.response_headers = response_headers or {}
|
||||
self.is_function_call = check_is_function_call(logging_obj)
|
||||
self.cumulative_tool_call_index: int = 0
|
||||
self.has_seen_tool_calls: bool = False
|
||||
|
||||
def _apply_stream_candidates(
|
||||
self,
|
||||
_candidates: List[Candidates],
|
||||
model_response: Any,
|
||||
) -> Tuple[List[dict], List[dict], List[dict], List[dict]]:
|
||||
(
|
||||
grounding_metadata,
|
||||
url_context_metadata,
|
||||
safety_ratings,
|
||||
citation_metadata,
|
||||
self.cumulative_tool_call_index,
|
||||
) = VertexGeminiConfig._process_candidates(
|
||||
_candidates,
|
||||
model_response,
|
||||
self.logging_obj.optional_params,
|
||||
cumulative_tool_call_index=self.cumulative_tool_call_index,
|
||||
)
|
||||
|
||||
# Track whether tool_calls have been seen across streaming chunks.
|
||||
# Gemini sends tool_calls and finishReason in separate chunks,
|
||||
# so we need to remember if earlier chunks contained tool_calls
|
||||
# to correctly set finish_reason="tool_calls" per the OpenAI spec.
|
||||
if not self.has_seen_tool_calls:
|
||||
for choice in model_response.choices:
|
||||
if (
|
||||
hasattr(choice, "delta")
|
||||
and choice.delta
|
||||
and choice.delta.tool_calls
|
||||
):
|
||||
self.has_seen_tool_calls = True
|
||||
break
|
||||
|
||||
# Handle final chunk with finishReason but no content.
|
||||
# _process_candidates skips candidates without "content",
|
||||
# so the finish_reason from the final chunk is lost.
|
||||
if not model_response.choices and _candidates:
|
||||
from litellm.types.utils import Delta, StreamingChoices
|
||||
|
||||
for candidate in _candidates:
|
||||
finish_reason_str = candidate.get("finishReason")
|
||||
if finish_reason_str is not None:
|
||||
if self.has_seen_tool_calls:
|
||||
mapped_finish_reason = "tool_calls"
|
||||
else:
|
||||
mapped_finish_reason = VertexGeminiConfig._check_finish_reason(
|
||||
None, finish_reason_str
|
||||
)
|
||||
choice = StreamingChoices(
|
||||
finish_reason=mapped_finish_reason,
|
||||
index=candidate.get("index", 0),
|
||||
delta=Delta(content=None, role=None),
|
||||
logprobs=None,
|
||||
enhancements=None,
|
||||
)
|
||||
model_response.choices.append(choice)
|
||||
|
||||
# Also handle the case where the final chunk has empty
|
||||
# content (e.g. text:"") WITH finishReason. In this case
|
||||
# _process_candidates DOES create a choice, but maps
|
||||
# finishReason="STOP" to "stop" because the current chunk
|
||||
# has no tool_calls. Override if we saw tool_calls earlier.
|
||||
if self.has_seen_tool_calls:
|
||||
for choice in model_response.choices:
|
||||
if choice.finish_reason == "stop":
|
||||
choice.finish_reason = "tool_calls"
|
||||
|
||||
setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore
|
||||
setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore
|
||||
setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore
|
||||
setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore
|
||||
|
||||
return grounding_metadata, url_context_metadata, safety_ratings, citation_metadata
|
||||
|
||||
def _apply_stream_usage_metadata(
|
||||
self,
|
||||
processed_chunk: Any,
|
||||
model_response: Any,
|
||||
grounding_metadata: List[dict],
|
||||
) -> Optional[Usage]:
|
||||
if "usageMetadata" not in processed_chunk:
|
||||
return None
|
||||
|
||||
usage = VertexGeminiConfig._calculate_usage(
|
||||
completion_response=processed_chunk,
|
||||
)
|
||||
|
||||
web_search_requests = VertexGeminiConfig._calculate_web_search_requests(
|
||||
grounding_metadata
|
||||
)
|
||||
if web_search_requests is not None:
|
||||
cast(
|
||||
PromptTokensDetailsWrapper, usage.prompt_tokens_details
|
||||
).web_search_requests = web_search_requests
|
||||
|
||||
traffic_type = processed_chunk.get("usageMetadata", {}).get("trafficType")
|
||||
if traffic_type:
|
||||
model_response._hidden_params.setdefault(
|
||||
"provider_specific_fields", {}
|
||||
)["traffic_type"] = traffic_type
|
||||
|
||||
service_tier = self.response_headers.get("x-gemini-service-tier")
|
||||
if service_tier:
|
||||
if service_tier.lower() == "standard":
|
||||
setattr(model_response, "service_tier", "default")
|
||||
else:
|
||||
setattr(model_response, "service_tier", service_tier.lower())
|
||||
|
||||
return usage
|
||||
|
||||
def chunk_parser(self, chunk: dict) -> Optional["ModelResponseStream"]:
|
||||
try:
|
||||
verbose_logger.debug(f"RAW GEMINI CHUNK: {chunk}")
|
||||
|
|
@ -3043,101 +3181,23 @@ class ModelResponseIterator:
|
|||
if blocked_response is not None:
|
||||
model_response = blocked_response
|
||||
|
||||
usage: Optional[Usage] = None
|
||||
_candidates: Optional[List[Candidates]] = processed_chunk.get("candidates")
|
||||
grounding_metadata: List[dict] = []
|
||||
url_context_metadata: List[dict] = []
|
||||
safety_ratings: List[dict] = []
|
||||
citation_metadata: List[dict] = []
|
||||
|
||||
_candidates: Optional[List[Candidates]] = processed_chunk.get("candidates")
|
||||
if _candidates:
|
||||
(
|
||||
grounding_metadata,
|
||||
url_context_metadata,
|
||||
safety_ratings,
|
||||
citation_metadata,
|
||||
self.cumulative_tool_call_index,
|
||||
) = VertexGeminiConfig._process_candidates(
|
||||
_candidates,
|
||||
model_response,
|
||||
self.logging_obj.optional_params,
|
||||
cumulative_tool_call_index=self.cumulative_tool_call_index,
|
||||
)
|
||||
) = self._apply_stream_candidates(_candidates, model_response)
|
||||
|
||||
# Track whether tool_calls have been seen across streaming chunks.
|
||||
# Gemini sends tool_calls and finishReason in separate chunks,
|
||||
# so we need to remember if earlier chunks contained tool_calls
|
||||
# to correctly set finish_reason="tool_calls" per the OpenAI spec.
|
||||
if not self.has_seen_tool_calls:
|
||||
for choice in model_response.choices:
|
||||
if (
|
||||
hasattr(choice, "delta")
|
||||
and choice.delta
|
||||
and choice.delta.tool_calls
|
||||
):
|
||||
self.has_seen_tool_calls = True
|
||||
break
|
||||
|
||||
# Handle final chunk with finishReason but no content.
|
||||
# _process_candidates skips candidates without "content",
|
||||
# so the finish_reason from the final chunk is lost.
|
||||
if not model_response.choices and _candidates:
|
||||
from litellm.types.utils import Delta, StreamingChoices
|
||||
|
||||
for candidate in _candidates:
|
||||
finish_reason_str = candidate.get("finishReason")
|
||||
if finish_reason_str is not None:
|
||||
if self.has_seen_tool_calls:
|
||||
mapped_finish_reason = "tool_calls"
|
||||
else:
|
||||
mapped_finish_reason = (
|
||||
VertexGeminiConfig._check_finish_reason(
|
||||
None, finish_reason_str
|
||||
)
|
||||
)
|
||||
choice = StreamingChoices(
|
||||
finish_reason=mapped_finish_reason,
|
||||
index=candidate.get("index", 0),
|
||||
delta=Delta(content=None, role=None),
|
||||
logprobs=None,
|
||||
enhancements=None,
|
||||
)
|
||||
model_response.choices.append(choice)
|
||||
|
||||
# Also handle the case where the final chunk has empty
|
||||
# content (e.g. text:"") WITH finishReason. In this case
|
||||
# _process_candidates DOES create a choice, but maps
|
||||
# finishReason="STOP" to "stop" because the current chunk
|
||||
# has no tool_calls. Override if we saw tool_calls earlier.
|
||||
if self.has_seen_tool_calls:
|
||||
for choice in model_response.choices:
|
||||
if choice.finish_reason == "stop":
|
||||
choice.finish_reason = "tool_calls"
|
||||
|
||||
setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore
|
||||
setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore
|
||||
setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore
|
||||
setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore
|
||||
|
||||
if "usageMetadata" in processed_chunk:
|
||||
usage = VertexGeminiConfig._calculate_usage(
|
||||
completion_response=processed_chunk,
|
||||
)
|
||||
|
||||
web_search_requests = VertexGeminiConfig._calculate_web_search_requests(
|
||||
grounding_metadata
|
||||
)
|
||||
if web_search_requests is not None:
|
||||
cast(
|
||||
PromptTokensDetailsWrapper, usage.prompt_tokens_details
|
||||
).web_search_requests = web_search_requests
|
||||
|
||||
traffic_type = processed_chunk.get("usageMetadata", {}).get(
|
||||
"trafficType"
|
||||
)
|
||||
if traffic_type:
|
||||
model_response._hidden_params.setdefault(
|
||||
"provider_specific_fields", {}
|
||||
)["traffic_type"] = traffic_type
|
||||
usage = self._apply_stream_usage_metadata(
|
||||
processed_chunk, model_response, grounding_metadata
|
||||
)
|
||||
|
||||
setattr(model_response, "usage", usage) # type: ignore
|
||||
|
||||
|
|
|
|||
|
|
@ -136,6 +136,11 @@ class VertexBase:
|
|||
json_obj,
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
)
|
||||
elif isinstance(credential_source, dict) and "executable" in credential_source:
|
||||
creds = self._credentials_from_pluggable(
|
||||
json_obj,
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
)
|
||||
else:
|
||||
creds = self._credentials_from_identity_pool(
|
||||
json_obj,
|
||||
|
|
@ -190,6 +195,17 @@ class VertexBase:
|
|||
creds = creds.with_scopes(scopes)
|
||||
return creds
|
||||
|
||||
def _credentials_from_pluggable(self, json_obj, scopes):
|
||||
try:
|
||||
from google.auth import pluggable
|
||||
except ImportError:
|
||||
raise ImportError(GOOGLE_IMPORT_ERROR_MESSAGE)
|
||||
|
||||
creds = pluggable.Credentials.from_info(json_obj)
|
||||
if scopes and hasattr(creds, "requires_scopes") and creds.requires_scopes:
|
||||
creds = creds.with_scopes(scopes)
|
||||
return creds
|
||||
|
||||
def _credentials_from_identity_pool_with_aws(self, json_obj, scopes):
|
||||
try:
|
||||
from google.auth import aws
|
||||
|
|
|
|||
|
|
@ -7819,26 +7819,6 @@
|
|||
"cache_read_input_token_cost": 3e-08,
|
||||
"cache_creation_input_token_cost": 3.75e-07
|
||||
},
|
||||
"bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0": {
|
||||
"input_cost_per_token": 3.6e-06,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.8e-05,
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"cache_read_input_token_cost": 3.6e-07,
|
||||
"cache_creation_input_token_cost": 4.5e-06
|
||||
},
|
||||
"bedrock/us-gov-east-1/anthropic.claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
"cache_read_input_token_cost": 3.3e-07,
|
||||
|
|
@ -7857,7 +7837,29 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"supports_native_structured_output": true
|
||||
},
|
||||
"bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
"cache_read_input_token_cost": 3.3e-07,
|
||||
"input_cost_per_token": 3.3e-06,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.65e-05,
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_native_structured_output": true
|
||||
},
|
||||
"bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0": {
|
||||
"input_cost_per_token": 2.65e-06,
|
||||
|
|
@ -7990,26 +7992,6 @@
|
|||
"cache_read_input_token_cost": 3e-08,
|
||||
"cache_creation_input_token_cost": 3.75e-07
|
||||
},
|
||||
"bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0": {
|
||||
"input_cost_per_token": 3.6e-06,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.8e-05,
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"cache_read_input_token_cost": 3.6e-07,
|
||||
"cache_creation_input_token_cost": 4.5e-06
|
||||
},
|
||||
"bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
"cache_read_input_token_cost": 3.3e-07,
|
||||
|
|
@ -8028,7 +8010,29 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"supports_native_structured_output": true
|
||||
},
|
||||
"bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
"cache_read_input_token_cost": 3.3e-07,
|
||||
"input_cost_per_token": 3.3e-06,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.65e-05,
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_native_structured_output": true
|
||||
},
|
||||
"bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": {
|
||||
"input_cost_per_token": 2.65e-06,
|
||||
|
|
@ -13740,7 +13744,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
"supports_web_search": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini-2.5-flash-image": {
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
|
|
@ -13789,7 +13794,8 @@
|
|||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": false,
|
||||
"tpm": 8000000
|
||||
"tpm": 8000000,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini-3-pro-image-preview": {
|
||||
"input_cost_per_image": 0.0011,
|
||||
|
|
@ -13823,7 +13829,8 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
"supports_web_search": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini-3.1-flash-image-preview": {
|
||||
"input_cost_per_image": 0.00056,
|
||||
|
|
@ -13906,7 +13913,8 @@
|
|||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_native_streaming": true
|
||||
"supports_native_streaming": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"deep-research-pro-preview-12-2025": {
|
||||
"input_cost_per_image": 0.0011,
|
||||
|
|
@ -13985,7 +13993,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
"supports_web_search": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini-2.5-flash-lite-preview-09-2025": {
|
||||
"cache_read_input_token_cost": 1e-08,
|
||||
|
|
@ -14256,7 +14265,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
"supports_web_search": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini-3-pro-preview": {
|
||||
"deprecation_date": "2026-03-26",
|
||||
|
|
@ -15038,7 +15048,8 @@
|
|||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"tpm": 8000000
|
||||
"tpm": 8000000,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini/gemini-2.5-flash-image": {
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
|
|
@ -15088,7 +15099,8 @@
|
|||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"tpm": 8000000
|
||||
"tpm": 8000000,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini/gemini-3-pro-image-preview": {
|
||||
"input_cost_per_image": 0.0011,
|
||||
|
|
@ -15124,7 +15136,8 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
"supports_web_search": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini/gemini-3.1-flash-image-preview": {
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
|
|
@ -15243,7 +15256,8 @@
|
|||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"tpm": 250000
|
||||
"tpm": 250000,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini/gemini-2.5-flash-lite-preview-09-2025": {
|
||||
"cache_read_input_token_cost": 1e-08,
|
||||
|
|
@ -15683,7 +15697,8 @@
|
|||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_native_streaming": true,
|
||||
"tpm": 250000
|
||||
"tpm": 250000,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini/gemini-3-flash-preview": {
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
|
|
@ -16924,6 +16939,72 @@
|
|||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-06
|
||||
},
|
||||
"baseten/MiniMaxAI/MiniMax-M2.5": {
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-06
|
||||
},
|
||||
"baseten/nvidia/Nemotron-120B-A12B": {
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-07
|
||||
},
|
||||
"baseten/zai-org/GLM-5": {
|
||||
"input_cost_per_token": 9.5e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3.15e-06
|
||||
},
|
||||
"baseten/zai-org/GLM-4.7": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.2e-06
|
||||
},
|
||||
"baseten/zai-org/GLM-4.6": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.2e-06
|
||||
},
|
||||
"baseten/moonshotai/Kimi-K2.5": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-06
|
||||
},
|
||||
"baseten/moonshotai/Kimi-K2-Thinking": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.5e-06
|
||||
},
|
||||
"baseten/moonshotai/Kimi-K2-Instruct-0905": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.5e-06
|
||||
},
|
||||
"baseten/openai/gpt-oss-120b": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-07
|
||||
},
|
||||
"baseten/deepseek-ai/DeepSeek-V3.1": {
|
||||
"input_cost_per_token": 5e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-06
|
||||
},
|
||||
"baseten/deepseek-ai/DeepSeek-V3-0324": {
|
||||
"input_cost_per_token": 7.7e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.7e-07
|
||||
},
|
||||
"gmi/Qwen/Qwen3-VL-235B-A22B-Instruct-FP8": {
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "gmi",
|
||||
|
|
@ -38149,4 +38230,4 @@
|
|||
"supports_native_structured_output": true,
|
||||
"supports_pdf_input": true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -21,7 +21,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from litellm.types.mcp import MCPCredentials
|
||||
|
||||
|
||||
|
|
@ -583,6 +585,7 @@ async def store_user_oauth_credential(
|
|||
refresh_token: Optional[str] = None,
|
||||
expires_in: Optional[int] = None,
|
||||
scopes: Optional[List[str]] = None,
|
||||
skip_byok_guard: bool = False,
|
||||
) -> None:
|
||||
"""Persist an OAuth2 access token for a user+server pair.
|
||||
|
||||
|
|
@ -611,21 +614,26 @@ async def store_user_oauth_credential(
|
|||
|
||||
# Guard against silently overwriting a BYOK credential with an OAuth token.
|
||||
# BYOK credentials lack a "type" field (or use a non-"oauth2" type).
|
||||
existing = await prisma_client.db.litellm_mcpusercredentials.find_unique(
|
||||
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
|
||||
)
|
||||
if existing is not None:
|
||||
_byok_error = ValueError(
|
||||
f"A non-OAuth2 credential already exists for user {user_id} "
|
||||
f"and server {server_id}. Refusing to overwrite."
|
||||
# Skip the guard when the caller knows the row is already an OAuth2 credential
|
||||
# (e.g. during token refresh), saving an extra DB round-trip.
|
||||
if not skip_byok_guard:
|
||||
existing = await prisma_client.db.litellm_mcpusercredentials.find_unique(
|
||||
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
|
||||
)
|
||||
try:
|
||||
raw = json.loads(base64.urlsafe_b64decode(existing.credential_b64).decode())
|
||||
except Exception:
|
||||
# Credential is not base64+JSON — it's a plain-text BYOK key.
|
||||
raise _byok_error
|
||||
if raw.get("type") != "oauth2":
|
||||
raise _byok_error
|
||||
if existing is not None:
|
||||
_byok_error = ValueError(
|
||||
f"A non-OAuth2 credential already exists for user {user_id} "
|
||||
f"and server {server_id}. Refusing to overwrite."
|
||||
)
|
||||
try:
|
||||
raw = json.loads(
|
||||
base64.urlsafe_b64decode(existing.credential_b64).decode()
|
||||
)
|
||||
except Exception:
|
||||
# Credential is not base64+JSON — it's a plain-text BYOK key.
|
||||
raise _byok_error
|
||||
if raw.get("type") != "oauth2":
|
||||
raise _byok_error
|
||||
|
||||
encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode()
|
||||
await prisma_client.db.litellm_mcpusercredentials.upsert(
|
||||
|
|
@ -704,6 +712,115 @@ async def list_user_oauth_credentials(
|
|||
return results
|
||||
|
||||
|
||||
async def refresh_user_oauth_token(
|
||||
prisma_client: PrismaClient,
|
||||
user_id: str,
|
||||
server: Any,
|
||||
cred: Dict[str, Any],
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Attempt to refresh a per-user OAuth2 token using its stored refresh_token.
|
||||
|
||||
POSTs to ``server.token_url`` with ``grant_type=refresh_token``.
|
||||
|
||||
On success: persists the new credential via ``store_user_oauth_credential``
|
||||
and returns the updated payload dict.
|
||||
On failure (network error, invalid_grant, missing refresh_token, …): logs a
|
||||
warning and returns ``None`` — the caller is responsible for clearing the
|
||||
stale credential and triggering re-authentication.
|
||||
"""
|
||||
refresh_token: Optional[str] = cred.get("refresh_token")
|
||||
token_url: Optional[str] = getattr(server, "token_url", None)
|
||||
server_id: str = getattr(server, "server_id", "")
|
||||
client_id: Optional[str] = getattr(server, "client_id", None)
|
||||
client_secret: Optional[str] = getattr(server, "client_secret", None)
|
||||
|
||||
if not refresh_token:
|
||||
verbose_proxy_logger.debug(
|
||||
"refresh_user_oauth_token: no refresh_token stored for user=%s server=%s",
|
||||
user_id,
|
||||
server_id,
|
||||
)
|
||||
return None
|
||||
if not token_url:
|
||||
verbose_proxy_logger.debug(
|
||||
"refresh_user_oauth_token: server=%s has no token_url configured",
|
||||
server_id,
|
||||
)
|
||||
return None
|
||||
|
||||
token_data: Dict[str, str] = {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refresh_token,
|
||||
}
|
||||
if client_id:
|
||||
token_data["client_id"] = client_id
|
||||
if client_secret:
|
||||
token_data["client_secret"] = client_secret
|
||||
|
||||
try:
|
||||
async_client = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.Oauth2Check
|
||||
)
|
||||
response = await async_client.post(
|
||||
token_url,
|
||||
headers={"Accept": "application/json"},
|
||||
data=token_data,
|
||||
)
|
||||
response.raise_for_status()
|
||||
body: Dict[str, Any] = response.json()
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.warning(
|
||||
"refresh_user_oauth_token: refresh request failed for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
access_token: Optional[str] = body.get("access_token")
|
||||
if not access_token:
|
||||
verbose_proxy_logger.warning(
|
||||
"refresh_user_oauth_token: token response missing access_token for "
|
||||
"user=%s server=%s",
|
||||
user_id,
|
||||
server_id,
|
||||
)
|
||||
return None
|
||||
|
||||
expires_in: Optional[int] = None
|
||||
raw_expires = body.get("expires_in")
|
||||
try:
|
||||
expires_in = int(raw_expires) if raw_expires is not None else None
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
|
||||
# Rotate refresh token when the provider returns a new one
|
||||
new_refresh_token: Optional[str] = body.get("refresh_token") or refresh_token
|
||||
|
||||
raw_scope = body.get("scope")
|
||||
scopes: Optional[List[str]] = (
|
||||
raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None
|
||||
) or cred.get("scopes")
|
||||
|
||||
await store_user_oauth_credential(
|
||||
prisma_client=prisma_client,
|
||||
user_id=user_id,
|
||||
server_id=server_id,
|
||||
access_token=access_token,
|
||||
refresh_token=new_refresh_token,
|
||||
expires_in=expires_in,
|
||||
scopes=scopes,
|
||||
skip_byok_guard=True, # Row is already OAuth2; skip the extra find_unique check
|
||||
)
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"refresh_user_oauth_token: refreshed token for user=%s server=%s",
|
||||
user_id,
|
||||
server_id,
|
||||
)
|
||||
return await get_user_oauth_credential(prisma_client, user_id, server_id)
|
||||
|
||||
|
||||
async def approve_mcp_server(
|
||||
prisma_client: PrismaClient,
|
||||
server_id: str,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import json
|
||||
from typing import Optional, cast
|
||||
from typing import Any, Dict, Optional, cast
|
||||
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
|
||||
|
||||
from fastapi import APIRouter, Form, HTTPException, Request
|
||||
|
|
@ -148,6 +148,160 @@ def _resolve_oauth2_server_for_root_endpoints(
|
|||
return None
|
||||
|
||||
|
||||
def _validate_token_response(
|
||||
token_response: Dict[str, Any],
|
||||
validation_rules: Dict[str, Any],
|
||||
server_id: str,
|
||||
) -> None:
|
||||
"""Raise HTTPException 403 if any validation rule doesn't match the token response.
|
||||
|
||||
Supports dot-notation for nested fields (e.g. ``"team.enterprise_id"`` checks
|
||||
``token_response["team"]["enterprise_id"]``). Top-level keys are tried first,
|
||||
then dot-split traversal. All comparisons are string-coerced so that numeric
|
||||
values in the response (e.g. ``"org_id": 12345``) match string rules
|
||||
(``"org_id": "12345"``).
|
||||
"""
|
||||
for key, expected in validation_rules.items():
|
||||
actual: Any = token_response.get(key)
|
||||
# Try dot-notation traversal when top-level lookup returns None
|
||||
if actual is None and "." in key:
|
||||
obj: Any = token_response
|
||||
for part in key.split("."):
|
||||
if isinstance(obj, dict):
|
||||
obj = obj.get(part)
|
||||
else:
|
||||
obj = None
|
||||
break
|
||||
actual = obj
|
||||
# Treat absent fields as a distinct failure from a mismatched value
|
||||
if actual is None:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "token_validation_failed",
|
||||
"server_id": server_id,
|
||||
"field": key,
|
||||
"message": (
|
||||
f"OAuth token rejected: required field '{key}' is absent"
|
||||
),
|
||||
},
|
||||
)
|
||||
if str(actual) != str(expected):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "token_validation_failed",
|
||||
"server_id": server_id,
|
||||
"field": key,
|
||||
"message": (
|
||||
f"OAuth token rejected: '{key}' = '{actual}', "
|
||||
f"expected '{expected}'"
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def _extract_user_id_from_request(request: Request) -> Optional[str]:
|
||||
"""Best-effort extraction of LiteLLM user_id from the request's Authorization header.
|
||||
|
||||
Called at the OAuth token endpoint so that per-user tokens can be stored
|
||||
server-side. Uses a read-only cache lookup to avoid re-running the full
|
||||
auth pipeline (which has side effects such as rate-limit increments and
|
||||
spend logging). Returns ``None`` if no cached credential is found.
|
||||
"""
|
||||
auth_header = request.headers.get("Authorization") or request.headers.get(
|
||||
"authorization"
|
||||
)
|
||||
if not auth_header:
|
||||
return None
|
||||
lower = auth_header.lower()
|
||||
if not lower.startswith("bearer "):
|
||||
return None
|
||||
token = auth_header[7:].strip()
|
||||
try:
|
||||
from litellm.proxy._types import hash_token # noqa: PLC0415
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
cached = await user_api_key_cache.async_get_cache(hash_token(token))
|
||||
return getattr(cached, "user_id", None)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def _store_per_user_token_server_side(
|
||||
server: MCPServer,
|
||||
user_id: str,
|
||||
token_response: Dict[str, Any],
|
||||
) -> None:
|
||||
"""Persist the OAuth token server-side and warm the Redis cache.
|
||||
|
||||
Called from the token endpoint after a successful code exchange or refresh.
|
||||
Errors are logged but NOT re-raised — the token is always returned to the
|
||||
client even when server-side storage fails.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415
|
||||
_compute_per_user_token_ttl,
|
||||
mcp_per_user_token_cache,
|
||||
)
|
||||
from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415
|
||||
|
||||
access_token: Optional[str] = token_response.get("access_token")
|
||||
if not access_token:
|
||||
return
|
||||
|
||||
raw_expires = token_response.get("expires_in")
|
||||
try:
|
||||
expires_in: Optional[int] = int(raw_expires) if raw_expires is not None else None
|
||||
except (TypeError, ValueError):
|
||||
expires_in = None
|
||||
|
||||
refresh_token: Optional[str] = token_response.get("refresh_token") or None
|
||||
raw_scope = token_response.get("scope")
|
||||
scopes: Optional[list] = (
|
||||
raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None
|
||||
)
|
||||
|
||||
try:
|
||||
prisma_client = get_prisma_client_or_throw(
|
||||
"Database not connected. Cannot store per-user OAuth token."
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
|
||||
store_user_oauth_credential,
|
||||
)
|
||||
|
||||
await store_user_oauth_credential(
|
||||
prisma_client=prisma_client,
|
||||
user_id=user_id,
|
||||
server_id=server.server_id,
|
||||
access_token=access_token,
|
||||
refresh_token=refresh_token,
|
||||
expires_in=expires_in,
|
||||
scopes=scopes,
|
||||
)
|
||||
verbose_logger.info(
|
||||
"_store_per_user_token_server_side: stored token for user=%s server=%s",
|
||||
user_id,
|
||||
server.server_id,
|
||||
)
|
||||
except Exception as exc:
|
||||
verbose_logger.warning(
|
||||
"_store_per_user_token_server_side: DB storage failed for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server.server_id,
|
||||
exc,
|
||||
)
|
||||
return # Don't warm Redis if DB write failed
|
||||
|
||||
# Warm the Redis cache so the first subsequent MCP call is a cache hit
|
||||
ttl = _compute_per_user_token_ttl(server, expires_in)
|
||||
await mcp_per_user_token_cache.set(
|
||||
user_id=user_id,
|
||||
server_id=server.server_id,
|
||||
access_token=access_token,
|
||||
ttl=ttl,
|
||||
)
|
||||
|
||||
|
||||
async def authorize_with_server(
|
||||
request: Request,
|
||||
mcp_server: MCPServer,
|
||||
|
|
@ -267,6 +421,44 @@ async def exchange_token_with_server(
|
|||
token_response = response.json()
|
||||
access_token = token_response["access_token"]
|
||||
|
||||
# Validate token response against server-configured rules before any storage.
|
||||
# This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc.
|
||||
if mcp_server.token_validation and isinstance(mcp_server.token_validation, dict):
|
||||
_validate_token_response(
|
||||
token_response=token_response,
|
||||
validation_rules=mcp_server.token_validation,
|
||||
server_id=mcp_server.server_id,
|
||||
)
|
||||
|
||||
# Store server-side when the server is configured for per-user OAuth and
|
||||
# the calling client has provided a valid LiteLLM identity.
|
||||
# Errors are non-fatal: the token is still returned to the client.
|
||||
if mcp_server.needs_user_oauth_token:
|
||||
user_id = await _extract_user_id_from_request(request)
|
||||
if user_id:
|
||||
try:
|
||||
await _store_per_user_token_server_side(
|
||||
server=mcp_server,
|
||||
user_id=user_id,
|
||||
token_response=token_response,
|
||||
)
|
||||
except Exception as exc:
|
||||
verbose_logger.warning(
|
||||
"exchange_token_with_server: server-side storage failed "
|
||||
"for user=%s server=%s: %s",
|
||||
user_id,
|
||||
mcp_server.server_id,
|
||||
exc,
|
||||
)
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
"exchange_token_with_server: no LiteLLM user_id found in request; "
|
||||
"per-user token for server=%s will not be stored server-side. "
|
||||
"The client should call POST /mcp/server/{id}/oauth-user-credential "
|
||||
"to store it manually.",
|
||||
mcp_server.server_id,
|
||||
)
|
||||
|
||||
result = {
|
||||
"access_token": access_token,
|
||||
"token_type": token_response.get("token_type", "Bearer"),
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import asyncio
|
|||
import datetime
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from typing import Any, Callable, Dict, List, Literal, Optional, Set, Tuple, Union, cast
|
||||
from urllib.parse import urlparse
|
||||
|
|
@ -35,6 +36,8 @@ from litellm.constants import (
|
|||
MCP_CLIENT_TIMEOUT,
|
||||
MCP_HEALTH_CHECK_TIMEOUT,
|
||||
MCP_METADATA_TIMEOUT,
|
||||
MCP_NPM_CACHE_DIR,
|
||||
MCP_STDIO_ALLOWED_COMMANDS,
|
||||
MCP_TOOL_LISTING_TIMEOUT,
|
||||
)
|
||||
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
|
||||
|
|
@ -1119,9 +1122,19 @@ class MCPServerManager:
|
|||
# In containers the default (~/.npm or /app/.npm) may not exist
|
||||
# or be read-only, causing npx to fail with ENOENT.
|
||||
if "NPM_CONFIG_CACHE" not in resolved_env:
|
||||
from litellm.constants import MCP_NPM_CACHE_DIR
|
||||
|
||||
resolved_env["NPM_CONFIG_CACHE"] = MCP_NPM_CACHE_DIR
|
||||
# Defense-in-depth: block commands not in the allowlist.
|
||||
# The Pydantic validator blocks new servers; this catches legacy
|
||||
# config/DB records predating the allowlist.
|
||||
if server.command:
|
||||
base_command = os.path.basename(server.command)
|
||||
if base_command not in MCP_STDIO_ALLOWED_COMMANDS:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"MCP stdio command '{server.command}' is not in the allowlist ({sorted(MCP_STDIO_ALLOWED_COMMANDS)}). "
|
||||
f"Add it to LITELLM_MCP_STDIO_EXTRA_COMMANDS to allow this command.",
|
||||
)
|
||||
|
||||
stdio_config: Optional[MCPStdioConfig] = None
|
||||
if server.command and server.args is not None:
|
||||
stdio_config = MCPStdioConfig(
|
||||
|
|
@ -2442,6 +2455,37 @@ class MCPServerManager:
|
|||
)
|
||||
tasks.append(during_hook_task)
|
||||
|
||||
# For per-user OAuth servers: if the client didn't supply a token in
|
||||
# oauth2_headers, look up the stored token from Redis / DB. This is the
|
||||
# call_tool equivalent of _get_user_oauth_extra_headers_from_db used in
|
||||
# list_tools.
|
||||
if (
|
||||
mcp_server.needs_user_oauth_token
|
||||
and not oauth2_headers
|
||||
and user_api_key_auth is not None
|
||||
):
|
||||
user_id = getattr(user_api_key_auth, "user_id", None)
|
||||
if user_id:
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415
|
||||
_get_user_oauth_extra_headers_from_db,
|
||||
)
|
||||
|
||||
stored_headers = await _get_user_oauth_extra_headers_from_db(
|
||||
server=mcp_server,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
)
|
||||
if stored_headers:
|
||||
oauth2_headers = stored_headers
|
||||
except Exception as _lookup_exc:
|
||||
verbose_logger.debug(
|
||||
"call_tool: per-user token lookup failed for "
|
||||
"user=%s server=%s: %s",
|
||||
user_id,
|
||||
mcp_server.server_id,
|
||||
_lookup_exc,
|
||||
)
|
||||
|
||||
# For OpenAPI servers, call the tool handler directly instead of via MCP client
|
||||
if mcp_server.spec_path:
|
||||
verbose_logger.debug(
|
||||
|
|
|
|||
|
|
@ -17,8 +17,15 @@ from litellm.constants import (
|
|||
MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE,
|
||||
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL,
|
||||
MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS,
|
||||
MCP_PER_USER_TOKEN_DEFAULT_TTL,
|
||||
MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS,
|
||||
MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -152,6 +159,107 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
mcp_oauth2_token_cache = MCPOAuth2TokenCache()
|
||||
|
||||
|
||||
def _compute_per_user_token_ttl(server: "MCPServer", expires_in: Optional[int]) -> int:
|
||||
"""Compute Redis TTL for a per-user token.
|
||||
|
||||
Uses server.token_storage_ttl_seconds when configured; otherwise derives
|
||||
TTL from expires_in minus the expiry buffer; falls back to the default TTL.
|
||||
"""
|
||||
if server.token_storage_ttl_seconds is not None:
|
||||
return max(server.token_storage_ttl_seconds, 1)
|
||||
if expires_in is not None:
|
||||
return max(
|
||||
expires_in - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS,
|
||||
1,
|
||||
)
|
||||
return MCP_PER_USER_TOKEN_DEFAULT_TTL
|
||||
|
||||
|
||||
class MCPPerUserTokenCache:
|
||||
"""Redis-backed cache for per-user OAuth2 access tokens.
|
||||
|
||||
Uses LiteLLM's existing ``user_api_key_cache`` (DualCache with optional
|
||||
Redis backend). Tokens are NaCl-encrypted with ``encrypt_value_helper``
|
||||
before storage so they are safe at rest in Redis.
|
||||
|
||||
Redis key format: ``mcp:per_user_token:{user_id}:{server_id}``
|
||||
Redis value: ``encrypt_value_helper(access_token)`` — URL-safe base64
|
||||
"""
|
||||
|
||||
def _cache_key(self, user_id: str, server_id: str) -> str:
|
||||
return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}"
|
||||
|
||||
async def get(self, user_id: str, server_id: str) -> Optional[str]:
|
||||
"""Return the plaintext access_token, or None on miss/error."""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
key = self._cache_key(user_id, server_id)
|
||||
encrypted = await user_api_key_cache.async_get_cache(key)
|
||||
if encrypted is None:
|
||||
return None
|
||||
plaintext = decrypt_value_helper(
|
||||
encrypted,
|
||||
key="mcp_per_user_token",
|
||||
exception_type="debug",
|
||||
)
|
||||
return plaintext or None
|
||||
except Exception as exc:
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.get failed for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
async def set(
|
||||
self,
|
||||
user_id: str,
|
||||
server_id: str,
|
||||
access_token: str,
|
||||
ttl: int,
|
||||
) -> None:
|
||||
"""Store NaCl-encrypted access_token in Redis with the given TTL."""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
key = self._cache_key(user_id, server_id)
|
||||
encrypted = encrypt_value_helper(access_token)
|
||||
await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl)
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds",
|
||||
user_id,
|
||||
server_id,
|
||||
ttl,
|
||||
)
|
||||
except Exception as exc:
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.set failed for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
exc,
|
||||
)
|
||||
|
||||
async def delete(self, user_id: str, server_id: str) -> None:
|
||||
"""Invalidate the cached token (removes from both in-memory and Redis layers)."""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
key = self._cache_key(user_id, server_id)
|
||||
await user_api_key_cache.async_delete_cache(key)
|
||||
except Exception as exc:
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.delete failed for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
exc,
|
||||
)
|
||||
|
||||
|
||||
mcp_per_user_token_cache = MCPPerUserTokenCache()
|
||||
|
||||
|
||||
async def resolve_mcp_auth(
|
||||
server: "MCPServer",
|
||||
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
|
||||
|
|
|
|||
|
|
@ -2,14 +2,14 @@ import importlib
|
|||
from datetime import datetime
|
||||
from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Set, Union
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
|
||||
build_effective_auth_contexts,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.utils import merge_mcp_headers
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
|
||||
|
|
@ -1027,6 +1027,13 @@ if MCP_AVAILABLE:
|
|||
"""
|
||||
Test if we can connect to the provided MCP server before adding it
|
||||
"""
|
||||
if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={
|
||||
"error": "User does not have permission to test MCP server connections. Only PROXY_ADMIN users can perform this action."
|
||||
},
|
||||
)
|
||||
|
||||
async def _test_connection_operation(client):
|
||||
async def _noop(session):
|
||||
|
|
@ -1041,7 +1048,7 @@ if MCP_AVAILABLE:
|
|||
raw_headers=_safe_get_request_headers(request),
|
||||
)
|
||||
|
||||
@router.post("/test/tools/list")
|
||||
@router.post("/test/tools/list", dependencies=[Depends(user_api_key_auth)])
|
||||
async def test_tools_list(
|
||||
request: Request,
|
||||
new_mcp_server_request: NewMCPServerRequest,
|
||||
|
|
@ -1050,6 +1057,14 @@ if MCP_AVAILABLE:
|
|||
"""
|
||||
Preview tools available from MCP server before adding it
|
||||
"""
|
||||
if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={
|
||||
"error": "User does not have permission to test MCP server tools. Only PROXY_ADMIN users can perform this action."
|
||||
},
|
||||
)
|
||||
|
||||
# For OpenAPI spec servers, generate tools from the spec directly
|
||||
if new_mcp_server_request.spec_path:
|
||||
return await _preview_openapi_tools(new_mcp_server_request.spec_path)
|
||||
|
|
|
|||
|
|
@ -896,11 +896,17 @@ if MCP_AVAILABLE:
|
|||
user_api_key_auth: Optional[UserAPIKeyAuth],
|
||||
prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None,
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""Look up stored OAuth2 token for (user, server) from DB and return as extra_headers dict.
|
||||
"""Look up stored OAuth2 token for (user, server) and return as extra_headers dict.
|
||||
|
||||
Lookup order:
|
||||
1. Redis cache (fast path, NaCl-decrypted) — skipped when prefetched_creds supplied
|
||||
2. prefetched_creds dict (pre-fetched batch DB query) or fresh DB query
|
||||
3. Auto-refresh when the stored token is expired and a refresh_token exists
|
||||
|
||||
Args:
|
||||
prefetched_creds: Optional dict keyed by server_id with credential payloads.
|
||||
When provided, avoids a per-server DB round-trip.
|
||||
When provided, the Redis and individual DB lookups are
|
||||
skipped in favour of the pre-fetched batch result.
|
||||
"""
|
||||
if server.auth_type != MCPAuth.oauth2:
|
||||
return None
|
||||
|
|
@ -914,8 +920,27 @@ if MCP_AVAILABLE:
|
|||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
|
||||
get_user_oauth_credential,
|
||||
is_oauth_credential_expired,
|
||||
refresh_user_oauth_token,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415
|
||||
_compute_per_user_token_ttl,
|
||||
mcp_per_user_token_cache,
|
||||
)
|
||||
|
||||
# ── Fast path: Redis cache ────────────────────────────────────────
|
||||
# Only used when prefetched_creds is not supplied (individual lookup).
|
||||
if prefetched_creds is None:
|
||||
cached_token = await mcp_per_user_token_cache.get(user_id, server_id)
|
||||
if cached_token is not None:
|
||||
verbose_logger.debug(
|
||||
"_get_user_oauth_extra_headers_from_db: Redis hit for "
|
||||
"user=%s server=%s",
|
||||
user_id,
|
||||
server_id,
|
||||
)
|
||||
return {"Authorization": f"Bearer {cached_token}"}
|
||||
|
||||
# ── Slow path: DB lookup ──────────────────────────────────────────
|
||||
if prefetched_creds is not None:
|
||||
cred = prefetched_creds.get(server_id)
|
||||
else:
|
||||
|
|
@ -929,18 +954,83 @@ if MCP_AVAILABLE:
|
|||
cred = await get_user_oauth_credential(
|
||||
prisma_client, user_id, server_id
|
||||
)
|
||||
if cred and cred.get("access_token"):
|
||||
if is_oauth_credential_expired(cred):
|
||||
verbose_logger.debug(
|
||||
f"_get_user_oauth_extra_headers_from_db: token expired for "
|
||||
f"user={user_id} server={server_id}"
|
||||
)
|
||||
|
||||
if not cred or not cred.get("access_token"):
|
||||
return None
|
||||
|
||||
if is_oauth_credential_expired(cred):
|
||||
verbose_logger.debug(
|
||||
"_get_user_oauth_extra_headers_from_db: token expired for "
|
||||
"user=%s server=%s — attempting refresh",
|
||||
user_id,
|
||||
server_id,
|
||||
)
|
||||
# Attempt token refresh; requires a DB client (not available from prefetch)
|
||||
if cred.get("refresh_token"):
|
||||
try:
|
||||
from litellm.proxy.utils import ( # noqa: PLC0415
|
||||
get_prisma_client_or_throw,
|
||||
)
|
||||
|
||||
prisma_client = get_prisma_client_or_throw(
|
||||
"Database not connected. Cannot refresh OAuth token."
|
||||
)
|
||||
cred = await refresh_user_oauth_token(
|
||||
prisma_client=prisma_client,
|
||||
user_id=user_id,
|
||||
server=server,
|
||||
cred=cred,
|
||||
)
|
||||
except Exception as refresh_exc:
|
||||
verbose_logger.warning(
|
||||
"_get_user_oauth_extra_headers_from_db: refresh failed "
|
||||
"for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
refresh_exc,
|
||||
)
|
||||
cred = None
|
||||
|
||||
if not cred or not cred.get("access_token"):
|
||||
# Clear stale Redis/cache entry so we don't serve it again.
|
||||
# Do this for both the individual and prefetch paths so the
|
||||
# next request doesn't get a stale cache hit.
|
||||
await mcp_per_user_token_cache.delete(user_id, server_id)
|
||||
return None
|
||||
return {"Authorization": f"Bearer {cred['access_token']}"}
|
||||
|
||||
access_token: str = cred["access_token"]
|
||||
|
||||
# Warm (or re-warm) the Redis cache from the DB result.
|
||||
# Always write regardless of whether expires_at is present — tokens
|
||||
# without an expiry are still valid and should be cached using the
|
||||
# server/default TTL so subsequent requests are fast.
|
||||
if prefetched_creds is None:
|
||||
raw_expires = None
|
||||
expires_at = cred.get("expires_at")
|
||||
if expires_at:
|
||||
from datetime import datetime, timezone # noqa: PLC0415
|
||||
|
||||
try:
|
||||
exp_dt = datetime.fromisoformat(expires_at)
|
||||
if exp_dt.tzinfo is None:
|
||||
exp_dt = exp_dt.replace(tzinfo=timezone.utc)
|
||||
remaining = int(
|
||||
(exp_dt - datetime.now(timezone.utc)).total_seconds()
|
||||
)
|
||||
raw_expires = max(remaining, 0) if remaining > 0 else None
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
ttl = _compute_per_user_token_ttl(server, raw_expires)
|
||||
await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl)
|
||||
|
||||
return {"Authorization": f"Bearer {access_token}"}
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for "
|
||||
f"user={user_id} server={server_id}: {e}"
|
||||
"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for "
|
||||
"user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
|
||||
|
|
@ -2504,6 +2594,14 @@ if MCP_AVAILABLE:
|
|||
server_name, client_ip=_client_ip
|
||||
)
|
||||
if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers:
|
||||
# For servers that store per-user tokens server-side, skip the
|
||||
# pre-emptive 401 — the call_tool / list_tools dispatch will look
|
||||
# up the stored token from Redis / DB and only fail at the MCP
|
||||
# protocol level if none is found, giving the client a proper
|
||||
# tool-execution error rather than an HTTP 401.
|
||||
if server.needs_user_oauth_token:
|
||||
continue
|
||||
|
||||
request = StarletteRequest(scope)
|
||||
base_url = get_request_base_url(request)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import enum
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union
|
||||
|
||||
|
|
@ -15,6 +16,7 @@ from pydantic import (
|
|||
from typing_extensions import Required, TypedDict
|
||||
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import MCP_STDIO_ALLOWED_COMMANDS
|
||||
from litellm.types.integrations.slack_alerting import AlertType
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
|
|
@ -492,10 +494,12 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/v2/key/info",
|
||||
"/model_group/info",
|
||||
"/health",
|
||||
"/health/services",
|
||||
"/key/list",
|
||||
"/user/filter/ui",
|
||||
"/models",
|
||||
"/v1/models",
|
||||
"/sso/get/ui_settings",
|
||||
]
|
||||
|
||||
# NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend
|
||||
|
|
@ -564,6 +568,8 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/spend/tags",
|
||||
"/spend/calculate",
|
||||
"/spend/logs",
|
||||
"/spend/logs/ui",
|
||||
"/spend/logs/session/ui",
|
||||
"/cost/estimate",
|
||||
]
|
||||
|
||||
|
|
@ -579,6 +585,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/global/spend/report",
|
||||
"/global/spend/provider",
|
||||
"/global/spend/tags",
|
||||
"/global/spend/all_tag_names",
|
||||
]
|
||||
|
||||
public_routes = set(
|
||||
|
|
@ -600,6 +607,9 @@ class LiteLLMRoutes(enum.Enum):
|
|||
]
|
||||
)
|
||||
|
||||
# Retained for backwards compatibility with JWT auth configs that reference
|
||||
# "ui_routes" in admin_allowed_routes. Not used by the proxy's own route
|
||||
# authorization — UI tokens now go through the same RBAC path as API tokens.
|
||||
ui_routes = [
|
||||
"/sso",
|
||||
"/sso/get/ui_settings",
|
||||
|
|
@ -625,19 +635,16 @@ class LiteLLMRoutes(enum.Enum):
|
|||
|
||||
internal_user_routes = (
|
||||
[
|
||||
"/global/spend/tags",
|
||||
"/global/spend/keys",
|
||||
"/global/spend/models",
|
||||
"/global/spend/provider",
|
||||
"/global/spend/end_users",
|
||||
"/global/activity",
|
||||
"/global/activity/model",
|
||||
"/global/activity/cache_hits",
|
||||
"/v1/models/{model_id}",
|
||||
"/models/{model_id}",
|
||||
"/guardrails/list",
|
||||
"/v2/guardrails/list",
|
||||
]
|
||||
+ spend_tracking_routes
|
||||
+ global_spend_tracking_routes
|
||||
+ key_management_routes
|
||||
)
|
||||
|
||||
|
|
@ -692,6 +699,9 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/tag/list",
|
||||
"/audit",
|
||||
"/audit/{id}",
|
||||
"/global/activity",
|
||||
"/global/activity/model",
|
||||
"/global/activity/cache_hits",
|
||||
] + info_routes
|
||||
|
||||
# All routes accesible by an Org Admin
|
||||
|
|
@ -890,9 +900,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase):
|
|||
allowed_cache_controls: Optional[list] = []
|
||||
config: Optional[dict] = {}
|
||||
permissions: Optional[dict] = {}
|
||||
model_max_budget: Optional[dict] = (
|
||||
{}
|
||||
) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {}
|
||||
model_max_budget: Optional[
|
||||
dict
|
||||
] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {}
|
||||
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
model_rpm_limit: Optional[dict] = None
|
||||
|
|
@ -1034,9 +1044,9 @@ class RegenerateKeyRequest(GenerateKeyRequest):
|
|||
spend: Optional[float] = None
|
||||
metadata: Optional[dict] = None
|
||||
new_master_key: Optional[str] = None
|
||||
grace_period: Optional[str] = (
|
||||
None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke
|
||||
)
|
||||
grace_period: Optional[
|
||||
str
|
||||
] = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke
|
||||
|
||||
|
||||
class ResetSpendRequest(LiteLLMPydanticObjectBase):
|
||||
|
|
@ -1162,6 +1172,13 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
|
|||
raise ValueError("command is required for stdio transport")
|
||||
if not values.get("args"):
|
||||
raise ValueError("args is required for stdio transport")
|
||||
# Validate command against allowlist to prevent arbitrary execution
|
||||
base_command = os.path.basename(values["command"])
|
||||
if base_command not in MCP_STDIO_ALLOWED_COMMANDS:
|
||||
raise ValueError(
|
||||
f"Command '{values['command']}' is not in the allowed commands list "
|
||||
f"for stdio transport. Allowed commands: {sorted(MCP_STDIO_ALLOWED_COMMANDS)}"
|
||||
)
|
||||
elif transport in [MCPTransport.http, MCPTransport.sse]:
|
||||
if not values.get("url") and not values.get("spec_path"):
|
||||
raise ValueError(
|
||||
|
|
@ -1222,6 +1239,13 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
|
|||
raise ValueError("command is required for stdio transport")
|
||||
if not values.get("args"):
|
||||
raise ValueError("args is required for stdio transport")
|
||||
# Validate command against allowlist to prevent arbitrary execution
|
||||
base_command = os.path.basename(values["command"])
|
||||
if base_command not in MCP_STDIO_ALLOWED_COMMANDS:
|
||||
raise ValueError(
|
||||
f"Command '{values['command']}' is not in the allowed commands list "
|
||||
f"for stdio transport. Allowed commands: {sorted(MCP_STDIO_ALLOWED_COMMANDS)}"
|
||||
)
|
||||
elif transport in [MCPTransport.http, MCPTransport.sse]:
|
||||
if not values.get("url") and not values.get("spec_path"):
|
||||
raise ValueError(
|
||||
|
|
@ -1546,12 +1570,12 @@ class NewCustomerRequest(BudgetNewRequest):
|
|||
blocked: bool = False # allow/disallow requests for this end-user
|
||||
budget_id: Optional[str] = None # give either a budget_id or max_budget
|
||||
spend: Optional[float] = None
|
||||
allowed_model_region: Optional[AllowedModelRegion] = (
|
||||
None # require all user requests to use models in this specific region
|
||||
)
|
||||
default_model: Optional[str] = (
|
||||
None # if no equivalent model in allowed region - default all requests to this model
|
||||
)
|
||||
allowed_model_region: Optional[
|
||||
AllowedModelRegion
|
||||
] = None # require all user requests to use models in this specific region
|
||||
default_model: Optional[
|
||||
str
|
||||
] = None # if no equivalent model in allowed region - default all requests to this model
|
||||
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
|
||||
|
||||
@model_validator(mode="before")
|
||||
|
|
@ -1574,12 +1598,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase):
|
|||
blocked: bool = False # allow/disallow requests for this end-user
|
||||
max_budget: Optional[float] = None
|
||||
budget_id: Optional[str] = None # give either a budget_id or max_budget
|
||||
allowed_model_region: Optional[AllowedModelRegion] = (
|
||||
None # require all user requests to use models in this specific region
|
||||
)
|
||||
default_model: Optional[str] = (
|
||||
None # if no equivalent model in allowed region - default all requests to this model
|
||||
)
|
||||
allowed_model_region: Optional[
|
||||
AllowedModelRegion
|
||||
] = None # require all user requests to use models in this specific region
|
||||
default_model: Optional[
|
||||
str
|
||||
] = None # if no equivalent model in allowed region - default all requests to this model
|
||||
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
|
||||
|
||||
|
||||
|
|
@ -1669,15 +1693,15 @@ class NewTeamRequest(TeamBase):
|
|||
] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm
|
||||
|
||||
model_tpm_limit: Optional[Dict[str, int]] = None
|
||||
team_member_budget: Optional[float] = (
|
||||
None # allow user to set a budget for all team members
|
||||
)
|
||||
team_member_rpm_limit: Optional[int] = (
|
||||
None # allow user to set RPM limit for all team members
|
||||
)
|
||||
team_member_tpm_limit: Optional[int] = (
|
||||
None # allow user to set TPM limit for all team members
|
||||
)
|
||||
team_member_budget: Optional[
|
||||
float
|
||||
] = None # allow user to set a budget for all team members
|
||||
team_member_rpm_limit: Optional[
|
||||
int
|
||||
] = None # allow user to set RPM limit for all team members
|
||||
team_member_tpm_limit: Optional[
|
||||
int
|
||||
] = None # allow user to set TPM limit for all team members
|
||||
team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m"
|
||||
team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo"
|
||||
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
|
||||
|
|
@ -1774,9 +1798,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase):
|
|||
|
||||
class AddTeamCallback(LiteLLMPydanticObjectBase):
|
||||
callback_name: str
|
||||
callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = (
|
||||
"success_and_failure"
|
||||
)
|
||||
callback_type: Optional[
|
||||
Literal["success", "failure", "success_and_failure"]
|
||||
] = "success_and_failure"
|
||||
callback_vars: Dict[str, str]
|
||||
|
||||
@model_validator(mode="before")
|
||||
|
|
@ -2118,9 +2142,9 @@ class ConfigList(LiteLLMPydanticObjectBase):
|
|||
stored_in_db: Optional[bool]
|
||||
field_default_value: Any
|
||||
premium_field: bool = False
|
||||
nested_fields: Optional[List[FieldDetail]] = (
|
||||
None # For nested dictionary or Pydantic fields
|
||||
)
|
||||
nested_fields: Optional[
|
||||
List[FieldDetail]
|
||||
] = None # For nested dictionary or Pydantic fields
|
||||
|
||||
|
||||
class UserHeaderMapping(LiteLLMPydanticObjectBase):
|
||||
|
|
@ -2419,6 +2443,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
|
|||
end_user_model_max_budget: Optional[dict] = None
|
||||
|
||||
# Organization Params
|
||||
organization_alias: Optional[str] = None
|
||||
organization_max_budget: Optional[float] = None
|
||||
organization_tpm_limit: Optional[int] = None
|
||||
organization_rpm_limit: Optional[int] = None
|
||||
|
|
@ -2478,9 +2503,9 @@ class UserAPIKeyAuth(
|
|||
user_max_budget: Optional[float] = None
|
||||
request_route: Optional[str] = None
|
||||
user: Optional[Any] = None # Expanded user object when expand=user is used
|
||||
created_by_user: Optional[Any] = (
|
||||
None # Expanded created_by user when expand=user is used
|
||||
)
|
||||
created_by_user: Optional[
|
||||
Any
|
||||
] = None # Expanded created_by user when expand=user is used
|
||||
end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
|
||||
# Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery
|
||||
# and forwarded into outbound tokens by guardrails such as MCPJWTSigner.
|
||||
|
|
@ -2619,9 +2644,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase):
|
|||
budget_id: Optional[str] = None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
user: Optional[Any] = (
|
||||
None # You might want to replace 'Any' with a more specific type if available
|
||||
)
|
||||
user: Optional[
|
||||
Any
|
||||
] = None # You might want to replace 'Any' with a more specific type if available
|
||||
litellm_budget_table: Optional[LiteLLM_BudgetTable] = None
|
||||
user_email: Optional[str] = None
|
||||
|
||||
|
|
@ -3776,9 +3801,9 @@ class TeamModelDeleteRequest(BaseModel):
|
|||
# Organization Member Requests
|
||||
class OrganizationMemberAddRequest(OrgMemberAddRequest):
|
||||
organization_id: str
|
||||
max_budget_in_organization: Optional[float] = (
|
||||
None # Users max budget within the organization
|
||||
)
|
||||
max_budget_in_organization: Optional[
|
||||
float
|
||||
] = None # Users max budget within the organization
|
||||
|
||||
|
||||
class OrganizationMemberDeleteRequest(MemberDeleteRequest):
|
||||
|
|
@ -4033,9 +4058,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase):
|
|||
Maps provider names to their budget configs.
|
||||
"""
|
||||
|
||||
providers: Dict[str, ProviderBudgetResponseObject] = (
|
||||
{}
|
||||
) # Dictionary mapping provider names to their budget configurations
|
||||
providers: Dict[
|
||||
str, ProviderBudgetResponseObject
|
||||
] = {} # Dictionary mapping provider names to their budget configurations
|
||||
|
||||
|
||||
class ProxyStateVariables(TypedDict):
|
||||
|
|
@ -4197,9 +4222,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
|
|||
enforce_rbac: bool = False
|
||||
roles_jwt_field: Optional[str] = None # v2 on role mappings
|
||||
role_mappings: Optional[List[RoleMapping]] = None
|
||||
object_id_jwt_field: Optional[str] = (
|
||||
None # can be either user / team, inferred from the role mapping
|
||||
)
|
||||
object_id_jwt_field: Optional[
|
||||
str
|
||||
] = None # can be either user / team, inferred from the role mapping
|
||||
scope_mappings: Optional[List[ScopeMapping]] = None
|
||||
enforce_scope_based_access: bool = False
|
||||
enforce_team_based_model_access: bool = False
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from litellm.types.agents import (
|
|||
MakeAgentsPublicRequest,
|
||||
PatchAgentRequest,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
SpendAnalyticsPaginatedResponse,
|
||||
|
|
@ -36,6 +37,28 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
|||
router = APIRouter()
|
||||
|
||||
|
||||
def _redact_sensitive_agent_fields(
|
||||
agents: List[AgentResponse],
|
||||
) -> List[AgentResponse]:
|
||||
"""
|
||||
Return copies of the given agents with sensitive configuration fields
|
||||
redacted. The original objects are not modified.
|
||||
"""
|
||||
redacted: List[AgentResponse] = []
|
||||
for agent in agents:
|
||||
copy = agent.model_copy(deep=True)
|
||||
copy.static_headers = None
|
||||
copy.extra_headers = None
|
||||
if copy.litellm_params:
|
||||
copy.litellm_params = _get_masked_values(
|
||||
copy.litellm_params,
|
||||
unmasked_length=4,
|
||||
number_of_asterisks=4,
|
||||
)
|
||||
redacted.append(copy)
|
||||
return redacted
|
||||
|
||||
|
||||
def _check_agent_management_permission(user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
"""
|
||||
Raises HTTP 403 if the caller does not have permission to create, update,
|
||||
|
|
@ -183,6 +206,14 @@ async def get_agents(
|
|||
agent.agent_id in litellm.public_agent_groups
|
||||
)
|
||||
|
||||
# Redact sensitive fields for non-admin users
|
||||
is_admin = (
|
||||
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
if not is_admin:
|
||||
returned_agents = _redact_sensitive_agent_fields(returned_agents)
|
||||
|
||||
if health_check:
|
||||
agents_with_url = [
|
||||
agent
|
||||
|
|
@ -399,6 +430,14 @@ async def get_agent_by_id(
|
|||
status_code=404, detail=f"Agent with ID {agent_id} not found"
|
||||
)
|
||||
|
||||
# Redact sensitive fields for non-admin users
|
||||
is_admin = (
|
||||
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
if not is_admin:
|
||||
agent = _redact_sensitive_agent_fields([agent])[0]
|
||||
|
||||
return agent
|
||||
except HTTPException:
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
#### Analytics Endpoints #####
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
from typing import List, Optional
|
||||
|
||||
import fastapi
|
||||
|
|
@ -58,8 +58,10 @@ async def get_global_activity(
|
|||
detail={"error": "Please provide start_date and end_date"},
|
||||
)
|
||||
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d")
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d")
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(
|
||||
tzinfo=timezone.utc
|
||||
)
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
|
|
@ -83,8 +85,9 @@ async def get_global_activity(
|
|||
SUM(CASE WHEN sl."cache_hit" != 'True' THEN sl."completion_tokens" ELSE 0 END) AS generated_completion_tokens
|
||||
FROM "LiteLLM_SpendLogs" sl
|
||||
LEFT JOIN "LiteLLM_VerificationToken" vt ON sl."api_key" = vt."token"
|
||||
WHERE
|
||||
sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
WHERE
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
GROUP BY
|
||||
vt."key_alias",
|
||||
sl."call_type",
|
||||
|
|
|
|||
|
|
@ -196,9 +196,7 @@ def _is_model_cost_zero(
|
|||
return True
|
||||
|
||||
|
||||
def _is_cost_explicitly_configured(
|
||||
model: str, llm_router: "Router"
|
||||
) -> bool:
|
||||
def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool:
|
||||
"""
|
||||
Check if any deployment in the model group has cost fields explicitly
|
||||
set in its litellm.model_cost entry.
|
||||
|
|
@ -215,10 +213,7 @@ def _is_cost_explicitly_configured(
|
|||
if model_id is None:
|
||||
continue
|
||||
raw_entry = litellm.model_cost.get(model_id, {})
|
||||
if (
|
||||
"input_cost_per_token" in raw_entry
|
||||
or "output_cost_per_token" in raw_entry
|
||||
):
|
||||
if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
|
@ -596,17 +591,12 @@ async def common_checks( # noqa: PLR0915
|
|||
user_object=user_object, route=route, request_body=request_body
|
||||
)
|
||||
|
||||
token_team = getattr(valid_token, "team_id", None)
|
||||
token_type: Literal["ui", "api"] = (
|
||||
"ui" if token_team is not None and token_team == "litellm-dashboard" else "api"
|
||||
)
|
||||
_is_route_allowed = _is_allowed_route(
|
||||
_is_route_allowed = _is_api_route_allowed(
|
||||
route=route,
|
||||
token_type=token_type,
|
||||
user_obj=user_object,
|
||||
request=request,
|
||||
request_data=request_body,
|
||||
valid_token=valid_token,
|
||||
user_obj=user_object,
|
||||
)
|
||||
|
||||
# 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store
|
||||
|
|
@ -629,31 +619,6 @@ async def common_checks( # noqa: PLR0915
|
|||
return True
|
||||
|
||||
|
||||
def _is_ui_route(
|
||||
route: str,
|
||||
user_obj: Optional[LiteLLM_UserTable] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
- Check if the route is a UI used route
|
||||
"""
|
||||
# this token is only used for managing the ui
|
||||
allowed_routes = LiteLLMRoutes.ui_routes.value
|
||||
# check if the current route startswith any of the allowed routes
|
||||
if (
|
||||
route is not None
|
||||
and isinstance(route, str)
|
||||
and any(route.startswith(allowed_route) for allowed_route in allowed_routes)
|
||||
):
|
||||
# Do something if the current route starts with any of the allowed routes
|
||||
return True
|
||||
elif any(
|
||||
RouteChecks._route_matches_pattern(route=route, pattern=allowed_route)
|
||||
for allowed_route in allowed_routes
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _get_user_role(
|
||||
user_obj: Optional[LiteLLM_UserTable],
|
||||
) -> Optional[LitellmUserRoles]:
|
||||
|
|
@ -717,30 +682,6 @@ def _is_user_proxy_admin(user_obj: Optional[LiteLLM_UserTable]):
|
|||
return False
|
||||
|
||||
|
||||
def _is_allowed_route(
|
||||
route: str,
|
||||
token_type: Literal["ui", "api"],
|
||||
request: Request,
|
||||
request_data: dict,
|
||||
valid_token: Optional[UserAPIKeyAuth],
|
||||
user_obj: Optional[LiteLLM_UserTable] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
- Route b/w ui token check and normal token check
|
||||
"""
|
||||
|
||||
if token_type == "ui" and _is_ui_route(route=route, user_obj=user_obj):
|
||||
return True
|
||||
else:
|
||||
return _is_api_route_allowed(
|
||||
route=route,
|
||||
request=request,
|
||||
request_data=request_data,
|
||||
valid_token=valid_token,
|
||||
user_obj=user_obj,
|
||||
)
|
||||
|
||||
|
||||
def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool:
|
||||
"""
|
||||
Return if a user is allowed to access route. Helper function for `allowed_routes_check`.
|
||||
|
|
|
|||
|
|
@ -60,12 +60,20 @@ def _get_guardrails_list_response(
|
|||
"""
|
||||
Helper function to get the guardrails list response
|
||||
"""
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
|
||||
|
||||
guardrail_configs: List[GuardrailInfoResponse] = []
|
||||
for guardrail in guardrails_config:
|
||||
litellm_params = guardrail.get("litellm_params") or {}
|
||||
masked_params = _get_masked_values(
|
||||
litellm_params,
|
||||
unmasked_length=4,
|
||||
number_of_asterisks=4,
|
||||
)
|
||||
guardrail_configs.append(
|
||||
GuardrailInfoResponse(
|
||||
guardrail_name=guardrail.get("guardrail_name"),
|
||||
litellm_params=guardrail.get("litellm_params"),
|
||||
litellm_params=masked_params,
|
||||
guardrail_info=guardrail.get("guardrail_info"),
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from starlette.datastructures import Headers
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
from litellm._service_logger import ServiceLogging
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.proxy._types import (
|
||||
AddTeamCallback,
|
||||
|
|
@ -684,6 +685,7 @@ class LiteLLMProxyRequestSetup:
|
|||
user_api_key_project_alias=user_api_key_dict.project_alias,
|
||||
user_api_key_user_id=user_api_key_dict.user_id,
|
||||
user_api_key_org_id=user_api_key_dict.org_id,
|
||||
user_api_key_org_alias=user_api_key_dict.organization_alias,
|
||||
user_api_key_team_alias=user_api_key_dict.team_alias,
|
||||
user_api_key_end_user_id=user_api_key_dict.end_user_id,
|
||||
user_api_key_user_email=user_api_key_dict.user_email,
|
||||
|
|
@ -1263,6 +1265,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# Save pre-alias model name for credential override lookup
|
||||
_pre_alias_model = data.get("model")
|
||||
|
||||
# Team Model Aliases
|
||||
_update_model_if_team_alias_exists(
|
||||
data=data,
|
||||
|
|
@ -1279,6 +1284,14 @@ async def add_litellm_data_to_request( # noqa: PLR0915
|
|||
"[PROXY] returned data from litellm_pre_call_utils: %s", data
|
||||
)
|
||||
|
||||
# Team/Project credential overrides from model_config
|
||||
# Placed after the debug log to avoid leaking credential secrets in logs
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
pre_alias_model_name=_pre_alias_model,
|
||||
)
|
||||
|
||||
## ENFORCED PARAMS CHECK
|
||||
# loop through each enforced param
|
||||
# example enforced_params ['user', 'metadata', 'metadata.generation_name']
|
||||
|
|
@ -1406,6 +1419,175 @@ def _update_model_if_key_alias_exists(
|
|||
return
|
||||
|
||||
|
||||
def _apply_credential_overrides_from_model_config(
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
pre_alias_model_name: Optional[str] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Walk the model_config precedence chain in team/project metadata.
|
||||
If a matching credential is found, set api_base/api_key/api_version on data
|
||||
so they override deployment defaults in the router.
|
||||
|
||||
Precedence (highest to lowest):
|
||||
1. Clientside credentials (already in data — skip if present)
|
||||
2. Project model-specific override
|
||||
3. Project default override (defaultconfig)
|
||||
4. Team model-specific override
|
||||
5. Team default override (defaultconfig)
|
||||
6. Deployment default (no action needed)
|
||||
"""
|
||||
# Feature flag gate — disabled by default, opt in with litellm.enable_model_config_credential_overrides = True
|
||||
if not litellm.enable_model_config_credential_overrides:
|
||||
return
|
||||
|
||||
# Respect clientside credentials — highest precedence
|
||||
if data.get("api_base") is not None or data.get("api_key") is not None:
|
||||
return
|
||||
|
||||
model_name = data.get("model")
|
||||
if not model_name:
|
||||
return
|
||||
|
||||
project_metadata = user_api_key_dict.project_metadata or {}
|
||||
team_metadata = user_api_key_dict.team_metadata or {}
|
||||
|
||||
project_model_config = project_metadata.get("model_config")
|
||||
team_model_config = team_metadata.get("model_config")
|
||||
|
||||
if not project_model_config and not team_model_config:
|
||||
return
|
||||
|
||||
# Extract provider hint from model name (e.g. "azure/gpt-4" -> "azure")
|
||||
provider: Optional[str] = None
|
||||
if "/" in model_name:
|
||||
provider = model_name.split("/", 1)[0]
|
||||
|
||||
credential_name = _resolve_credential_from_model_config(
|
||||
model_name=model_name,
|
||||
project_model_config=project_model_config,
|
||||
team_model_config=team_model_config,
|
||||
pre_alias_model_name=pre_alias_model_name,
|
||||
provider=provider,
|
||||
)
|
||||
|
||||
if not credential_name:
|
||||
return
|
||||
|
||||
credential_values = CredentialAccessor.get_credential_values(credential_name)
|
||||
if not credential_values:
|
||||
_safe_cred = str(credential_name).replace("\n", "").replace("\r", "")
|
||||
verbose_proxy_logger.warning(
|
||||
"model_config references credential '%s' but it was not found or has no values",
|
||||
_safe_cred,
|
||||
)
|
||||
return
|
||||
|
||||
# Apply credential overrides only for keys not already in the request
|
||||
for key in ("api_base", "api_key", "api_version"):
|
||||
if key in credential_values and key not in data:
|
||||
data[key] = credential_values[key]
|
||||
|
||||
_safe_model = str(model_name).replace("\n", "").replace("\r", "")
|
||||
_safe_cred = str(credential_name).replace("\n", "").replace("\r", "")
|
||||
verbose_proxy_logger.debug(
|
||||
"Applied credential override '%s' for model '%s'",
|
||||
_safe_cred,
|
||||
_safe_model,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_credential_from_model_config(
|
||||
model_name: str,
|
||||
project_model_config: Optional[dict],
|
||||
team_model_config: Optional[dict],
|
||||
pre_alias_model_name: Optional[str] = None,
|
||||
provider: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Walk the precedence chain and return the first matching credential name.
|
||||
|
||||
Checks (in order):
|
||||
1. project_model_config[model_name][provider] — project model-specific
|
||||
2. project_model_config[pre_alias_model_name][provider] — project pre-alias
|
||||
3. project_model_config["defaultconfig"][provider] — project default
|
||||
4. team_model_config[model_name][provider] — team model-specific
|
||||
5. team_model_config[pre_alias_model_name][provider] — team pre-alias
|
||||
6. team_model_config["defaultconfig"][provider] — team default
|
||||
|
||||
When a model-specific entry exists but contains no litellm_credentials,
|
||||
the function falls through to defaultconfig. This is intentional —
|
||||
an entry without litellm_credentials is treated as incomplete config,
|
||||
not as an explicit "no override" signal.
|
||||
"""
|
||||
# Build the list of model names to try (post-alias first, then pre-alias)
|
||||
model_names_to_try = [model_name]
|
||||
if pre_alias_model_name and pre_alias_model_name != model_name:
|
||||
model_names_to_try.append(pre_alias_model_name)
|
||||
|
||||
for model_config in (project_model_config, team_model_config):
|
||||
if not model_config or not isinstance(model_config, dict):
|
||||
continue
|
||||
|
||||
# Model-specific check (try resolved name, then pre-alias name)
|
||||
for name in model_names_to_try:
|
||||
model_entry = model_config.get(name)
|
||||
if model_entry:
|
||||
credential_name = _extract_credential_from_entry(
|
||||
model_entry, provider=provider
|
||||
)
|
||||
if credential_name:
|
||||
return credential_name
|
||||
_safe_name = str(name).replace("\n", "").replace("\r", "")
|
||||
verbose_proxy_logger.debug(
|
||||
"model_config entry '%s' found but has no litellm_credentials, "
|
||||
"trying next candidate",
|
||||
_safe_name,
|
||||
)
|
||||
|
||||
# Default check
|
||||
default_entry = model_config.get("defaultconfig")
|
||||
if default_entry:
|
||||
credential_name = _extract_credential_from_entry(
|
||||
default_entry, provider=provider
|
||||
)
|
||||
if credential_name:
|
||||
return credential_name
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _extract_credential_from_entry(
|
||||
entry: dict, provider: Optional[str] = None
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Extract litellm_credentials from a model_config entry.
|
||||
|
||||
Entry structure: {"azure": {"litellm_credentials": "name"}, ...}
|
||||
|
||||
When provider is given (e.g. "azure"), tries an exact provider match first.
|
||||
Falls back to the first credential found across all provider keys.
|
||||
"""
|
||||
if not isinstance(entry, dict):
|
||||
return None
|
||||
|
||||
# Prefer exact provider match when provider hint is available
|
||||
if provider and provider in entry:
|
||||
provider_config = entry[provider]
|
||||
if isinstance(provider_config, dict):
|
||||
credential_name = provider_config.get("litellm_credentials")
|
||||
if credential_name:
|
||||
return credential_name
|
||||
|
||||
# Fall back to first available provider
|
||||
for provider_config in entry.values():
|
||||
if isinstance(provider_config, dict):
|
||||
credential_name = provider_config.get("litellm_credentials")
|
||||
if credential_name:
|
||||
return credential_name
|
||||
return None
|
||||
|
||||
|
||||
def _get_enforced_params(
|
||||
general_settings: Optional[dict], user_api_key_dict: UserAPIKeyAuth
|
||||
) -> Optional[list]:
|
||||
|
|
|
|||
|
|
@ -456,6 +456,34 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict:
|
|||
return data_json
|
||||
|
||||
|
||||
def _check_allowed_routes_caller_permission(
|
||||
allowed_routes: Optional[list],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
"""
|
||||
Only proxy admins may set `allowed_routes` on a key.
|
||||
|
||||
`allowed_routes` bypasses the standard role-based route gate in
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check, so if a non-admin is
|
||||
allowed to set it they can grant themselves access to any endpoint.
|
||||
Non-admins should use `key_type` to pick a preset route bucket instead.
|
||||
"""
|
||||
# Empty list is the default on GenerateKeyRequest — treat as "not set".
|
||||
if not allowed_routes:
|
||||
return
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
"Only proxy admins can set `allowed_routes` on a key. "
|
||||
"Use `key_type` to pick a preset route bucket instead."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def validate_team_id_used_in_service_account_request(
|
||||
team_id: Optional[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
|
|
@ -740,9 +768,9 @@ async def _common_key_generation_helper( # noqa: PLR0915
|
|||
request_type="key", **data_json, table_name="key"
|
||||
)
|
||||
|
||||
response[
|
||||
"soft_budget"
|
||||
] = data.soft_budget # include the user-input soft budget in the response
|
||||
response["soft_budget"] = (
|
||||
data.soft_budget
|
||||
) # include the user-input soft budget in the response
|
||||
|
||||
response = GenerateKeyResponse(**response)
|
||||
|
||||
|
|
@ -1254,6 +1282,12 @@ async def generate_key_fn(
|
|||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN, detail=message
|
||||
)
|
||||
|
||||
_check_allowed_routes_caller_permission(
|
||||
allowed_routes=data.allowed_routes,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# For non-admin internal users: auto-assign caller's user_id if not provided
|
||||
# This prevents creating unbound keys with no user association (LIT-1884)
|
||||
_is_proxy_admin = (
|
||||
|
|
@ -1888,6 +1922,11 @@ async def _validate_update_key_data(
|
|||
"""Validate permissions and constraints for key update."""
|
||||
_is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
|
||||
_check_allowed_routes_caller_permission(
|
||||
allowed_routes=data.allowed_routes,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# Prevent non-admin from removing user_id (setting to empty string) (LIT-1884)
|
||||
if data.user_id is not None and data.user_id == "" and not _is_proxy_admin:
|
||||
raise HTTPException(
|
||||
|
|
@ -3233,10 +3272,10 @@ async def delete_verification_tokens(
|
|||
try:
|
||||
if prisma_client:
|
||||
tokens = [_hash_token_if_needed(token=key) for key in tokens]
|
||||
_keys_being_deleted: List[
|
||||
LiteLLM_VerificationToken
|
||||
] = await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"token": {"in": tokens}}
|
||||
_keys_being_deleted: List[LiteLLM_VerificationToken] = (
|
||||
await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"token": {"in": tokens}}
|
||||
)
|
||||
)
|
||||
|
||||
if len(_keys_being_deleted) == 0:
|
||||
|
|
@ -3436,9 +3475,9 @@ async def _rotate_master_key( # noqa: PLR0915
|
|||
from litellm.proxy.proxy_server import proxy_config
|
||||
|
||||
try:
|
||||
models: Optional[
|
||||
List
|
||||
] = await prisma_client.db.litellm_proxymodeltable.find_many()
|
||||
models: Optional[List] = (
|
||||
await prisma_client.db.litellm_proxymodeltable.find_many()
|
||||
)
|
||||
except Exception:
|
||||
models = None
|
||||
# 2. process model table
|
||||
|
|
@ -4078,11 +4117,11 @@ async def validate_key_list_check(
|
|||
param="user_id",
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
)
|
||||
complete_user_info_db_obj: Optional[
|
||||
BaseModel
|
||||
] = await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_api_key_dict.user_id},
|
||||
include={"organization_memberships": True},
|
||||
complete_user_info_db_obj: Optional[BaseModel] = (
|
||||
await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_api_key_dict.user_id},
|
||||
include={"organization_memberships": True},
|
||||
)
|
||||
)
|
||||
|
||||
if complete_user_info_db_obj is None:
|
||||
|
|
@ -4165,10 +4204,10 @@ async def _fetch_user_team_objects(
|
|||
if complete_user_info is None or not complete_user_info.teams:
|
||||
return []
|
||||
|
||||
teams: Optional[
|
||||
List[BaseModel]
|
||||
] = await prisma_client.db.litellm_teamtable.find_many(
|
||||
where={"team_id": {"in": complete_user_info.teams}}
|
||||
teams: Optional[List[BaseModel]] = (
|
||||
await prisma_client.db.litellm_teamtable.find_many(
|
||||
where={"team_id": {"in": complete_user_info.teams}}
|
||||
)
|
||||
)
|
||||
if teams is None:
|
||||
return []
|
||||
|
|
|
|||
|
|
@ -1115,7 +1115,10 @@ async def delete_file(
|
|||
file_id=original_file_id,
|
||||
)
|
||||
|
||||
response = await litellm.afile_delete(**data) # type: ignore
|
||||
response = await litellm.afile_delete(
|
||||
custom_llm_provider=credentials["custom_llm_provider"], # type: ignore
|
||||
**data,
|
||||
) # type: ignore
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"Deleted file using model: {model_used}"
|
||||
|
|
|
|||
|
|
@ -7128,6 +7128,13 @@ async def chat_completion( # noqa: PLR0915
|
|||
and user_api_key_dict.org_id is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id
|
||||
if (
|
||||
hasattr(user_api_key_dict, "organization_alias")
|
||||
and user_api_key_dict.organization_alias is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_alias"] = (
|
||||
user_api_key_dict.organization_alias
|
||||
)
|
||||
if (
|
||||
hasattr(user_api_key_dict, "agent_id")
|
||||
and user_api_key_dict.agent_id is not None
|
||||
|
|
@ -7302,6 +7309,13 @@ async def completion( # noqa: PLR0915
|
|||
and user_api_key_dict.org_id is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id
|
||||
if (
|
||||
hasattr(user_api_key_dict, "organization_alias")
|
||||
and user_api_key_dict.organization_alias is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_alias"] = (
|
||||
user_api_key_dict.organization_alias
|
||||
)
|
||||
if (
|
||||
hasattr(user_api_key_dict, "agent_id")
|
||||
and user_api_key_dict.agent_id is not None
|
||||
|
|
@ -7544,6 +7558,13 @@ async def embeddings( # noqa: PLR0915
|
|||
and user_api_key_dict.org_id is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id
|
||||
if (
|
||||
hasattr(user_api_key_dict, "organization_alias")
|
||||
and user_api_key_dict.organization_alias is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_alias"] = (
|
||||
user_api_key_dict.organization_alias
|
||||
)
|
||||
if (
|
||||
hasattr(user_api_key_dict, "agent_id")
|
||||
and user_api_key_dict.agent_id is not None
|
||||
|
|
|
|||
|
|
@ -223,7 +223,8 @@ async def get_global_activity_internal_user(
|
|||
COUNT(*) AS api_requests,
|
||||
SUM(total_tokens) AS total_tokens
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND "user" = $3
|
||||
GROUP BY date_trunc('day', "startTime")
|
||||
"""
|
||||
|
|
@ -282,8 +283,10 @@ async def get_global_activity(
|
|||
detail={"error": "Please provide start_date and end_date"},
|
||||
)
|
||||
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d")
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d")
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(
|
||||
tzinfo=timezone.utc
|
||||
)
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
|
|
@ -307,7 +310,8 @@ async def get_global_activity(
|
|||
COUNT(*) AS api_requests,
|
||||
SUM(total_tokens) AS total_tokens
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
GROUP BY date_trunc('day', "startTime")
|
||||
"""
|
||||
db_response = await prisma_client.db.query_raw(
|
||||
|
|
@ -366,7 +370,8 @@ async def get_global_activity_model_internal_user(
|
|||
COUNT(*) AS api_requests,
|
||||
SUM(total_tokens) AS total_tokens
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND "user" = $3
|
||||
GROUP BY model_group, date_trunc('day', "startTime")
|
||||
"""
|
||||
|
|
@ -448,8 +453,10 @@ async def get_global_activity_model(
|
|||
detail={"error": "Please provide start_date and end_date"},
|
||||
)
|
||||
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d")
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d")
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(
|
||||
tzinfo=timezone.utc
|
||||
)
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
|
|
@ -474,7 +481,8 @@ async def get_global_activity_model(
|
|||
COUNT(*) AS api_requests,
|
||||
SUM(total_tokens) AS total_tokens
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
GROUP BY model_group, date_trunc('day', "startTime")
|
||||
"""
|
||||
db_response = await prisma_client.db.query_raw(
|
||||
|
|
@ -600,8 +608,10 @@ async def get_global_activity_exceptions_per_deployment(
|
|||
detail={"error": "Please provide start_date and end_date"},
|
||||
)
|
||||
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d")
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d")
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(
|
||||
tzinfo=timezone.utc
|
||||
)
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
|
|
@ -619,7 +629,8 @@ async def get_global_activity_exceptions_per_deployment(
|
|||
FROM
|
||||
"LiteLLM_ErrorLogs"
|
||||
WHERE
|
||||
"startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
"startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND model_group = $3
|
||||
AND status_code = '429'
|
||||
GROUP BY
|
||||
|
|
@ -732,8 +743,10 @@ async def get_global_activity_exceptions(
|
|||
detail={"error": "Please provide start_date and end_date"},
|
||||
)
|
||||
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d")
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d")
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(
|
||||
tzinfo=timezone.utc
|
||||
)
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
|
|
@ -750,7 +763,8 @@ async def get_global_activity_exceptions(
|
|||
FROM
|
||||
"LiteLLM_ErrorLogs"
|
||||
WHERE
|
||||
"startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
"startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND model_group = $3
|
||||
AND status_code = '429'
|
||||
GROUP BY
|
||||
|
|
@ -837,8 +851,10 @@ async def get_global_spend_provider(
|
|||
detail={"error": "Please provide start_date and end_date"},
|
||||
)
|
||||
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d")
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d")
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(
|
||||
tzinfo=timezone.utc
|
||||
)
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
|
||||
from litellm.proxy.proxy_server import llm_router, prisma_client
|
||||
|
||||
|
|
@ -863,7 +879,8 @@ async def get_global_spend_provider(
|
|||
model_id,
|
||||
SUM(spend) AS spend
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND length(model_id) > 0
|
||||
AND "user" = $3
|
||||
GROUP BY model_id
|
||||
|
|
@ -877,7 +894,9 @@ async def get_global_spend_provider(
|
|||
model_id,
|
||||
SUM(spend) AS spend
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND length(model_id) > 0
|
||||
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND length(model_id) > 0
|
||||
GROUP BY model_id
|
||||
"""
|
||||
db_response = await prisma_client.db.query_raw(
|
||||
|
|
@ -996,8 +1015,10 @@ async def get_global_spend_report(
|
|||
detail={"error": "Please provide start_date and end_date"},
|
||||
)
|
||||
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d")
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d")
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(
|
||||
tzinfo=timezone.utc
|
||||
)
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client
|
||||
|
||||
|
|
@ -1029,7 +1050,9 @@ async def get_global_spend_report(
|
|||
FROM
|
||||
"LiteLLM_SpendLogs" sl
|
||||
WHERE
|
||||
sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND sl.api_key = $3
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND sl.api_key = $3
|
||||
GROUP BY
|
||||
sl.api_key,
|
||||
sl.model
|
||||
|
|
@ -1074,7 +1097,9 @@ async def get_global_spend_report(
|
|||
FROM
|
||||
"LiteLLM_SpendLogs" sl
|
||||
WHERE
|
||||
sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\') AND sl.user = $3
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND sl.user = $3
|
||||
GROUP BY
|
||||
sl.api_key,
|
||||
sl.model
|
||||
|
|
@ -1128,7 +1153,8 @@ async def get_global_spend_report(
|
|||
ON
|
||||
sl.team_id = tt.team_id
|
||||
WHERE
|
||||
sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
GROUP BY
|
||||
date_trunc('day', sl."startTime"),
|
||||
tt.team_alias,
|
||||
|
|
@ -1187,7 +1213,8 @@ async def get_global_spend_report(
|
|||
FROM
|
||||
"LiteLLM_SpendLogs" sl
|
||||
WHERE
|
||||
sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
GROUP BY
|
||||
date_trunc('day', sl."startTime"),
|
||||
customer,
|
||||
|
|
@ -1244,7 +1271,8 @@ async def get_global_spend_report(
|
|||
FROM
|
||||
"LiteLLM_SpendLogs" sl
|
||||
WHERE
|
||||
sl."startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
GROUP BY
|
||||
sl.api_key,
|
||||
sl.model
|
||||
|
|
@ -1428,6 +1456,16 @@ async def _get_spend_report_for_time_range(
|
|||
)
|
||||
return None
|
||||
|
||||
# Normalize string inputs to tz-aware UTC datetimes so Prisma serializes
|
||||
# them with an explicit +00:00 suffix. Raw strings get bound as untyped
|
||||
# text, which forces Postgres to parse `::timestamptz` using the DB
|
||||
# session timezone and drifts the window by the offset even with the
|
||||
# AT TIME ZONE 'UTC' wrap below.
|
||||
start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(
|
||||
tzinfo=timezone.utc
|
||||
)
|
||||
end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||||
|
||||
try:
|
||||
sql_query = """
|
||||
SELECT
|
||||
|
|
@ -1438,27 +1476,31 @@ async def _get_spend_report_for_time_range(
|
|||
LEFT JOIN
|
||||
"LiteLLM_TeamTable" t ON s.team_id = t.team_id
|
||||
WHERE
|
||||
s."startTime" >= $1::date AND s."startTime" < ($2::date + INTERVAL '1 day')
|
||||
s."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND s."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
GROUP BY
|
||||
t.team_alias
|
||||
ORDER BY
|
||||
total_spend DESC;
|
||||
"""
|
||||
response = await prisma_client.db.query_raw(sql_query, start_date, end_date)
|
||||
response = await prisma_client.db.query_raw(
|
||||
sql_query, start_date_obj, end_date_obj
|
||||
)
|
||||
|
||||
# get spend per tag for today
|
||||
sql_query = """
|
||||
SELECT
|
||||
SELECT
|
||||
jsonb_array_elements_text(request_tags) AS individual_request_tag,
|
||||
SUM(spend) AS total_spend
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE "startTime" >= $1::timestamptz AND "startTime" < ($2::timestamptz + INTERVAL \'1 day\')
|
||||
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
GROUP BY individual_request_tag
|
||||
ORDER BY total_spend DESC;
|
||||
"""
|
||||
|
||||
spend_per_tag = await prisma_client.db.query_raw(
|
||||
sql_query, start_date, end_date
|
||||
sql_query, start_date_obj, end_date_obj
|
||||
)
|
||||
|
||||
return response, spend_per_tag
|
||||
|
|
@ -1910,11 +1952,17 @@ async def ui_view_spend_logs( # noqa: PLR0915
|
|||
sql_params: List[Any] = []
|
||||
p = 1 # parameter index counter
|
||||
|
||||
# Date range (always present)
|
||||
sql_conditions.append(f'"startTime" >= ${p}::timestamptz')
|
||||
# Date range (always present). Wrap the param side with
|
||||
# `AT TIME ZONE 'UTC'` so comparison against the plain `timestamp`
|
||||
# column does not depend on the DB session timezone (see #22529).
|
||||
sql_conditions.append(
|
||||
f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')"
|
||||
)
|
||||
sql_params.append(start_date_obj)
|
||||
p += 1
|
||||
sql_conditions.append(f'"startTime" <= ${p}::timestamptz')
|
||||
sql_conditions.append(
|
||||
f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')"
|
||||
)
|
||||
sql_params.append(end_date_obj)
|
||||
p += 1
|
||||
|
||||
|
|
@ -2897,8 +2945,8 @@ async def global_spend_end_users(data: Optional[GlobalEndUsersSpend] = None):
|
|||
sql_query = """
|
||||
SELECT end_user, COUNT(*) AS total_count, SUM(spend) AS total_spend
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE "startTime" >= $1::timestamptz
|
||||
AND "startTime" < $2::timestamptz
|
||||
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND "startTime" < ($2::timestamptz AT TIME ZONE 'UTC')
|
||||
AND (
|
||||
CASE
|
||||
WHEN $3::TEXT IS NULL THEN TRUE
|
||||
|
|
|
|||
|
|
@ -559,7 +559,8 @@ async def get_spend_by_team_and_customer(
|
|||
ON
|
||||
sl.team_id = tt.team_id
|
||||
WHERE
|
||||
sl."startTime" >= $1::timestamptz AND sl."startTime" < ($2::timestamptz + INTERVAL '1 day')
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND sl.team_id = $3
|
||||
AND sl.end_user = $4
|
||||
GROUP BY
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
#### CRUD ENDPOINTS for UI Settings #####
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from fastapi import APIRouter, Depends, File, HTTPException, UploadFile
|
||||
|
||||
|
|
@ -817,6 +818,29 @@ async def get_ui_theme_settings():
|
|||
)
|
||||
|
||||
|
||||
def _validate_public_image_url(value: Optional[str], field_name: str) -> None:
|
||||
"""
|
||||
Reject anything that isn't a plain http(s) URL with a host. This value is
|
||||
later served via the unauthenticated /get_image endpoint, so local paths
|
||||
like "/etc/passwd" or "file://..." must not be accepted.
|
||||
"""
|
||||
if value is None:
|
||||
return
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
return
|
||||
parsed = urlparse(value.strip())
|
||||
if parsed.scheme not in ("http", "https") or not parsed.netloc:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": (
|
||||
f"Invalid {field_name}: must be an http(s) URL with a host. "
|
||||
"Local filesystem paths and non-http schemes are not allowed."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/update/ui_theme_settings",
|
||||
tags=["UI Theme Settings"],
|
||||
|
|
@ -831,6 +855,9 @@ async def update_ui_theme_settings(theme_config: UIThemeConfig):
|
|||
|
||||
from litellm.proxy.proxy_server import proxy_config, store_model_in_db
|
||||
|
||||
_validate_public_image_url(theme_config.logo_url, "logo_url")
|
||||
_validate_public_image_url(theme_config.favicon_url, "favicon_url")
|
||||
|
||||
if store_model_in_db is not True:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
|
|
|
|||
|
|
@ -2645,7 +2645,7 @@ class PrismaClient:
|
|||
raise e
|
||||
|
||||
async def _query_first_with_cached_plan_fallback(
|
||||
self, sql_query: str
|
||||
self, sql_query: str, *args
|
||||
) -> Optional[dict]:
|
||||
"""
|
||||
Execute a query with automatic fallback for PostgreSQL cached plan errors.
|
||||
|
|
@ -2664,7 +2664,7 @@ class PrismaClient:
|
|||
Original exception if not a cached plan error
|
||||
"""
|
||||
try:
|
||||
return await self.db.query_first(query=sql_query)
|
||||
return await self.db.query_first(sql_query, *args)
|
||||
except Exception as e:
|
||||
error_str = str(e)
|
||||
if "cached plan must not change result type" in error_str:
|
||||
|
|
@ -2679,7 +2679,7 @@ class PrismaClient:
|
|||
"retrying with fresh plan. This may occur during rolling deployments "
|
||||
"when schema changes are applied."
|
||||
)
|
||||
return await self.db.query_first(query=sql_query_retry)
|
||||
return await self.db.query_first(sql_query_retry, *args)
|
||||
else:
|
||||
raise
|
||||
|
||||
|
|
@ -2978,7 +2978,7 @@ class PrismaClient:
|
|||
detail={"error": f"No token passed in. Token={token}"},
|
||||
)
|
||||
|
||||
sql_query = f"""
|
||||
sql_query = """
|
||||
SELECT
|
||||
v.*,
|
||||
t.spend AS team_spend,
|
||||
|
|
@ -3004,6 +3004,7 @@ class PrismaClient:
|
|||
b.model_max_budget as litellm_budget_table_model_max_budget,
|
||||
b.soft_budget as litellm_budget_table_soft_budget,
|
||||
o.metadata as organization_metadata,
|
||||
o.organization_alias as organization_alias,
|
||||
b2.max_budget as organization_max_budget,
|
||||
b2.tpm_limit as organization_tpm_limit,
|
||||
b2.rpm_limit as organization_rpm_limit
|
||||
|
|
@ -3015,11 +3016,11 @@ class PrismaClient:
|
|||
LEFT JOIN "LiteLLM_ProjectTable" AS p ON v.project_id = p.project_id
|
||||
LEFT JOIN "LiteLLM_OrganizationTable" AS o ON v.organization_id = o.organization_id
|
||||
LEFT JOIN "LiteLLM_BudgetTable" AS b2 ON o.budget_id = b2.budget_id
|
||||
WHERE v.token = '{token}'
|
||||
WHERE v.token = $1
|
||||
"""
|
||||
|
||||
response = await self._query_first_with_cached_plan_fallback(
|
||||
sql_query
|
||||
sql_query, hashed_token
|
||||
)
|
||||
|
||||
# If not found in main table, check deprecated keys (grace period)
|
||||
|
|
@ -5293,11 +5294,12 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException:
|
|||
)
|
||||
elif isinstance(e, ProxyException):
|
||||
return e
|
||||
_status_code = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||||
return ProxyException(
|
||||
message="Internal Server Error, " + str(e),
|
||||
message=str(e),
|
||||
type=ProxyErrorTypes.internal_server_error,
|
||||
param=getattr(e, "param", "None"),
|
||||
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
code=_status_code,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1967,11 +1967,18 @@ async def _aresponses_websocket(
|
|||
)
|
||||
|
||||
# Extract params that we're passing explicitly to avoid duplicates in **kwargs
|
||||
remaining_kwargs = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if k not in {"user_api_key_dict", "litellm_metadata"}
|
||||
_explicit_keys = {
|
||||
"user_api_key_dict",
|
||||
"litellm_metadata",
|
||||
"custom_llm_provider",
|
||||
"model",
|
||||
"websocket",
|
||||
"litellm_logging_obj",
|
||||
"api_base",
|
||||
"api_key",
|
||||
"timeout",
|
||||
}
|
||||
remaining_kwargs = {k: v for k, v in kwargs.items() if k not in _explicit_keys}
|
||||
|
||||
await base_llm_http_handler.async_responses_websocket(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -3864,14 +3864,29 @@ class Router:
|
|||
self._add_deployment_model_to_endpoint_for_llm_passthrough_route(
|
||||
kwargs=kwargs, model=model, model_name=model_name
|
||||
)
|
||||
### get custom
|
||||
response = original_generic_function(
|
||||
**{
|
||||
**data,
|
||||
"caching": self.cache_responses,
|
||||
**kwargs,
|
||||
}
|
||||
)
|
||||
|
||||
# Get custom_llm_provider from deployment params
|
||||
try:
|
||||
custom_llm_provider = data.get("custom_llm_provider")
|
||||
_, inferred_custom_llm_provider, _, _ = get_llm_provider(
|
||||
model=data["model"],
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
|
||||
except Exception:
|
||||
custom_llm_provider = None
|
||||
|
||||
# Build response kwargs
|
||||
response_kwargs = {
|
||||
**data,
|
||||
"caching": self.cache_responses,
|
||||
**kwargs,
|
||||
}
|
||||
# Only set custom_llm_provider if it's not None
|
||||
if custom_llm_provider is not None:
|
||||
response_kwargs["custom_llm_provider"] = custom_llm_provider
|
||||
|
||||
response = original_generic_function(**response_kwargs)
|
||||
|
||||
rpm_semaphore = self._get_client(
|
||||
deployment=deployment,
|
||||
|
|
@ -3961,7 +3976,12 @@ class Router:
|
|||
self.routing_strategy_pre_call_checks(deployment=deployment)
|
||||
|
||||
try:
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model=data["model"])
|
||||
custom_llm_provider = data.get("custom_llm_provider")
|
||||
_, inferred_custom_llm_provider, _, _ = get_llm_provider(
|
||||
model=data["model"],
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
|
||||
except Exception:
|
||||
custom_llm_provider = None
|
||||
|
||||
|
|
@ -4219,9 +4239,14 @@ class Router:
|
|||
self.total_calls[model_name] += 1
|
||||
|
||||
## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ##
|
||||
stripped_model, custom_llm_provider, _, _ = get_llm_provider(
|
||||
model=data["model"]
|
||||
# For DB/config deployments, use provider from deployment params
|
||||
custom_llm_provider = data.get("custom_llm_provider")
|
||||
stripped_model, inferred_custom_llm_provider, _, _ = get_llm_provider(
|
||||
model=data["model"],
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
# Preserve explicitly stored provider, fallback to inferred
|
||||
custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
|
||||
|
||||
## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ##
|
||||
purpose = cast(Optional[OpenAIFilesPurpose], kwargs.get("purpose"))
|
||||
|
|
@ -4367,8 +4392,13 @@ class Router:
|
|||
)
|
||||
self.total_calls[model_name] += 1
|
||||
|
||||
# Get custom provider
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model=data["model"])
|
||||
# Get custom provider from deployment params
|
||||
custom_llm_provider = data.get("custom_llm_provider")
|
||||
_, inferred_custom_llm_provider, _, _ = get_llm_provider(
|
||||
model=data["model"],
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
|
||||
|
||||
response = avector_store_create_sdk(
|
||||
**{
|
||||
|
|
@ -4486,7 +4516,12 @@ class Router:
|
|||
self.total_calls[model_name] += 1
|
||||
|
||||
## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ##
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model=data["model"])
|
||||
custom_llm_provider = data.get("custom_llm_provider")
|
||||
_, inferred_custom_llm_provider, _, _ = get_llm_provider(
|
||||
model=data["model"],
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
|
||||
|
||||
response = litellm.acreate_batch(
|
||||
**{
|
||||
|
|
@ -4720,7 +4755,12 @@ class Router:
|
|||
self.total_calls[model_name] += 1
|
||||
|
||||
## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ##
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model=data["model"])
|
||||
custom_llm_provider = data.get("custom_llm_provider")
|
||||
_, inferred_custom_llm_provider, _, _ = get_llm_provider(
|
||||
model=data["model"],
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
|
||||
|
||||
response = litellm.acancel_batch(
|
||||
**{
|
||||
|
|
|
|||
|
|
@ -139,9 +139,16 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
typed_healthy_deployments = cast(List[dict], healthy_deployments)
|
||||
|
||||
# Signal to the response post-processor that encrypted item IDs should be
|
||||
# encoded in the output of this request.
|
||||
litellm_metadata = request_kwargs.setdefault("litellm_metadata", {})
|
||||
litellm_metadata["encrypted_content_affinity_enabled"] = True
|
||||
# encoded in the output of this request. Only set the flag when
|
||||
# litellm_metadata already exists (Responses API path). Using
|
||||
# setdefault would create an empty litellm_metadata dict for chat
|
||||
# completions / embeddings, which breaks tag-based routing because
|
||||
# _get_metadata_variable_name_from_kwargs would pick "litellm_metadata"
|
||||
# over "metadata" where tags are actually stored.
|
||||
if "litellm_metadata" in request_kwargs:
|
||||
request_kwargs["litellm_metadata"][
|
||||
"encrypted_content_affinity_enabled"
|
||||
] = True
|
||||
|
||||
request_input = request_kwargs.get("input")
|
||||
model_id = self._extract_model_id_from_input(request_input)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import re
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Literal, Optional, Tuple
|
||||
from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
from typing_extensions import Annotated
|
||||
|
|
@ -665,6 +665,24 @@ class PrometheusMetricLabels:
|
|||
litellm_cache_misses_metric = _cache_metric_labels
|
||||
litellm_cached_tokens_metric = _cache_metric_labels
|
||||
|
||||
# Metrics whose emission paths supply org context (used by get_labels)
|
||||
_org_label_metrics: ClassVar[frozenset] = frozenset(
|
||||
{
|
||||
"litellm_llm_api_latency_metric",
|
||||
"litellm_llm_api_time_to_first_token_metric",
|
||||
"litellm_request_total_latency_metric",
|
||||
"litellm_request_queue_time_seconds",
|
||||
"litellm_proxy_total_requests_metric",
|
||||
"litellm_proxy_failed_requests_metric",
|
||||
"litellm_deployment_latency_per_output_token",
|
||||
"litellm_requests_metric",
|
||||
"litellm_spend_metric",
|
||||
"litellm_input_tokens_metric",
|
||||
"litellm_total_tokens_metric",
|
||||
"litellm_output_tokens_metric",
|
||||
}
|
||||
)
|
||||
|
||||
# Managed batch metrics
|
||||
_batch_user_labels = [
|
||||
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
|
||||
|
|
@ -731,6 +749,14 @@ class PrometheusMetricLabels:
|
|||
):
|
||||
custom_labels.append(UserAPIKeyLabelNames.STREAM.value)
|
||||
|
||||
if label_name in PrometheusMetricLabels._org_label_metrics:
|
||||
for label in [
|
||||
UserAPIKeyLabelNames.ORG_ID.value,
|
||||
UserAPIKeyLabelNames.ORG_ALIAS.value,
|
||||
]:
|
||||
if label not in default_labels and label not in custom_labels:
|
||||
custom_labels.append(label)
|
||||
|
||||
return default_labels + custom_labels
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -325,6 +325,7 @@ class RequestBody(TypedDict, total=False):
|
|||
generationConfig: GenerationConfig
|
||||
cachedContent: str
|
||||
labels: Dict[str, str]
|
||||
serviceTier: str
|
||||
|
||||
|
||||
class CachedContentRequestBody(TypedDict, total=False):
|
||||
|
|
|
|||
|
|
@ -71,6 +71,15 @@ class MCPServer(BaseModel):
|
|||
# OAuth2 flow type. Defaults to None (interactive / authorization_code).
|
||||
# Set to "client_credentials" to enable M2M token fetching.
|
||||
oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None
|
||||
# Per-user OAuth server-side storage config.
|
||||
# token_validation: key-value pairs that must match fields in the OAuth token
|
||||
# response (supports dot-notation for nested fields, e.g. "team.enterprise_id").
|
||||
# Tokens that fail validation are rejected before storage.
|
||||
token_validation: Optional[Dict[str, Any]] = None
|
||||
# Optional TTL override (seconds) for the Redis per-user token cache.
|
||||
# Defaults to the token's expires_in minus the expiry buffer, or
|
||||
# MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent.
|
||||
token_storage_ttl_seconds: Optional[int] = None
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
@property
|
||||
|
|
|
|||
|
|
@ -2507,6 +2507,7 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict):
|
|||
user_api_key_max_budget: Optional[float]
|
||||
user_api_key_budget_reset_at: Optional[str]
|
||||
user_api_key_org_id: Optional[str]
|
||||
user_api_key_org_alias: Optional[str]
|
||||
user_api_key_team_id: Optional[str]
|
||||
user_api_key_project_id: Optional[str]
|
||||
user_api_key_project_alias: Optional[str]
|
||||
|
|
|
|||
|
|
@ -7818,26 +7818,6 @@
|
|||
"cache_read_input_token_cost": 3e-08,
|
||||
"cache_creation_input_token_cost": 3.75e-07
|
||||
},
|
||||
"bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0": {
|
||||
"input_cost_per_token": 3.6e-06,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.8e-05,
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"cache_read_input_token_cost": 3.6e-07,
|
||||
"cache_creation_input_token_cost": 4.5e-06
|
||||
},
|
||||
"bedrock/us-gov-east-1/anthropic.claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
"cache_read_input_token_cost": 3.3e-07,
|
||||
|
|
@ -7856,7 +7836,29 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"supports_native_structured_output": true
|
||||
},
|
||||
"bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
"cache_read_input_token_cost": 3.3e-07,
|
||||
"input_cost_per_token": 3.3e-06,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.65e-05,
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_native_structured_output": true
|
||||
},
|
||||
"bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0": {
|
||||
"input_cost_per_token": 2.65e-06,
|
||||
|
|
@ -7989,26 +7991,6 @@
|
|||
"cache_read_input_token_cost": 3e-08,
|
||||
"cache_creation_input_token_cost": 3.75e-07
|
||||
},
|
||||
"bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0": {
|
||||
"input_cost_per_token": 3.6e-06,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.8e-05,
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"cache_read_input_token_cost": 3.6e-07,
|
||||
"cache_creation_input_token_cost": 4.5e-06
|
||||
},
|
||||
"bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
"cache_read_input_token_cost": 3.3e-07,
|
||||
|
|
@ -8027,7 +8009,29 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"supports_native_structured_output": true
|
||||
},
|
||||
"bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
"cache_read_input_token_cost": 3.3e-07,
|
||||
"input_cost_per_token": 3.3e-06,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.65e-05,
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_native_structured_output": true
|
||||
},
|
||||
"bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": {
|
||||
"input_cost_per_token": 2.65e-06,
|
||||
|
|
@ -13735,7 +13739,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
"supports_web_search": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini-2.5-flash-image": {
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
|
|
@ -13784,7 +13789,8 @@
|
|||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": false,
|
||||
"tpm": 8000000
|
||||
"tpm": 8000000,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini-3-pro-image-preview": {
|
||||
"input_cost_per_image": 0.0011,
|
||||
|
|
@ -13818,7 +13824,8 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
"supports_web_search": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini-3.1-flash-image-preview": {
|
||||
"input_cost_per_image": 0.00056,
|
||||
|
|
@ -13901,7 +13908,8 @@
|
|||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_native_streaming": true
|
||||
"supports_native_streaming": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"deep-research-pro-preview-12-2025": {
|
||||
"input_cost_per_image": 0.0011,
|
||||
|
|
@ -13980,7 +13988,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
"supports_web_search": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini-2.5-flash-lite-preview-09-2025": {
|
||||
"cache_read_input_token_cost": 1e-08,
|
||||
|
|
@ -14251,7 +14260,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
"supports_web_search": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini-3-pro-preview": {
|
||||
"deprecation_date": "2026-03-26",
|
||||
|
|
@ -15033,7 +15043,8 @@
|
|||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"tpm": 8000000
|
||||
"tpm": 8000000,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini/gemini-2.5-flash-image": {
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
|
|
@ -15083,7 +15094,8 @@
|
|||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"tpm": 8000000
|
||||
"tpm": 8000000,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini/gemini-3-pro-image-preview": {
|
||||
"input_cost_per_image": 0.0011,
|
||||
|
|
@ -15119,7 +15131,8 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
"supports_web_search": true,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini/gemini-3.1-flash-image-preview": {
|
||||
"input_cost_per_token": 2.5e-07,
|
||||
|
|
@ -15238,7 +15251,8 @@
|
|||
"supports_url_context": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"tpm": 250000
|
||||
"tpm": 250000,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini/gemini-2.5-flash-lite-preview-09-2025": {
|
||||
"cache_read_input_token_cost": 1e-08,
|
||||
|
|
@ -15678,7 +15692,8 @@
|
|||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_native_streaming": true,
|
||||
"tpm": 250000
|
||||
"tpm": 250000,
|
||||
"supports_service_tier": true
|
||||
},
|
||||
"gemini/gemini-3-flash-preview": {
|
||||
"cache_read_input_token_cost": 5e-08,
|
||||
|
|
@ -16919,6 +16934,72 @@
|
|||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-06
|
||||
},
|
||||
"baseten/MiniMaxAI/MiniMax-M2.5": {
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-06
|
||||
},
|
||||
"baseten/nvidia/Nemotron-120B-A12B": {
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-07
|
||||
},
|
||||
"baseten/zai-org/GLM-5": {
|
||||
"input_cost_per_token": 9.5e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3.15e-06
|
||||
},
|
||||
"baseten/zai-org/GLM-4.7": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.2e-06
|
||||
},
|
||||
"baseten/zai-org/GLM-4.6": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.2e-06
|
||||
},
|
||||
"baseten/moonshotai/Kimi-K2.5": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-06
|
||||
},
|
||||
"baseten/moonshotai/Kimi-K2-Thinking": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.5e-06
|
||||
},
|
||||
"baseten/moonshotai/Kimi-K2-Instruct-0905": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.5e-06
|
||||
},
|
||||
"baseten/openai/gpt-oss-120b": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-07
|
||||
},
|
||||
"baseten/deepseek-ai/DeepSeek-V3.1": {
|
||||
"input_cost_per_token": 5e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-06
|
||||
},
|
||||
"baseten/deepseek-ai/DeepSeek-V3-0324": {
|
||||
"input_cost_per_token": 7.7e-07,
|
||||
"litellm_provider": "baseten",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.7e-07
|
||||
},
|
||||
"gmi/Qwen/Qwen3-VL-235B-A22B-Instruct-FP8": {
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "gmi",
|
||||
|
|
@ -38108,4 +38189,4 @@
|
|||
"supports_native_structured_output": true,
|
||||
"supports_pdf_input": true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
[tool.poetry]
|
||||
name = "litellm"
|
||||
version = "1.83.5"
|
||||
version = "1.83.6"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
authors = ["BerriAI"]
|
||||
license = "MIT"
|
||||
|
|
@ -181,7 +181,7 @@ requires = ["poetry-core", "wheel"]
|
|||
build-backend = "poetry.core.masonry.api"
|
||||
|
||||
[tool.commitizen]
|
||||
version = "1.83.5"
|
||||
version = "1.83.6"
|
||||
version_files = [
|
||||
"pyproject.toml:^version"
|
||||
]
|
||||
|
|
|
|||
|
|
@ -223,6 +223,8 @@ def test_increment_token_metrics(prometheus_logger):
|
|||
api_key_alias="test_alias",
|
||||
team="test_team",
|
||||
team_alias="test_team_alias",
|
||||
org_id=None,
|
||||
org_alias=None,
|
||||
requested_model=None,
|
||||
model="gpt-3.5-turbo",
|
||||
model_id="model-123",
|
||||
|
|
@ -237,6 +239,8 @@ def test_increment_token_metrics(prometheus_logger):
|
|||
api_key_alias="test_alias",
|
||||
team="test_team",
|
||||
team_alias="test_team_alias",
|
||||
org_id=None,
|
||||
org_alias=None,
|
||||
requested_model=None,
|
||||
model="gpt-3.5-turbo",
|
||||
model_id="model-123",
|
||||
|
|
@ -253,6 +257,8 @@ def test_increment_token_metrics(prometheus_logger):
|
|||
api_key_alias="test_alias",
|
||||
team="test_team",
|
||||
team_alias="test_team_alias",
|
||||
org_id=None,
|
||||
org_alias=None,
|
||||
requested_model=None,
|
||||
model="gpt-3.5-turbo",
|
||||
model_id="model-123",
|
||||
|
|
@ -414,6 +420,8 @@ def test_set_latency_metrics(prometheus_logger):
|
|||
api_key_alias="test_alias",
|
||||
team="test_team",
|
||||
team_alias="test_team_alias",
|
||||
org_id=None,
|
||||
org_alias=None,
|
||||
requested_model="openai-gpt",
|
||||
model="gpt-3.5-turbo",
|
||||
model_id="model-123",
|
||||
|
|
@ -430,6 +438,8 @@ def test_set_latency_metrics(prometheus_logger):
|
|||
api_key_alias="test_alias",
|
||||
team="test_team",
|
||||
team_alias="test_team_alias",
|
||||
org_id=None,
|
||||
org_alias=None,
|
||||
requested_model="openai-gpt",
|
||||
model="gpt-3.5-turbo",
|
||||
model_id="model-123",
|
||||
|
|
@ -446,6 +456,8 @@ def test_set_latency_metrics(prometheus_logger):
|
|||
api_key_alias="test_alias",
|
||||
team="test_team",
|
||||
team_alias="test_team_alias",
|
||||
org_id=None,
|
||||
org_alias=None,
|
||||
requested_model="openai-gpt",
|
||||
model="gpt-3.5-turbo",
|
||||
model_id="model-123",
|
||||
|
|
@ -589,6 +601,8 @@ def test_increment_top_level_request_and_spend_metrics(prometheus_logger):
|
|||
api_key_alias="test_alias",
|
||||
team="test_team",
|
||||
team_alias="test_team_alias",
|
||||
org_id=None,
|
||||
org_alias=None,
|
||||
model="gpt-3.5-turbo",
|
||||
model_id="model-123",
|
||||
client_ip=None,
|
||||
|
|
@ -605,6 +619,8 @@ def test_increment_top_level_request_and_spend_metrics(prometheus_logger):
|
|||
api_key_alias="test_alias",
|
||||
team="test_team",
|
||||
team_alias="test_team_alias",
|
||||
org_id=None,
|
||||
org_alias=None,
|
||||
model="gpt-3.5-turbo",
|
||||
model_id="model-123",
|
||||
client_ip=None,
|
||||
|
|
@ -758,6 +774,8 @@ async def test_async_post_call_failure_hook(prometheus_logger):
|
|||
api_key_alias="test_alias",
|
||||
team="test_team",
|
||||
team_alias="test_team_alias",
|
||||
org_id=None,
|
||||
org_alias=None,
|
||||
requested_model="gpt-3.5-turbo",
|
||||
exception_status="429",
|
||||
exception_class="Openai.RateLimitError",
|
||||
|
|
@ -776,6 +794,8 @@ async def test_async_post_call_failure_hook(prometheus_logger):
|
|||
requested_model="gpt-3.5-turbo",
|
||||
team="test_team",
|
||||
team_alias="test_team_alias",
|
||||
org_id=None,
|
||||
org_alias=None,
|
||||
user="test_user",
|
||||
status_code="429",
|
||||
user_email=None,
|
||||
|
|
@ -955,6 +975,8 @@ def test_set_llm_deployment_success_metrics(prometheus_logger):
|
|||
api_key_alias=standard_logging_payload["metadata"]["user_api_key_alias"],
|
||||
team=standard_logging_payload["metadata"]["user_api_key_team_id"],
|
||||
team_alias=standard_logging_payload["metadata"]["user_api_key_team_alias"],
|
||||
org_id=None,
|
||||
org_alias=None,
|
||||
)
|
||||
prometheus_logger.litellm_overhead_latency_metric.labels.assert_called_once_with(
|
||||
api_base="https://api.openai.com",
|
||||
|
|
|
|||
|
|
@ -50,6 +50,138 @@ def test_split_embedding_by_shape_fails_with_shape_value_error():
|
|||
)
|
||||
|
||||
|
||||
def test_triton_embedding_response_sets_usage_with_token_counter():
|
||||
config = TritonEmbeddingConfig()
|
||||
mock_http_response = MagicMock()
|
||||
mock_http_response.status_code = 200
|
||||
mock_http_response.json.return_value = {
|
||||
"model_name": "gte-base-en-v1",
|
||||
"outputs": [
|
||||
{
|
||||
"name": "embedding",
|
||||
"shape": [1, 2],
|
||||
"data": [0.1, 0.2],
|
||||
}
|
||||
],
|
||||
}
|
||||
model_response = litellm.EmbeddingResponse()
|
||||
request_data = {
|
||||
"inputs": [
|
||||
{
|
||||
"name": "input_text",
|
||||
"shape": [1],
|
||||
"datatype": "BYTES",
|
||||
"data": ["hello from triton"],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.llms.triton.embedding.transformation.token_counter",
|
||||
return_value=7,
|
||||
):
|
||||
transformed = config.transform_embedding_response(
|
||||
model="triton/gte-base-en-v1",
|
||||
raw_response=mock_http_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data=request_data,
|
||||
)
|
||||
|
||||
assert transformed.usage is not None
|
||||
assert transformed.usage.prompt_tokens == 7
|
||||
assert transformed.usage.completion_tokens == 0
|
||||
assert transformed.usage.total_tokens == 7
|
||||
|
||||
|
||||
def test_triton_embedding_response_sets_usage_with_word_count_fallback():
|
||||
config = TritonEmbeddingConfig()
|
||||
mock_http_response = MagicMock()
|
||||
mock_http_response.status_code = 200
|
||||
mock_http_response.json.return_value = {
|
||||
"model_name": "gte-base-en-v1",
|
||||
"outputs": [
|
||||
{
|
||||
"name": "embedding",
|
||||
"shape": [1, 2],
|
||||
"data": [0.1, 0.2],
|
||||
}
|
||||
],
|
||||
}
|
||||
model_response = litellm.EmbeddingResponse()
|
||||
request_data = {
|
||||
"inputs": [
|
||||
{
|
||||
"name": "input_text",
|
||||
"shape": [1],
|
||||
"datatype": "BYTES",
|
||||
"data": ["hello from triton"],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.llms.triton.embedding.transformation.token_counter",
|
||||
side_effect=Exception("tokenizer error"),
|
||||
):
|
||||
transformed = config.transform_embedding_response(
|
||||
model="triton/gte-base-en-v1",
|
||||
raw_response=mock_http_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data=request_data,
|
||||
)
|
||||
|
||||
assert transformed.usage is not None
|
||||
assert transformed.usage.prompt_tokens == 3
|
||||
assert transformed.usage.completion_tokens == 0
|
||||
assert transformed.usage.total_tokens == 3
|
||||
|
||||
|
||||
def test_triton_embedding_batch_usage_sums_per_input_token_counts():
|
||||
"""Batch inputs must not be joined before token counting (avoids extra newline tokens)."""
|
||||
config = TritonEmbeddingConfig()
|
||||
mock_http_response = MagicMock()
|
||||
mock_http_response.status_code = 200
|
||||
mock_http_response.json.return_value = {
|
||||
"model_name": "gte-base-en-v1",
|
||||
"outputs": [
|
||||
{
|
||||
"name": "embedding",
|
||||
"shape": [2, 2],
|
||||
"data": [0.1, 0.2, 0.3, 0.4],
|
||||
}
|
||||
],
|
||||
}
|
||||
model_response = litellm.EmbeddingResponse()
|
||||
request_data = {
|
||||
"inputs": [
|
||||
{
|
||||
"name": "input_text",
|
||||
"shape": [2],
|
||||
"datatype": "BYTES",
|
||||
"data": ["first input", "second input"],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.llms.triton.embedding.transformation.token_counter",
|
||||
side_effect=[5, 7],
|
||||
):
|
||||
transformed = config.transform_embedding_response(
|
||||
model="triton/gte-base-en-v1",
|
||||
raw_response=mock_http_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data=request_data,
|
||||
)
|
||||
|
||||
assert transformed.usage is not None
|
||||
assert transformed.usage.prompt_tokens == 12
|
||||
assert transformed.usage.total_tokens == 12
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [True, False])
|
||||
def test_completion_triton_generate_api(stream):
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -178,7 +178,6 @@ async def test_get_response():
|
|||
async def test_aavertex_ai_anthropic_async():
|
||||
# load_vertex_ai_credentials()
|
||||
try:
|
||||
|
||||
model = "claude-3-5-sonnet@20240620"
|
||||
|
||||
vertex_ai_project = "pathrise-convert-1606954137718"
|
||||
|
|
@ -351,7 +350,6 @@ def test_avertex_ai_stream():
|
|||
@pytest.mark.flaky(retries=3, delay=1)
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_vertexai_response_basic():
|
||||
|
||||
load_vertex_ai_credentials()
|
||||
try:
|
||||
user_message = "Hello, how are you?"
|
||||
|
|
@ -1382,7 +1380,6 @@ async def test_gemini_pro_json_schema_args_sent_httpx(
|
|||
]
|
||||
)
|
||||
elif resp is not None:
|
||||
|
||||
assert resp.model == model.split("/")[1]
|
||||
|
||||
|
||||
|
|
@ -2291,6 +2288,8 @@ def test_prompt_factory_nested():
|
|||
async def test_completion_fine_tuned_model():
|
||||
load_vertex_ai_credentials()
|
||||
mock_response = AsyncMock()
|
||||
mock_response.headers = {}
|
||||
mock_response.status_code = 200
|
||||
|
||||
def return_val():
|
||||
return {
|
||||
|
|
@ -2326,7 +2325,6 @@ async def test_completion_fine_tuned_model():
|
|||
}
|
||||
|
||||
mock_response.json = return_val
|
||||
mock_response.status_code = 200
|
||||
|
||||
expected_payload = {
|
||||
"contents": [
|
||||
|
|
|
|||
527
tests/mcp_tests/test_per_user_oauth_cache.py
Normal file
527
tests/mcp_tests/test_per_user_oauth_cache.py
Normal file
|
|
@ -0,0 +1,527 @@
|
|||
"""
|
||||
Unit tests for per-user MCP OAuth token storage:
|
||||
- MCPPerUserTokenCache (NaCl-encrypted Redis cache)
|
||||
- _validate_token_response (token validation rules)
|
||||
- _compute_per_user_token_ttl (TTL computation)
|
||||
- refresh_user_oauth_token (token refresh flow)
|
||||
"""
|
||||
|
||||
import sys
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# Stub out modules that aren't available in the unit-test environment
|
||||
# so we can import the targets without a full proxy stack.
|
||||
for _mod in ("orjson",):
|
||||
if _mod not in sys.modules:
|
||||
sys.modules[_mod] = MagicMock()
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: E402
|
||||
MCPPerUserTokenCache,
|
||||
_compute_per_user_token_ttl,
|
||||
mcp_per_user_token_cache,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth, MCPTransport # noqa: E402
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer # noqa: E402
|
||||
|
||||
|
||||
def _import_validate():
|
||||
"""Lazy import to avoid pulling orjson at collection time."""
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
_validate_token_response,
|
||||
)
|
||||
|
||||
return _validate_token_response
|
||||
|
||||
|
||||
# ── Fixtures ─────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_server(**kwargs) -> MCPServer:
|
||||
defaults: Dict[str, Any] = {
|
||||
"server_id": "slack-test",
|
||||
"name": "Slack",
|
||||
"server_name": "slack",
|
||||
"url": "https://slack-mcp.example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
"auth_type": MCPAuth.oauth2,
|
||||
"client_id": "SLACK_CLIENT_ID",
|
||||
"client_secret": "SLACK_CLIENT_SECRET",
|
||||
"token_url": "https://slack.com/api/oauth.v2.access",
|
||||
"authorization_url": "https://slack.com/oauth/v2/authorize",
|
||||
}
|
||||
defaults.update(kwargs)
|
||||
return MCPServer(**defaults)
|
||||
|
||||
|
||||
# ── _validate_token_response ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestValidateTokenResponse:
|
||||
def test_passes_when_all_rules_match(self):
|
||||
_validate_token_response = _import_validate()
|
||||
token_response = {
|
||||
"access_token": "xoxb-123",
|
||||
"enterprise_id": "E04XXXXXX",
|
||||
"team": {"id": "T123", "name": "Acme"},
|
||||
}
|
||||
# Should not raise
|
||||
_validate_token_response(
|
||||
token_response=token_response,
|
||||
validation_rules={"enterprise_id": "E04XXXXXX"},
|
||||
server_id="slack-test",
|
||||
)
|
||||
|
||||
def test_raises_on_mismatch(self):
|
||||
from fastapi import HTTPException
|
||||
|
||||
_validate_token_response = _import_validate()
|
||||
token_response = {"access_token": "xoxb-123", "enterprise_id": "E99999999"}
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
_validate_token_response(
|
||||
token_response=token_response,
|
||||
validation_rules={"enterprise_id": "E04XXXXXX"},
|
||||
server_id="slack-test",
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
detail = exc_info.value.detail
|
||||
assert detail["error"] == "token_validation_failed"
|
||||
assert detail["field"] == "enterprise_id"
|
||||
|
||||
def test_raises_when_field_absent(self):
|
||||
from fastapi import HTTPException
|
||||
|
||||
_validate_token_response = _import_validate()
|
||||
token_response = {"access_token": "xoxb-123"}
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
_validate_token_response(
|
||||
token_response=token_response,
|
||||
validation_rules={"enterprise_id": "E04XXXXXX"},
|
||||
server_id="slack-test",
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
# Absent field should produce a distinct "absent" message, not str(None)
|
||||
assert "absent" in exc_info.value.detail["message"]
|
||||
|
||||
def test_absent_field_does_not_match_string_none(self):
|
||||
"""str(None)='None' must NOT match the string rule value 'None'."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
_validate_token_response = _import_validate()
|
||||
token_response = {"access_token": "tok"} # enterprise_id absent
|
||||
# Even if admin writes validation_rules={"enterprise_id": "None"}, absent
|
||||
# field should raise, not pass.
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
_validate_token_response(
|
||||
token_response=token_response,
|
||||
validation_rules={"enterprise_id": "None"},
|
||||
server_id="slack-test",
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "absent" in exc_info.value.detail["message"]
|
||||
|
||||
def test_dot_notation_nested_field(self):
|
||||
_validate_token_response = _import_validate()
|
||||
token_response = {
|
||||
"access_token": "xoxb-123",
|
||||
"team": {"enterprise_id": "E04XXXXXX"},
|
||||
}
|
||||
# Should not raise — dot-notation traverses nested dict
|
||||
_validate_token_response(
|
||||
token_response=token_response,
|
||||
validation_rules={"team.enterprise_id": "E04XXXXXX"},
|
||||
server_id="slack-test",
|
||||
)
|
||||
|
||||
def test_dot_notation_mismatch(self):
|
||||
from fastapi import HTTPException
|
||||
|
||||
_validate_token_response = _import_validate()
|
||||
token_response = {
|
||||
"access_token": "xoxb-123",
|
||||
"team": {"enterprise_id": "WRONG"},
|
||||
}
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
_validate_token_response(
|
||||
token_response=token_response,
|
||||
validation_rules={"team.enterprise_id": "E04XXXXXX"},
|
||||
server_id="slack-test",
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail["field"] == "team.enterprise_id"
|
||||
|
||||
def test_numeric_value_string_coercion(self):
|
||||
"""Numeric values in token response should match string rules."""
|
||||
_validate_token_response = _import_validate()
|
||||
token_response = {"access_token": "tok", "org_id": 12345}
|
||||
# Should not raise — str(12345) == "12345"
|
||||
_validate_token_response(
|
||||
token_response=token_response,
|
||||
validation_rules={"org_id": "12345"},
|
||||
server_id="test",
|
||||
)
|
||||
|
||||
def test_multiple_rules_all_must_match(self):
|
||||
from fastapi import HTTPException
|
||||
|
||||
_validate_token_response = _import_validate()
|
||||
token_response = {
|
||||
"access_token": "tok",
|
||||
"enterprise_id": "E04XXXXXX",
|
||||
"cloud_id": "WRONG_CLOUD",
|
||||
}
|
||||
with pytest.raises(HTTPException):
|
||||
_validate_token_response(
|
||||
token_response=token_response,
|
||||
validation_rules={
|
||||
"enterprise_id": "E04XXXXXX",
|
||||
"cloud_id": "abc-123",
|
||||
},
|
||||
server_id="atlassian",
|
||||
)
|
||||
|
||||
|
||||
# ── _compute_per_user_token_ttl ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestComputePerUserTokenTtl:
|
||||
def test_uses_server_override_when_set(self):
|
||||
server = _make_server(token_storage_ttl_seconds=7200)
|
||||
assert _compute_per_user_token_ttl(server, expires_in=99999) == 7200
|
||||
|
||||
def test_uses_expires_in_minus_buffer(self):
|
||||
server = _make_server()
|
||||
# Default buffer is 60s
|
||||
ttl = _compute_per_user_token_ttl(server, expires_in=3600)
|
||||
assert ttl == 3600 - 60
|
||||
|
||||
def test_minimum_ttl_is_1(self):
|
||||
server = _make_server()
|
||||
# expires_in smaller than buffer → clamp to 1
|
||||
ttl = _compute_per_user_token_ttl(server, expires_in=30)
|
||||
assert ttl == 1
|
||||
|
||||
def test_default_ttl_when_expires_in_none(self):
|
||||
from litellm.constants import MCP_PER_USER_TOKEN_DEFAULT_TTL
|
||||
|
||||
server = _make_server()
|
||||
ttl = _compute_per_user_token_ttl(server, expires_in=None)
|
||||
assert ttl == MCP_PER_USER_TOKEN_DEFAULT_TTL
|
||||
|
||||
|
||||
# ── MCPPerUserTokenCache ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestMCPPerUserTokenCache:
|
||||
"""Tests for Redis-backed per-user token cache.
|
||||
|
||||
Patches ``user_api_key_cache`` to avoid needing a real Redis instance.
|
||||
Patches ``encrypt_value_helper`` / ``decrypt_value_helper`` to verify
|
||||
encryption is applied before Redis writes and decryption after reads.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def cache(self):
|
||||
return MCPPerUserTokenCache()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_dual_cache(self):
|
||||
dc = MagicMock()
|
||||
dc.async_get_cache = AsyncMock(return_value=None)
|
||||
dc.async_set_cache = AsyncMock()
|
||||
return dc
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_returns_none_on_miss(self, cache, mock_dual_cache):
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper"
|
||||
) as mock_decrypt, patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache
|
||||
):
|
||||
mock_dual_cache.async_get_cache.return_value = None
|
||||
result = await cache.get("alice", "slack-test")
|
||||
assert result is None
|
||||
mock_decrypt.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_decrypts_cached_value(self, cache, mock_dual_cache):
|
||||
fake_encrypted = "encrypted_blob_abc123"
|
||||
fake_plaintext = "xoxb-slack-token"
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper",
|
||||
return_value=fake_plaintext,
|
||||
) as mock_decrypt, patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache
|
||||
):
|
||||
mock_dual_cache.async_get_cache.return_value = fake_encrypted
|
||||
result = await cache.get("alice", "slack-test")
|
||||
|
||||
assert result == fake_plaintext
|
||||
mock_decrypt.assert_called_once_with(
|
||||
fake_encrypted,
|
||||
key="mcp_per_user_token",
|
||||
exception_type="debug",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_encrypts_before_storing(self, cache, mock_dual_cache):
|
||||
fake_encrypted = "encrypted_blob_xyz"
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper",
|
||||
return_value=fake_encrypted,
|
||||
) as mock_encrypt, patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache
|
||||
):
|
||||
await cache.set("alice", "slack-test", "xoxb-token", ttl=3540)
|
||||
|
||||
mock_encrypt.assert_called_once_with("xoxb-token")
|
||||
mock_dual_cache.async_set_cache.assert_called_once()
|
||||
call_kwargs = mock_dual_cache.async_set_cache.call_args
|
||||
assert call_kwargs[0][1] == fake_encrypted # encrypted value stored
|
||||
assert call_kwargs[1]["ttl"] == 3540
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_uses_correct_cache_key(self, cache, mock_dual_cache):
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper",
|
||||
return_value="enc",
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache
|
||||
):
|
||||
await cache.set("bob", "github-server", "ghp_token", ttl=3600)
|
||||
|
||||
key_used = mock_dual_cache.async_set_cache.call_args[0][0]
|
||||
assert key_used == "mcp:per_user_token:bob:github-server"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_calls_async_delete_cache(self, cache, mock_dual_cache):
|
||||
mock_dual_cache.async_delete_cache = AsyncMock()
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache
|
||||
):
|
||||
await cache.delete("alice", "slack-test")
|
||||
|
||||
mock_dual_cache.async_delete_cache.assert_called_once_with(
|
||||
"mcp:per_user_token:alice:slack-test"
|
||||
)
|
||||
mock_dual_cache.async_set_cache.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_returns_none_on_decrypt_failure(self, cache, mock_dual_cache):
|
||||
"""Cache misses and decrypt errors should both return None without raising."""
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper",
|
||||
return_value=None, # decrypt returns None on failure
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache
|
||||
):
|
||||
mock_dual_cache.async_get_cache.return_value = "bad_encrypted_data"
|
||||
result = await cache.get("alice", "slack-test")
|
||||
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_is_noop_on_cache_error(self, cache, mock_dual_cache):
|
||||
"""Errors in the cache layer must not propagate to the caller."""
|
||||
mock_dual_cache.async_set_cache.side_effect = RuntimeError("Redis down")
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper",
|
||||
return_value="enc",
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache
|
||||
):
|
||||
# Should not raise
|
||||
await cache.set("alice", "slack-test", "token", ttl=3600)
|
||||
|
||||
|
||||
# ── refresh_user_oauth_token ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestRefreshUserOauthToken:
|
||||
"""Tests for the DB-level token refresh helper."""
|
||||
|
||||
@pytest.fixture
|
||||
def server(self):
|
||||
return _make_server()
|
||||
|
||||
@pytest.fixture
|
||||
def cred(self):
|
||||
return {
|
||||
"type": "oauth2",
|
||||
"access_token": "OLD_TOKEN",
|
||||
"refresh_token": "REFRESH_TOKEN_123",
|
||||
"expires_at": (
|
||||
datetime.now(timezone.utc) - timedelta(hours=1)
|
||||
).isoformat(),
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_none_when_no_refresh_token(self, server):
|
||||
from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token
|
||||
|
||||
cred = {"type": "oauth2", "access_token": "OLD"} # no refresh_token
|
||||
result = await refresh_user_oauth_token(
|
||||
prisma_client=MagicMock(),
|
||||
user_id="alice",
|
||||
server=server,
|
||||
cred=cred,
|
||||
)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_none_when_no_token_url(self, cred):
|
||||
from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token
|
||||
|
||||
server = _make_server(token_url=None)
|
||||
result = await refresh_user_oauth_token(
|
||||
prisma_client=MagicMock(),
|
||||
user_id="alice",
|
||||
server=server,
|
||||
cred=cred,
|
||||
)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_none_on_http_error(self, server, cred):
|
||||
from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.side_effect = Exception("Connection refused")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
):
|
||||
result = await refresh_user_oauth_token(
|
||||
prisma_client=MagicMock(),
|
||||
user_id="alice",
|
||||
server=server,
|
||||
cred=cred,
|
||||
)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stores_and_returns_new_credential(self, server, cred):
|
||||
from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token
|
||||
|
||||
new_token_response = MagicMock()
|
||||
new_token_response.json.return_value = {
|
||||
"access_token": "NEW_TOKEN",
|
||||
"expires_in": 3600,
|
||||
"refresh_token": "NEW_REFRESH",
|
||||
"scope": "channels:read chat:write",
|
||||
}
|
||||
new_token_response.raise_for_status = MagicMock()
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.return_value = new_token_response
|
||||
|
||||
stored_cred = {
|
||||
"type": "oauth2",
|
||||
"access_token": "NEW_TOKEN",
|
||||
"refresh_token": "NEW_REFRESH",
|
||||
}
|
||||
mock_prisma = AsyncMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
), patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.store_user_oauth_credential",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_store, patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential",
|
||||
new_callable=AsyncMock,
|
||||
return_value=stored_cred,
|
||||
):
|
||||
result = await refresh_user_oauth_token(
|
||||
prisma_client=mock_prisma,
|
||||
user_id="alice",
|
||||
server=server,
|
||||
cred=cred,
|
||||
)
|
||||
|
||||
assert result == stored_cred
|
||||
mock_store.assert_called_once()
|
||||
call_kwargs = mock_store.call_args[1]
|
||||
assert call_kwargs["access_token"] == "NEW_TOKEN"
|
||||
assert call_kwargs["refresh_token"] == "NEW_REFRESH"
|
||||
assert call_kwargs["expires_in"] == 3600
|
||||
assert call_kwargs["scopes"] == ["channels:read", "chat:write"]
|
||||
# Refresh path must skip the BYOK guard (row is already OAuth2)
|
||||
assert call_kwargs.get("skip_byok_guard") is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_falls_back_to_old_refresh_token_when_not_rotated(
|
||||
self, server, cred
|
||||
):
|
||||
"""When provider doesn't return a new refresh_token, keep the old one."""
|
||||
from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token
|
||||
|
||||
new_token_response = MagicMock()
|
||||
new_token_response.json.return_value = {
|
||||
"access_token": "NEW_TOKEN",
|
||||
"expires_in": 3600,
|
||||
# No refresh_token in response
|
||||
}
|
||||
new_token_response.raise_for_status = MagicMock()
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.return_value = new_token_response
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
), patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.store_user_oauth_credential",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_store, patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"type": "oauth2", "access_token": "NEW_TOKEN"},
|
||||
):
|
||||
await refresh_user_oauth_token(
|
||||
prisma_client=AsyncMock(),
|
||||
user_id="alice",
|
||||
server=server,
|
||||
cred=cred,
|
||||
)
|
||||
|
||||
call_kwargs = mock_store.call_args[1]
|
||||
# Old refresh_token preserved when provider doesn't rotate
|
||||
assert call_kwargs["refresh_token"] == "REFRESH_TOKEN_123"
|
||||
|
||||
|
||||
# ── MCPServer new fields ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestMCPServerNewFields:
|
||||
def test_token_validation_default_none(self):
|
||||
server = _make_server()
|
||||
assert server.token_validation is None
|
||||
|
||||
def test_token_validation_set(self):
|
||||
server = _make_server(token_validation={"enterprise_id": "E04XXXXXX"})
|
||||
assert server.token_validation == {"enterprise_id": "E04XXXXXX"}
|
||||
|
||||
def test_token_storage_ttl_default_none(self):
|
||||
server = _make_server()
|
||||
assert server.token_storage_ttl_seconds is None
|
||||
|
||||
def test_token_storage_ttl_set(self):
|
||||
server = _make_server(token_storage_ttl_seconds=7200)
|
||||
assert server.token_storage_ttl_seconds == 7200
|
||||
|
||||
def test_needs_user_oauth_token_true_for_oauth2_without_m2m(self):
|
||||
server = _make_server(auth_type=MCPAuth.oauth2)
|
||||
assert server.needs_user_oauth_token is True
|
||||
|
||||
def test_needs_user_oauth_token_false_for_m2m(self):
|
||||
server = _make_server(
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth2_flow="client_credentials",
|
||||
)
|
||||
assert server.needs_user_oauth_token is False
|
||||
|
|
@ -934,10 +934,7 @@ async def mock_user_object(*args, **kwargs):
|
|||
user_id = kwargs.get("user_id")
|
||||
user_email = kwargs.get("user_email")
|
||||
return LiteLLM_UserTable(
|
||||
spend=0,
|
||||
user_id=user_id,
|
||||
max_budget=None,
|
||||
user_email=user_email
|
||||
spend=0, user_id=user_id, max_budget=None, user_email=user_email
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1170,15 +1167,13 @@ async def test_end_user_jwt_auth(monkeypatch):
|
|||
# use generated key to auth in
|
||||
from litellm import Router
|
||||
from litellm.types.router import RouterGeneralSettings
|
||||
|
||||
|
||||
# Create a router with pass_through_all_models enabled
|
||||
router = Router(
|
||||
model_list=[],
|
||||
router_general_settings=RouterGeneralSettings(
|
||||
pass_through_all_models=True
|
||||
),
|
||||
router_general_settings=RouterGeneralSettings(pass_through_all_models=True),
|
||||
)
|
||||
|
||||
|
||||
setattr(litellm.proxy.proxy_server, "premium_user", True)
|
||||
setattr(
|
||||
litellm.proxy.proxy_server,
|
||||
|
|
@ -1196,7 +1191,7 @@ async def test_end_user_jwt_auth(monkeypatch):
|
|||
|
||||
cost_tracking()
|
||||
result = await user_api_key_auth(request=request, api_key=bearer_token)
|
||||
|
||||
|
||||
# Assert that end_user_id is correctly extracted from JWT token's 'sub' field
|
||||
assert result.end_user_id == "81b3e52a-67a6-4efb-9645-70527e101479"
|
||||
|
||||
|
|
@ -1228,7 +1223,9 @@ async def test_end_user_jwt_auth(monkeypatch):
|
|||
),
|
||||
)
|
||||
|
||||
with patch("litellm.acompletion", new=AsyncMock(return_value=mock_response)) as mock_completion:
|
||||
with patch(
|
||||
"litellm.acompletion", new=AsyncMock(return_value=mock_response)
|
||||
) as mock_completion:
|
||||
resp = await chat_completion(
|
||||
request=request,
|
||||
fastapi_response=temp_response,
|
||||
|
|
@ -1243,10 +1240,13 @@ async def test_end_user_jwt_auth(monkeypatch):
|
|||
# Verify the completion was called with correct end_user_id
|
||||
mock_completion.assert_called_once()
|
||||
call_kwargs = mock_completion.call_args.kwargs
|
||||
|
||||
|
||||
# end_user_id is passed in metadata as 'user_api_key_end_user_id'
|
||||
metadata = call_kwargs.get("metadata", {})
|
||||
assert metadata.get("user_api_key_end_user_id") == "81b3e52a-67a6-4efb-9645-70527e101479"
|
||||
assert (
|
||||
metadata.get("user_api_key_end_user_id")
|
||||
== "81b3e52a-67a6-4efb-9645-70527e101479"
|
||||
)
|
||||
|
||||
|
||||
def test_can_rbac_role_call_route():
|
||||
|
|
@ -1278,13 +1278,13 @@ def test_user_api_key_auth_jwt_hashing():
|
|||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
|
||||
|
||||
# Test with a JWT token (3 parts separated by dots)
|
||||
jwt_token = "test-jwt-token-header.payload.signature"
|
||||
|
||||
|
||||
# Create UserAPIKeyAuth instance with JWT
|
||||
user_auth = UserAPIKeyAuth(api_key=jwt_token)
|
||||
|
||||
|
||||
# Verify that the API key is hashed with "hashed-jwt-" prefix
|
||||
# critical - the raw JWT token should not be in the api_key or token
|
||||
assert user_auth.api_key.startswith("hashed-jwt-")
|
||||
|
|
@ -1292,19 +1292,18 @@ def test_user_api_key_auth_jwt_hashing():
|
|||
assert jwt_token not in user_auth.api_key
|
||||
assert jwt_token not in user_auth.token
|
||||
|
||||
|
||||
# Test with a regular API key (should not be hashed)
|
||||
regular_api_key = "sk-1234567890abcdef"
|
||||
user_auth_regular = UserAPIKeyAuth(api_key=regular_api_key)
|
||||
|
||||
|
||||
# Verify that regular API key is hashed normally (without "hashed-jwt-" prefix)
|
||||
assert not user_auth_regular.api_key.startswith("hashed-jwt-")
|
||||
assert not user_auth_regular.token.startswith("hashed-jwt-")
|
||||
|
||||
|
||||
# Test with a non-JWT, non-sk string (should not be hashed)
|
||||
non_jwt_key = "some-random-key"
|
||||
user_auth_non_jwt = UserAPIKeyAuth(api_key=non_jwt_key)
|
||||
|
||||
|
||||
# Verify that non-JWT key is not hashed
|
||||
assert user_auth_non_jwt.api_key == non_jwt_key
|
||||
assert user_auth_non_jwt.token == non_jwt_key
|
||||
|
|
@ -1315,19 +1314,19 @@ def test_jwt_handler_is_jwt_static_method():
|
|||
Test that JWTHandler.is_jwt is a static method and works correctly
|
||||
"""
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
|
||||
|
||||
# Test with valid JWT format
|
||||
valid_jwt = "test-jwt-token-header.payload.signature"
|
||||
assert JWTHandler.is_jwt(valid_jwt) == True
|
||||
|
||||
|
||||
# Test with invalid JWT format (only 2 parts)
|
||||
invalid_jwt = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ"
|
||||
assert JWTHandler.is_jwt(invalid_jwt) == False
|
||||
|
||||
|
||||
# Test with regular API key
|
||||
regular_key = "sk-1234567890abcdef"
|
||||
assert JWTHandler.is_jwt(regular_key) == False
|
||||
|
||||
|
||||
# Test with empty string
|
||||
assert JWTHandler.is_jwt("") == False
|
||||
|
||||
|
|
@ -1461,7 +1460,13 @@ async def test_auth_jwt_es256_jwk_path(monkeypatch):
|
|||
|
||||
now = int(time.time())
|
||||
token = jwt.encode(
|
||||
{"sub": "alice", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300},
|
||||
{
|
||||
"sub": "alice",
|
||||
"aud": "litellm-proxy",
|
||||
"iss": "http://example",
|
||||
"iat": now,
|
||||
"exp": now + 300,
|
||||
},
|
||||
ec_priv_pem,
|
||||
algorithm="ES256",
|
||||
headers={"kid": "ec1"},
|
||||
|
|
@ -1508,7 +1513,13 @@ async def test_auth_jwt_rs256_regression(monkeypatch):
|
|||
|
||||
now = int(time.time())
|
||||
token = jwt.encode(
|
||||
{"sub": "bob", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300},
|
||||
{
|
||||
"sub": "bob",
|
||||
"aud": "litellm-proxy",
|
||||
"iss": "http://example",
|
||||
"iat": now,
|
||||
"exp": now + 300,
|
||||
},
|
||||
rsa_priv_pem,
|
||||
algorithm="RS256",
|
||||
headers={"kid": "rsa1"},
|
||||
|
|
@ -1540,7 +1551,13 @@ async def test_auth_jwt_mismatched_key_fails(monkeypatch):
|
|||
)
|
||||
now = int(time.time())
|
||||
token = jwt.encode(
|
||||
{"sub": "mallory", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300},
|
||||
{
|
||||
"sub": "mallory",
|
||||
"aud": "litellm-proxy",
|
||||
"iss": "http://example",
|
||||
"iat": now,
|
||||
"exp": now + 300,
|
||||
},
|
||||
ec_priv_pem,
|
||||
algorithm="ES256",
|
||||
headers={"kid": "ec1"},
|
||||
|
|
@ -1566,4 +1583,4 @@ async def test_auth_jwt_mismatched_key_fails(monkeypatch):
|
|||
with patch.object(h, "get_public_key", new=AsyncMock(return_value=rsa_jwk)):
|
||||
with pytest.raises(Exception) as exc:
|
||||
await h.auth_jwt(token)
|
||||
assert "Validation fails" in str(exc.value)
|
||||
assert "Validation fails" in str(exc.value)
|
||||
|
|
|
|||
|
|
@ -2637,3 +2637,50 @@ async def test_handle_logging_proxy_only_error_skips_handlers_for_pass_through()
|
|||
mock_async.assert_not_called()
|
||||
mock_sync.assert_not_called()
|
||||
assert logging_obj.call_type == CallTypes.pass_through.value
|
||||
|
||||
|
||||
def test_handle_exception_on_proxy_preserves_status_code():
|
||||
"""
|
||||
OpenAI batch creation returns 429 for rate limits. LiteLLM wraps this as a
|
||||
RateLimitError with status_code=429. handle_exception_on_proxy must pass
|
||||
that status code through instead of hardcoding 500.
|
||||
"""
|
||||
from litellm.proxy.utils import handle_exception_on_proxy
|
||||
|
||||
rate_limit_error = litellm.RateLimitError(
|
||||
message="Rate limit exceeded: batch creation limit of 2000/hour hit",
|
||||
llm_provider="openai",
|
||||
model="gpt-4o",
|
||||
)
|
||||
|
||||
result = handle_exception_on_proxy(rate_limit_error)
|
||||
|
||||
assert int(result.code) == 429, f"Expected 429, got {result.code}"
|
||||
|
||||
|
||||
def test_handle_exception_on_proxy_defaults_to_500_for_unknown_exceptions():
|
||||
"""
|
||||
Generic exceptions with no status_code should still return 500.
|
||||
"""
|
||||
from litellm.proxy.utils import handle_exception_on_proxy
|
||||
|
||||
result = handle_exception_on_proxy(Exception("something went wrong"))
|
||||
|
||||
assert int(result.code) == 500, f"Expected 500, got {result.code}"
|
||||
|
||||
|
||||
def test_handle_exception_on_proxy_preserves_auth_error_status_code():
|
||||
"""
|
||||
AuthenticationError (401) should also pass through correctly.
|
||||
"""
|
||||
from litellm.proxy.utils import handle_exception_on_proxy
|
||||
|
||||
auth_error = litellm.AuthenticationError(
|
||||
message="Invalid API key",
|
||||
llm_provider="openai",
|
||||
model="gpt-4o",
|
||||
)
|
||||
|
||||
result = handle_exception_on_proxy(auth_error)
|
||||
|
||||
assert int(result.code) == 401, f"Expected 401, got {result.code}"
|
||||
|
|
|
|||
|
|
@ -359,27 +359,38 @@ async def test_auth_with_allowed_routes(route, should_raise_error):
|
|||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route, user_role, expected_result",
|
||||
"route, user_role, should_be_allowed",
|
||||
[
|
||||
# Proxy Admin checks
|
||||
# Admin can access everything
|
||||
("/config/update", "proxy_admin", True),
|
||||
("/global/spend/logs", "proxy_admin", True),
|
||||
("/key/delete", "proxy_admin", False),
|
||||
("/key/generate", "proxy_admin", False),
|
||||
("/key/regenerate", "proxy_admin", False),
|
||||
# Internal User checks - allowed routes
|
||||
("/global/activity/cache_hits", "proxy_admin", True),
|
||||
# Internal User - allowed read-only routes
|
||||
("/global/spend/logs", "internal_user", True),
|
||||
("/key/delete", "internal_user", False),
|
||||
("/key/generate", "internal_user", False),
|
||||
("/key/82akk800000000jjsk/regenerate", "internal_user", False),
|
||||
# Internal User Viewer
|
||||
("/key/generate", "internal_user_viewer", False),
|
||||
# Internal User checks - disallowed routes
|
||||
("/spend/logs/ui", "internal_user", True),
|
||||
("/global/activity/cache_hits", "internal_user", True),
|
||||
("/health/services", "internal_user", True),
|
||||
# Internal User - BLOCKED from admin routes (security fix)
|
||||
("/config/update", "internal_user", False),
|
||||
("/config/pass_through_endpoint", "internal_user", False),
|
||||
("/config/field/update", "internal_user", False),
|
||||
("/organization/member_add", "internal_user", False),
|
||||
# Internal User Viewer - allowed spend routes only
|
||||
("/spend/logs/ui", "internal_user_viewer", True),
|
||||
("/global/spend/all_tag_names", "internal_user_viewer", True),
|
||||
# Internal User Viewer - blocked from admin routes
|
||||
("/config/update", "internal_user_viewer", False),
|
||||
("/key/generate", "internal_user_viewer", False),
|
||||
],
|
||||
)
|
||||
def test_is_ui_route_allowed(route, user_role, expected_result):
|
||||
from litellm.proxy.auth.auth_checks import _is_ui_route
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
def test_ui_token_route_access(route, user_role, should_be_allowed):
|
||||
"""
|
||||
Verify that UI tokens (team_id=litellm-dashboard) go through the same
|
||||
RBAC checks as API tokens. Non-admin dashboard users must not be able
|
||||
to access admin-only routes like /config/update.
|
||||
"""
|
||||
from litellm.proxy.auth.auth_checks import _is_api_route_allowed
|
||||
from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth
|
||||
|
||||
user_obj = LiteLLM_UserTable(
|
||||
user_id="3b803c0e-666e-4e99-bd5c-6e534c07e297",
|
||||
|
|
@ -395,18 +406,36 @@ def test_is_ui_route_allowed(route, user_role, expected_result):
|
|||
organization_memberships=[],
|
||||
)
|
||||
|
||||
received_args: dict = {
|
||||
"route": route,
|
||||
"user_obj": user_obj,
|
||||
}
|
||||
try:
|
||||
assert _is_ui_route(**received_args) == expected_result
|
||||
except Exception as e:
|
||||
# If expected result is False, we expect an error
|
||||
if expected_result is False:
|
||||
pass
|
||||
else:
|
||||
raise e
|
||||
valid_token = UserAPIKeyAuth(
|
||||
user_id="3b803c0e-666e-4e99-bd5c-6e534c07e297",
|
||||
team_id="litellm-dashboard",
|
||||
user_role=user_role,
|
||||
)
|
||||
|
||||
from starlette.datastructures import URL
|
||||
from fastapi import Request
|
||||
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url=route)
|
||||
|
||||
if should_be_allowed:
|
||||
result = _is_api_route_allowed(
|
||||
route=route,
|
||||
request=request,
|
||||
request_data={},
|
||||
valid_token=valid_token,
|
||||
user_obj=user_obj,
|
||||
)
|
||||
assert result is True
|
||||
else:
|
||||
with pytest.raises(Exception):
|
||||
_is_api_route_allowed(
|
||||
route=route,
|
||||
request=request,
|
||||
request_data={},
|
||||
valid_token=valid_token,
|
||||
user_obj=user_obj,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -684,7 +713,7 @@ async def test_soft_budget_alert():
|
|||
|
||||
|
||||
def test_is_allowed_route():
|
||||
from litellm.proxy.auth.auth_checks import _is_allowed_route
|
||||
from litellm.proxy.auth.auth_checks import _is_api_route_allowed
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
import datetime
|
||||
|
||||
|
|
@ -692,7 +721,6 @@ def test_is_allowed_route():
|
|||
|
||||
args = {
|
||||
"route": "/embeddings",
|
||||
"token_type": "api",
|
||||
"request": request,
|
||||
"request_data": {"input": ["hello world"], "model": "embedding-small"},
|
||||
"valid_token": UserAPIKeyAuth(
|
||||
|
|
@ -752,7 +780,7 @@ def test_is_allowed_route():
|
|||
"user_obj": None,
|
||||
}
|
||||
|
||||
assert _is_allowed_route(**args)
|
||||
assert _is_api_route_allowed(**args)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -836,7 +864,6 @@ async def test_user_api_key_auth_websocket():
|
|||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True
|
||||
) as mock_user_api_key_auth:
|
||||
|
||||
# Make the call to the WebSocket function
|
||||
await user_api_key_auth_websocket(mock_websocket)
|
||||
|
||||
|
|
@ -845,10 +872,14 @@ async def test_user_api_key_auth_websocket():
|
|||
|
||||
# Get the request object that was passed to user_api_key_auth
|
||||
request_arg = mock_user_api_key_auth.call_args.kwargs["request"]
|
||||
|
||||
|
||||
# Verify that the request has headers set
|
||||
assert hasattr(request_arg, "headers"), "Request object should have headers attribute"
|
||||
assert "authorization" in request_arg.headers, "Request headers should contain authorization"
|
||||
assert hasattr(
|
||||
request_arg, "headers"
|
||||
), "Request object should have headers attribute"
|
||||
assert (
|
||||
"authorization" in request_arg.headers
|
||||
), "Request headers should contain authorization"
|
||||
assert request_arg.headers["authorization"] == "Bearer some_api_key"
|
||||
|
||||
assert (
|
||||
|
|
@ -1036,7 +1067,10 @@ async def test_jwt_non_admin_team_route_access(monkeypatch):
|
|||
|
||||
# Create request
|
||||
request = Request(
|
||||
scope={"type": "http", "headers": [(b"authorization", b"Bearer fake.jwt.token")]}
|
||||
scope={
|
||||
"type": "http",
|
||||
"headers": [(b"authorization", b"Bearer fake.jwt.token")],
|
||||
}
|
||||
)
|
||||
request._url = URL(url="/team/new")
|
||||
|
||||
|
|
@ -1101,14 +1135,14 @@ async def test_x_litellm_api_key():
|
|||
ignored_key = "aj12445"
|
||||
|
||||
# Create request with headers as bytes
|
||||
request = Request(
|
||||
scope={
|
||||
"type": "http"
|
||||
}
|
||||
)
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
|
||||
valid_token = await user_api_key_auth(request=request, api_key="Bearer " + ignored_key, custom_litellm_key_header=master_key)
|
||||
valid_token = await user_api_key_auth(
|
||||
request=request,
|
||||
api_key="Bearer " + ignored_key,
|
||||
custom_litellm_key_header=master_key,
|
||||
)
|
||||
assert valid_token.token == hash_token(master_key)
|
||||
|
||||
|
||||
|
|
@ -1123,7 +1157,9 @@ async def test_user_api_key_from_query_param():
|
|||
from litellm.proxy.proxy_server import hash_token, user_api_key_cache
|
||||
|
||||
user_key = "sk-query-1234"
|
||||
user_api_key_cache.set_cache(key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key)))
|
||||
user_api_key_cache.set_cache(
|
||||
key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key))
|
||||
)
|
||||
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
||||
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
|
|
@ -1136,7 +1172,9 @@ async def test_user_api_key_from_query_param():
|
|||
"query_string": f"alt=sse&key={user_key}".encode(),
|
||||
}
|
||||
)
|
||||
request._url = URL(url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}")
|
||||
request._url = URL(
|
||||
url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}"
|
||||
)
|
||||
|
||||
async def return_body():
|
||||
return b"{}"
|
||||
|
|
@ -1145,4 +1183,3 @@ async def test_user_api_key_from_query_param():
|
|||
|
||||
valid_token = await user_api_key_auth(request=request, api_key="")
|
||||
assert valid_token.token == hash_token(user_key)
|
||||
|
||||
|
|
|
|||
|
|
@ -118,6 +118,8 @@ async def test_async_post_call_success_hook_includes_client_ip_user_agent():
|
|||
"user_api_key_alias": "alias_1",
|
||||
"user_api_key_team_id": "team_1",
|
||||
"user_api_key_team_alias": "team_alias_1",
|
||||
"user_api_key_org_id": None,
|
||||
"user_api_key_org_alias": None,
|
||||
"user_api_key_user_email": "test@example.com",
|
||||
"user_api_key_request_route": "/chat/completions",
|
||||
"requester_ip_address": "192.168.1.1",
|
||||
|
|
|
|||
|
|
@ -525,6 +525,61 @@ async def test_set_user_budget_metrics_after_api_request_inf_when_genuinely_no_b
|
|||
)
|
||||
|
||||
|
||||
def test_per_request_metrics_emit_all_identity_labels(prometheus_logger):
|
||||
"""Verify org labels appear when flag is on and are absent when flag is off."""
|
||||
import litellm
|
||||
from litellm.types.integrations.prometheus import UserAPIKeyLabelValues
|
||||
|
||||
prometheus_logger.litellm_requests_metric = MagicMock()
|
||||
prometheus_logger.litellm_spend_metric = MagicMock()
|
||||
|
||||
enum_values = UserAPIKeyLabelValues(
|
||||
hashed_api_key="hashed-key",
|
||||
api_key_alias="my-key",
|
||||
model="gpt-4",
|
||||
team="team-abc",
|
||||
team_alias="my-team",
|
||||
org_id="org-abc",
|
||||
org_alias="my-org",
|
||||
user="user-1",
|
||||
)
|
||||
|
||||
common_kwargs = dict(
|
||||
end_user_id=None,
|
||||
user_api_key="hashed-key",
|
||||
user_api_key_alias="my-key",
|
||||
model="gpt-4",
|
||||
user_api_team="team-abc",
|
||||
user_api_team_alias="my-team",
|
||||
user_id="user-1",
|
||||
response_cost=0.001,
|
||||
enum_values=enum_values,
|
||||
)
|
||||
|
||||
try:
|
||||
# org labels are always included in per-request metrics
|
||||
prometheus_logger._increment_top_level_request_and_spend_metrics(**common_kwargs)
|
||||
label_kwargs = prometheus_logger.litellm_requests_metric.labels.call_args.kwargs
|
||||
assert label_kwargs["org_id"] == "org-abc"
|
||||
assert label_kwargs["org_alias"] == "my-org"
|
||||
assert label_kwargs["team"] == "team-abc"
|
||||
assert label_kwargs["user"] == "user-1"
|
||||
|
||||
# Metrics not in the org-emission list must NOT get org labels
|
||||
from litellm.types.integrations.prometheus import PrometheusMetricLabels
|
||||
for metric in ("litellm_remaining_api_key_budget_metric", "litellm_remaining_team_budget_metric"):
|
||||
labels = PrometheusMetricLabels.get_labels(metric)
|
||||
assert "org_id" not in labels, f"{metric} should not have org_id"
|
||||
assert "org_alias" not in labels, f"{metric} should not have org_alias"
|
||||
|
||||
# org_id in custom_prometheus_metadata_labels must not produce duplicate labels
|
||||
litellm.custom_prometheus_metadata_labels = ["org_id"]
|
||||
labels = PrometheusMetricLabels.get_labels("litellm_requests_metric")
|
||||
assert labels.count("org_id") == 1
|
||||
finally:
|
||||
litellm.custom_prometheus_metadata_labels = []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Org budget metric tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
from unittest.mock import MagicMock
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.llms.anthropic.chat.handler import ModelResponseIterator
|
||||
from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call
|
||||
from litellm.types.llms.openai import (
|
||||
ChatCompletionToolCallChunk,
|
||||
ChatCompletionToolCallFunctionChunk,
|
||||
|
|
@ -9,6 +11,33 @@ from litellm.types.llms.openai import (
|
|||
from litellm.types.responses.main import OutputCodeInterpreterCall
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_make_call_passes_logging_obj_to_client_post():
|
||||
"""make_call must pass logging_obj to client.post so track_llm_api_timing can set llm_api_duration_ms for litellm_overhead_time_ms."""
|
||||
mock_client = AsyncMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.aiter_lines = MagicMock(return_value=iter([b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n']))
|
||||
mock_client.post.return_value = mock_response
|
||||
|
||||
logging_obj = MagicMock()
|
||||
|
||||
await make_call(
|
||||
client=mock_client,
|
||||
api_base="https://api.anthropic.com/v1/messages",
|
||||
headers={},
|
||||
data="{}",
|
||||
model="claude-3-5-haiku",
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
logging_obj=logging_obj,
|
||||
timeout=60.0,
|
||||
json_mode=False,
|
||||
)
|
||||
|
||||
mock_client.post.assert_called_once()
|
||||
call_kwargs = mock_client.post.call_args[1]
|
||||
assert call_kwargs.get("logging_obj") is logging_obj
|
||||
|
||||
|
||||
def test_redacted_thinking_content_block_delta():
|
||||
chunk = {
|
||||
"type": "content_block_start",
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import json
|
|||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -125,7 +126,8 @@ async def test_bedrock_sse_wrapper_keeps_usage_in_message_start_and_message_delt
|
|||
assert "usage" in delta_json
|
||||
assert delta_json["usage"]["cache_creation_input_tokens"] == 1562
|
||||
assert delta_json["usage"]["cache_read_input_tokens"] == 32392
|
||||
assert delta_json["usage"]["input_tokens"] == 3 + 1562 + 32392
|
||||
assert delta_json["usage"]["input_tokens"] == 3
|
||||
assert delta_json["usage"]["output_tokens"] == 8
|
||||
|
||||
|
||||
def test_chunk_parser_usage_transformation():
|
||||
|
|
@ -402,3 +404,111 @@ def test_bedrock_messages_strips_output_config_with_output_format():
|
|||
|
||||
assert "output_config" not in result
|
||||
assert "output_format" not in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_promote_message_stop_usage_preserves_message_delta_output_tokens():
|
||||
"""
|
||||
Bedrock unified /messages streaming can send full usage on message_delta and a
|
||||
conflicting smaller usage on message_stop (e.g. output_tokens 9 vs 12).
|
||||
_promote_message_stop_usage must not replace message_delta output_tokens.
|
||||
"""
|
||||
cfg = AmazonAnthropicClaudeMessagesConfig()
|
||||
|
||||
async def _stream(): # type: ignore[return-type]
|
||||
yield {
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {
|
||||
"input_tokens": 3,
|
||||
"cache_creation_input_tokens": 10553,
|
||||
"cache_read_input_tokens": 25490,
|
||||
"output_tokens": 12,
|
||||
},
|
||||
}
|
||||
yield {
|
||||
"type": "message_stop",
|
||||
"usage": {"input_tokens": 3, "output_tokens": 9},
|
||||
}
|
||||
|
||||
merged: list[dict] = []
|
||||
async for chunk in cfg._promote_message_stop_usage(_stream()):
|
||||
if isinstance(chunk, dict):
|
||||
merged.append(chunk)
|
||||
|
||||
assert len(merged) >= 1
|
||||
delta_out = merged[0]
|
||||
assert delta_out["type"] == "message_delta"
|
||||
assert delta_out["usage"]["output_tokens"] == 12
|
||||
assert delta_out["usage"]["cache_creation_input_tokens"] == 10553
|
||||
assert delta_out["usage"]["cache_read_input_tokens"] == 25490
|
||||
assert delta_out["usage"]["input_tokens"] == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unified_bedrock_messages_sse_usage_and_cost_claude_sonnet_46():
|
||||
"""
|
||||
End-to-end for Bedrock Invoke Anthropic Messages (unified) streaming path:
|
||||
dict chunks -> _promote_message_stop_usage -> bedrock_sse_wrapper SSE bytes ->
|
||||
same logging reconstruction as Anthropic /messages. Ensures token counts and
|
||||
completion_cost match model_prices for us.anthropic.claude-sonnet-4-6.
|
||||
"""
|
||||
from litellm import completion_cost
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
|
||||
AnthropicPassthroughLoggingHandler,
|
||||
)
|
||||
|
||||
cfg = AmazonAnthropicClaudeMessagesConfig()
|
||||
|
||||
async def _stream(): # type: ignore[return-type]
|
||||
yield {
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {
|
||||
"input_tokens": 3,
|
||||
"cache_creation_input_tokens": 10553,
|
||||
"cache_read_input_tokens": 25490,
|
||||
"output_tokens": 12,
|
||||
},
|
||||
}
|
||||
yield {
|
||||
"type": "message_stop",
|
||||
"usage": {"input_tokens": 3, "output_tokens": 9},
|
||||
}
|
||||
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-6",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
stream=True,
|
||||
call_type="chat",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test_unified_bedrock_messages_sse_cost",
|
||||
function_id="test_unified_bedrock_messages_sse_cost",
|
||||
)
|
||||
|
||||
collected: list[bytes] = []
|
||||
async for sse in cfg.bedrock_sse_wrapper(
|
||||
completion_stream=_stream(),
|
||||
litellm_logging_obj=logging_obj,
|
||||
request_body={"model": "us.anthropic.claude-sonnet-4-6"},
|
||||
):
|
||||
collected.append(sse)
|
||||
|
||||
built = AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
|
||||
all_chunks=collected,
|
||||
model="us.anthropic.claude-sonnet-4-6",
|
||||
litellm_logging_obj=Mock(),
|
||||
)
|
||||
assert built.usage is not None
|
||||
assert built.usage.completion_tokens == 12
|
||||
assert built.usage.prompt_tokens == 36046
|
||||
assert built.usage.total_tokens == 36058
|
||||
assert built.usage.cache_creation_input_tokens == 10553
|
||||
assert built.usage.cache_read_input_tokens == 25490
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=built,
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-6",
|
||||
custom_llm_provider="bedrock",
|
||||
)
|
||||
assert cost == pytest.approx(0.052150725, rel=0, abs=1e-9)
|
||||
|
|
|
|||
|
|
@ -127,11 +127,12 @@ class TestToolTransformationIntegration:
|
|||
}
|
||||
|
||||
validated_tool = validate_dict(openai_tool, ChatCompletionTool)
|
||||
|
||||
|
||||
# After validation, parameters should have type='object'
|
||||
assert validated_tool["function"]["parameters"]["type"] == "object"
|
||||
assert "properties" in validated_tool["function"]["parameters"]
|
||||
|
||||
|
||||
def test_should_transform_tool_with_existing_parameters(self):
|
||||
"""Tool with parameters should preserve them while ensuring type='object'."""
|
||||
from litellm.llms.sap.chat.transformation import validate_dict
|
||||
|
|
|
|||
564
tests/test_litellm/llms/sap/chat/test_sap_transformation.py
Normal file
564
tests/test_litellm/llms/sap/chat/test_sap_transformation.py
Normal file
|
|
@ -0,0 +1,564 @@
|
|||
import warnings
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
class TestSAPTransformationIntegration:
|
||||
"""Integration tests for SAP transformation."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_config(self):
|
||||
from litellm.llms.sap.chat.transformation import GenAIHubOrchestrationConfig
|
||||
|
||||
config = GenAIHubOrchestrationConfig()
|
||||
config.token_creator = lambda: "Bearer TEST_TOKEN"
|
||||
config._base_url = "https://api.test-sap.com"
|
||||
config._resource_group = "test-group"
|
||||
|
||||
return config
|
||||
|
||||
def test_parameter_classification_in_transform_request(self, mock_config):
|
||||
"""Test parameter classification within the actual transform_request method."""
|
||||
|
||||
model = "gpt-4o"
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
|
||||
optional_params = {
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100,
|
||||
"deployment_url": "https://custom.sap.com/deployment/123",
|
||||
"model_version": "v1.5",
|
||||
"tools": [{"type": "function", "function": {"name": "calculator"}}],
|
||||
"frequency_penalty": 0.1
|
||||
}
|
||||
|
||||
result = mock_config.transform_request(
|
||||
model, messages, optional_params, {}, {}
|
||||
)
|
||||
|
||||
model_params = result["config"]["modules"]["prompt_templating"]["model"]["params"]
|
||||
|
||||
assert "temperature" in model_params
|
||||
assert "frequency_penalty" in model_params
|
||||
assert "deployment_url" not in model_params
|
||||
assert "model_version" not in model_params
|
||||
assert "tools" not in model_params
|
||||
|
||||
model_version = result["config"]["modules"]["prompt_templating"]["model"]["version"]
|
||||
assert model_version == "v1.5"
|
||||
|
||||
prompt = result["config"]["modules"]["prompt_templating"]["prompt"]
|
||||
if "tools" in prompt:
|
||||
assert isinstance(prompt["tools"], list)
|
||||
for tool in prompt["tools"]:
|
||||
assert tool["function"]["parameters"]["type"] == "object", (
|
||||
"SAP API requires parameters.type == 'object'"
|
||||
)
|
||||
assert "properties" in tool["function"]["parameters"]
|
||||
|
||||
def test_transform_request_parameter_handling_robustness(self, mock_config):
|
||||
"""Test transform_request method handles various parameter combinations correctly."""
|
||||
|
||||
model = "gpt-4o"
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
|
||||
test_cases = [
|
||||
# Case 1: Basic parameters only
|
||||
{
|
||||
"params": {"temperature": 0.7, "max_tokens": 100},
|
||||
"expected_in_model": {"temperature", "max_tokens"},
|
||||
"expected_excluded": set()
|
||||
},
|
||||
# Case 2: Parameters with auth/infrastructure components
|
||||
{
|
||||
"params": {
|
||||
"temperature": 0.8,
|
||||
"deployment_url": "https://api.sap.com/deployments/test",
|
||||
"max_tokens": 150
|
||||
},
|
||||
"expected_in_model": {"temperature", "max_tokens"},
|
||||
"expected_excluded": {"deployment_url"}
|
||||
},
|
||||
# Case 3: Parameters with framework components
|
||||
{
|
||||
"params": {
|
||||
"temperature": 0.6,
|
||||
"model_version": "v2.0",
|
||||
"tools": [{"function": {"name": "test"}}],
|
||||
"frequency_penalty": 0.1
|
||||
},
|
||||
"expected_in_model": {"temperature", "frequency_penalty"},
|
||||
"expected_excluded": {"model_version", "tools"}
|
||||
}
|
||||
]
|
||||
|
||||
for i, test_case in enumerate(test_cases):
|
||||
filtered_params = {
|
||||
k: v for k, v in test_case["params"].items()
|
||||
if k not in {"tools", "model_version", "deployment_url"}
|
||||
}
|
||||
|
||||
for expected_param in test_case["expected_in_model"]:
|
||||
assert expected_param in filtered_params, f"Case {i + 1}: {expected_param} should be in model params"
|
||||
|
||||
for excluded_param in test_case["expected_excluded"]:
|
||||
assert excluded_param not in filtered_params, f"Case {i + 1}: {excluded_param} should be excluded from model params"
|
||||
|
||||
result = mock_config.transform_request(
|
||||
model, messages, test_case["params"], {}, {}
|
||||
)
|
||||
if result and "config" in result:
|
||||
model_params = result["config"]["modules"]["prompt_templating"]["model"]["params"]
|
||||
|
||||
for excluded_param in test_case["expected_excluded"]:
|
||||
assert excluded_param not in model_params, (
|
||||
f"Case {i + 1}: {excluded_param} should not be in actual model params"
|
||||
)
|
||||
|
||||
def test_config_transform_with_response_format_json_object(self, mock_config):
|
||||
expected_dict = {'config':
|
||||
{'modules':
|
||||
{'prompt_templating':
|
||||
{'prompt':
|
||||
{'template':
|
||||
[{'role': 'user', 'content': 'First man on the moon, answer in json'}],
|
||||
'response_format': {'type': 'json_object'}},
|
||||
'model': {'name': 'gpt-4o', 'params': {}, 'version': 'latest'}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
config = mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{'role': 'user', 'content': 'First man on the moon, answer in json'}],
|
||||
optional_params={'response_format': {'type': 'json_object'},
|
||||
'deployment_url': "shouldn't be in results"},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
assert config == expected_dict
|
||||
|
||||
def test_config_transform_with_response_format_json_schema(self, mock_config):
|
||||
|
||||
expected_response_format = {
|
||||
'type': 'json_schema',
|
||||
'json_schema': {
|
||||
'description': 'Schema for person information',
|
||||
'name': 'person_info',
|
||||
'schema': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'name': {
|
||||
'type': 'string',
|
||||
'description': "The person's full name"
|
||||
},
|
||||
'age': {
|
||||
'type': 'integer',
|
||||
'description': "The person's age in years"
|
||||
},
|
||||
'occupation': {
|
||||
'type': 'string',
|
||||
'description': "The person's job title"
|
||||
}
|
||||
},
|
||||
'required': ['name', 'age', 'occupation'],
|
||||
'additionalProperties': False
|
||||
},
|
||||
'strict': True
|
||||
}
|
||||
}
|
||||
|
||||
config = mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{'role': 'user', 'content': 'First man on the moon, answer in json'}],
|
||||
optional_params={'response_format': expected_response_format,
|
||||
'deployment_url': "shouldn't be in results"},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
assert config["config"]["modules"]["prompt_templating"]["prompt"]["response_format"] == expected_response_format
|
||||
assert len(config["config"]["modules"]["prompt_templating"]["model"]["params"]) == 0
|
||||
|
||||
def test_config_transform_with_stream(self, mock_config):
|
||||
expected_dict = {
|
||||
'config': {
|
||||
'modules': {
|
||||
'prompt_templating': {
|
||||
'prompt': {
|
||||
'template': [{'role': 'user', 'content': 'Hello, how are you?'}]
|
||||
},
|
||||
'model': {
|
||||
'name': 'anthropic--claude-4-sonnet',
|
||||
'params': {},
|
||||
'version': 'latest'
|
||||
}
|
||||
}
|
||||
},
|
||||
'stream': {'chunk_size': 10}
|
||||
}
|
||||
}
|
||||
config = mock_config.transform_request(
|
||||
model="anthropic--claude-4-sonnet",
|
||||
messages=[{'content': 'Hello, how are you?', 'role': 'user'}],
|
||||
optional_params={'stream': True,
|
||||
'stream_options': {'chunk_size': 10},
|
||||
'model_version': 'latest',
|
||||
'deployment_url': "shouldn't be in results"},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
|
||||
assert config == expected_dict
|
||||
|
||||
def test_sap_placeholder_defaults(self, mock_config):
|
||||
config = mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[
|
||||
{"role": "user", "content": "Hello. Answer {{ ?user_query }}"}
|
||||
],
|
||||
optional_params={'deployment_url': "shouldn't be in results",
|
||||
"placeholder_defaults": {"user_query": "default value"}},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
|
||||
assert config["config"]["modules"]["prompt_templating"]["prompt"]["defaults"] == {
|
||||
"user_query": "default value"}
|
||||
assert config["config"]["modules"]["prompt_templating"]["model"]["params"] == {}
|
||||
|
||||
def test_sap_placeholder_values(self, mock_config):
|
||||
placeholder_values = {"user_query": "Some text"}
|
||||
config = mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[
|
||||
{"role": "user", "content": "Hello. Answer {{ ?user_query }}"}
|
||||
],
|
||||
optional_params={'deployment_url': "shouldn't be in results",
|
||||
"placeholder_values": placeholder_values},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
|
||||
assert config["placeholder_values"] == placeholder_values
|
||||
assert config["config"]["modules"]["prompt_templating"]["model"]["params"] == {}
|
||||
|
||||
def test_sap_grounding(self, mock_config):
|
||||
grounding_config = {
|
||||
'type': 'document_grounding_service',
|
||||
'config': {
|
||||
'filters': [
|
||||
{'id': 's3-docs',
|
||||
'data_repository_type': 'vector',
|
||||
'search_config': {'max_chunk_count': 2},
|
||||
'data_repositories': ['123456890-test']
|
||||
}
|
||||
],
|
||||
'placeholders': {'input': ['user_query'], 'output': 'grounding_response'},
|
||||
'metadata_params': ['source', 'webUrl', 'title', 'mimeType', 'fileSuffix']
|
||||
}
|
||||
}
|
||||
placeholder_values = {"user_query": "Some text"}
|
||||
config = mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[
|
||||
{"role": "user", "content": "Hello. Answer {{ ?user_query }} using context: {{ ?grounding_response }}"}
|
||||
],
|
||||
optional_params={'deployment_url': "shouldn't be in results",
|
||||
"grounding": grounding_config,
|
||||
"placeholder_values": placeholder_values},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
assert config["placeholder_values"] == placeholder_values
|
||||
modules = config["config"]["modules"]
|
||||
assert modules["grounding"]["type"] == "document_grounding_service"
|
||||
assert modules["grounding"]["config"]["placeholders"]["output"] == "grounding_response"
|
||||
assert modules["grounding"]["config"]["filters"][0]["data_repository_type"] == "vector"
|
||||
assert modules["prompt_templating"]["model"]["params"] == {}
|
||||
|
||||
def test_grounding_search_config_rejects_both_count_fields(self, mock_config):
|
||||
with pytest.raises(ValidationError):
|
||||
mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
optional_params={
|
||||
"grounding": {
|
||||
"type": "document_grounding_service",
|
||||
"config": {
|
||||
"filters": [{"data_repository_type": "vector",
|
||||
"search_config": {"max_chunk_count": 2,
|
||||
"max_document_count": 5}}],
|
||||
"placeholders": {"input": ["q"], "output": "r"},
|
||||
}
|
||||
}
|
||||
},
|
||||
litellm_params={}, headers={}
|
||||
)
|
||||
|
||||
def test_sap_filtering(self, mock_config):
|
||||
filtering_config_azure = {
|
||||
'input':
|
||||
{
|
||||
'filters':
|
||||
[
|
||||
{'type': 'azure_content_safety',
|
||||
'config':
|
||||
{'hate': 0,
|
||||
'sexual': 0,
|
||||
'violence': 0,
|
||||
'self_harm': 0
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
'output':
|
||||
{
|
||||
'filters':
|
||||
[
|
||||
{'type': 'azure_content_safety',
|
||||
'config': {'hate': 0,
|
||||
'sexual': 0,
|
||||
'violence': 0,
|
||||
'self_harm': 0
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
filtering_config_llama = {
|
||||
'input':
|
||||
{
|
||||
'filters':
|
||||
[
|
||||
{
|
||||
'type': 'llama_guard_3_8b',
|
||||
'config': {'hate': True,
|
||||
"elections": True}
|
||||
}
|
||||
]
|
||||
},
|
||||
'output':
|
||||
{
|
||||
'filters':
|
||||
[
|
||||
{
|
||||
'type': 'llama_guard_3_8b',
|
||||
'config': {'hate': True, "elections": True}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
config = mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello."}],
|
||||
optional_params={'deployment_url': "shouldn't be in results",
|
||||
"filtering": filtering_config_azure},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
assert config["config"]["modules"]["filtering"] == filtering_config_azure
|
||||
assert config["config"]["modules"]["prompt_templating"]["model"]["params"] == {}
|
||||
|
||||
config = mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello."}],
|
||||
optional_params={'deployment_url': "shouldn't be in results",
|
||||
"filtering": filtering_config_llama},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
assert config["config"]["modules"]["filtering"] == filtering_config_llama
|
||||
assert config["config"]["modules"]["prompt_templating"]["model"]["params"] == {}
|
||||
|
||||
def test_filtering_config_requires_at_least_one_property(self, mock_config):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
optional_params={
|
||||
"filtering": {}
|
||||
},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
|
||||
assert "For using SAP Filtering Module you must provide at least one property" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_sap_masking(self, mock_config):
|
||||
masking_config = {
|
||||
'providers':
|
||||
[
|
||||
{
|
||||
'type': 'sap_data_privacy_integration',
|
||||
'method': 'anonymization',
|
||||
'entities': [
|
||||
{'type': 'profile-address'},
|
||||
{'type': 'profile-email'},
|
||||
{'type': 'profile-phone'},
|
||||
{'type': 'profile-person'},
|
||||
{'type': 'profile-location'}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
config = mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello."}],
|
||||
optional_params={'deployment_url': "shouldn't be in results",
|
||||
"masking": masking_config},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
assert config["config"]["modules"]["masking"] == masking_config
|
||||
assert config["config"]["modules"]["prompt_templating"]["model"]["params"] == {}
|
||||
|
||||
def test_masking_config_requires_exactly_one_provider_list(self, mock_config):
|
||||
masking_config = {
|
||||
'providers':
|
||||
[
|
||||
{
|
||||
'type': 'sap_data_privacy_integration',
|
||||
'method': 'anonymization',
|
||||
'entities': [
|
||||
{'type': 'profile-address'},
|
||||
{'type': 'profile-email'},
|
||||
{'type': 'profile-phone'},
|
||||
{'type': 'profile-person'},
|
||||
{'type': 'profile-location'}
|
||||
]
|
||||
}
|
||||
],
|
||||
'masking_providers':
|
||||
[
|
||||
{
|
||||
'type': 'sap_data_privacy_integration',
|
||||
'method': 'anonymization',
|
||||
'entities': [
|
||||
{'type': 'profile-address'}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
optional_params={
|
||||
"masking": masking_config
|
||||
},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
|
||||
assert "must set exactly one of: 'providers' or 'masking_providers'" in str(exc_info.value)
|
||||
|
||||
def test_masking_providers_deprecated_emits_warning(self, mock_config):
|
||||
masking_config = {
|
||||
'masking_providers':
|
||||
[
|
||||
{
|
||||
'type': 'sap_data_privacy_integration',
|
||||
'method': 'anonymization',
|
||||
'entities': [
|
||||
{'type': 'profile-address'}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
warnings.simplefilter("always")
|
||||
mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
optional_params={"masking": masking_config},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
assert any(
|
||||
issubclass(warning.category, DeprecationWarning)
|
||||
and "masking_providers" in str(warning.message)
|
||||
for warning in w
|
||||
), "Expected DeprecationWarning for 'masking_providers'"
|
||||
|
||||
def test_sap_translation(self, mock_config):
|
||||
translation_config = {
|
||||
'input':
|
||||
{'type': 'sap_document_translation',
|
||||
'config':
|
||||
{'source_language': 'en-US',
|
||||
'target_language': 'de-DE'}
|
||||
},
|
||||
'output':
|
||||
{'type': 'sap_document_translation',
|
||||
'config':
|
||||
{'source_language': 'de-DE',
|
||||
'target_language': 'fr-FR'}
|
||||
}
|
||||
}
|
||||
|
||||
config = mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello."}],
|
||||
optional_params={'deployment_url': "shouldn't be in results",
|
||||
"translation": translation_config},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
assert config["config"]["modules"]["translation"] == translation_config
|
||||
assert config["config"]["modules"]["prompt_templating"]["model"]["params"] == {}
|
||||
|
||||
def test_translation_config_requires_at_least_one_property(self, mock_config):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
optional_params={
|
||||
"translation": {}
|
||||
},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
|
||||
assert "TranslationModuleConfig requires at least one of 'input' or 'output'" in str(exc_info.value)
|
||||
|
||||
def test_sap_multiple_modules(self, mock_config):
|
||||
translation_config = {
|
||||
'input':
|
||||
{'type': 'sap_document_translation',
|
||||
'config':
|
||||
{'source_language': 'en-US',
|
||||
'target_language': 'de-DE'}
|
||||
},
|
||||
'output':
|
||||
{'type': 'sap_document_translation',
|
||||
'config':
|
||||
{'source_language': 'de-DE',
|
||||
'target_language': 'fr-FR'}
|
||||
}
|
||||
}
|
||||
for model in ["sap/gpt-5", "gpt-5"]:
|
||||
config = mock_config.transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hello."}],
|
||||
optional_params={'deployment_url': "shouldn't be in results",
|
||||
"fallback_sap_modules": [{"model": model,
|
||||
"messages": [{"role": "user", "content": "Hello world!"}],
|
||||
"translation": translation_config
|
||||
}]
|
||||
,
|
||||
},
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
assert "translation" not in config["config"]["modules"][0]
|
||||
translation = config["config"]["modules"][1]["translation"]
|
||||
assert translation["input"]["config"]["source_language"] == "en-US"
|
||||
assert translation["input"]["config"]["target_language"] == "de-DE"
|
||||
assert translation["output"]["config"]["target_language"] == "fr-FR"
|
||||
assert config["config"]["modules"][1]["prompt_templating"]["model"]["name"] == "gpt-5"
|
||||
assert config["config"]["modules"][0]["prompt_templating"]["model"]["name"] == "gpt-4o"
|
||||
assert config["config"]["modules"][0]["prompt_templating"]["model"]["params"] == {}
|
||||
assert config["config"]["modules"][1]["prompt_templating"]["prompt"]["template"][0]["content"] == "Hello world!"
|
||||
assert config["config"]["modules"][0]["prompt_templating"]["prompt"]["template"][0]["content"] == "Hello."
|
||||
assert config["config"]["modules"][1]["translation"]["input"]["type"] == "sap_document_translation"
|
||||
|
|
@ -0,0 +1,97 @@
|
|||
from unittest.mock import patch, PropertyMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.llms.sap.embed.transformation import GenAIHubEmbeddingConfig
|
||||
|
||||
@pytest.fixture
|
||||
def fake_token_creator():
|
||||
return (lambda: "Bearer FAKE_TOKEN", "https://api.ai.moke-sap.com", "fake-group")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_deployment_url():
|
||||
return "https://api.ai.moke-sap.com/v2/inference/deployments/mokeid"
|
||||
|
||||
def test_basic_config_transform(fake_token_creator, fake_deployment_url):
|
||||
expected_dict = {
|
||||
'config': {
|
||||
'modules': {
|
||||
'embeddings': {
|
||||
'model': {
|
||||
'name': 'text-embedding-3-small',
|
||||
'version': 'latest',
|
||||
'params': {}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
'input': {
|
||||
'text': 'Hi'
|
||||
}
|
||||
}
|
||||
with patch(
|
||||
"litellm.llms.sap.embed.transformation.GenAIHubEmbeddingConfig.deployment_url",
|
||||
new_callable=PropertyMock,
|
||||
return_value=fake_deployment_url,
|
||||
), patch(
|
||||
"litellm.llms.sap.embed.transformation.get_token_creator",
|
||||
return_value=fake_token_creator,
|
||||
):
|
||||
body = GenAIHubEmbeddingConfig().transform_embedding_request(
|
||||
model="text-embedding-3-small",
|
||||
input="Hi",
|
||||
optional_params={},
|
||||
headers={}
|
||||
)
|
||||
assert body == expected_dict
|
||||
|
||||
def test_model_params(fake_token_creator, fake_deployment_url):
|
||||
with patch(
|
||||
"litellm.llms.sap.embed.transformation.GenAIHubEmbeddingConfig.deployment_url",
|
||||
new_callable=PropertyMock,
|
||||
return_value=fake_deployment_url,
|
||||
), patch(
|
||||
"litellm.llms.sap.embed.transformation.get_token_creator",
|
||||
return_value=fake_token_creator,
|
||||
):
|
||||
body = GenAIHubEmbeddingConfig().transform_embedding_request(
|
||||
model="text-embedding-3-small",
|
||||
input="Hi",
|
||||
optional_params={"parameters": {"truncate": "END"}},
|
||||
headers={}
|
||||
)
|
||||
assert body["config"]["modules"]["embeddings"]["model"]["params"] == {"truncate": "END"}
|
||||
|
||||
def test_embed_with_masking(fake_token_creator, fake_deployment_url):
|
||||
masking_config = {
|
||||
'providers':
|
||||
[
|
||||
{
|
||||
'type': 'sap_data_privacy_integration',
|
||||
'method': 'anonymization',
|
||||
'entities': [
|
||||
{'type': 'profile-address'},
|
||||
{'type': 'profile-phone'},
|
||||
{'type': 'profile-person'},
|
||||
{'type': 'profile-location'}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
with patch(
|
||||
"litellm.llms.sap.embed.transformation.GenAIHubEmbeddingConfig.deployment_url",
|
||||
new_callable=PropertyMock,
|
||||
return_value=fake_deployment_url,
|
||||
), patch(
|
||||
"litellm.llms.sap.embed.transformation.get_token_creator",
|
||||
return_value=fake_token_creator,
|
||||
):
|
||||
body = GenAIHubEmbeddingConfig().transform_embedding_request(
|
||||
model="text-embedding-3-small",
|
||||
input="Hi",
|
||||
optional_params={"parameters": {"truncate": "END"},
|
||||
"masking": masking_config},
|
||||
headers={}
|
||||
)
|
||||
assert body["config"]["modules"]["masking"] == masking_config
|
||||
142
tests/test_litellm/llms/sap/test_sap_fetch_creds.py
Normal file
142
tests/test_litellm/llms/sap/test_sap_fetch_creds.py
Normal file
|
|
@ -0,0 +1,142 @@
|
|||
import json
|
||||
import pytest
|
||||
import litellm.llms.sap.credentials as sap_credentials
|
||||
|
||||
mock_sap_service_key_dict = {
|
||||
"serviceurls":
|
||||
{
|
||||
"AI_API_URL":"https://testurl.hana.ondemand.com/"
|
||||
},
|
||||
"clientid":"mockclientid",
|
||||
"clientsecret":"mockclientsecret",
|
||||
"url":"https://test.sap.hana.ondemand.com/"
|
||||
}
|
||||
|
||||
mock_wrapped_sap_service_key_dict = {
|
||||
"credentials": {
|
||||
"serviceurls":
|
||||
{
|
||||
"AI_API_URL":"https://testurl.hana.ondemand.com/"
|
||||
},
|
||||
"clientid":"mockclientid",
|
||||
"clientsecret":"mockclientsecret",
|
||||
"url":"https://test.sap.hana.ondemand.com/"
|
||||
}
|
||||
}
|
||||
|
||||
expected_creds = {'client_id': "mockclientid",
|
||||
'client_secret': "mockclientsecret",
|
||||
'auth_url': 'https://test.sap.hana.ondemand.com/oauth/token',
|
||||
'base_url': 'https://testurl.hana.ondemand.com/v2',
|
||||
'resource_group': 'default'}
|
||||
|
||||
mock_sap_vcap_service_key_dict = {
|
||||
'aicore': [{
|
||||
'label': 'aicore',
|
||||
'name': 'aicore-instance',
|
||||
'instance_guid': '53ad5b47-a49a-4fec-9f0b-cd921c00b828',
|
||||
'credentials': {
|
||||
'serviceurls': {
|
||||
'AI_API_URL': 'vcap-api-url'
|
||||
},
|
||||
'url': 'vcap-auth-url',
|
||||
'clientid': 'vcap-clientid',
|
||||
'clientsecret': 'vcap-clientsecret'
|
||||
}
|
||||
}]
|
||||
}
|
||||
def _prep_env(monkeypatch):
|
||||
for var in ("AICORE_CLIENT_ID", "AICORE_CLIENT_SECRET", "AICORE_AUTH_URL", "AICORE_RESOURCE_GROUP",
|
||||
"AICORE_BASE_URL", "AICORE_CERT_URL", "AICORE_SERVICE_KEY", "VCAP_SERVICES"):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
monkeypatch.setenv("AICORE_HOME", 'notexist')
|
||||
monkeypatch.setattr('litellm.sap_service_key', None)
|
||||
|
||||
def test_sap_fetch_creds_from_env_service_key(monkeypatch):
|
||||
_prep_env(monkeypatch)
|
||||
monkeypatch.setenv("AICORE_SERVICE_KEY", json.dumps(mock_sap_service_key_dict))
|
||||
creds = sap_credentials.fetch_credentials()
|
||||
assert creds == expected_creds
|
||||
|
||||
def test_sap_fetch_creds_from_env_wrapped_service_key(monkeypatch):
|
||||
_prep_env(monkeypatch)
|
||||
monkeypatch.setenv("AICORE_SERVICE_KEY", json.dumps(mock_wrapped_sap_service_key_dict))
|
||||
creds = sap_credentials.fetch_credentials()
|
||||
assert creds == expected_creds
|
||||
|
||||
def test_sap_fetch_creds_from_arg_service_key(monkeypatch):
|
||||
_prep_env(monkeypatch)
|
||||
creds = sap_credentials.fetch_credentials(service_key=json.dumps(mock_sap_service_key_dict))
|
||||
assert creds == expected_creds
|
||||
|
||||
def test_fetch_creds_from_env_vcap_service(monkeypatch):
|
||||
_prep_env(monkeypatch)
|
||||
monkeypatch.setenv("VCAP_SERVICES", json.dumps(mock_sap_vcap_service_key_dict))
|
||||
creds = sap_credentials.fetch_credentials()
|
||||
assert creds['client_id'] == "vcap-clientid"
|
||||
assert creds['client_secret'] == "vcap-clientsecret"
|
||||
assert creds['auth_url'] == "vcap-auth-url/oauth/token"
|
||||
assert creds['base_url'] == "vcap-api-url/v2"
|
||||
assert creds['resource_group'] == "default"
|
||||
|
||||
def test_fetch_creds_from_env(monkeypatch):
|
||||
_prep_env(monkeypatch)
|
||||
monkeypatch.setenv("AICORE_CLIENT_ID", "env-client-id")
|
||||
monkeypatch.setenv("AICORE_CLIENT_SECRET", "env-client-secret")
|
||||
monkeypatch.setenv("AICORE_AUTH_URL", "env-auth-url")
|
||||
monkeypatch.setenv("AICORE_BASE_URL", "env-base-url")
|
||||
monkeypatch.setenv("AICORE_RESOURCE_GROUP", "env-resource-group")
|
||||
|
||||
creds = sap_credentials.fetch_credentials()
|
||||
|
||||
assert creds['client_id'] == "env-client-id"
|
||||
assert creds['client_secret'] == "env-client-secret"
|
||||
assert creds['auth_url'] == "env-auth-url/oauth/token"
|
||||
assert creds['base_url'] == "env-base-url/v2"
|
||||
assert creds['resource_group'] == "env-resource-group"
|
||||
|
||||
def test_creds_priority_order(monkeypatch):
|
||||
_prep_env(monkeypatch)
|
||||
monkeypatch.setenv("AICORE_CLIENT_ID", "env-client-id")
|
||||
monkeypatch.setenv("AICORE_CLIENT_SECRET", "env-client-secret")
|
||||
monkeypatch.setenv("AICORE_AUTH_URL", "env-auth-url")
|
||||
monkeypatch.setenv("AICORE_BASE_URL", "env-base-url")
|
||||
monkeypatch.setenv("AICORE_RESOURCE_GROUP", "env-resource-group")
|
||||
creds = sap_credentials.fetch_credentials(service_key=json.dumps(mock_sap_service_key_dict))
|
||||
assert creds['client_id'] == "mockclientid"
|
||||
assert creds['resource_group'] == "env-resource-group"
|
||||
|
||||
def test_no_credentials_configured(monkeypatch):
|
||||
_prep_env(monkeypatch)
|
||||
with pytest.raises(ValueError, match="No credentials found in any source"):
|
||||
sap_credentials.fetch_credentials()
|
||||
|
||||
|
||||
def test_partial_credentials_missing_auth_url(monkeypatch):
|
||||
_prep_env(monkeypatch)
|
||||
|
||||
# Set only client_id and base_url, missing auth_url
|
||||
monkeypatch.setenv("AICORE_CLIENT_ID", "test-client-id")
|
||||
monkeypatch.setenv("AICORE_BASE_URL", "test-base-url")
|
||||
|
||||
# fetch_credentials should succeed (it returns whatever it finds)
|
||||
creds = sap_credentials.fetch_credentials()
|
||||
creds.pop('resource_group')
|
||||
|
||||
with pytest.raises(ValueError, match="SAP AI Core credentials not found"):
|
||||
sap_credentials.validate_credentials(**creds)
|
||||
|
||||
def test_credentials_without_authentication_mode(monkeypatch):
|
||||
_prep_env(monkeypatch)
|
||||
|
||||
# Set all required fields but no authentication mode (no client_secret, no certs)
|
||||
monkeypatch.setenv("AICORE_CLIENT_ID", "test-client-id")
|
||||
monkeypatch.setenv("AICORE_AUTH_URL", "test-auth-url")
|
||||
monkeypatch.setenv("AICORE_BASE_URL", "test-base-url")
|
||||
|
||||
creds = sap_credentials.fetch_credentials()
|
||||
creds.pop('resource_group')
|
||||
|
||||
# validate_credentials should raise because no authentication mode is provided
|
||||
with pytest.raises(ValueError, match="SAP AI Core credentials are incomplete"):
|
||||
sap_credentials.validate_credentials(**creds)
|
||||
|
|
@ -127,6 +127,24 @@ def test_vertex_ai_includes_labels():
|
|||
assert result["labels"] == {"project": "test", "team": "ai"}
|
||||
|
||||
|
||||
def test_service_tier_forwarded_to_vertex_ai():
|
||||
"""Test that service_tier in optional_params is mapped to serviceTier in request body."""
|
||||
messages = [{"role": "user", "content": "test"}]
|
||||
optional_params = {"service_tier": "flex"}
|
||||
litellm_params = {}
|
||||
|
||||
result = _transform_request_body(
|
||||
messages=messages,
|
||||
model="gemini-2.5-pro",
|
||||
optional_params=optional_params,
|
||||
custom_llm_provider="vertex_ai",
|
||||
litellm_params=litellm_params,
|
||||
cached_content=None,
|
||||
)
|
||||
|
||||
assert "serviceTier" in result
|
||||
assert result["serviceTier"] == "flex"
|
||||
|
||||
|
||||
def test_extra_body_cache_not_forwarded_to_vertex_ai():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3504,6 +3504,73 @@ def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming():
|
|||
assert result._hidden_params["provider_specific_fields"]["traffic_type"] == "PROVISIONED_THROUGHPUT"
|
||||
|
||||
|
||||
def test_vertex_ai_service_tier_streaming():
|
||||
"""Test service_tier is preserved in model_response from headers for streaming."""
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator,
|
||||
)
|
||||
|
||||
chunk = {
|
||||
"candidates": [{"content": {"parts": [{"text": "Hello"}]}}],
|
||||
}
|
||||
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=[],
|
||||
sync_stream=True,
|
||||
logging_obj=MagicMock(),
|
||||
response_headers={"x-gemini-service-tier": "FLEX"},
|
||||
)
|
||||
# Undefined when usageMetadata is missing
|
||||
result = iterator.chunk_parser(chunk)
|
||||
|
||||
# But definitely set when usageMetadata is present
|
||||
chunk_with_usage = {
|
||||
"candidates": [{"content": {"parts": [{"text": "hi"}]}}],
|
||||
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}
|
||||
}
|
||||
result_with_usage = iterator.chunk_parser(chunk_with_usage)
|
||||
assert result_with_usage.service_tier == "flex"
|
||||
|
||||
|
||||
def test_vertex_ai_service_tier_non_streaming():
|
||||
"""Test service_tier is preserved in model_response from headers for non-streaming."""
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
|
||||
completion_response = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": [{"text": "Hello"}], "role": "model"},
|
||||
"finishReason": "STOP",
|
||||
}
|
||||
],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 50,
|
||||
"candidatesTokenCount": 100,
|
||||
"totalTokenCount": 150,
|
||||
},
|
||||
}
|
||||
|
||||
raw_response = MagicMock()
|
||||
raw_response.json.return_value = completion_response
|
||||
raw_response.headers = {"x-gemini-service-tier": "FLEX"}
|
||||
|
||||
result = VertexGeminiConfig().transform_response(
|
||||
model="gemini-pro",
|
||||
raw_response=raw_response,
|
||||
model_response=ModelResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert result.service_tier == "flex"
|
||||
|
||||
|
||||
def test_vertex_ai_traffic_type_surfaced_in_responses_api():
|
||||
"""Test trafficType is surfaced as provider_specific_fields in ResponsesAPIResponse."""
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
|
|
@ -3609,6 +3676,54 @@ def test_vertex_ai_web_search_options_in_map_openai_params():
|
|||
assert "web_search_options" not in optional_params, "web_search_options should be removed after transformation"
|
||||
|
||||
|
||||
def test_vertex_ai_service_tier_in_map_openai_params():
|
||||
"""Test that service_tier is correctly mapped to optional_params."""
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
|
||||
v = VertexGeminiConfig()
|
||||
|
||||
# Test pass-through
|
||||
optional_params = {}
|
||||
non_default_params = {"service_tier": "FLEX"}
|
||||
|
||||
result = v.map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params=optional_params,
|
||||
model="gemini-3-pro-preview",
|
||||
drop_params=True,
|
||||
)
|
||||
|
||||
assert result["service_tier"] == "flex"
|
||||
|
||||
# Test auto -> priority
|
||||
optional_params_auto = {}
|
||||
non_default_params_auto = {"service_tier": "auto"}
|
||||
|
||||
result_auto = v.map_openai_params(
|
||||
non_default_params=non_default_params_auto,
|
||||
optional_params=optional_params_auto,
|
||||
model="gemini-3-pro-preview",
|
||||
drop_params=True,
|
||||
)
|
||||
|
||||
assert result_auto["service_tier"] == "priority"
|
||||
|
||||
# Test AUTO (uppercase) -> priority
|
||||
optional_params_auto_upper = {}
|
||||
non_default_params_auto_upper = {"service_tier": "AUTO"}
|
||||
|
||||
result_auto_upper = v.map_openai_params(
|
||||
non_default_params=non_default_params_auto_upper,
|
||||
optional_params=optional_params_auto_upper,
|
||||
model="gemini-3-pro-preview",
|
||||
drop_params=True,
|
||||
)
|
||||
|
||||
assert result_auto_upper["service_tier"] == "priority"
|
||||
|
||||
|
||||
def test_vertex_ai_usage_metadata_with_video_tokens_in_prompt():
|
||||
"""Test promptTokensDetails with VIDEO modality for video inputs.
|
||||
|
||||
|
|
|
|||
|
|
@ -1050,6 +1050,85 @@ class TestVertexBase:
|
|||
mock_creds.with_scopes.assert_called_once_with(scopes)
|
||||
assert result == "scoped_creds"
|
||||
|
||||
def test_credentials_from_pluggable_implementation(self):
|
||||
"""Test _credentials_from_pluggable dispatches to pluggable.Credentials"""
|
||||
vertex_base = VertexBase()
|
||||
json_obj = {
|
||||
"type": "external_account",
|
||||
"credential_source": {
|
||||
"executable": {"command": "/path/to/executable", "timeout_millis": 5000}
|
||||
},
|
||||
}
|
||||
scopes = ["https://www.googleapis.com/auth/cloud-platform"]
|
||||
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.requires_scopes = True
|
||||
mock_creds.with_scopes.return_value = "scoped_creds"
|
||||
|
||||
with patch("google.auth.pluggable.Credentials") as MockCredentials:
|
||||
MockCredentials.from_info.return_value = mock_creds
|
||||
|
||||
result = vertex_base._credentials_from_pluggable(json_obj, scopes)
|
||||
|
||||
MockCredentials.from_info.assert_called_once_with(json_obj)
|
||||
mock_creds.with_scopes.assert_called_once_with(scopes)
|
||||
assert result == "scoped_creds"
|
||||
|
||||
def test_credentials_from_pluggable_no_scopes_needed(self):
|
||||
"""Test _credentials_from_pluggable when scopes are not needed"""
|
||||
vertex_base = VertexBase()
|
||||
json_obj = {
|
||||
"type": "external_account",
|
||||
"credential_source": {
|
||||
"executable": {"command": "/path/to/executable"}
|
||||
},
|
||||
}
|
||||
scopes = ["https://www.googleapis.com/auth/cloud-platform"]
|
||||
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.requires_scopes = False
|
||||
|
||||
with patch("google.auth.pluggable.Credentials") as MockCredentials:
|
||||
MockCredentials.from_info.return_value = mock_creds
|
||||
|
||||
result = vertex_base._credentials_from_pluggable(json_obj, scopes)
|
||||
|
||||
MockCredentials.from_info.assert_called_once_with(json_obj)
|
||||
mock_creds.with_scopes.assert_not_called()
|
||||
assert result == mock_creds
|
||||
|
||||
def test_load_auth_dispatches_to_pluggable_for_executable(self):
|
||||
"""Test that load_auth routes executable credential_source to _credentials_from_pluggable"""
|
||||
vertex_base = VertexBase()
|
||||
json_obj = {
|
||||
"type": "external_account",
|
||||
"credential_source": {
|
||||
"executable": {"command": "/path/to/executable", "timeout_millis": 5000}
|
||||
},
|
||||
}
|
||||
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.project_id = "test-project"
|
||||
|
||||
with patch.object(
|
||||
vertex_base, "_credentials_from_pluggable", return_value=mock_creds
|
||||
) as mock_pluggable, patch.object(
|
||||
vertex_base, "_credentials_from_identity_pool"
|
||||
) as mock_identity_pool, patch.object(
|
||||
vertex_base, "refresh_auth"
|
||||
):
|
||||
creds, project_id = vertex_base.load_auth(
|
||||
credentials=json.dumps(json_obj), project_id=None
|
||||
)
|
||||
|
||||
mock_pluggable.assert_called_once_with(
|
||||
json_obj,
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
)
|
||||
mock_identity_pool.assert_not_called()
|
||||
assert creds == mock_creds
|
||||
assert project_id == "test-project"
|
||||
|
||||
def test_extract_aws_params(self):
|
||||
"""Test _extract_aws_params: extraction, empty case, and unrecognized keys."""
|
||||
# Case 1: Extracts recognized aws_* keys, ignores GCP-standard fields
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from litellm.proxy._experimental.mcp_server import rest_endpoints
|
|||
from litellm.proxy._experimental.mcp_server.auth import (
|
||||
user_api_key_auth_mcp as auth_mcp,
|
||||
)
|
||||
from litellm.proxy._types import NewMCPServerRequest, UserAPIKeyAuth
|
||||
from litellm.proxy._types import NewMCPServerRequest, UpdateMCPServerRequest, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
|
|
@ -156,7 +156,6 @@ class TestExecuteWithMcpClient:
|
|||
"Authorization": "STATIC token",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_m2m_credentials_forwarded_to_server_model(self, monkeypatch):
|
||||
"""M2M OAuth credentials (client_id, client_secret) from the nested
|
||||
|
|
@ -199,9 +198,7 @@ class TestExecuteWithMcpClient:
|
|||
},
|
||||
)
|
||||
|
||||
result = await rest_endpoints._execute_with_mcp_client(
|
||||
payload, ok_operation
|
||||
)
|
||||
result = await rest_endpoints._execute_with_mcp_client(payload, ok_operation)
|
||||
|
||||
assert result["status"] == "ok"
|
||||
server = captured["server"]
|
||||
|
|
@ -262,7 +259,10 @@ class TestExecuteWithMcpClient:
|
|||
assert result["status"] == "ok"
|
||||
# The incoming Authorization must be dropped — extra_headers should
|
||||
# contain no oauth2 headers (only static_headers, which are None here).
|
||||
assert captured["extra_headers"] is None or "Authorization" not in captured["extra_headers"]
|
||||
assert (
|
||||
captured["extra_headers"] is None
|
||||
or "Authorization" not in captured["extra_headers"]
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_catches_exception_group(self, monkeypatch):
|
||||
|
|
@ -300,9 +300,7 @@ class TestExecuteWithMcpClient:
|
|||
auth_type=MCPAuth.none,
|
||||
)
|
||||
|
||||
result = await rest_endpoints._execute_with_mcp_client(
|
||||
payload, ok_operation
|
||||
)
|
||||
result = await rest_endpoints._execute_with_mcp_client(payload, ok_operation)
|
||||
|
||||
assert result["status"] == "error"
|
||||
assert result["error"] is True
|
||||
|
|
@ -365,8 +363,12 @@ class TestTestToolsList:
|
|||
credentials={"auth_value": "secret-key"},
|
||||
)
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
result = await rest_endpoints.test_tools_list(
|
||||
request, payload, user_api_key_dict=UserAPIKeyAuth()
|
||||
request,
|
||||
payload,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
assert result["message"] == "Successfully retrieved tools"
|
||||
|
|
@ -419,8 +421,12 @@ class TestTestToolsList:
|
|||
auth_type=MCPAuth.oauth2,
|
||||
)
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
result = await rest_endpoints.test_tools_list(
|
||||
request, payload, user_api_key_dict=UserAPIKeyAuth()
|
||||
request,
|
||||
payload,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
assert result["message"] == "Successfully retrieved tools"
|
||||
|
|
@ -484,7 +490,11 @@ class TestListToolsRestAPI:
|
|||
captured = {"called": False}
|
||||
|
||||
async def fake_get_tools(
|
||||
server, server_auth_header, raw_headers=None, user_api_key_auth=None, extra_headers=None
|
||||
server,
|
||||
server_auth_header,
|
||||
raw_headers=None,
|
||||
user_api_key_auth=None,
|
||||
extra_headers=None,
|
||||
):
|
||||
captured["called"] = True
|
||||
captured["server"] = server
|
||||
|
|
@ -555,27 +565,47 @@ class TestListToolsRestAPI:
|
|||
|
||||
captured = {"called": False, "server_arg": None}
|
||||
|
||||
async def fake_get_tools(server, server_auth_header, raw_headers=None, user_api_key_auth=None, extra_headers=None):
|
||||
async def fake_get_tools(
|
||||
server,
|
||||
server_auth_header,
|
||||
raw_headers=None,
|
||||
user_api_key_auth=None,
|
||||
extra_headers=None,
|
||||
):
|
||||
captured["called"] = True
|
||||
captured["server_arg"] = server
|
||||
return ["tool-x"]
|
||||
|
||||
monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager, "get_allowed_mcp_servers",
|
||||
fake_get_allowed_mcp_servers, raising=False,
|
||||
rest_endpoints,
|
||||
"build_effective_auth_contexts",
|
||||
fake_contexts,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_name",
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
fake_get_allowed_mcp_servers,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_mcp_server_by_name",
|
||||
lambda name: stub_server if name == "my-server" else None,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_id",
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_mcp_server_by_id",
|
||||
lambda sid: stub_server if sid == "uuid-abc-123" else None,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(rest_endpoints, "_get_tools_for_single_server", fake_get_tools, raising=False)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints,
|
||||
"_get_tools_for_single_server",
|
||||
fake_get_tools,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
request = _build_request(path="/mcp-rest/tools/list", method="GET")
|
||||
result = await rest_endpoints.list_tool_rest_api(
|
||||
|
|
@ -609,18 +639,27 @@ class TestListToolsRestAPI:
|
|||
async def fake_get_allowed_mcp_servers(*args, **kwargs):
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager, "get_allowed_mcp_servers",
|
||||
fake_get_allowed_mcp_servers, raising=False,
|
||||
rest_endpoints,
|
||||
"build_effective_auth_contexts",
|
||||
fake_contexts,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_name",
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
fake_get_allowed_mcp_servers,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_mcp_server_by_name",
|
||||
lambda name: stub_server if name == "restricted-server" else None,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_id",
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_mcp_server_by_id",
|
||||
lambda sid: stub_server if sid == "uuid-xyz-999" else None,
|
||||
raising=False,
|
||||
)
|
||||
|
|
@ -662,31 +701,54 @@ class TestListToolsRestAPI:
|
|||
|
||||
oauth_headers = {"Authorization": "Bearer user-oauth-token"}
|
||||
|
||||
async def fake_get_user_oauth_extra_headers(server, user_api_key_dict, prefetched_creds=None):
|
||||
async def fake_get_user_oauth_extra_headers(
|
||||
server, user_api_key_dict, prefetched_creds=None
|
||||
):
|
||||
return oauth_headers
|
||||
|
||||
captured = {}
|
||||
|
||||
async def fake_get_tools(server, server_auth_header, raw_headers=None, user_api_key_auth=None, extra_headers=None):
|
||||
async def fake_get_tools(
|
||||
server,
|
||||
server_auth_header,
|
||||
raw_headers=None,
|
||||
user_api_key_auth=None,
|
||||
extra_headers=None,
|
||||
):
|
||||
captured["server"] = server
|
||||
captured["auth_header"] = server_auth_header
|
||||
return ["oauth-tool"]
|
||||
|
||||
monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager, "get_allowed_mcp_servers",
|
||||
fake_get_allowed_mcp_servers, raising=False,
|
||||
rest_endpoints,
|
||||
"build_effective_auth_contexts",
|
||||
fake_contexts,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_id",
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
fake_get_allowed_mcp_servers,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_mcp_server_by_id",
|
||||
lambda sid: stub_server if sid == "oauth-server-id" else None,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints, "_get_user_oauth_extra_headers",
|
||||
fake_get_user_oauth_extra_headers, raising=False,
|
||||
rest_endpoints,
|
||||
"_get_user_oauth_extra_headers",
|
||||
fake_get_user_oauth_extra_headers,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints,
|
||||
"_get_tools_for_single_server",
|
||||
fake_get_tools,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(rest_endpoints, "_get_tools_for_single_server", fake_get_tools, raising=False)
|
||||
|
||||
request = _build_request(path="/mcp-rest/tools/list", method="GET")
|
||||
result = await rest_endpoints.list_tool_rest_api(
|
||||
|
|
@ -1124,3 +1186,189 @@ class TestGetToolsForSingleServer:
|
|||
assert "tool3" in tool_names
|
||||
assert "tool1" not in tool_names
|
||||
assert "tool4" not in tool_names
|
||||
|
||||
|
||||
class TestStdioCommandAllowlist:
|
||||
"""Tests for MCP stdio command allowlist validation."""
|
||||
|
||||
def test_allowed_command_passes_validation(self):
|
||||
"""npx, uvx, python, etc. should be accepted."""
|
||||
req = NewMCPServerRequest(
|
||||
server_name="test",
|
||||
transport="stdio",
|
||||
command="npx",
|
||||
args=["-y", "@modelcontextprotocol/server-filesystem"],
|
||||
)
|
||||
assert req.command == "npx"
|
||||
|
||||
def test_disallowed_command_raises(self):
|
||||
"""Arbitrary commands like bash should be rejected."""
|
||||
with pytest.raises(ValueError, match="not in the allowed commands list"):
|
||||
NewMCPServerRequest(
|
||||
server_name="test",
|
||||
transport="stdio",
|
||||
command="bash",
|
||||
args=["-c", "echo pwned"],
|
||||
)
|
||||
|
||||
def test_sh_command_raises(self):
|
||||
"""sh should be rejected."""
|
||||
with pytest.raises(ValueError, match="not in the allowed commands list"):
|
||||
NewMCPServerRequest(
|
||||
server_name="test",
|
||||
transport="stdio",
|
||||
command="sh",
|
||||
args=["-c", "id > /tmp/output.txt"],
|
||||
)
|
||||
|
||||
def test_absolute_path_bypass_blocked(self):
|
||||
"""/bin/bash should be blocked (basename is 'bash')."""
|
||||
with pytest.raises(ValueError, match="not in the allowed commands list"):
|
||||
NewMCPServerRequest(
|
||||
server_name="test",
|
||||
transport="stdio",
|
||||
command="/bin/bash",
|
||||
args=["-c", "echo pwned"],
|
||||
)
|
||||
|
||||
def test_absolute_path_to_allowed_command_works(self):
|
||||
"""/usr/bin/python3 should pass (basename is 'python3')."""
|
||||
req = NewMCPServerRequest(
|
||||
server_name="test",
|
||||
transport="stdio",
|
||||
command="/usr/bin/python3",
|
||||
args=["-m", "some_module"],
|
||||
)
|
||||
assert req.command == "/usr/bin/python3"
|
||||
|
||||
def test_http_transport_ignores_allowlist(self):
|
||||
"""HTTP/SSE transport should not trigger command validation."""
|
||||
req = NewMCPServerRequest(
|
||||
server_name="test",
|
||||
transport="sse",
|
||||
url="https://example.com/mcp",
|
||||
)
|
||||
assert req.transport == "sse"
|
||||
|
||||
def test_uvx_command_passes(self):
|
||||
req = NewMCPServerRequest(
|
||||
server_name="test",
|
||||
transport="stdio",
|
||||
command="uvx",
|
||||
args=["mcp-server-sqlite"],
|
||||
)
|
||||
assert req.command == "uvx"
|
||||
|
||||
def test_node_command_passes(self):
|
||||
req = NewMCPServerRequest(
|
||||
server_name="test",
|
||||
transport="stdio",
|
||||
command="node",
|
||||
args=["server.js"],
|
||||
)
|
||||
assert req.command == "node"
|
||||
|
||||
def test_update_request_disallowed_command_raises(self):
|
||||
"""UpdateMCPServerRequest should also block non-allowlisted commands."""
|
||||
with pytest.raises(ValueError, match="not in the allowed commands list"):
|
||||
UpdateMCPServerRequest(
|
||||
server_id="some-id",
|
||||
transport="stdio",
|
||||
command="bash",
|
||||
args=["-c", "echo pwned"],
|
||||
)
|
||||
|
||||
|
||||
class TestEndpointRoleChecks:
|
||||
"""Tests for PROXY_ADMIN role checks on MCP test endpoints."""
|
||||
|
||||
def test_test_connection_has_auth_dependency(self):
|
||||
route = _get_route("/mcp-rest/test/connection", "POST")
|
||||
assert _route_has_dependency(route, user_api_key_auth)
|
||||
|
||||
def test_test_tools_list_has_auth_dependency(self):
|
||||
route = _get_route("/mcp-rest/test/tools/list", "POST")
|
||||
assert _route_has_dependency(route, user_api_key_auth)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_test_connection_rejects_non_admin(self):
|
||||
"""Non-admin users should get 403 from test_connection."""
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
payload = NewMCPServerRequest(
|
||||
server_name="test",
|
||||
url="https://example.com/mcp",
|
||||
auth_type=MCPAuth.none,
|
||||
)
|
||||
user_key = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="non_admin",
|
||||
api_key="sk-test",
|
||||
)
|
||||
request = _build_request()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await rest_endpoints.test_connection(
|
||||
request=request,
|
||||
new_mcp_server_request=payload,
|
||||
user_api_key_dict=user_key,
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_test_tools_list_rejects_non_admin(self):
|
||||
"""Non-admin users should get 403 from test_tools_list."""
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
payload = NewMCPServerRequest(
|
||||
server_name="test",
|
||||
url="https://example.com/mcp",
|
||||
auth_type=MCPAuth.none,
|
||||
)
|
||||
user_key = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="non_admin",
|
||||
api_key="sk-test",
|
||||
)
|
||||
request = _build_request()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await rest_endpoints.test_tools_list(
|
||||
request=request,
|
||||
new_mcp_server_request=payload,
|
||||
user_api_key_dict=user_key,
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_test_connection_allows_admin(self, monkeypatch):
|
||||
"""PROXY_ADMIN should pass the role check."""
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
async def fake_execute(*args, **kwargs):
|
||||
return {"status": "ok"}
|
||||
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints,
|
||||
"_execute_with_mcp_client",
|
||||
fake_execute,
|
||||
)
|
||||
|
||||
payload = NewMCPServerRequest(
|
||||
server_name="test",
|
||||
url="https://example.com/mcp",
|
||||
auth_type=MCPAuth.none,
|
||||
)
|
||||
user_key = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="admin",
|
||||
api_key="sk-admin",
|
||||
)
|
||||
request = _build_request()
|
||||
|
||||
result = await rest_endpoints.test_connection(
|
||||
request=request,
|
||||
new_mcp_server_request=payload,
|
||||
user_api_key_dict=user_key,
|
||||
)
|
||||
assert result["status"] == "ok"
|
||||
|
|
|
|||
|
|
@ -8579,3 +8579,170 @@ def test_enforce_upperbound_no_config_is_noop():
|
|||
assert data.tpm_limit == 999999
|
||||
finally:
|
||||
litellm.upperbound_key_generate_params = original
|
||||
|
||||
|
||||
class TestAllowedRoutesCallerPermission:
|
||||
"""
|
||||
Non-admins must not be able to set `allowed_routes` on a key. The field
|
||||
bypasses the role-based route gate in
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check, so allowing a non-admin
|
||||
to populate it grants them arbitrary endpoint access.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_admin_generate_key_with_allowed_routes_rejected(self):
|
||||
data = GenerateKeyRequest(
|
||||
key_alias="escalate",
|
||||
allowed_routes=["/*"],
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user-123",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", MagicMock()
|
||||
), patch("litellm.proxy.proxy_server.user_custom_key_generate", None), patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||
new_callable=AsyncMock,
|
||||
return_value=MagicMock(),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await generate_key_fn(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
assert str(exc_info.value.code) == "403"
|
||||
assert "allowed_routes" in str(exc_info.value.message)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_generate_key_with_allowed_routes_allowed(self):
|
||||
data = GenerateKeyRequest(
|
||||
key_alias="admin-key",
|
||||
allowed_routes=["/chat/completions"],
|
||||
user_id="admin-user",
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="admin-user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
mock_prisma_client = AsyncMock()
|
||||
stub_response = MagicMock()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", MagicMock()
|
||||
), patch("litellm.proxy.proxy_server.user_custom_key_generate", None), patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||
new_callable=AsyncMock,
|
||||
return_value=stub_response,
|
||||
):
|
||||
result = await generate_key_fn(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
assert result is stub_response
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_admin_generate_key_default_empty_allowed_routes_ok(self):
|
||||
"""
|
||||
Regression guard: GenerateKeyRequest.allowed_routes defaults to [], so
|
||||
the helper must treat empty-list as "not set" or every non-admin key
|
||||
creation breaks.
|
||||
"""
|
||||
data = GenerateKeyRequest(key_alias="plain-key")
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user-123",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
mock_prisma_client = AsyncMock()
|
||||
stub_response = MagicMock()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", MagicMock()
|
||||
), patch("litellm.proxy.proxy_server.user_custom_key_generate", None), patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||
new_callable=AsyncMock,
|
||||
return_value=stub_response,
|
||||
):
|
||||
result = await generate_key_fn(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
assert result is stub_response
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_admin_update_key_with_allowed_routes_rejected(self):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
update_key_fn,
|
||||
)
|
||||
|
||||
data = UpdateKeyRequest(key="sk-test", allowed_routes=["/*"])
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user-123",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", MagicMock()
|
||||
), patch("litellm.proxy.proxy_server.user_custom_key_update", None), patch(
|
||||
"litellm.proxy.proxy_server.llm_router", None
|
||||
), patch("litellm.proxy.proxy_server.premium_user", True), patch(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()
|
||||
), patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._get_and_validate_existing_key",
|
||||
new_callable=AsyncMock,
|
||||
return_value=MagicMock(),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await update_key_fn(
|
||||
request=MagicMock(),
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
assert str(exc_info.value.code) == "403"
|
||||
assert "allowed_routes" in str(exc_info.value.message)
|
||||
|
||||
|
||||
def test_jinja_prompt_manager_is_sandboxed():
|
||||
"""
|
||||
PromptManager renders user-supplied templates via /prompts/test, so its
|
||||
jinja env must reject access to unsafe Python attributes like
|
||||
``__class__`` and ``__mro__``.
|
||||
"""
|
||||
from jinja2.exceptions import SecurityError
|
||||
|
||||
from litellm.integrations.dotprompt.prompt_manager import PromptManager
|
||||
|
||||
pm = PromptManager()
|
||||
template = pm.jinja_env.from_string("{{ ''.__class__.__mro__ }}")
|
||||
with pytest.raises(SecurityError):
|
||||
template.render()
|
||||
|
||||
|
||||
def test_validate_public_image_url_rejects_local_paths():
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
_validate_public_image_url,
|
||||
)
|
||||
|
||||
for bad in ("/etc/passwd", "file:///etc/passwd", "../../etc/passwd"):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
_validate_public_image_url(bad, "logo_url")
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
|
||||
def test_validate_public_image_url_accepts_http_and_noop_empty():
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
_validate_public_image_url,
|
||||
)
|
||||
|
||||
_validate_public_image_url("https://example.com/logo.png", "logo_url")
|
||||
_validate_public_image_url("http://cdn.internal/logo.svg", "logo_url")
|
||||
_validate_public_image_url(None, "logo_url")
|
||||
_validate_public_image_url("", "logo_url")
|
||||
_validate_public_image_url(" ", "logo_url")
|
||||
|
|
|
|||
|
|
@ -60,25 +60,165 @@ async def test_spend_query_uses_timestamp_filtering():
|
|||
params = call_args[1:]
|
||||
|
||||
# 1) SQL should NOT cast the startTime column to DATE (prevents index usage)
|
||||
assert "::date" not in sql.lower(), \
|
||||
"SQL should not use '::date' casting which prevents index usage"
|
||||
assert "date(" not in sql.lower(), \
|
||||
"SQL should not use DATE() function which prevents index usage"
|
||||
assert (
|
||||
"::date" not in sql.lower()
|
||||
), "SQL should not use '::date' casting which prevents index usage"
|
||||
assert (
|
||||
"date(" not in sql.lower()
|
||||
), "SQL should not use DATE() function which prevents index usage"
|
||||
|
||||
# 2) SQL should use timestamp-range filtering pattern for index optimization
|
||||
assert '"startTime" >=' in sql or '"startTime">=' in sql, \
|
||||
"SQL should use >= operator for lower bound"
|
||||
assert '"startTime" <' in sql or '"startTime"<' in sql, \
|
||||
"SQL should use < operator for upper bound"
|
||||
assert "interval '1 day'" in sql.lower(), \
|
||||
"SQL should use INTERVAL for date arithmetic"
|
||||
assert (
|
||||
'"startTime" >=' in sql or '"startTime">=' in sql
|
||||
), "SQL should use >= operator for lower bound"
|
||||
assert (
|
||||
'"startTime" <' in sql or '"startTime"<' in sql
|
||||
), "SQL should use < operator for upper bound"
|
||||
assert (
|
||||
"interval '1 day'" in sql.lower()
|
||||
), "SQL should use INTERVAL for date arithmetic"
|
||||
|
||||
# 3) Parameters should be datetime objects (not date objects)
|
||||
assert isinstance(params[0], datetime.datetime), \
|
||||
"First parameter (start_date) should be datetime object"
|
||||
assert isinstance(params[1], datetime.datetime), \
|
||||
"Second parameter (end_date) should be datetime object"
|
||||
assert params[0].tzinfo is not None, \
|
||||
"start_date should be timezone-aware"
|
||||
assert params[1].tzinfo is not None, \
|
||||
"end_date should be timezone-aware"
|
||||
assert isinstance(
|
||||
params[0], datetime.datetime
|
||||
), "First parameter (start_date) should be datetime object"
|
||||
assert isinstance(
|
||||
params[1], datetime.datetime
|
||||
), "Second parameter (end_date) should be datetime object"
|
||||
assert params[0].tzinfo is not None, "start_date should be timezone-aware"
|
||||
assert params[1].tzinfo is not None, "end_date should be timezone-aware"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_activity_wraps_params_in_at_time_zone_utc(monkeypatch):
|
||||
"""
|
||||
/global/activity must emit `AT TIME ZONE 'UTC'` around its date params
|
||||
so the date window and `date_trunc` bucketing do not depend on the DB
|
||||
session timezone. Regression guard for Issue 1.
|
||||
"""
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
||||
get_global_activity,
|
||||
)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(return_value=[])
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
|
||||
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin")
|
||||
|
||||
await get_global_activity(
|
||||
start_date="2026-02-16",
|
||||
end_date="2026-02-16",
|
||||
user_api_key_dict=auth,
|
||||
)
|
||||
|
||||
assert mock_prisma.db.query_raw.called, "query_raw should have been called"
|
||||
call_args = mock_prisma.db.query_raw.call_args[0]
|
||||
sql = call_args[0]
|
||||
params = call_args[1:]
|
||||
|
||||
# 1) SQL must wrap both bounds in `AT TIME ZONE 'UTC'`.
|
||||
assert sql.count("AT TIME ZONE 'UTC'") >= 2, (
|
||||
"Both date bounds must be wrapped with `AT TIME ZONE 'UTC'` so that "
|
||||
"comparison against the plain-timestamp column is session-TZ-independent. "
|
||||
f"SQL was:\n{sql}"
|
||||
)
|
||||
|
||||
# 2) Params must still be tz-aware UTC datetimes (preserves existing contract).
|
||||
assert isinstance(params[0], datetime.datetime)
|
||||
assert isinstance(params[1], datetime.datetime)
|
||||
assert params[0].tzinfo is not None and params[0].utcoffset() == datetime.timedelta(
|
||||
0
|
||||
)
|
||||
assert params[1].tzinfo is not None and params[1].utcoffset() == datetime.timedelta(
|
||||
0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_activity_internal_user_wraps_params_in_at_time_zone_utc(
|
||||
monkeypatch,
|
||||
):
|
||||
"""
|
||||
The internal-user branch of /global/activity goes through a different
|
||||
helper (`get_global_activity_internal_user`) and has its own SQL string.
|
||||
Both branches must carry the fix. Regression guard for Issue 1.
|
||||
"""
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
||||
get_global_activity,
|
||||
)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(return_value=[])
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
|
||||
auth = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user_1"
|
||||
)
|
||||
|
||||
await get_global_activity(
|
||||
start_date="2026-02-16",
|
||||
end_date="2026-02-16",
|
||||
user_api_key_dict=auth,
|
||||
)
|
||||
|
||||
assert mock_prisma.db.query_raw.called
|
||||
sql = mock_prisma.db.query_raw.call_args[0][0]
|
||||
assert sql.count("AT TIME ZONE 'UTC'") >= 2, (
|
||||
"Internal-user branch must also wrap date bounds with "
|
||||
f"`AT TIME ZONE 'UTC'`. SQL was:\n{sql}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch):
|
||||
"""
|
||||
/spend/logs/ui builds its WHERE clause dynamically. The date-range
|
||||
conditions must wrap the param side with `AT TIME ZONE 'UTC'` so the
|
||||
log filter window doesn't drift with the DB session TZ. Regression
|
||||
guard for GH #22529.
|
||||
"""
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
||||
ui_view_spend_logs,
|
||||
)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_spendlogs = MagicMock()
|
||||
mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
|
||||
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin")
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = "/spend/logs/ui"
|
||||
|
||||
await ui_view_spend_logs(
|
||||
request=mock_request,
|
||||
api_key=None,
|
||||
user_id=None,
|
||||
request_id=None,
|
||||
start_date="2026-02-16 00:00:00",
|
||||
end_date="2026-02-16 23:59:59",
|
||||
page=1,
|
||||
page_size=50,
|
||||
sort_by="startTime",
|
||||
sort_order="desc",
|
||||
user_api_key_dict=auth,
|
||||
)
|
||||
|
||||
assert mock_prisma.db.query_raw.called, "query_raw should have been called"
|
||||
sql = mock_prisma.db.query_raw.call_args[0][0]
|
||||
assert sql.count("AT TIME ZONE 'UTC'") >= 2, (
|
||||
"/spend/logs/ui must wrap both `startTime` bounds with "
|
||||
f"`AT TIME ZONE 'UTC'`. SQL was:\n{sql}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -13,14 +13,18 @@ from litellm.proxy._types import TeamCallbackMetadata, UserAPIKeyAuth
|
|||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
KeyAndTeamLoggingSettings,
|
||||
LiteLLMProxyRequestSetup,
|
||||
_apply_credential_overrides_from_model_config,
|
||||
_extract_credential_from_entry,
|
||||
_get_dynamic_logging_metadata,
|
||||
_get_enforced_params,
|
||||
_get_metadata_variable_name,
|
||||
_resolve_credential_from_model_config,
|
||||
_update_model_if_key_alias_exists,
|
||||
add_guardrails_from_policy_engine,
|
||||
add_litellm_data_to_request,
|
||||
check_if_token_is_service_account,
|
||||
)
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../..")
|
||||
|
|
@ -1912,3 +1916,542 @@ async def test_bearer_token_not_in_debug_logs():
|
|||
f"Bearer token leaked in debug logs. "
|
||||
f"Found token in log output:\n{log_output[:500]}"
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tests for credential overrides from model_config (team/project metadata)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def setup_test_credentials():
|
||||
"""Populate litellm.credential_list with test credentials and enable feature flag, clean up after."""
|
||||
original = litellm.credential_list[:]
|
||||
original_flag = litellm.enable_model_config_credential_overrides
|
||||
litellm.enable_model_config_credential_overrides = True
|
||||
litellm.credential_list.extend(
|
||||
[
|
||||
CredentialItem(
|
||||
credential_name="hotel-azure-eastus",
|
||||
credential_info={},
|
||||
credential_values={
|
||||
"api_base": "https://hotel-eastus.openai.azure.com/",
|
||||
"api_key": "key-hotel-eastus",
|
||||
},
|
||||
),
|
||||
CredentialItem(
|
||||
credential_name="hotel-azure-westus",
|
||||
credential_info={},
|
||||
credential_values={
|
||||
"api_base": "https://hotel-westus.openai.azure.com/",
|
||||
"api_key": "key-hotel-westus",
|
||||
},
|
||||
),
|
||||
CredentialItem(
|
||||
credential_name="hotel-rec-azure",
|
||||
credential_info={},
|
||||
credential_values={
|
||||
"api_base": "https://hotel-rec-app.openai.azure.com/",
|
||||
"api_key": "key-hotel-rec",
|
||||
},
|
||||
),
|
||||
CredentialItem(
|
||||
credential_name="hotel-rec-vision",
|
||||
credential_info={},
|
||||
credential_values={
|
||||
"api_base": "https://hotel-rec-vision.openai.azure.com/",
|
||||
"api_key": "key-hotel-rec-vision",
|
||||
"api_version": "2024-06-01",
|
||||
},
|
||||
),
|
||||
CredentialItem(
|
||||
credential_name="flight-azure-centralus",
|
||||
credential_info={},
|
||||
credential_values={
|
||||
"api_base": "https://flight-centralus.openai.azure.com/",
|
||||
"api_key": "key-flight-centralus",
|
||||
},
|
||||
),
|
||||
]
|
||||
)
|
||||
yield
|
||||
litellm.credential_list[:] = original
|
||||
litellm.enable_model_config_credential_overrides = original_flag
|
||||
|
||||
|
||||
# --- Unit tests for _extract_credential_from_entry ---
|
||||
|
||||
|
||||
def test_extract_credential_from_entry_azure():
|
||||
entry = {"azure": {"litellm_credentials": "my-cred"}}
|
||||
assert _extract_credential_from_entry(entry) == "my-cred"
|
||||
|
||||
|
||||
def test_extract_credential_from_entry_no_credential():
|
||||
entry = {"azure": {"some_other_key": "value"}}
|
||||
assert _extract_credential_from_entry(entry) is None
|
||||
|
||||
|
||||
def test_extract_credential_from_entry_empty():
|
||||
assert _extract_credential_from_entry({}) is None
|
||||
|
||||
|
||||
def test_extract_credential_from_entry_non_dict_value():
|
||||
entry = {"azure": "not-a-dict"}
|
||||
assert _extract_credential_from_entry(entry) is None
|
||||
|
||||
|
||||
def test_extract_credential_from_entry_non_dict_entry():
|
||||
"""Non-dict entry (e.g. string) should return None, not crash."""
|
||||
assert _extract_credential_from_entry("my-cred-name") is None
|
||||
assert _extract_credential_from_entry(["a", "list"]) is None
|
||||
assert _extract_credential_from_entry(42) is None
|
||||
|
||||
|
||||
# --- Unit tests for _resolve_credential_from_model_config ---
|
||||
|
||||
|
||||
def test_resolve_project_model_specific_wins():
|
||||
project_config = {
|
||||
"gpt-4": {"azure": {"litellm_credentials": "proj-gpt4"}},
|
||||
"defaultconfig": {"azure": {"litellm_credentials": "proj-default"}},
|
||||
}
|
||||
team_config = {
|
||||
"gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}},
|
||||
"defaultconfig": {"azure": {"litellm_credentials": "team-default"}},
|
||||
}
|
||||
result = _resolve_credential_from_model_config(
|
||||
"gpt-4", project_config, team_config
|
||||
)
|
||||
assert result == "proj-gpt4"
|
||||
|
||||
|
||||
def test_resolve_project_default_wins_over_team():
|
||||
project_config = {
|
||||
"defaultconfig": {"azure": {"litellm_credentials": "proj-default"}},
|
||||
}
|
||||
team_config = {
|
||||
"gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}},
|
||||
"defaultconfig": {"azure": {"litellm_credentials": "team-default"}},
|
||||
}
|
||||
result = _resolve_credential_from_model_config(
|
||||
"gpt-4", project_config, team_config
|
||||
)
|
||||
assert result == "proj-default"
|
||||
|
||||
|
||||
def test_resolve_team_model_specific_wins_over_team_default():
|
||||
team_config = {
|
||||
"gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}},
|
||||
"defaultconfig": {"azure": {"litellm_credentials": "team-default"}},
|
||||
}
|
||||
result = _resolve_credential_from_model_config("gpt-4", None, team_config)
|
||||
assert result == "team-gpt4"
|
||||
|
||||
|
||||
def test_resolve_team_default_used_as_fallback():
|
||||
team_config = {
|
||||
"defaultconfig": {"azure": {"litellm_credentials": "team-default"}},
|
||||
}
|
||||
result = _resolve_credential_from_model_config("gpt-3.5", None, team_config)
|
||||
assert result == "team-default"
|
||||
|
||||
|
||||
def test_resolve_no_match_returns_none():
|
||||
result = _resolve_credential_from_model_config("gpt-4", None, None)
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_resolve_empty_configs_returns_none():
|
||||
result = _resolve_credential_from_model_config("gpt-4", {}, {})
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_resolve_model_not_in_any_config():
|
||||
project_config = {"gpt-4": {"azure": {"litellm_credentials": "x"}}}
|
||||
result = _resolve_credential_from_model_config("gpt-3.5", project_config, None)
|
||||
assert result is None
|
||||
|
||||
|
||||
# --- Integration tests for _apply_credential_overrides_from_model_config ---
|
||||
|
||||
|
||||
def test_apply_overrides_project_model_specific(setup_test_credentials):
|
||||
"""Scenario 2: Hotel Rec App -> gpt-4-vision -> project model-specific."""
|
||||
data = {"model": "gpt-4-vision"}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {"litellm_credentials": "hotel-azure-eastus"}
|
||||
},
|
||||
"gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}},
|
||||
}
|
||||
},
|
||||
project_metadata={
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {"litellm_credentials": "hotel-rec-azure"}
|
||||
},
|
||||
"gpt-4-vision": {
|
||||
"azure": {"litellm_credentials": "hotel-rec-vision"}
|
||||
},
|
||||
}
|
||||
},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
assert data["api_base"] == "https://hotel-rec-vision.openai.azure.com/"
|
||||
assert data["api_key"] == "key-hotel-rec-vision"
|
||||
assert data["api_version"] == "2024-06-01"
|
||||
|
||||
|
||||
def test_apply_overrides_project_default(setup_test_credentials):
|
||||
"""Scenario 1: Hotel Rec App -> gpt-4 -> project default."""
|
||||
data = {"model": "gpt-4"}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {"litellm_credentials": "hotel-azure-eastus"}
|
||||
},
|
||||
"gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}},
|
||||
}
|
||||
},
|
||||
project_metadata={
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {"litellm_credentials": "hotel-rec-azure"}
|
||||
},
|
||||
"gpt-4-vision": {
|
||||
"azure": {"litellm_credentials": "hotel-rec-vision"}
|
||||
},
|
||||
}
|
||||
},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
assert data["api_base"] == "https://hotel-rec-app.openai.azure.com/"
|
||||
assert data["api_key"] == "key-hotel-rec"
|
||||
|
||||
|
||||
def test_apply_overrides_team_model_specific(setup_test_credentials):
|
||||
"""Scenario 4: Hotel Review App -> gpt-4 -> team model-specific."""
|
||||
data = {"model": "gpt-4"}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {"litellm_credentials": "hotel-azure-eastus"}
|
||||
},
|
||||
"gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}},
|
||||
}
|
||||
},
|
||||
project_metadata={},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
assert data["api_base"] == "https://hotel-westus.openai.azure.com/"
|
||||
assert data["api_key"] == "key-hotel-westus"
|
||||
|
||||
|
||||
def test_apply_overrides_team_default(setup_test_credentials):
|
||||
"""Scenario 3: Hotel Review App -> gpt-3.5 -> team default."""
|
||||
data = {"model": "gpt-3.5"}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {"litellm_credentials": "hotel-azure-eastus"}
|
||||
},
|
||||
"gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}},
|
||||
}
|
||||
},
|
||||
project_metadata={},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
assert data["api_base"] == "https://hotel-eastus.openai.azure.com/"
|
||||
assert data["api_key"] == "key-hotel-eastus"
|
||||
|
||||
|
||||
def test_apply_overrides_no_config(setup_test_credentials):
|
||||
"""Scenario 6: No model_config anywhere -> data unchanged."""
|
||||
data = {"model": "gpt-4"}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={},
|
||||
project_metadata={},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
assert "api_base" not in data
|
||||
assert "api_key" not in data
|
||||
|
||||
|
||||
def test_apply_overrides_clientside_credentials_take_precedence(
|
||||
setup_test_credentials,
|
||||
):
|
||||
"""Clientside api_base/api_key in data should block model_config override."""
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"api_base": "https://my-custom-endpoint.openai.azure.com/",
|
||||
"api_key": "my-custom-key",
|
||||
}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {"litellm_credentials": "hotel-azure-eastus"}
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
assert data["api_base"] == "https://my-custom-endpoint.openai.azure.com/"
|
||||
assert data["api_key"] == "my-custom-key"
|
||||
|
||||
|
||||
def test_apply_overrides_missing_credential_name(setup_test_credentials):
|
||||
"""model_config references a credential that doesn't exist -> no override."""
|
||||
data = {"model": "gpt-4"}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={
|
||||
"model_config": {
|
||||
"gpt-4": {
|
||||
"azure": {"litellm_credentials": "nonexistent-credential"}
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
assert "api_base" not in data
|
||||
assert "api_key" not in data
|
||||
|
||||
|
||||
def test_apply_overrides_api_version_only_if_present(setup_test_credentials):
|
||||
"""api_version should only be set if the credential contains it."""
|
||||
data = {"model": "gpt-3.5"}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {"litellm_credentials": "hotel-azure-eastus"}
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
assert data["api_base"] == "https://hotel-eastus.openai.azure.com/"
|
||||
assert data["api_key"] == "key-hotel-eastus"
|
||||
assert "api_version" not in data
|
||||
|
||||
|
||||
def test_apply_overrides_no_model_in_data(setup_test_credentials):
|
||||
"""No model in request data -> skip override."""
|
||||
data = {"messages": [{"role": "user", "content": "hello"}]}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={
|
||||
"model_config": {
|
||||
"defaultconfig": {
|
||||
"azure": {"litellm_credentials": "some-cred"}
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
assert "api_base" not in data
|
||||
|
||||
|
||||
def test_apply_overrides_none_metadata(setup_test_credentials):
|
||||
"""None metadata on both team and project -> skip override."""
|
||||
data = {"model": "gpt-4"}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata=None,
|
||||
project_metadata=None,
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
assert "api_base" not in data
|
||||
|
||||
|
||||
def test_apply_overrides_clientside_api_version_preserved(setup_test_credentials):
|
||||
"""Clientside api_version should not be overwritten by credential."""
|
||||
data = {"model": "gpt-4-vision", "api_version": "2025-01-01"}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={
|
||||
"model_config": {
|
||||
"gpt-4-vision": {
|
||||
"azure": {"litellm_credentials": "hotel-rec-vision"}
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
# api_base and api_key should be set from credential
|
||||
assert data["api_base"] == "https://hotel-rec-vision.openai.azure.com/"
|
||||
assert data["api_key"] == "key-hotel-rec-vision"
|
||||
# api_version should be preserved from the request, not overwritten
|
||||
assert data["api_version"] == "2025-01-01"
|
||||
|
||||
|
||||
def test_resolve_non_dict_model_config_ignored():
|
||||
"""Non-dict model_config (e.g. string) should be safely skipped."""
|
||||
result = _resolve_credential_from_model_config("gpt-4", "not-a-dict", None)
|
||||
assert result is None
|
||||
|
||||
result = _resolve_credential_from_model_config(
|
||||
"gpt-4", None, ["also", "not", "a", "dict"]
|
||||
)
|
||||
assert result is None
|
||||
|
||||
# Valid config still works alongside invalid one
|
||||
result = _resolve_credential_from_model_config(
|
||||
"gpt-4",
|
||||
"invalid",
|
||||
{"gpt-4": {"azure": {"litellm_credentials": "valid-cred"}}},
|
||||
)
|
||||
assert result == "valid-cred"
|
||||
|
||||
|
||||
def test_resolve_pre_alias_model_name_fallback():
|
||||
"""model_config keyed on pre-alias name should match after alias resolution."""
|
||||
team_config = {
|
||||
"gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}},
|
||||
}
|
||||
# Post-alias name doesn't match, but pre-alias does (team scope)
|
||||
result = _resolve_credential_from_model_config(
|
||||
"azure/gpt-4-0613", None, team_config, pre_alias_model_name="gpt-4"
|
||||
)
|
||||
assert result == "team-gpt4"
|
||||
|
||||
# Same test for project scope
|
||||
project_config = {
|
||||
"gpt-4": {"azure": {"litellm_credentials": "proj-gpt4"}},
|
||||
}
|
||||
result = _resolve_credential_from_model_config(
|
||||
"azure/gpt-4-0613", project_config, None, pre_alias_model_name="gpt-4"
|
||||
)
|
||||
assert result == "proj-gpt4"
|
||||
|
||||
|
||||
def test_resolve_post_alias_name_takes_priority():
|
||||
"""Post-alias (resolved) name should be tried before pre-alias name."""
|
||||
team_config = {
|
||||
"gpt-4": {"azure": {"litellm_credentials": "pre-alias-cred"}},
|
||||
"gpt-4o-team-1": {"azure": {"litellm_credentials": "post-alias-cred"}},
|
||||
}
|
||||
# Team scope
|
||||
result = _resolve_credential_from_model_config(
|
||||
"gpt-4o-team-1", None, team_config, pre_alias_model_name="gpt-4"
|
||||
)
|
||||
assert result == "post-alias-cred"
|
||||
|
||||
# Project scope
|
||||
result = _resolve_credential_from_model_config(
|
||||
"gpt-4o-team-1", team_config, None, pre_alias_model_name="gpt-4"
|
||||
)
|
||||
assert result == "post-alias-cred"
|
||||
|
||||
|
||||
def test_apply_overrides_with_alias(setup_test_credentials):
|
||||
"""Credential override should work when model name was changed by alias."""
|
||||
# Simulate: user called "my-gpt4", alias resolved to "azure/gpt-4-custom"
|
||||
# model_config is keyed on "my-gpt4" (the pre-alias name)
|
||||
data = {"model": "azure/gpt-4-custom"}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={
|
||||
"model_config": {
|
||||
"my-gpt4": {"azure": {"litellm_credentials": "hotel-azure-eastus"}},
|
||||
}
|
||||
},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
pre_alias_model_name="my-gpt4",
|
||||
)
|
||||
assert data["api_base"] == "https://hotel-eastus.openai.azure.com/"
|
||||
assert data["api_key"] == "key-hotel-eastus"
|
||||
|
||||
|
||||
def test_apply_overrides_feature_flag_disabled_by_default():
|
||||
"""Feature flag defaults to False — credential overrides are inert until explicitly enabled."""
|
||||
assert litellm.enable_model_config_credential_overrides is False
|
||||
data = {"model": "gpt-4"}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_metadata={
|
||||
"model_config": {
|
||||
"gpt-4": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}
|
||||
}
|
||||
},
|
||||
)
|
||||
_apply_credential_overrides_from_model_config(
|
||||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
assert "api_base" not in data
|
||||
assert "api_key" not in data
|
||||
|
||||
|
||||
def test_extract_credential_provider_hint_prefers_exact_match():
|
||||
"""Provider hint selects the correct provider in a multi-provider entry."""
|
||||
entry = {
|
||||
"openai": {"litellm_credentials": "openai-cred"},
|
||||
"azure": {"litellm_credentials": "azure-cred"},
|
||||
}
|
||||
# With provider hint, should pick the exact match
|
||||
assert _extract_credential_from_entry(entry, provider="azure") == "azure-cred"
|
||||
assert _extract_credential_from_entry(entry, provider="openai") == "openai-cred"
|
||||
|
||||
# Without provider hint, falls back to first key (insertion order)
|
||||
result = _extract_credential_from_entry(entry)
|
||||
assert result in ("openai-cred", "azure-cred")
|
||||
|
||||
# Unknown provider falls back to first available
|
||||
result = _extract_credential_from_entry(entry, provider="bedrock")
|
||||
assert result in ("openai-cred", "azure-cred")
|
||||
|
||||
|
||||
def test_resolve_provider_hint_from_model_name():
|
||||
"""Provider prefix in model name (e.g. azure/gpt-4) threads through to entry extraction."""
|
||||
config = {
|
||||
"gpt-4": {
|
||||
"openai": {"litellm_credentials": "openai-cred"},
|
||||
"azure": {"litellm_credentials": "azure-cred"},
|
||||
},
|
||||
}
|
||||
# Model name "azure/gpt-4" -> provider="azure" -> should prefer azure-cred
|
||||
# But _resolve_credential_from_model_config tries "azure/gpt-4" first (no match),
|
||||
# then falls to defaultconfig (no match). So we need to use pre_alias_model_name.
|
||||
result = _resolve_credential_from_model_config(
|
||||
"azure/gpt-4", config, None, pre_alias_model_name="gpt-4", provider="azure"
|
||||
)
|
||||
assert result == "azure-cred"
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue