tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 02:21:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge main and resolve Snowflake test conflict

Browse Source 
Main rewrote the same tests we moved. Resolution: keep the tests only
in the unit test directory, adopting main's improved patterns (AsyncMock,
assert_called_once, stronger content assertions on streaming).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yuneng Jiang
2026-03-30 18:06:37 -07:00
co-authored by Claude Opus 4.6
parent 19f8b58046 949a395b01
commit 2b374a2abf
141 changed files with 4689 additions and 6245 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.circleci/config.yml
-625
View File
@@ -406,119 +406,6 @@ jobs:
# Store test results
- store_test_results:
path: test-results
caching_unit_tests:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
resource_class: large
working_directory: ~/project
parallelism: 2
steps:
- checkout
- setup_google_dns
- run:
name: DNS lookup for Redis host
command: |
sudo apt-get update
sudo apt-get install -y dnsutils
dig redis-19899.c239.us-east-1-2.ec2.redns.redis-cloud.com +short
- run:
name: Show git commit hash
command: |
echo "Git commit hash: $CIRCLE_SHA1"
- restore_cache:
keys:
- v2-caching-deps-{{ checksum ".circleci/requirements.txt" }}
- v2-caching-deps-
- run:
name: Install Dependencies
command: |
python -m pip install --upgrade pip
python -m pip install -r .circleci/requirements.txt
pip install "pytest==7.3.1"
pip install "pytest-retry==1.6.3"
pip install "pytest-asyncio==0.21.1"
pip install "pytest-cov==5.0.0"
pip install "mypy==1.18.2"
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install pyarrow
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install lunary==0.2.5
pip install "azure-identity==1.16.1"
pip install "langfuse==2.59.7"
pip install "logfire==0.29.0"
pip install numpydoc
pip install traceloop-sdk==0.21.1
pip install opentelemetry-api==1.25.0
pip install opentelemetry-sdk==1.25.0
pip install opentelemetry-exporter-otlp==1.25.0
pip install openai==1.100.1
pip install prisma==0.11.0
pip install "detect_secrets==1.5.0"
pip install "httpx==0.24.1"
pip install "respx==0.22.0"
pip install fastapi
pip install "gunicorn==21.2.0"
pip install "anyio==4.2.0"
pip install "aiodynamo==23.10.1"
pip install "asyncio==3.4.3"
pip install "apscheduler==3.10.4"
pip install "PyGithub==1.59.1"
pip install argon2-cffi
pip install "pytest-mock==3.12.0"
pip install python-multipart
pip install google-cloud-aiplatform
pip install prometheus-client==0.20.0
pip install "pydantic==2.10.2"
pip install "diskcache==5.6.1"
pip install "Pillow==10.3.0"
pip install "jsonschema==4.22.0"
pip install "websockets==13.1.0"
pip install "pytest-xdist==3.6.1"
- setup_litellm_enterprise_pip
- save_cache:
paths:
- /home/circleci/.pyenv/versions
- /home/circleci/.local
key: v2-caching-deps-{{ checksum ".circleci/requirements.txt" }}
- run:
name: Run prisma ./docker/entrypoint.sh
command: |
set +e
chmod +x docker/entrypoint.sh
./docker/entrypoint.sh
set -e
# Run pytest and generate JUnit XML report
- run:
name: Run tests
command: |
pwd
ls
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--split-by=timings \
--verbose \
--command="xargs python -m pytest \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-k 'caching or cache'"
no_output_timeout: 15m
# Store test results
- store_test_results:
path: test-results
auth_ui_unit_tests:
docker:
- image: cimg/python:3.11
@@ -664,376 +551,6 @@ jobs:
# Store test results
- store_test_results:
path: test-results
litellm_security_tests:
docker:
- image: cimg/python:3.13
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
- image: cimg/postgres:14.0
environment:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: circle_test
resource_class: xlarge
working_directory: ~/project
environment:
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/circle_test"
steps:
- checkout
- setup_google_dns
- run:
name: Show git commit hash
command: |
echo "Git commit hash: $CIRCLE_SHA1"
- setup_remote_docker:
docker_layer_caching: true
- restore_cache:
keys:
- v3-litellm-uv-deps-{{ checksum "requirements.txt" }}-{{ checksum ".circleci/config.yml" }}
- 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" "pytest-retry==1.6.3" "pytest-mock==3.12.0" \
"pytest-asyncio==0.21.1" "pytest-cov==5.0.0"
- save_cache:
paths:
- ~/.local/lib
- ~/.local/bin
- ~/.cache/uv
key: v3-litellm-uv-deps-{{ checksum "requirements.txt" }}-{{ checksum ".circleci/config.yml" }}
- run:
name: Install dockerize
command: |
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
rm dockerize-linux-amd64-v0.6.1.tar.gz
- run:
name: Wait for PostgreSQL to be ready
command: dockerize -wait tcp://localhost:5432 -timeout 1m
- run:
name: Run Security Scans
command: |
chmod +x ci_cd/security_scans.sh
./ci_cd/security_scans.sh
- run:
name: Run prisma ./docker/entrypoint.sh
command: |
set +e
chmod +x docker/entrypoint.sh
./docker/entrypoint.sh
set -e
# Run pytest and generate JUnit XML report
- run:
name: Run tests
command: |
python -m pytest tests/proxy_security_tests -v -x --junitxml=test-results/junit.xml --durations=5
no_output_timeout: 15m
# Store test results
- store_test_results:
path: test-results
# Split proxy unit tests into 3 jobs for faster execution and better debugging
# test_key_generate_prisma runs separately without parallel execution to avoid event loop issues with logging worker
litellm_proxy_unit_testing_key_generation:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: medium
steps:
- checkout
- setup_google_dns
- run:
name: Show git commit hash
command: |
echo "Git commit hash: $CIRCLE_SHA1"
- run:
name: Install PostgreSQL
command: |
sudo apt-get update
sudo apt-get install -y postgresql-14 postgresql-contrib-14
- restore_cache:
keys:
- v1-dependencies-{{ checksum ".circleci/requirements.txt" }}
- run:
name: Install Dependencies
command: |
python -m pip install --upgrade pip
python -m pip install -r .circleci/requirements.txt
pip install "pytest==7.3.1"
pip install "pytest-retry==1.6.3"
pip install "pytest-asyncio==0.21.1"
pip install "pytest-cov==5.0.0"
pip install "pytest-timeout==2.2.0"
pip install "pytest-forked==1.6.0"
pip install "mypy==1.18.2"
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install "google-genai==1.22.0"
pip install pyarrow
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install lunary==0.2.5
pip install "azure-identity==1.16.1"
pip install "langfuse==2.59.7"
pip install "logfire==0.29.0"
pip install numpydoc
pip install traceloop-sdk==0.21.1
pip install opentelemetry-api==1.25.0
pip install opentelemetry-sdk==1.25.0
pip install opentelemetry-exporter-otlp==1.25.0
pip install openai==1.100.1
pip install prisma==0.11.0
pip install "detect_secrets==1.5.0"
pip install "httpx==0.24.1"
pip install "respx==0.22.0"
pip install fastapi
pip install "gunicorn==21.2.0"
pip install "anyio==4.2.0"
pip install "aiodynamo==23.10.1"
pip install "asyncio==3.4.3"
pip install "apscheduler==3.10.4"
pip install "PyGithub==1.59.1"
pip install argon2-cffi
pip install "pytest-mock==3.12.0"
pip install python-multipart
pip install google-cloud-aiplatform
pip install prometheus-client==0.20.0
pip install "pydantic==2.10.2"
pip install "diskcache==5.6.1"
pip install "Pillow==10.3.0"
pip install "jsonschema==4.22.0"
pip install "pytest-postgresql==7.0.1"
pip install "fakeredis==2.28.1"
- setup_litellm_enterprise_pip
- save_cache:
paths:
- ./venv
key: v1-dependencies-{{ checksum ".circleci/requirements.txt" }}
- run:
name: Run prisma ./docker/entrypoint.sh
command: |
set +e
chmod +x docker/entrypoint.sh
./docker/entrypoint.sh
set -e
- run:
name: Run key generation tests (no parallel execution to avoid event loop issues)
command: |
pwd
ls
# Run without -n flag to avoid pytest-xdist event loop conflicts with logging worker
python -m pytest tests/proxy_unit_tests/test_key_generate_prisma.py --cov=litellm --cov-report=xml --junitxml=test-results/junit-key-generation.xml --durations=10 --timeout=300 -vv --log-cli-level=INFO
no_output_timeout: 15m
- run:
name: Rename the coverage files
command: |
mv coverage.xml litellm_proxy_unit_tests_key_generation_coverage.xml
mv .coverage litellm_proxy_unit_tests_key_generation_coverage
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- litellm_proxy_unit_tests_key_generation_coverage.xml
- litellm_proxy_unit_tests_key_generation_coverage
litellm_proxy_unit_testing_part1:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: xlarge
steps:
- checkout
- setup_google_dns
- run:
name: Show git commit hash
command: |
echo "Git commit hash: $CIRCLE_SHA1"
- run:
name: Install PostgreSQL
command: |
sudo apt-get update
sudo apt-get install -y postgresql-14 postgresql-contrib-14
- restore_cache:
keys:
- v1-dependencies-{{ checksum ".circleci/requirements.txt" }}
- run:
name: Install Dependencies
command: |
python -m pip install --upgrade pip
python -m pip install -r .circleci/requirements.txt
pip install "pytest==7.3.1"
pip install "pytest-retry==1.6.3"
pip install "pytest-asyncio==0.21.1"
pip install "pytest-cov==5.0.0"
pip install "pytest-timeout==2.2.0"
pip install "pytest-forked==1.6.0"
pip install "mypy==1.18.2"
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install "google-genai==1.22.0"
pip install pyarrow
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install lunary==0.2.5
pip install "azure-identity==1.16.1"
pip install "langfuse==2.59.7"
pip install "logfire==0.29.0"
pip install numpydoc
pip install traceloop-sdk==0.21.1
pip install opentelemetry-api==1.25.0
pip install opentelemetry-sdk==1.25.0
pip install opentelemetry-exporter-otlp==1.25.0
pip install openai==1.100.1
pip install prisma==0.11.0
pip install "detect_secrets==1.5.0"
pip install "httpx==0.24.1"
pip install "respx==0.22.0"
pip install fastapi
pip install "gunicorn==21.2.0"
pip install "anyio==4.2.0"
pip install "aiodynamo==23.10.1"
pip install "asyncio==3.4.3"
pip install "apscheduler==3.10.4"
pip install "PyGithub==1.59.1"
pip install argon2-cffi
pip install "pytest-mock==3.12.0"
pip install python-multipart
pip install google-cloud-aiplatform
pip install prometheus-client==0.20.0
pip install "pydantic==2.10.2"
pip install "diskcache==5.6.1"
pip install "Pillow==10.3.0"
pip install "jsonschema==4.22.0"
pip install "pytest-postgresql==7.0.1"
pip install "fakeredis==2.28.1"
pip install "pytest-xdist==3.6.1"
- setup_litellm_enterprise_pip
- save_cache:
paths:
- ./venv
key: v1-dependencies-{{ checksum ".circleci/requirements.txt" }}
- run:
name: Run prisma ./docker/entrypoint.sh
command: |
set +e
chmod +x docker/entrypoint.sh
./docker/entrypoint.sh
set -e
- run:
name: Run proxy unit tests (part 1 - auth checks)
command: |
pwd
ls
python -m pytest tests/proxy_unit_tests/test_auth_checks.py tests/proxy_unit_tests/test_user_api_key_auth.py --junitxml=test-results/junit-part1.xml --durations=10 -n 8 --timeout=300 -v
no_output_timeout: 15m
- store_test_results:
path: test-results
litellm_proxy_unit_testing_part2:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: xlarge
steps:
- checkout
- setup_google_dns
- run:
name: Show git commit hash
command: |
echo "Git commit hash: $CIRCLE_SHA1"
- run:
name: Install PostgreSQL
command: |
sudo apt-get update
sudo apt-get install -y postgresql-14 postgresql-contrib-14
- restore_cache:
keys:
- v1-dependencies-{{ checksum ".circleci/requirements.txt" }}
- run:
name: Install Dependencies
command: |
python -m pip install --upgrade pip
python -m pip install -r .circleci/requirements.txt
pip install "pytest==7.3.1"
pip install "pytest-retry==1.6.3"
pip install "pytest-asyncio==0.21.1"
pip install "pytest-cov==5.0.0"
pip install "pytest-timeout==2.2.0"
pip install "pytest-forked==1.6.0"
pip install "mypy==1.18.2"
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install "google-genai==1.22.0"
pip install pyarrow
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install lunary==0.2.5
pip install "azure-identity==1.16.1"
pip install "langfuse==2.59.7"
pip install "logfire==0.29.0"
pip install numpydoc
pip install traceloop-sdk==0.21.1
pip install opentelemetry-api==1.25.0
pip install opentelemetry-sdk==1.25.0
pip install opentelemetry-exporter-otlp==1.25.0
pip install openai==1.100.1
pip install prisma==0.11.0
pip install "detect_secrets==1.5.0"
pip install "httpx==0.24.1"
pip install "respx==0.22.0"
pip install fastapi
pip install "gunicorn==21.2.0"
pip install "anyio==4.2.0"
pip install "aiodynamo==23.10.1"
pip install "asyncio==3.4.3"
pip install "apscheduler==3.10.4"
pip install "PyGithub==1.59.1"
pip install argon2-cffi
pip install "pytest-mock==3.12.0"
pip install python-multipart
pip install google-cloud-aiplatform
pip install prometheus-client==0.20.0
pip install "pydantic==2.10.2"
pip install "diskcache==5.6.1"
pip install "Pillow==10.3.0"
pip install "jsonschema==4.22.0"
pip install "pytest-postgresql==7.0.1"
pip install "fakeredis==2.28.1"
pip install "pytest-xdist==3.6.1"
- setup_litellm_enterprise_pip
- save_cache:
paths:
- ./venv
key: v1-dependencies-{{ checksum ".circleci/requirements.txt" }}
- run:
name: Run prisma ./docker/entrypoint.sh
command: |
set +e
chmod +x docker/entrypoint.sh
./docker/entrypoint.sh
set -e
- run:
name: Run proxy unit tests (part 2 - remaining tests)
command: |
pwd
ls
python -m pytest tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py --junitxml=test-results/junit-part2.xml --durations=10 -n 8 --timeout=300 -v
no_output_timeout: 15m
- store_test_results:
path: test-results
litellm_assistants_api_testing: # Runs all tests with the "assistants" keyword
docker:
- image: cimg/python:3.13.1
@@ -1507,101 +1024,6 @@ jobs:
no_output_timeout: 15m
- store_test_results:
path: test-results
litellm_mapped_tests_llms:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: large
steps:
- setup_litellm_test_deps
- run:
name: Run LLM provider tests
command: |
python -m pytest tests/test_litellm/llms --junitxml=test-results/junit-llms.xml --durations=10 -n 4 --maxfail=5 --timeout=300 -vv --log-cli-level=WARNING
no_output_timeout: 15m
- store_test_results:
path: test-results
litellm_mapped_tests_core:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: large
steps:
- setup_litellm_test_deps
- run:
name: Run core tests
command: |
python -m pytest tests/test_litellm --ignore=tests/test_litellm/proxy --ignore=tests/test_litellm/llms --ignore=tests/test_litellm/integrations --ignore=tests/test_litellm/litellm_core_utils --ignore=tests/test_litellm/experimental_mcp_client --junitxml=test-results/junit-core.xml --durations=10 -n 4 --maxfail=5 --timeout=300 -vv --log-cli-level=WARNING
no_output_timeout: 15m
- store_test_results:
path: test-results
litellm_mapped_tests_litellm_core_utils:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: large
steps:
- setup_litellm_test_deps
- run:
name: Run litellm_core_utils tests
command: |
python -m pytest tests/test_litellm/litellm_core_utils --junitxml=test-results/junit-litellm-core-utils.xml --durations=10 -n 4 --maxfail=5 --timeout=300 -vv --log-cli-level=WARNING
no_output_timeout: 15m
- store_test_results:
path: test-results
litellm_mapped_tests_mcps:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: medium
steps:
- setup_litellm_test_deps
- run:
name: Run MCP client tests
command: |
python -m pytest tests/test_litellm/experimental_mcp_client --cov=litellm --cov-report=xml --junitxml=test-results/junit-mcps.xml --durations=10 -n 2 --maxfail=5 --timeout=300 -vv --log-cli-level=WARNING
no_output_timeout: 15m
- run:
name: Rename the coverage files
command: |
mv coverage.xml litellm_mcps_tests_coverage.xml
mv .coverage litellm_mcps_tests_coverage
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- litellm_mcps_tests_coverage.xml
- litellm_mcps_tests_coverage
litellm_mapped_tests_integrations:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: large
steps:
- setup_litellm_test_deps
- run:
name: Run integrations tests
command: |
python -m pytest tests/test_litellm/integrations --junitxml=test-results/junit-integrations.xml --durations=10 -n 4 --maxfail=5 --timeout=300 -vv --log-cli-level=WARNING
no_output_timeout: 15m
- store_test_results:
path: test-results
litellm_mapped_enterprise_tests:
docker:
- image: cimg/python:3.11
@@ -4002,34 +3424,14 @@ workflows:
only:
- main
- /litellm_.*/
- caching_unit_tests:
filters:
branches:
only:
- main
- /litellm_.*/
- litellm_proxy_unit_testing_key_generation:
filters:
branches:
only:
- main
- /litellm_.*/
- litellm_proxy_unit_testing_part1:
filters:
branches:
only:
- main
- /litellm_.*/
- litellm_proxy_unit_testing_part2:
filters:
branches:
only:
- main
- /litellm_.*/
- litellm_security_tests:
filters:
branches:
only:
- main
- /litellm_.*/
- litellm_assistants_api_testing:
@@ -4265,34 +3667,14 @@ workflows:
only:
- main
- /litellm_.*/
- litellm_mapped_tests_llms:
filters:
branches:
only:
- main
- /litellm_.*/
- litellm_mapped_tests_core:
filters:
branches:
only:
- main
- /litellm_.*/
- litellm_mapped_tests_mcps:
filters:
branches:
only:
- main
- /litellm_.*/
- litellm_mapped_tests_integrations:
filters:
branches:
only:
- main
- /litellm_.*/
- litellm_mapped_tests_litellm_core_utils:
filters:
branches:
only:
- main
- /litellm_.*/
- batches_testing:
@@ -4342,11 +3724,6 @@ workflows:
- search_testing
- litellm_mapped_tests_proxy_part1
- litellm_mapped_tests_proxy_part2
- litellm_mapped_tests_llms
- litellm_mapped_tests_core
- litellm_mapped_tests_mcps
- litellm_mapped_tests_integrations
- litellm_mapped_tests_litellm_core_utils
- litellm_mapped_enterprise_tests
- batches_testing
- litellm_utils_testing
@@ -4354,8 +3731,6 @@ workflows:
- image_gen_testing
- logging_testing
- audio_testing
- caching_unit_tests
- litellm_proxy_unit_testing_key_generation
- langfuse_logging_unit_tests
- local_testing_part1
- local_testing_part2
.github/workflows/scorecard.yml
+47
View File
@@ -0,0 +1,47 @@
name: Scorecard supply-chain security
on:
branch_protection_rule:
schedule:
- cron: '27 12 * * 4'
push:
branches: ["main"]
permissions: read-all
jobs:
analysis:
name: Scorecard analysis
runs-on: ubuntu-latest
if: github.event.repository.default_branch == github.ref_name
permissions:
security-events: write
id-token: write
# Uncomment for private repos if needed:
# contents: read
# actions: read
steps:
- name: Checkout code
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
persist-credentials: false
- name: Run analysis
uses: ossf/scorecard-action@f49aabe0b5af0936a0987cfb85d86b75731b0186 # v2.4.1
with:
results_file: results.sarif
results_format: sarif
publish_results: true
- name: Upload artifact
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: SARIF file
path: results.sarif
retention-days: 5
- name: Upload to code scanning
uses: github/codeql-action/upload-sarif@c10b806170c8ee63ea24152429041b5624f0baf5 # v4.35.1
with:
sarif_file: results.sarif
docs/my-website/blog/security_townhall_updates/index.md
+18 -6
View File
@@ -24,12 +24,6 @@ On March 24, 2026 at 10:39 UTC, LiteLLM v1.82.7 was pushed to PyPI. Version v1.8
At this point, our understanding is that this was a supply-chain incident affecting those two published versions.
### Q. Were older packages impacted?
Our current findings show no indicators of compromise in the last 20 versions of LiteLLM. This was manually verified by our team and independently reviewed by Veria Labs.
We have also published the verified versions for users to use. [Check Security Blog for release verification.](https://docs.litellm.ai/blog/security-update-march-2026#verified-safe-versions)
## How did this happen?
Our understanding is that the issue came from the [compromised Trivy security scanner](https://www.aquasec.com/blog/trivy-supply-chain-attack-what-you-need-to-know/) dependency in our CI/CD pipeline.
@@ -160,6 +154,24 @@ We believe that [Cosign](https://github.com/sigstore/cosign) is a good fit for t
We've added zizmor to help us catch issues such as unpinned dependencies and credential leakage. [commit](https://github.com/BerriAI/litellm/commit/a671275f5c5b0e1fb1adacdf3b6ef779aaa5d56c).
## Frequently Asked Questions
**Q: Did you observe any lateral movement into your corporate environment during this incident?**
A: No. Our investigation to date, conducted in coordination with external security experts, has found no evidence of lateral movement into our internal corporate systems. The incident was isolated to the CI/CD pipeline and the release path for specific versions (v1.82.7 and v1.82.8). As a proactive measure, we have rotated all potentially impacted or adjacent secrets—including PyPI, GitHub, and Docker credentials—and updated maintainer account security to ensure continued isolation.
**Q: Do you expect delays in future product releases due to these new security measures?**
A: We are committed to balancing security with speed. While we have temporarily paused releases to implement stronger controls, we are moving quickly to automate our new security protocols. We are currently implementing isolated CI/CD environments, ephemeral credentials (via Trusted Publishers), and release auditing with Cosign. These improvements are designed to be integrated into our automated pipeline, allowing us to maintain a fast release cadence while ensuring every package is verified and secure.
**Q: Were older packages impacted?**
Our current findings show no indicators of compromise in the last 20 versions of LiteLLM. This was manually verified by our team and independently reviewed by Veria Labs.
We have also published the verified versions for users to use. [Check Security Blog for release verification.](https://docs.litellm.ai/blog/security-update-march-2026#verified-safe-versions)
## Questions & Support
If you believe your systems may be affected, contact us immediately:
docs/my-website/blog/vanta_compliance_recertification/index.md
+18
View File
@@ -0,0 +1,18 @@
---
slug: vanta-compliance-recertification
title: "LiteLLM + Vanta: SOC 2 Type 2 and ISO 27001 Recertification"
date: 2026-03-30T10:00:00
authors:
- krrish
description: "LiteLLM is partnering with Vanta on SOC 2 Type 2 and ISO 27001 recertification and engaging independent auditors for verification."
tags: [security, compliance]
hide_table_of_contents: true
---
![LiteLLM x Vanta SOC-2 Recertification](/img/blog/vanta_soc2_recertification.png)
We are partnering with [Vanta](https://www.vanta.com/) to recertify LiteLLM's compliance for SOC 2 Type 2 and ISO 27001.
As part of this process, we are also identifying independent auditors to validate and verify our compliance posture.
This is part of our commitment to being the most secure and transparent AI Gateway possible.
docs/my-website/docs/learn/sdk_quickstart.md
+1 -1
View File
@@ -11,7 +11,7 @@ Use this path if you are integrating LiteLLM directly into application code.
## 1. Install LiteLLM
```bash
pip install litellm
pip install litellm==1.82.6
```
## 2. Set Provider Credentials
docs/my-website/static/img/blog/vanta_soc2_recertification.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 39 KiB

enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py
+2
View File
@@ -116,6 +116,7 @@ class PagerDutyAlerting(SlackAlerting):
user_api_key_org_id=_meta.get("user_api_key_org_id"),
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"),
user_api_key_user_id=_meta.get("user_api_key_user_id"),
user_api_key_team_alias=_meta.get("user_api_key_team_alias"),
user_api_key_end_user_id=_meta.get("user_api_key_end_user_id"),
@@ -197,6 +198,7 @@ class PagerDutyAlerting(SlackAlerting):
user_api_key_org_id=user_api_key_dict.org_id,
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,
user_api_key_user_id=user_api_key_dict.user_id,
user_api_key_team_alias=user_api_key_dict.team_alias,
user_api_key_end_user_id=user_api_key_dict.end_user_id,
litellm/completion_extras/litellm_responses_transformation/transformation.py
+13 -3
View File
@@ -56,6 +56,16 @@ if TYPE_CHECKING:
)
def _get_reasoning_items(
msg: "AllMessageValues",
) -> List[ChatCompletionReasoningItem]:
"""Extract reasoning_items from a message dict with proper typing."""
items = msg.get("reasoning_items") # type: ignore[union-attr]
if items:
return items # type: ignore[return-value]
return []
def _build_reasoning_item(
item_id: str,
encrypted_content: Optional[str],
@@ -86,7 +96,7 @@ def _build_reasoning_item(
}
def _reasoning_item_to_response_input(r_item: Dict[str, Any]) -> Dict[str, Any]:
def _reasoning_item_to_response_input(r_item: Union[ChatCompletionReasoningItem, Dict[str, Any]]) -> Dict[str, Any]:
"""Convert a stored ChatCompletionReasoningItem back to a Responses API input item."""
r_input: Dict[str, Any] = {
"type": "reasoning",
@@ -246,7 +256,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
}
)
elif role == "assistant" and tool_calls and isinstance(tool_calls, list):
for r_item in msg.get("reasoning_items") or []:
for r_item in _get_reasoning_items(msg):
input_items.append(_reasoning_item_to_response_input(r_item))
for tool_call in tool_calls:
function = tool_call.get("function")
@@ -264,7 +274,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
raise ValueError(f"tool call not supported: {tool_call}")
elif content is not None:
if role == "assistant":
for r_item in msg.get("reasoning_items") or []:
for r_item in _get_reasoning_items(msg):
input_items.append(_reasoning_item_to_response_input(r_item))
input_items.append(
{
litellm/integrations/prometheus.py
+196 -2
View File
@@ -191,6 +191,31 @@ class PrometheusLogger(CustomLogger):
),
)
# Remaining Budget for Org
self.litellm_remaining_org_budget_metric = self._gauge_factory(
"litellm_remaining_org_budget_metric",
"Remaining budget for org",
labelnames=self.get_labels_for_metric(
"litellm_remaining_org_budget_metric"
),
)
# Max Budget for Org
self.litellm_org_max_budget_metric = self._gauge_factory(
"litellm_org_max_budget_metric",
"Maximum budget set for org",
labelnames=self.get_labels_for_metric("litellm_org_max_budget_metric"),
)
# Org Budget Reset At
self.litellm_org_budget_remaining_hours_metric = self._gauge_factory(
"litellm_org_budget_remaining_hours_metric",
"Remaining hours for org budget to be reset",
labelnames=self.get_labels_for_metric(
"litellm_org_budget_remaining_hours_metric"
),
)
# Remaining Budget for API Key
self.litellm_remaining_api_key_budget_metric = self._gauge_factory(
"litellm_remaining_api_key_budget_metric",
@@ -1003,6 +1028,9 @@ class PrometheusLogger(CustomLogger):
user_api_team_alias = standard_logging_payload["metadata"][
"user_api_key_team_alias"
]
user_api_key_org_id = standard_logging_payload["metadata"].get(
"user_api_key_org_id"
)
output_tokens = standard_logging_payload["completion_tokens"]
tokens_used = standard_logging_payload["total_tokens"]
response_cost = standard_logging_payload["response_cost"]
@@ -1111,6 +1139,7 @@ class PrometheusLogger(CustomLogger):
litellm_params=litellm_params,
response_cost=response_cost,
user_id=user_id,
user_api_key_org_id=user_api_key_org_id,
)
# set proxy virtual key rpm/tpm metrics
@@ -1266,6 +1295,7 @@ class PrometheusLogger(CustomLogger):
litellm_params: dict,
response_cost: float,
user_id: Optional[str] = None,
user_api_key_org_id: Optional[str] = None,
):
_metadata = litellm_params.get("metadata") or {}
_team_spend = _metadata.get("user_api_key_team_spend", None)
@@ -1298,12 +1328,16 @@ class PrometheusLogger(CustomLogger):
user_max_budget=_user_max_budget,
response_cost=response_cost,
),
self._set_org_budget_metrics_after_api_request(
org_id=user_api_key_org_id,
response_cost=response_cost,
),
return_exceptions=True,
)
for i, r in enumerate(results):
if isinstance(r, Exception):
verbose_logger.debug(
f"[Non-Blocking] Prometheus: Budget metric lookup {['key', 'team', 'user'][i]} failed: {r}"
f"[Non-Blocking] Prometheus: Budget metric lookup {['key', 'team', 'user', 'org'][i]} failed: {r}"
)
def _increment_top_level_request_and_spend_metrics(
@@ -1492,6 +1526,9 @@ class PrometheusLogger(CustomLogger):
user_api_team_alias = standard_logging_payload["metadata"][
"user_api_key_team_alias"
]
user_api_key_org_id = standard_logging_payload["metadata"].get(
"user_api_key_org_id"
)
try:
self.litellm_llm_api_failed_requests_metric.labels(
@@ -1507,6 +1544,10 @@ class PrometheusLogger(CustomLogger):
),
).inc()
self.set_llm_deployment_failure_metrics(kwargs)
await self._set_org_budget_metrics_after_api_request(
org_id=user_api_key_org_id,
response_cost=0,
)
except Exception as e:
verbose_logger.exception(
"prometheus Layer Error(): Exception occured - {}".format(str(e))
@@ -2736,6 +2777,37 @@ class PrometheusLogger(CustomLogger):
data_type="users",
)
async def _initialize_org_budget_metrics(self):
"""
Initialize org budget metrics by reusing the generic pagination logic.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
verbose_logger.debug(
"Prometheus: skipping org metrics initialization, DB not initialized"
)
return
async def fetch_orgs(
page_size: int, page: int
) -> Tuple[list, Optional[int]]:
skip = (page - 1) * page_size
orgs = await prisma_client.db.litellm_organizationtable.find_many(
skip=skip,
take=page_size,
order={"created_at": "desc"},
include={"litellm_budget_table": True},
)
total_count = await prisma_client.db.litellm_organizationtable.count()
return orgs, total_count
await self._initialize_budget_metrics(
data_fetch_function=fetch_orgs,
set_metrics_function=self._set_org_list_budget_metrics,
data_type="orgs",
)
async def initialize_remaining_budget_metrics(self):
"""
Handler for initializing remaining budget metrics for all teams to avoid metric discrepancies.
@@ -2770,10 +2842,11 @@ class PrometheusLogger(CustomLogger):
"""
Helper to initialize remaining budget metrics for all teams, API keys, and users.
"""
verbose_logger.debug("Emitting key, team, user budget metrics....")
verbose_logger.debug("Emitting key, team, user, org budget metrics....")
await self._initialize_team_budget_metrics()
await self._initialize_api_key_budget_metrics()
await self._initialize_user_budget_metrics()
await self._initialize_org_budget_metrics()
await self._initialize_user_and_team_count_metrics()
async def _initialize_user_and_team_count_metrics(self):
@@ -2829,6 +2902,20 @@ class PrometheusLogger(CustomLogger):
for user in users:
self._set_user_budget_metrics(user)
async def _set_org_list_budget_metrics(self, orgs: list):
"""Helper function to set budget metrics for a list of orgs"""
for org in orgs:
budget_table = getattr(org, "litellm_budget_table", None)
self._set_org_budget_metrics(
org_id=org.organization_id or "",
org_alias=org.organization_alias or "",
spend=org.spend or 0.0,
max_budget=budget_table.max_budget if budget_table else None,
budget_reset_at=getattr(budget_table, "budget_reset_at", None)
if budget_table
else None,
)
async def _set_team_budget_metrics_after_api_request(
self,
user_api_team: Optional[str],
@@ -2950,6 +3037,113 @@ class PrometheusLogger(CustomLogger):
)
)
async def _set_org_budget_metrics_after_api_request(
self,
org_id: Optional[str],
response_cost: float,
):
"""
Set org budget metrics after an LLM API request
- Fetches org info via cache (get_org_object)
- Sets org budget metrics
"""
if not org_id:
return
from litellm.proxy.auth.auth_checks import get_org_object
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if prisma_client is None:
return
try:
org_info = await get_org_object(
org_id=org_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
include_budget_table=True,
)
except Exception as e:
verbose_logger.debug(
f"[Non-Blocking] Prometheus: Error getting org info: {str(e)}"
)
return
if org_info is None:
return
org_alias = org_info.organization_alias or ""
_total_org_spend = (org_info.spend or 0.0) + response_cost
budget_table = org_info.litellm_budget_table
max_budget = budget_table.max_budget if budget_table else None
budget_reset_at = (
getattr(budget_table, "budget_reset_at", None) if budget_table else None
)
self._set_org_budget_metrics(
org_id=org_id,
org_alias=org_alias,
spend=_total_org_spend,
max_budget=max_budget,
budget_reset_at=budget_reset_at,
)
def _set_org_budget_metrics(
self,
org_id: str,
org_alias: str,
spend: float,
max_budget: Optional[float],
budget_reset_at: Optional[datetime],
):
"""
Set org budget metrics for a single org
- Remaining Budget
- Max Budget
- Budget Reset At
"""
enum_values = UserAPIKeyLabelValues(
org_id=org_id,
org_alias=org_alias,
)
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_remaining_org_budget_metric"
),
enum_values=enum_values,
)
self.litellm_remaining_org_budget_metric.labels(**_labels).set(
self._safe_get_remaining_budget(
max_budget=max_budget,
spend=spend,
)
)
if max_budget is not None:
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_org_max_budget_metric"
),
enum_values=enum_values,
)
self.litellm_org_max_budget_metric.labels(**_labels).set(max_budget)
if budget_reset_at is not None:
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_org_budget_remaining_hours_metric"
),
enum_values=enum_values,
)
self.litellm_org_budget_remaining_hours_metric.labels(**_labels).set(
self._get_remaining_hours_for_budget_reset(
budget_reset_at=budget_reset_at
)
)
def _set_key_budget_metrics(self, user_api_key_dict: UserAPIKeyAuth):
"""
Set virtual key budget metrics
litellm/llms/a2a/chat/guardrail_translation/handler.py
+43 -28
View File
@@ -234,37 +234,13 @@ class A2AGuardrailHandler(BaseTranslation):
then the combined guardrailed text is written into the first chunk that had text
and all other text parts in other chunks are cleared (in-place).
"""
from litellm.llms.a2a.common_utils import extract_text_from_a2a_response
# Parse each item; keep alignment with responses_so_far (None where unparseable)
parsed: List[Optional[Dict[str, Any]]] = [None] * len(responses_so_far)
for i, item in enumerate(responses_so_far):
if isinstance(item, dict):
obj = item
elif isinstance(item, str):
try:
obj = json.loads(item.strip())
except (json.JSONDecodeError, TypeError):
continue
else:
continue
if isinstance(obj.get("result"), dict):
parsed[i] = obj
valid_parsed = [(i, obj) for i, obj in enumerate(parsed) if obj is not None]
parsed, valid_parsed = self._parse_streaming_responses(responses_so_far)
if not valid_parsed:
return responses_so_far
# Collect text from each chunk in order (by original index in responses_so_far)
text_parts: List[str] = []
chunk_indices_with_text: List[int] = [] # indices into valid_parsed
for idx, (orig_i, obj) in enumerate(valid_parsed):
t = extract_text_from_a2a_response(obj)
if t:
text_parts.append(t)
chunk_indices_with_text.append(orig_i)
combined_text = "".join(text_parts)
combined_text, chunk_indices_with_text = (
self._collect_text_from_parsed_chunks(valid_parsed)
)
if not combined_text:
return responses_so_far
@@ -337,6 +313,45 @@ class A2AGuardrailHandler(BaseTranslation):
return responses_so_far
def _parse_streaming_responses(
self,
responses_so_far: List[Any],
) -> Tuple[
List[Optional[Dict[str, Any]]], List[Tuple[int, Dict[str, Any]]]
]:
"""Parse JSON-RPC items, returning aligned parsed list and valid entries."""
parsed: List[Optional[Dict[str, Any]]] = [None] * len(responses_so_far)
for i, item in enumerate(responses_so_far):
if isinstance(item, dict):
obj = item
elif isinstance(item, str):
try:
obj = json.loads(item.strip())
except (json.JSONDecodeError, TypeError):
continue
else:
continue
if isinstance(obj.get("result"), dict):
parsed[i] = obj
valid_parsed = [(i, obj) for i, obj in enumerate(parsed) if obj is not None]
return parsed, valid_parsed
def _collect_text_from_parsed_chunks(
self,
valid_parsed: List[Tuple[int, Dict[str, Any]]],
) -> Tuple[str, List[int]]:
"""Collect text from parsed chunks, returning combined text and indices."""
from litellm.llms.a2a.common_utils import extract_text_from_a2a_response
text_parts: List[str] = []
chunk_indices_with_text: List[int] = []
for _idx, (orig_i, obj) in enumerate(valid_parsed):
t = extract_text_from_a2a_response(obj)
if t:
text_parts.append(t)
chunk_indices_with_text.append(orig_i)
return "".join(text_parts), chunk_indices_with_text
def _extract_texts_from_result(
self,
result: Dict[str, Any],
litellm/llms/anthropic/chat/guardrail_translation/handler.py
+100 -67
View File
@@ -277,82 +277,26 @@ class AnthropicMessagesHandler(BaseTranslation):
images_to_check: List[str] = []
tool_calls_to_check: List[ChatCompletionToolCallChunk] = []
task_mappings: List[Tuple[int, Optional[int]]] = []
# Track (content_index, None) for each text
# Handle both dict and object responses
response_content: List[Any] = []
if isinstance(response, dict):
response_content = response.get("content", []) or []
elif hasattr(response, "content"):
content = getattr(response, "content", None)
response_content = content or []
else:
response_content = []
response_content = self._get_response_content(response)
if not response_content:
return response
# Step 1: Extract all text content and tool calls from response
for content_idx, content_block in enumerate(response_content):
# Handle both dict and Pydantic object content blocks
block_dict: Dict[str, Any] = {}
if isinstance(content_block, dict):
block_type = content_block.get("type")
block_dict = cast(Dict[str, Any], content_block)
elif hasattr(content_block, "type"):
block_type = getattr(content_block, "type", None)
# Convert Pydantic object to dict for processing
if hasattr(content_block, "model_dump"):
block_dict = content_block.model_dump()
else:
block_dict = {
"type": block_type,
"text": getattr(content_block, "text", None),
}
else:
continue
if block_type in ["text", "tool_use"]:
self._extract_output_text_and_images(
content_block=block_dict,
content_idx=content_idx,
texts_to_check=texts_to_check,
images_to_check=images_to_check,
task_mappings=task_mappings,
tool_calls_to_check=tool_calls_to_check,
)
self._extract_from_content_blocks(
response_content, texts_to_check, images_to_check,
task_mappings, tool_calls_to_check,
)
# Step 2: Apply guardrail to all texts in batch
if texts_to_check or tool_calls_to_check:
# Use the real request_data if provided (proxy path), otherwise
# create a standalone dict (SDK / direct-call path).
if request_data is None:
request_data = {"response": response}
else:
if "response" not in request_data:
request_data["response"] = response
request_data = self._prepare_request_data(
request_data, response, user_api_key_dict, key="response",
)
# Add user API key metadata with prefixed keys
if "litellm_metadata" not in request_data:
user_metadata = self.transform_user_api_key_dict_to_metadata(
user_api_key_dict
)
if user_metadata:
request_data["litellm_metadata"] = user_metadata
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check
# Include model information from the response if available
response_model = None
if isinstance(response, dict):
response_model = response.get("model")
elif hasattr(response, "model"):
response_model = getattr(response, "model", None)
if response_model:
inputs["model"] = response_model
inputs = self._build_guardrail_inputs(
texts_to_check, images_to_check, tool_calls_to_check, response,
)
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
@@ -440,6 +384,95 @@ class AnthropicMessagesHandler(BaseTranslation):
)
return responses_so_far
def _prepare_request_data(
self,
request_data: Optional[dict],
response: Any,
user_api_key_dict: Optional[Any],
key: str,
) -> dict:
"""Ensure request_data has the response/responses_so_far key and metadata."""
if request_data is None:
request_data = {key: response}
else:
if key not in request_data:
request_data[key] = response
if "litellm_metadata" not in request_data:
user_metadata = self.transform_user_api_key_dict_to_metadata(
user_api_key_dict
)
if user_metadata:
request_data["litellm_metadata"] = user_metadata
return request_data
@staticmethod
def _get_response_content(response: Any) -> List[Any]:
"""Extract content list from a dict or object response."""
if isinstance(response, dict):
return response.get("content", []) or []
elif hasattr(response, "content"):
return getattr(response, "content", None) or []
return []
def _extract_from_content_blocks(
self,
response_content: List[Any],
texts_to_check: List[str],
images_to_check: List[str],
task_mappings: List[Tuple[int, Optional[int]]],
tool_calls_to_check: List["ChatCompletionToolCallChunk"],
) -> None:
"""Extract text, images, and tool calls from content blocks."""
for content_idx, content_block in enumerate(response_content):
block_dict: Dict[str, Any] = {}
if isinstance(content_block, dict):
block_type = content_block.get("type")
block_dict = cast(Dict[str, Any], content_block)
elif hasattr(content_block, "type"):
block_type = getattr(content_block, "type", None)
if hasattr(content_block, "model_dump"):
block_dict = content_block.model_dump()
else:
block_dict = {
"type": block_type,
"text": getattr(content_block, "text", None),
}
else:
continue
if block_type in ["text", "tool_use"]:
self._extract_output_text_and_images(
content_block=block_dict,
content_idx=content_idx,
texts_to_check=texts_to_check,
images_to_check=images_to_check,
task_mappings=task_mappings,
tool_calls_to_check=tool_calls_to_check,
)
@staticmethod
def _build_guardrail_inputs(
texts_to_check: List[str],
images_to_check: List[str],
tool_calls_to_check: List["ChatCompletionToolCallChunk"],
response: Any,
) -> "GenericGuardrailAPIInputs":
"""Build GenericGuardrailAPIInputs with optional images, tool calls, model."""
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check
response_model = None
if isinstance(response, dict):
response_model = response.get("model")
elif hasattr(response, "model"):
response_model = getattr(response, "model", None)
if response_model:
inputs["model"] = response_model
return inputs
def get_streaming_string_so_far(self, responses_so_far: List[Any]) -> str:
"""
Parse streaming responses and extract accumulated text content.
litellm/llms/bedrock/chat/converse_transformation.py
+94 -267
View File
@@ -91,34 +91,6 @@ UNSUPPORTED_BEDROCK_CONVERSE_BETA_PATTERNS = [
"compact-2026-01-12", # The compact beta feature is not currently supported on the Converse and ConverseStream APIs
]
# Models that support Bedrock's native structured outputs API (outputConfig.textFormat)
# Uses substring matching against the Bedrock model ID
# Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html
BEDROCK_NATIVE_STRUCTURED_OUTPUT_MODELS = {
# Anthropic Claude 4.5+
"claude-haiku-4-5",
"claude-sonnet-4-5",
"claude-opus-4-5",
"claude-opus-4-6",
# Qwen3
"qwen3",
# DeepSeek
"deepseek-v3.1",
# Gemma 3
"gemma-3",
# MiniMax
"minimax-m2",
# Mistral (magistral-small excluded: broken constrained decoding on Bedrock)
"ministral",
"mistral-large-3",
"voxtral",
# Moonshot
"kimi-k2",
# NVIDIA
"nemotron-nano",
# OpenAI (gpt-oss excluded: broken constrained decoding, works via tool-call fallback)
}
class AmazonConverseConfig(BaseConfig):
"""
@@ -188,8 +160,7 @@ class AmazonConverseConfig(BaseConfig):
if isinstance(content, list):
has_guarded_text = any(
isinstance(item, dict) and item.get("type") == "guarded_text"
for item in content
isinstance(item, dict) and item.get("type") == "guarded_text" for item in content
)
if has_guarded_text:
continue # Skip this message if it already has guarded_text
@@ -350,13 +321,9 @@ class AmazonConverseConfig(BaseConfig):
# Check if the model is a Nova 2 model (matches nova-2-lite, nova-2-pro, etc.)
# Also check for nova-2/ spec prefix for imported models
return model_without_region.startswith(
"amazon.nova-2-"
) or model_without_region.startswith("nova-2/")
return model_without_region.startswith("amazon.nova-2-") or model_without_region.startswith("nova-2/")
def _map_web_search_options(
self, web_search_options: dict, model: str
) -> Optional[BedrockToolBlock]:
def _map_web_search_options(self, web_search_options: dict, model: str) -> Optional[BedrockToolBlock]:
"""
Map web_search_options to Nova grounding systemTool.
@@ -385,9 +352,7 @@ class AmazonConverseConfig(BaseConfig):
# (unlike Anthropic), so we just enable grounding with no options
return BedrockToolBlock(systemTool={"name": "nova_grounding"})
def _transform_reasoning_effort_to_reasoning_config(
self, reasoning_effort: str
) -> dict:
def _transform_reasoning_effort_to_reasoning_config(self, reasoning_effort: str) -> dict:
"""
Transform reasoning_effort parameter to Nova 2 reasoningConfig structure.
@@ -432,9 +397,7 @@ class AmazonConverseConfig(BaseConfig):
}
}
def _handle_reasoning_effort_parameter(
self, model: str, reasoning_effort: str, optional_params: dict
) -> None:
def _handle_reasoning_effort_parameter(self, model: str, reasoning_effort: str, optional_params: dict) -> None:
"""
Handle the reasoning_effort parameter based on the model type.
@@ -471,9 +434,7 @@ class AmazonConverseConfig(BaseConfig):
optional_params["reasoning_effort"] = reasoning_effort
elif self._is_nova_2_model(model):
# Nova 2 models: transform to reasoningConfig
reasoning_config = self._transform_reasoning_effort_to_reasoning_config(
reasoning_effort
)
reasoning_config = self._transform_reasoning_effort_to_reasoning_config(reasoning_effort)
optional_params.update(reasoning_config)
else:
# Anthropic and other models: convert to thinking parameter
@@ -493,8 +454,7 @@ class AmazonConverseConfig(BaseConfig):
budget = thinking.get("budget_tokens")
if isinstance(budget, int) and budget < BEDROCK_MIN_THINKING_BUDGET_TOKENS:
verbose_logger.debug(
"Bedrock requires thinking.budget_tokens >= %d, got %d. "
"Clamping to minimum.",
"Bedrock requires thinking.budget_tokens >= %d, got %d. Clamping to minimum.",
BEDROCK_MIN_THINKING_BUDGET_TOKENS,
budget,
)
@@ -518,9 +478,7 @@ class AmazonConverseConfig(BaseConfig):
"parallel_tool_calls",
]
if (
"arn" in model
): # we can't infer the model from the arn, so just add all params
if "arn" in model: # we can't infer the model from the arn, so just add all params
supported_params.append("tools")
supported_params.append("tool_choice")
supported_params.append("thinking")
@@ -542,9 +500,7 @@ class AmazonConverseConfig(BaseConfig):
or base_model.startswith("meta.llama3-3")
or base_model.startswith("meta.llama4")
or base_model.startswith("amazon.nova")
or supports_function_calling(
model=model, custom_llm_provider=self.custom_llm_provider
)
or supports_function_calling(model=model, custom_llm_provider=self.custom_llm_provider)
):
supported_params.append("tools")
@@ -554,9 +510,7 @@ class AmazonConverseConfig(BaseConfig):
if litellm.utils.supports_tool_choice(
model=model, custom_llm_provider=self.custom_llm_provider
) or litellm.utils.supports_tool_choice(
model=base_model, custom_llm_provider=self.custom_llm_provider
):
) or litellm.utils.supports_tool_choice(model=base_model, custom_llm_provider=self.custom_llm_provider):
# only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
supported_params.append("tool_choice")
@@ -575,9 +529,7 @@ class AmazonConverseConfig(BaseConfig):
model=model,
custom_llm_provider=self.custom_llm_provider,
)
or supports_reasoning(
model=base_model, custom_llm_provider=self.custom_llm_provider
)
or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider)
):
supported_params.append("thinking")
supported_params.append("reasoning_effort")
@@ -602,9 +554,7 @@ class AmazonConverseConfig(BaseConfig):
return ToolChoiceValuesBlock(auto={})
elif isinstance(tool_choice, dict):
# only supported for anthropic + mistral models - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
specific_tool = SpecificToolChoiceBlock(
name=tool_choice.get("function", {}).get("name", "")
)
specific_tool = SpecificToolChoiceBlock(name=tool_choice.get("function", {}).get("name", ""))
return ToolChoiceValuesBlock(tool=specific_tool)
else:
raise litellm.utils.UnsupportedParamsError(
@@ -624,15 +574,9 @@ class AmazonConverseConfig(BaseConfig):
return ["mp4", "mov", "mkv", "webm", "flv", "mpeg", "mpg", "wmv", "3gp"]
def get_all_supported_content_types(self) -> List[str]:
return (
self.get_supported_image_types()
+ self.get_supported_document_types()
+ self.get_supported_video_types()
)
return self.get_supported_image_types() + self.get_supported_document_types() + self.get_supported_video_types()
def is_computer_use_tool_used(
self, tools: Optional[List[OpenAIChatCompletionToolParam]], model: str
) -> bool:
def is_computer_use_tool_used(self, tools: Optional[List[OpenAIChatCompletionToolParam]], model: str) -> bool:
"""Check if computer use tools are being used in the request."""
if tools is None:
return False
@@ -645,9 +589,7 @@ class AmazonConverseConfig(BaseConfig):
return True
return False
def _transform_computer_use_tools(
self, computer_use_tools: List[OpenAIChatCompletionToolParam]
) -> List[dict]:
def _transform_computer_use_tools(self, computer_use_tools: List[OpenAIChatCompletionToolParam]) -> List[dict]:
"""Transform computer use tools to Bedrock format."""
transformed_tools: List[dict] = []
@@ -689,9 +631,7 @@ class AmazonConverseConfig(BaseConfig):
def _separate_computer_use_tools(
self, tools: List[OpenAIChatCompletionToolParam], model: str
) -> Tuple[
List[OpenAIChatCompletionToolParam], List[OpenAIChatCompletionToolParam]
]:
) -> Tuple[List[OpenAIChatCompletionToolParam], List[OpenAIChatCompletionToolParam]]:
"""
Separate computer use tools from regular function tools.
@@ -763,10 +703,20 @@ class AmazonConverseConfig(BaseConfig):
return _tool
@staticmethod
def _supports_native_structured_outputs(model: str) -> bool:
"""Check if the Bedrock model supports native structured outputs (outputConfig.textFormat)."""
return any(
substring in model for substring in BEDROCK_NATIVE_STRUCTURED_OUTPUT_MODELS
def _supports_native_structured_outputs(
model: str, custom_llm_provider: Optional[str] = None
) -> bool:
"""Check if the Bedrock model supports native structured outputs (outputConfig.textFormat).
Delegates to the standard ``supports_native_structured_output`` utility
which looks up the flag in ``litellm.model_cost`` via
``_get_model_info_helper``.
Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html
"""
from litellm.utils import supports_native_structured_output
return supports_native_structured_output(
model=model, custom_llm_provider=custom_llm_provider
)
@staticmethod
@@ -790,25 +740,18 @@ class AmazonConverseConfig(BaseConfig):
# Recurse into nested schemas
if "properties" in result and isinstance(result["properties"], dict):
result["properties"] = {
k: AmazonConverseConfig._add_additional_properties_to_schema(v)
for k, v in result["properties"].items()
k: AmazonConverseConfig._add_additional_properties_to_schema(v) for k, v in result["properties"].items()
}
if "items" in result and isinstance(result["items"], dict):
result["items"] = AmazonConverseConfig._add_additional_properties_to_schema(
result["items"]
)
result["items"] = AmazonConverseConfig._add_additional_properties_to_schema(result["items"])
for defs_key in ("$defs", "definitions"):
if defs_key in result and isinstance(result[defs_key], dict):
result[defs_key] = {
k: AmazonConverseConfig._add_additional_properties_to_schema(v)
for k, v in result[defs_key].items()
k: AmazonConverseConfig._add_additional_properties_to_schema(v) for k, v in result[defs_key].items()
}
for key in ("anyOf", "allOf", "oneOf"):
if key in result and isinstance(result[key], list):
result[key] = [
AmazonConverseConfig._add_additional_properties_to_schema(item)
for item in result[key]
]
result[key] = [AmazonConverseConfig._add_additional_properties_to_schema(item) for item in result[key]]
return result
@@ -838,9 +781,7 @@ class AmazonConverseConfig(BaseConfig):
}
"""
if json_schema is not None:
json_schema = AmazonConverseConfig._add_additional_properties_to_schema(
json_schema
)
json_schema = AmazonConverseConfig._add_additional_properties_to_schema(json_schema)
schema_str = json.dumps(json_schema) if json_schema is not None else "{}"
json_schema_def: JsonSchemaDefinition = {"schema": schema_str}
if name is not None:
@@ -862,14 +803,9 @@ class AmazonConverseConfig(BaseConfig):
non_default_params: dict,
optional_params: dict,
):
optional_params = self._add_tools_to_optional_params(
optional_params=optional_params, tools=tools
)
optional_params = self._add_tools_to_optional_params(optional_params=optional_params, tools=tools)
if (
"meta.llama3-3-70b-instruct-v1:0" in model
and non_default_params.get("stream", False) is True
):
if "meta.llama3-3-70b-instruct-v1:0" in model and non_default_params.get("stream", False) is True:
optional_params["fake_stream"] = True
def map_openai_params(
@@ -913,7 +849,9 @@ class AmazonConverseConfig(BaseConfig):
)
if param == "tool_choice":
_tool_choice_value = self.map_tool_choice_values(
model=model, tool_choice=value, drop_params=drop_params # type: ignore
model=model,
tool_choice=value,
drop_params=drop_params, # type: ignore
)
if _tool_choice_value is not None:
optional_params["tool_choice"] = _tool_choice_value
@@ -1006,7 +944,7 @@ class AmazonConverseConfig(BaseConfig):
if "type" in value and value["type"] == "text":
return optional_params
if self._supports_native_structured_outputs(model) and json_schema is not None:
if self._supports_native_structured_outputs(model, self.custom_llm_provider) and json_schema is not None:
# Use Bedrock's native structured outputs API (outputConfig.textFormat)
# No synthetic tool injection, no fake_stream needed.
# Requires an explicit schema — json_object with no schema falls through
@@ -1024,14 +962,10 @@ class AmazonConverseConfig(BaseConfig):
json_schema=json_schema,
description=description,
)
optional_params = self._add_tools_to_optional_params(
optional_params=optional_params, tools=[_tool]
)
optional_params = self._add_tools_to_optional_params(optional_params=optional_params, tools=[_tool])
if (
litellm.utils.supports_tool_choice(
model=model, custom_llm_provider=self.custom_llm_provider
)
litellm.utils.supports_tool_choice(model=model, custom_llm_provider=self.custom_llm_provider)
and not is_thinking_enabled
):
optional_params["tool_choice"] = ToolChoiceValuesBlock(
@@ -1043,9 +977,7 @@ class AmazonConverseConfig(BaseConfig):
optional_params["json_mode"] = True
return optional_params
def update_optional_params_with_thinking_tokens(
self, non_default_params: dict, optional_params: dict
):
def update_optional_params_with_thinking_tokens(self, non_default_params: dict, optional_params: dict):
"""
Handles scenario where max tokens is not specified. For anthropic models (anthropic api/bedrock/vertex ai), this requires having the max tokens being set and being greater than the thinking token budget.
@@ -1063,13 +995,9 @@ class AmazonConverseConfig(BaseConfig):
is_thinking_enabled = self.is_thinking_enabled(optional_params)
is_max_tokens_in_request = self.is_max_tokens_in_request(non_default_params)
if is_thinking_enabled and not is_max_tokens_in_request:
thinking_token_budget = cast(dict, optional_params["thinking"]).get(
"budget_tokens", None
)
thinking_token_budget = cast(dict, optional_params["thinking"]).get("budget_tokens", None)
if thinking_token_budget is not None:
optional_params["maxTokens"] = (
thinking_token_budget + DEFAULT_MAX_TOKENS
)
optional_params["maxTokens"] = thinking_token_budget + DEFAULT_MAX_TOKENS
@overload
def _get_cache_point_block(
@@ -1135,23 +1063,15 @@ class AmazonConverseConfig(BaseConfig):
if message["role"] == "system":
system_prompt_indices.append(idx)
if isinstance(message["content"], str) and message["content"]:
system_content_blocks.append(
SystemContentBlock(text=message["content"])
)
cache_block = self._get_cache_point_block(
message, block_type="system", model=model
)
system_content_blocks.append(SystemContentBlock(text=message["content"]))
cache_block = self._get_cache_point_block(message, block_type="system", model=model)
if cache_block:
system_content_blocks.append(cache_block)
elif isinstance(message["content"], list):
for m in message["content"]:
if m.get("type") == "text" and m.get("text"):
system_content_blocks.append(
SystemContentBlock(text=m["text"])
)
cache_block = self._get_cache_point_block(
m, block_type="system", model=model
)
system_content_blocks.append(SystemContentBlock(text=m["text"]))
cache_block = self._get_cache_point_block(m, block_type="system", model=model)
if cache_block:
system_content_blocks.append(cache_block)
if len(system_prompt_indices) > 0:
@@ -1189,16 +1109,10 @@ class AmazonConverseConfig(BaseConfig):
# Exceptions should not be stored in optional_params (this is a defensive fix)
cleaned_params = filter_exceptions_from_params(optional_params)
inference_params = safe_deep_copy(cleaned_params)
supported_converse_params = list(
AmazonConverseConfig.__annotations__.keys()
) + ["top_k"]
supported_converse_params = list(AmazonConverseConfig.__annotations__.keys()) + ["top_k"]
supported_tool_call_params = ["tools", "tool_choice"]
supported_config_params = list(self.get_config_blocks().keys())
total_supported_params = (
supported_converse_params
+ supported_tool_call_params
+ supported_config_params
)
total_supported_params = supported_converse_params + supported_tool_call_params + supported_config_params
inference_params.pop("json_mode", None) # used for handling json_schema
# Anthropic-only key. Bedrock expects `outputConfig` (camelCase) and
# will reject `output_config` if it leaks through pass-through routes.
@@ -1209,25 +1123,15 @@ class AmazonConverseConfig(BaseConfig):
if request_metadata is not None:
self._validate_request_metadata(request_metadata)
output_config: Optional[OutputConfigBlock] = inference_params.pop(
"outputConfig", None
)
inference_params.pop(
"output_config", None
) # Bedrock Converse doesn't support it
output_config: Optional[OutputConfigBlock] = inference_params.pop("outputConfig", None)
inference_params.pop("output_config", None) # Bedrock Converse doesn't support it
# keep supported params in 'inference_params', and set all model-specific params in 'additional_request_params'
additional_request_params = {
k: v for k, v in inference_params.items() if k not in total_supported_params
}
inference_params = {
k: v for k, v in inference_params.items() if k in total_supported_params
}
additional_request_params = {k: v for k, v in inference_params.items() if k not in total_supported_params}
inference_params = {k: v for k, v in inference_params.items() if k in total_supported_params}
# Handle parallel_tool_calls configuration
parallel_tool_use_config = additional_request_params.pop(
"_parallel_tool_use_config", None
)
parallel_tool_use_config = additional_request_params.pop("_parallel_tool_use_config", None)
if parallel_tool_use_config is not None and is_claude_4_5_on_bedrock(model):
for key, value in parallel_tool_use_config.items():
if (
@@ -1242,9 +1146,7 @@ class AmazonConverseConfig(BaseConfig):
additional_request_params.pop("parallel_tool_calls", None)
# Only set the topK value in for models that support it
additional_request_params.update(
self._handle_top_k_value(model, inference_params)
)
additional_request_params.update(self._handle_top_k_value(model, inference_params))
# Filter out internal/MCP-related parameters that shouldn't be sent to the API
# These are LiteLLM internal parameters, not API parameters
@@ -1253,9 +1155,7 @@ class AmazonConverseConfig(BaseConfig):
# Filter out non-serializable objects (exceptions, callables, logging objects, etc.)
# from additional_request_params to prevent JSON serialization errors
# This filters: Exception objects, callable objects (functions), Logging objects, etc.
additional_request_params = filter_exceptions_from_params(
additional_request_params
)
additional_request_params = filter_exceptions_from_params(additional_request_params)
return (
inference_params,
@@ -1302,9 +1202,7 @@ class AmazonConverseConfig(BaseConfig):
# Only separate tools if computer use tools are actually present
if filtered_tools and self.is_computer_use_tool_used(filtered_tools, model):
# Separate computer use tools from regular function tools
computer_use_tools, regular_tools = self._separate_computer_use_tools(
filtered_tools, model
)
computer_use_tools, regular_tools = self._separate_computer_use_tools(filtered_tools, model)
# Process regular function tools using existing logic
bedrock_tools = _bedrock_tools_pt(regular_tools)
@@ -1365,9 +1263,7 @@ class AmazonConverseConfig(BaseConfig):
anthropic_beta_list.append(computer_use_header)
# Transform computer use tools to proper Bedrock format
transformed_computer_tools = self._transform_computer_use_tools(
computer_use_tools
)
transformed_computer_tools = self._transform_computer_use_tools(computer_use_tools)
additional_request_params["tools"] = transformed_computer_tools
else:
# No computer use tools, process all tools as regular tools
@@ -1396,15 +1292,9 @@ class AmazonConverseConfig(BaseConfig):
"""
Bedrock doesn't support tool calling without `tools=` param specified.
"""
if (
"tools" not in optional_params
and messages is not None
and has_tool_call_blocks(messages)
):
if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages):
if litellm.modify_params:
optional_params["tools"] = add_dummy_tool(
custom_llm_provider="bedrock_converse"
)
optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse")
else:
raise litellm.UnsupportedParamsError(
message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
@@ -1458,9 +1348,7 @@ class AmazonConverseConfig(BaseConfig):
bedrock_tool_config: Optional[ToolConfigBlock] = None
if len(bedrock_tools) > 0:
tool_choice_values: ToolChoiceValuesBlock = inference_params.pop(
"tool_choice", None
)
tool_choice_values: ToolChoiceValuesBlock = inference_params.pop("tool_choice", None)
bedrock_tool_config = ToolConfigBlock(
tools=bedrock_tools,
)
@@ -1470,9 +1358,7 @@ class AmazonConverseConfig(BaseConfig):
data: CommonRequestObject = {
"additionalModelRequestFields": additional_request_params,
"system": system_content_blocks,
"inferenceConfig": self._transform_inference_params(
inference_params=inference_params
),
"inferenceConfig": self._transform_inference_params(inference_params=inference_params),
}
# Handle all config blocks
@@ -1502,14 +1388,10 @@ class AmazonConverseConfig(BaseConfig):
litellm_params: dict,
headers: Optional[dict] = None,
) -> RequestObject:
messages, system_content_blocks = self._transform_system_message(
messages, model=model
)
messages, system_content_blocks = self._transform_system_message(messages, model=model)
# Convert last user message to guarded_text if guardrailConfig is present
messages = self._convert_consecutive_user_messages_to_guarded_text(
messages, optional_params
)
messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
## TRANSFORMATION ##
_data: CommonRequestObject = self._transform_request_helper(
@@ -1520,13 +1402,11 @@ class AmazonConverseConfig(BaseConfig):
headers=headers,
)
bedrock_messages = (
await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model=model,
llm_provider="bedrock_converse",
user_continue_message=litellm_params.pop("user_continue_message", None),
)
bedrock_messages = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model=model,
llm_provider="bedrock_converse",
user_continue_message=litellm_params.pop("user_continue_message", None),
)
data: RequestObject = {"messages": bedrock_messages, **_data}
@@ -1560,14 +1440,10 @@ class AmazonConverseConfig(BaseConfig):
litellm_params: dict,
headers: Optional[dict] = None,
) -> RequestObject:
messages, system_content_blocks = self._transform_system_message(
messages, model=model
)
messages, system_content_blocks = self._transform_system_message(messages, model=model)
# Convert last user message to guarded_text if guardrailConfig is present
messages = self._convert_consecutive_user_messages_to_guarded_text(
messages, optional_params
)
messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
_data: CommonRequestObject = self._transform_request_helper(
model=model,
@@ -1616,9 +1492,7 @@ class AmazonConverseConfig(BaseConfig):
encoding=encoding,
)
def _transform_reasoning_content(
self, reasoning_content_blocks: List[BedrockConverseReasoningContentBlock]
) -> str:
def _transform_reasoning_content(self, reasoning_content_blocks: List[BedrockConverseReasoningContentBlock]) -> str:
"""
Extract the reasoning text from the reasoning content blocks
@@ -1634,9 +1508,7 @@ class AmazonConverseConfig(BaseConfig):
self, thinking_blocks: List[BedrockConverseReasoningContentBlock]
) -> List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]:
"""Return a consistent format for thinking blocks between Anthropic and Bedrock."""
thinking_blocks_list: List[
Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]
] = []
thinking_blocks_list: List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] = []
for block in thinking_blocks:
if "reasoningText" in block:
_thinking_block = ChatCompletionThinkingBlock(type="thinking")
@@ -1672,21 +1544,11 @@ class AmazonConverseConfig(BaseConfig):
cache_creation_input_tokens = usage["cacheWriteInputTokens"]
input_tokens += cache_creation_input_tokens
prompt_tokens_details = PromptTokensDetailsWrapper(
cached_tokens=cache_read_input_tokens
)
reasoning_tokens = (
token_counter(text=reasoning_content, count_response_tokens=True)
if reasoning_content
else 0
)
prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=cache_read_input_tokens)
reasoning_tokens = token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0
completion_tokens_details = CompletionTokensDetailsWrapper(
reasoning_tokens=reasoning_tokens,
text_tokens=(
output_tokens - reasoning_tokens
if reasoning_tokens > 0
else output_tokens
),
text_tokens=(output_tokens - reasoning_tokens if reasoning_tokens > 0 else output_tokens),
)
openai_usage = Usage(
prompt_tokens=input_tokens,
@@ -1701,9 +1563,7 @@ class AmazonConverseConfig(BaseConfig):
def get_tool_call_names(
self,
tools: Optional[
Union[List[ToolBlock], List[OpenAIChatCompletionToolParam]]
] = None,
tools: Optional[Union[List[ToolBlock], List[OpenAIChatCompletionToolParam]]] = None,
) -> List[str]:
if tools is None:
return []
@@ -1742,13 +1602,8 @@ class AmazonConverseConfig(BaseConfig):
try:
tool_call_names = self.get_tool_call_names(tools)
json_content = json.loads(message.content)
if (
json_content.get("type") == "function"
and json_content.get("name") in tool_call_names
):
tool_calls = [
ChatCompletionMessageToolCall(function=Function(**json_content))
]
if json_content.get("type") == "function" and json_content.get("name") in tool_call_names:
tool_calls = [ChatCompletionMessageToolCall(function=Function(**json_content))]
message.tool_calls = tool_calls
message.content = None
@@ -1777,9 +1632,7 @@ class AmazonConverseConfig(BaseConfig):
"""
content_str = ""
tools: List[ChatCompletionToolCallChunk] = []
reasoningContentBlocks: Optional[
List[BedrockConverseReasoningContentBlock]
] = None
reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = None
citationsContentBlocks: Optional[List[CitationsContentBlock]] = None
for idx, content in enumerate(content_blocks):
"""
@@ -1796,9 +1649,7 @@ class AmazonConverseConfig(BaseConfig):
if "toolUse" in content:
## check tool name was formatted by litellm
_response_tool_name = content["toolUse"]["name"]
response_tool_name = get_bedrock_tool_name(
response_tool_name=_response_tool_name
)
response_tool_name = get_bedrock_tool_name(response_tool_name=_response_tool_name)
_function_chunk = ChatCompletionToolCallFunctionChunk(
name=response_tool_name,
arguments=json.dumps(content["toolUse"]["input"]),
@@ -1849,11 +1700,7 @@ class AmazonConverseConfig(BaseConfig):
"""
try:
response_data = json.loads(json_str)
if (
isinstance(response_data, dict)
and "properties" in response_data
and len(response_data) == 1
):
if isinstance(response_data, dict) and "properties" in response_data and len(response_data) == 1:
response_data = response_data["properties"]
return json.dumps(response_data)
except json.JSONDecodeError:
@@ -1877,11 +1724,7 @@ class AmazonConverseConfig(BaseConfig):
if not json_mode or not tools:
return tools if tools else None
json_tool_indices = [
i
for i, t in enumerate(tools)
if t["function"].get("name") == RESPONSE_FORMAT_TOOL_NAME
]
json_tool_indices = [i for i, t in enumerate(tools) if t["function"].get("name") == RESPONSE_FORMAT_TOOL_NAME]
if not json_tool_indices:
# No json_tool_call found, return tools unchanged
@@ -1889,14 +1732,10 @@ class AmazonConverseConfig(BaseConfig):
if len(json_tool_indices) == len(tools):
# All tools are json_tool_call — convert first one to content
verbose_logger.debug(
"Processing JSON tool call response for response_format"
)
verbose_logger.debug("Processing JSON tool call response for response_format")
json_mode_content_str: Optional[str] = tools[0]["function"].get("arguments")
if json_mode_content_str is not None:
json_mode_content_str = AmazonConverseConfig._unwrap_bedrock_properties(
json_mode_content_str
)
json_mode_content_str = AmazonConverseConfig._unwrap_bedrock_properties(json_mode_content_str)
chat_completion_message["content"] = json_mode_content_str
return None
@@ -1906,13 +1745,9 @@ class AmazonConverseConfig(BaseConfig):
first_idx = json_tool_indices[0]
json_mode_args = tools[first_idx]["function"].get("arguments")
if json_mode_args is not None:
json_mode_args = AmazonConverseConfig._unwrap_bedrock_properties(
json_mode_args
)
json_mode_args = AmazonConverseConfig._unwrap_bedrock_properties(json_mode_args)
existing = chat_completion_message.get("content") or ""
chat_completion_message["content"] = (
existing + json_mode_args if existing else json_mode_args
)
chat_completion_message["content"] = existing + json_mode_args if existing else json_mode_args
real_tools = [t for i, t in enumerate(tools) if i not in json_tool_indices]
return real_tools if real_tools else None
@@ -1990,9 +1825,7 @@ class AmazonConverseConfig(BaseConfig):
chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"}
content_str = ""
tools: List[ChatCompletionToolCallChunk] = []
reasoningContentBlocks: Optional[
List[BedrockConverseReasoningContentBlock]
] = None
reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = None
citationsContentBlocks: Optional[List[CitationsContentBlock]] = None
if message is not None:
@@ -2011,17 +1844,11 @@ class AmazonConverseConfig(BaseConfig):
provider_specific_fields["citationsContent"] = citationsContentBlocks
if provider_specific_fields:
chat_completion_message[
"provider_specific_fields"
] = provider_specific_fields
chat_completion_message["provider_specific_fields"] = provider_specific_fields
if reasoningContentBlocks is not None:
chat_completion_message[
"reasoning_content"
] = self._transform_reasoning_content(reasoningContentBlocks)
chat_completion_message[
"thinking_blocks"
] = self._transform_thinking_blocks(reasoningContentBlocks)
chat_completion_message["reasoning_content"] = self._transform_reasoning_content(reasoningContentBlocks)
chat_completion_message["thinking_blocks"] = self._transform_thinking_blocks(reasoningContentBlocks)
chat_completion_message["content"] = content_str
filtered_tools = self._filter_json_mode_tools(
json_mode=json_mode,
litellm/model_prices_and_context_window_backup.json
+88 -56
View File
@@ -722,7 +722,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"anthropic.claude-haiku-4-5@20251001": {
"cache_creation_input_token_cost": 1.25e-06,
@@ -745,7 +746,8 @@
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_streaming": true
"supports_native_streaming": true,
"supports_native_structured_output": true
},
"anthropic.claude-3-5-sonnet-20240620-v1:0": {
"input_cost_per_token": 3e-06,
@@ -967,7 +969,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_native_structured_output": true
},
"anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.25e-06,
@@ -993,7 +996,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"global.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.25e-06,
@@ -1019,7 +1023,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"us.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.875e-06,
@@ -1045,7 +1050,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"eu.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.875e-06,
@@ -1071,7 +1077,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"au.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.875e-06,
@@ -1097,7 +1104,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 3.75e-06,
@@ -1123,7 +1131,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"global.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 3.75e-06,
@@ -1149,7 +1158,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"us.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 4.125e-06,
@@ -1175,7 +1185,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"eu.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 4.125e-06,
@@ -1201,7 +1212,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"au.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 4.125e-06,
@@ -1227,7 +1239,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"anthropic.claude-sonnet-4-20250514-v1:0": {
"cache_creation_input_token_cost": 3.75e-06,
@@ -1287,7 +1300,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_native_structured_output": true
},
"anthropic.claude-v1": {
"input_cost_per_token": 8e-06,
@@ -1537,7 +1551,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"apac.anthropic.claude-3-sonnet-20240229-v1:0": {
"input_cost_per_token": 3e-06,
@@ -1625,7 +1640,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"azure/ada": {
"input_cost_per_token": 1e-07,
@@ -11690,7 +11706,8 @@
"output_cost_per_token": 1.68e-06,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_native_structured_output": true
},
"deepseek.v3.2": {
"input_cost_per_token": 6.2e-07,
@@ -12135,7 +12152,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"eu.anthropic.claude-3-5-sonnet-20240620-v1:0": {
"input_cost_per_token": 3e-06,
@@ -12349,7 +12367,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"eu.meta.llama3-2-1b-instruct-v1:0": {
"input_cost_per_token": 1.3e-07,
@@ -14579,18 +14598,6 @@
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
"uses_embed_content": true
},
"vertex_ai/gemini-embedding-2-preview": {
"input_cost_per_token": 1.5e-07,
"litellm_provider": "vertex_ai",
"max_input_tokens": 8192,
"max_tokens": 8192,
"mode": "embedding",
"output_cost_per_token": 0,
"output_vector_size": 3072,
"source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal",
"supports_multimodal": true,
"uses_embed_content": true
},
"gemini/gemini-embedding-001": {
"input_cost_per_token": 1.5e-07,
"litellm_provider": "gemini",
@@ -16767,7 +16774,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"global.anthropic.claude-sonnet-4-20250514-v1:0": {
"cache_creation_input_token_cost": 3.75e-06,
@@ -16819,7 +16827,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"global.amazon.nova-2-lite-v1:0": {
"cache_read_input_token_cost": 7.5e-08,
@@ -20548,7 +20557,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"jp.anthropic.claude-haiku-4-5-20251001-v1:0": {
"cache_creation_input_token_cost": 1.375e-06,
@@ -20570,7 +20580,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"lambda_ai/deepseek-llama3.3-70b": {
"input_cost_per_token": 2e-07,
@@ -21265,7 +21276,8 @@
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"minimax.minimax-m2.1": {
"input_cost_per_token": 3e-07,
@@ -21421,7 +21433,8 @@
"mode": "chat",
"output_cost_per_token": 2e-07,
"supports_function_calling": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral.ministral-3-3b-instruct": {
"input_cost_per_token": 1e-07,
@@ -21432,7 +21445,8 @@
"mode": "chat",
"output_cost_per_token": 1e-07,
"supports_function_calling": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral.ministral-3-8b-instruct": {
"input_cost_per_token": 1.5e-07,
@@ -21443,7 +21457,8 @@
"mode": "chat",
"output_cost_per_token": 1.5e-07,
"supports_function_calling": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral.mistral-7b-instruct-v0:2": {
"input_cost_per_token": 1.5e-07,
@@ -21485,7 +21500,8 @@
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"supports_function_calling": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral.mistral-small-2402-v1:0": {
"input_cost_per_token": 1e-06,
@@ -21516,7 +21532,8 @@
"mode": "chat",
"output_cost_per_token": 4e-08,
"supports_audio_input": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral.voxtral-small-24b-2507": {
"input_cost_per_token": 1e-07,
@@ -21527,7 +21544,8 @@
"mode": "chat",
"output_cost_per_token": 3e-07,
"supports_audio_input": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral/codestral-2405": {
"input_cost_per_token": 1e-06,
@@ -22214,7 +22232,8 @@
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"supports_reasoning": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"moonshotai.kimi-k2.5": {
"input_cost_per_token": 6e-07,
@@ -23089,7 +23108,8 @@
"supports_function_calling": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_native_structured_output": true
},
"o1": {
"cache_read_input_token_cost": 7.5e-06,
@@ -26173,7 +26193,8 @@
"output_cost_per_token": 1.8e-06,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_native_structured_output": true
},
"qwen.qwen3-235b-a22b-2507-v1:0": {
"input_cost_per_token": 2.2e-07,
@@ -26185,7 +26206,8 @@
"output_cost_per_token": 8.8e-07,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_native_structured_output": true
},
"qwen.qwen3-coder-30b-a3b-v1:0": {
"input_cost_per_token": 1.5e-07,
@@ -26197,7 +26219,8 @@
"output_cost_per_token": 6e-07,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_native_structured_output": true
},
"qwen.qwen3-32b-v1:0": {
"input_cost_per_token": 1.5e-07,
@@ -26209,7 +26232,8 @@
"output_cost_per_token": 6e-07,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_native_structured_output": true
},
"qwen.qwen3-next-80b-a3b": {
"input_cost_per_token": 1.5e-07,
@@ -26220,7 +26244,8 @@
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"supports_function_calling": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"qwen.qwen3-vl-235b-a22b": {
"input_cost_per_token": 5.3e-07,
@@ -26232,7 +26257,8 @@
"output_cost_per_token": 2.66e-06,
"supports_function_calling": true,
"supports_system_messages": true,
"supports_vision": true
"supports_vision": true,
"supports_native_structured_output": true
},
"qwen.qwen3-coder-next": {
"input_cost_per_token": 5e-07,
@@ -28257,7 +28283,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"us.anthropic.claude-3-5-sonnet-20240620-v1:0": {
"input_cost_per_token": 3e-06,
@@ -28415,7 +28442,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"au.anthropic.claude-haiku-4-5-20251001-v1:0": {
"cache_creation_input_token_cost": 1.375e-06,
@@ -28436,7 +28464,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"us.anthropic.claude-opus-4-20250514-v1:0": {
"cache_creation_input_token_cost": 1.875e-05,
@@ -28488,7 +28517,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_native_structured_output": true
},
"global.anthropic.claude-opus-4-5-20251101-v1:0": {
"cache_creation_input_token_cost": 6.25e-06,
@@ -28514,7 +28544,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_native_structured_output": true
},
"eu.anthropic.claude-opus-4-5-20251101-v1:0": {
"cache_creation_input_token_cost": 6.25e-06,
@@ -28540,7 +28571,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_native_structured_output": true
},
"us.anthropic.claude-sonnet-4-20250514-v1:0": {
"cache_creation_input_token_cost": 3.75e-06,
litellm/proxy/auth/handle_jwt.py
+6 -2
View File
@@ -1397,8 +1397,12 @@ class JWTAuthManager:
request_headers: Optional[dict] = None,
) -> JWTAuthBuilderResult:
"""Main authentication and authorization builder"""
# Check if OIDC UserInfo endpoint is enabled
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled:
# Check if OIDC UserInfo endpoint is enabled, but fall back to standard
# JWT auth if the token itself is a well-formed JWT (3-part structure).
if (
jwt_handler.litellm_jwtauth.oidc_userinfo_enabled
and not jwt_handler.is_jwt(token=api_key)
):
verbose_proxy_logger.debug(
"OIDC UserInfo is enabled. Fetching user info from UserInfo endpoint."
)
litellm/proxy/auth/user_api_key_auth.py
+4 -1
View File
@@ -686,7 +686,10 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if jwt_handler.litellm_jwtauth.virtual_key_claim_field is not None:
# Decode JWT to get claims without running full auth_builder
jwt_claims: Optional[dict]
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled:
if (
jwt_handler.litellm_jwtauth.oidc_userinfo_enabled
and not jwt_handler.is_jwt(token=api_key)
):
jwt_claims = await jwt_handler.get_oidc_userinfo(token=api_key)
else:
jwt_claims = await jwt_handler.auth_jwt(token=api_key)
litellm/proxy/common_request_processing.py
+22 -2
View File
@@ -293,19 +293,20 @@ def _override_openai_response_model(
we preserve the actual model that was used (the fallback model).
2. If the request was to an Azure Model Router, we preserve the actual model
that was used (e.g., gpt-5-nano-2025-08-07) instead of the router model.
3. If this was a fastest_response batch completion, use the winning model's
model group name instead of the comma-separated list the client sent.
"""
if not requested_model:
return
# Check if a fallback occurred - if so, preserve the actual model used
hidden_params = getattr(response_obj, "_hidden_params", {}) or {}
if isinstance(hidden_params, dict):
# Check if a fallback occurred - if so, preserve the actual model used
fallback_headers = hidden_params.get("additional_headers", {}) or {}
attempted_fallbacks = fallback_headers.get(
"x-litellm-attempted-fallbacks", None
)
if attempted_fallbacks is not None and attempted_fallbacks > 0:
# A fallback occurred - preserve the actual model that was used
verbose_proxy_logger.debug(
"%s: fallback detected (attempted_fallbacks=%d), preserving actual model used instead of overriding to requested model.",
log_context,
@@ -313,6 +314,25 @@ def _override_openai_response_model(
)
return
# For fastest_response batch completions, use the winning model's group
# name rather than the comma-separated list the client sent.
if hidden_params.get("fastest_response_batch_completion"):
winning_model = fallback_headers.get("x-litellm-model-group")
if winning_model:
verbose_proxy_logger.debug(
"%s: fastest_response detected, using winning model group=%r instead of requested=%r.",
log_context,
winning_model,
requested_model,
)
requested_model = winning_model
else:
verbose_proxy_logger.debug(
"%s: fastest_response detected but no model group header found, preserving actual model from response.",
log_context,
)
return
# Check if this is an Azure Model Router request - if so, preserve the actual model used
if _is_azure_model_router_request(requested_model):
verbose_proxy_logger.debug(
litellm/proxy/proxy_server.py
+5
View File
@@ -5721,6 +5721,11 @@ def _restamp_streaming_chunk_model(
if _is_azure_model_router_request(requested_model_from_client):
return chunk, model_mismatch_logged
# For fastest_response batch completions, preserve the winning model's name
# instead of stamping the comma-separated list the client sent.
if request_data.get("fastest_response", False):
return chunk, model_mismatch_logged
downstream_model = (
chunk.get("model") if isinstance(chunk, dict) else getattr(chunk, "model", None)
)
litellm/proxy/spend_tracking/spend_tracking_utils.py
+1
View File
@@ -90,6 +90,7 @@ def _get_spend_logs_metadata(
user_api_key_alias=None,
user_api_key_team_id=None,
user_api_key_project_id=None,
user_api_key_project_alias=None,
user_api_key_org_id=None,
user_api_key_user_id=None,
user_api_key_team_alias=None,
litellm/responses/litellm_completion_transformation/streaming_iterator.py
+9
View File
@@ -845,6 +845,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
)
event.__dict__["sequence_number"] = self._sequence_number
self._pending_response_events.append(event)
# Emit content_part.added immediately after output_item.added for message
# items. The OpenAI Responses spec requires this event before any
# output_text.delta events so downstream parsers can initialize the
# text part structure.
if not self.sent_content_part_added_event:
self.sent_content_part_added_event = True
content_part_event = self.create_content_part_added_event()
self._pending_response_events.append(content_part_event)
return
async def __anext__(
litellm/setup_wizard.py
+3 -3
View File
@@ -71,7 +71,7 @@ PROVIDERS: List[Dict] = [
"id": "azure",
"name": "Azure OpenAI",
"description": "GPT-4o via Azure",
"env_key": "AZURE_API_KEY",
"env_key": "AZURE_AI_API_KEY",
"key_hint": "your-azure-key",
"test_model": None, # needs deployment name — skip validation
"models": [],
@@ -428,7 +428,7 @@ class SetupWizard:
f" {blue('❯')} Azure endpoint URL {grey(p.get('api_base_hint', ''))}: "
)
if api_base:
env_vars[f"_LITELLM_AZURE_API_BASE_{p['id'].upper()}"] = api_base
env_vars[f"_LITELLM_AZURE_AI_API_BASE_{p['id'].upper()}"] = api_base
deployment = _styled_input(
f" {blue('❯')} Azure deployment name {grey('(e.g. my-gpt4o)')}: "
)
@@ -557,7 +557,7 @@ class SetupWizard:
f' api_base: "{_yaml_escape(str(p["api_base"]))}"'
)
elif p.get("needs_api_base"):
azure_base_key = f"_LITELLM_AZURE_API_BASE_{p['id'].upper()}"
azure_base_key = f"_LITELLM_AZURE_AI_API_BASE_{p['id'].upper()}"
if azure_base_key in env_copy:
lines.append(
f' api_base: "{_yaml_escape(env_copy.pop(azure_base_key))}"'
litellm/types/integrations/prometheus.py
+26
View File
@@ -202,6 +202,8 @@ class UserAPIKeyLabelNames(Enum):
USER_AGENT = "user_agent"
CALLBACK_NAME = "callback_name"
STREAM = "stream"
ORG_ID = "org_id"
ORG_ALIAS = "org_alias"
DEFINED_PROMETHEUS_METRICS = Literal[
@@ -224,6 +226,9 @@ DEFINED_PROMETHEUS_METRICS = Literal[
"litellm_remaining_team_budget_metric",
"litellm_team_max_budget_metric",
"litellm_team_budget_remaining_hours_metric",
"litellm_remaining_org_budget_metric",
"litellm_org_max_budget_metric",
"litellm_org_budget_remaining_hours_metric",
"litellm_remaining_api_key_budget_metric",
"litellm_api_key_max_budget_metric",
"litellm_api_key_budget_remaining_hours_metric",
@@ -517,6 +522,21 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.TEAM_ALIAS.value,
]
litellm_remaining_org_budget_metric = [
UserAPIKeyLabelNames.ORG_ID.value,
UserAPIKeyLabelNames.ORG_ALIAS.value,
]
litellm_org_max_budget_metric = [
UserAPIKeyLabelNames.ORG_ID.value,
UserAPIKeyLabelNames.ORG_ALIAS.value,
]
litellm_org_budget_remaining_hours_metric = [
UserAPIKeyLabelNames.ORG_ID.value,
UserAPIKeyLabelNames.ORG_ALIAS.value,
]
litellm_remaining_api_key_budget_metric = [
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
@@ -785,6 +805,12 @@ class UserAPIKeyLabelValues(BaseModel):
stream: Annotated[
Optional[str], Field(..., alias=UserAPIKeyLabelNames.STREAM.value)
] = None
org_id: Annotated[
Optional[str], Field(..., alias=UserAPIKeyLabelNames.ORG_ID.value)
] = None
org_alias: Annotated[
Optional[str], Field(..., alias=UserAPIKeyLabelNames.ORG_ALIAS.value)
] = None
@field_validator("stream", mode="before")
@classmethod
litellm/types/llms/openai.py
+1
View File
@@ -747,6 +747,7 @@ class ChatCompletionAssistantMessage(OpenAIChatCompletionAssistantMessage, total
thinking_blocks: Optional[
List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]
]
reasoning_items: Optional[List[ChatCompletionReasoningItem]]
class ChatCompletionToolMessage(TypedDict):
litellm/types/utils.py
+1
View File
@@ -133,6 +133,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
supports_audio_output: Optional[bool]
supports_pdf_input: Optional[bool]
supports_native_streaming: Optional[bool]
supports_native_structured_output: Optional[bool]
supports_parallel_function_calling: Optional[bool]
supports_web_search: Optional[bool]
supports_reasoning: Optional[bool]
litellm/utils.py
+16
View File
@@ -2735,6 +2735,19 @@ def supports_reasoning(model: str, custom_llm_provider: Optional[str] = None) ->
)
def supports_native_structured_output(
model: str, custom_llm_provider: Optional[str] = None
) -> bool:
"""
Check if the given model supports native structured outputs and return a boolean value.
"""
return _supports_factory(
model=model,
custom_llm_provider=custom_llm_provider,
key="supports_native_structured_output",
)
def get_supported_regions(
model: str, custom_llm_provider: Optional[str] = None
) -> Optional[List[str]]:
@@ -5845,6 +5858,9 @@ def _get_model_info_helper( # noqa: PLR0915
supports_native_streaming=_model_info.get(
"supports_native_streaming", None
),
supports_native_structured_output=_model_info.get(
"supports_native_structured_output", None
),
supports_web_search=_model_info.get("supports_web_search", None),
supports_url_context=_model_info.get("supports_url_context", None),
supports_reasoning=_model_info.get("supports_reasoning", None),
model_prices_and_context_window.json
+88 -56
View File
@@ -722,7 +722,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"anthropic.claude-haiku-4-5@20251001": {
"cache_creation_input_token_cost": 1.25e-06,
@@ -745,7 +746,8 @@
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_streaming": true
"supports_native_streaming": true,
"supports_native_structured_output": true
},
"anthropic.claude-3-5-sonnet-20240620-v1:0": {
"input_cost_per_token": 3e-06,
@@ -967,7 +969,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_native_structured_output": true
},
"anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.25e-06,
@@ -993,7 +996,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"global.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.25e-06,
@@ -1019,7 +1023,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"us.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.875e-06,
@@ -1045,7 +1050,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"eu.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.875e-06,
@@ -1071,7 +1077,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"au.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.875e-06,
@@ -1097,7 +1104,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 3.75e-06,
@@ -1123,7 +1131,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"global.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 3.75e-06,
@@ -1149,7 +1158,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"us.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 4.125e-06,
@@ -1175,7 +1185,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"eu.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 4.125e-06,
@@ -1201,7 +1212,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"au.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 4.125e-06,
@@ -1227,7 +1239,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"anthropic.claude-sonnet-4-20250514-v1:0": {
"cache_creation_input_token_cost": 3.75e-06,
@@ -1287,7 +1300,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_native_structured_output": true
},
"anthropic.claude-v1": {
"input_cost_per_token": 8e-06,
@@ -1537,7 +1551,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"apac.anthropic.claude-3-sonnet-20240229-v1:0": {
"input_cost_per_token": 3e-06,
@@ -1625,7 +1640,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"azure/ada": {
"input_cost_per_token": 1e-07,
@@ -11690,7 +11706,8 @@
"output_cost_per_token": 1.68e-06,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_native_structured_output": true
},
"deepseek.v3.2": {
"input_cost_per_token": 6.2e-07,
@@ -12135,7 +12152,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"eu.anthropic.claude-3-5-sonnet-20240620-v1:0": {
"input_cost_per_token": 3e-06,
@@ -12349,7 +12367,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"eu.meta.llama3-2-1b-instruct-v1:0": {
"input_cost_per_token": 1.3e-07,
@@ -14579,18 +14598,6 @@
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
"uses_embed_content": true
},
"vertex_ai/gemini-embedding-2-preview": {
"input_cost_per_token": 1.5e-07,
"litellm_provider": "vertex_ai",
"max_input_tokens": 8192,
"max_tokens": 8192,
"mode": "embedding",
"output_cost_per_token": 0,
"output_vector_size": 3072,
"source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal",
"supports_multimodal": true,
"uses_embed_content": true
},
"gemini/gemini-embedding-001": {
"input_cost_per_token": 1.5e-07,
"litellm_provider": "gemini",
@@ -16767,7 +16774,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"global.anthropic.claude-sonnet-4-20250514-v1:0": {
"cache_creation_input_token_cost": 3.75e-06,
@@ -16819,7 +16827,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"global.amazon.nova-2-lite-v1:0": {
"cache_read_input_token_cost": 7.5e-08,
@@ -20548,7 +20557,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"jp.anthropic.claude-haiku-4-5-20251001-v1:0": {
"cache_creation_input_token_cost": 1.375e-06,
@@ -20570,7 +20580,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"lambda_ai/deepseek-llama3.3-70b": {
"input_cost_per_token": 2e-07,
@@ -21265,7 +21276,8 @@
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"minimax.minimax-m2.1": {
"input_cost_per_token": 3e-07,
@@ -21421,7 +21433,8 @@
"mode": "chat",
"output_cost_per_token": 2e-07,
"supports_function_calling": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral.ministral-3-3b-instruct": {
"input_cost_per_token": 1e-07,
@@ -21432,7 +21445,8 @@
"mode": "chat",
"output_cost_per_token": 1e-07,
"supports_function_calling": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral.ministral-3-8b-instruct": {
"input_cost_per_token": 1.5e-07,
@@ -21443,7 +21457,8 @@
"mode": "chat",
"output_cost_per_token": 1.5e-07,
"supports_function_calling": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral.mistral-7b-instruct-v0:2": {
"input_cost_per_token": 1.5e-07,
@@ -21485,7 +21500,8 @@
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"supports_function_calling": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral.mistral-small-2402-v1:0": {
"input_cost_per_token": 1e-06,
@@ -21516,7 +21532,8 @@
"mode": "chat",
"output_cost_per_token": 4e-08,
"supports_audio_input": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral.voxtral-small-24b-2507": {
"input_cost_per_token": 1e-07,
@@ -21527,7 +21544,8 @@
"mode": "chat",
"output_cost_per_token": 3e-07,
"supports_audio_input": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"mistral/codestral-2405": {
"input_cost_per_token": 1e-06,
@@ -22214,7 +22232,8 @@
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"supports_reasoning": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"moonshotai.kimi-k2.5": {
"input_cost_per_token": 6e-07,
@@ -23089,7 +23108,8 @@
"supports_function_calling": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"source": "https://aws.amazon.com/bedrock/pricing/"
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_native_structured_output": true
},
"o1": {
"cache_read_input_token_cost": 7.5e-06,
@@ -26173,7 +26193,8 @@
"output_cost_per_token": 1.8e-06,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_native_structured_output": true
},
"qwen.qwen3-235b-a22b-2507-v1:0": {
"input_cost_per_token": 2.2e-07,
@@ -26185,7 +26206,8 @@
"output_cost_per_token": 8.8e-07,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_native_structured_output": true
},
"qwen.qwen3-coder-30b-a3b-v1:0": {
"input_cost_per_token": 1.5e-07,
@@ -26197,7 +26219,8 @@
"output_cost_per_token": 6e-07,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_native_structured_output": true
},
"qwen.qwen3-32b-v1:0": {
"input_cost_per_token": 1.5e-07,
@@ -26209,7 +26232,8 @@
"output_cost_per_token": 6e-07,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_native_structured_output": true
},
"qwen.qwen3-next-80b-a3b": {
"input_cost_per_token": 1.5e-07,
@@ -26220,7 +26244,8 @@
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"supports_function_calling": true,
"supports_system_messages": true
"supports_system_messages": true,
"supports_native_structured_output": true
},
"qwen.qwen3-vl-235b-a22b": {
"input_cost_per_token": 5.3e-07,
@@ -26232,7 +26257,8 @@
"output_cost_per_token": 2.66e-06,
"supports_function_calling": true,
"supports_system_messages": true,
"supports_vision": true
"supports_vision": true,
"supports_native_structured_output": true
},
"qwen.qwen3-coder-next": {
"input_cost_per_token": 5e-07,
@@ -28257,7 +28283,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"us.anthropic.claude-3-5-sonnet-20240620-v1:0": {
"input_cost_per_token": 3e-06,
@@ -28415,7 +28442,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"au.anthropic.claude-haiku-4-5-20251001-v1:0": {
"cache_creation_input_token_cost": 1.375e-06,
@@ -28436,7 +28464,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
},
"us.anthropic.claude-opus-4-20250514-v1:0": {
"cache_creation_input_token_cost": 1.875e-05,
@@ -28488,7 +28517,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_native_structured_output": true
},
"global.anthropic.claude-opus-4-5-20251101-v1:0": {
"cache_creation_input_token_cost": 6.25e-06,
@@ -28514,7 +28544,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_native_structured_output": true
},
"eu.anthropic.claude-opus-4-5-20251101-v1:0": {
"cache_creation_input_token_cost": 6.25e-06,
@@ -28540,7 +28571,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_native_structured_output": true
},
"us.anthropic.claude-sonnet-4-20250514-v1:0": {
"cache_creation_input_token_cost": 3.75e-06,
tests/agent_tests/test_a2a_agent.py
+57 -25
View File
@@ -1,38 +1,72 @@
"""
Simple A2A agent tests - non-streaming and streaming.
These tests validate the localhost URL retry logic: if an A2A agent's card
contains a localhost/internal URL (e.g., http://0.0.0.0:8001/), the request
will fail with a connection error. LiteLLM detects this and automatically
retries using the original api_base URL instead.
Requires A2A_AGENT_URL environment variable to be set.
Run with:
A2A_AGENT_URL=https://your-agent.example.com pytest tests/agent_tests/test_a2a_agent.py -v -s
These tests use a mocked A2A client to avoid network/env dependencies.
"""
import os
import pytest
from types import SimpleNamespace
from uuid import uuid4
import pytest
def get_a2a_agent_url():
"""Get A2A agent URL from environment, skip test if not set."""
url = os.environ.get("A2A_AGENT_URL")
return url
class MockA2AResponse:
def __init__(self, text: str):
self._payload = {
"id": str(uuid4()),
"jsonrpc": "2.0",
"result": {
"message": {
"role": "agent",
"parts": [{"kind": "text", "text": text}],
"messageId": uuid4().hex,
}
},
}
def model_dump(self, mode="json", exclude_none=True):
return self._payload
class MockA2AStreamingChunk(MockA2AResponse):
def __init__(self, text: str, state: str):
super().__init__(text=text)
self._payload["result"]["status"] = {"state": state}
class MockA2AClient:
def __init__(self):
self._litellm_agent_card = SimpleNamespace(
name="mock-agent", url="http://mock-agent.local"
)
async def send_message(self, request):
return MockA2AResponse(text="hello")
def send_message_streaming(self, request):
async def _stream():
yield MockA2AStreamingChunk(text="hel", state="in_progress")
yield MockA2AStreamingChunk(text="hello", state="completed")
return _stream()
@pytest.fixture
def mock_a2a_client(monkeypatch):
import litellm.a2a_protocol.main as a2a_main
async def _fake_create_a2a_client(base_url, timeout=60.0, extra_headers=None):
return MockA2AClient()
monkeypatch.setattr(a2a_main, "create_a2a_client", _fake_create_a2a_client)
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=5)
async def test_a2a_non_streaming():
async def test_a2a_non_streaming(mock_a2a_client):
"""Test non-streaming A2A request."""
from a2a.types import MessageSendParams, SendMessageRequest
from litellm.a2a_protocol import asend_message
api_base = get_a2a_agent_url()
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
@@ -46,7 +80,7 @@ async def test_a2a_non_streaming():
response = await asend_message(
request=request,
api_base=api_base,
api_base="http://mock",
)
assert response is not None
@@ -54,13 +88,11 @@ async def test_a2a_non_streaming():
@pytest.mark.asyncio
async def test_a2a_streaming():
async def test_a2a_streaming(mock_a2a_client):
"""Test streaming A2A request."""
from a2a.types import MessageSendParams, SendStreamingMessageRequest
from litellm.a2a_protocol import asend_message_streaming
api_base = get_a2a_agent_url()
request = SendStreamingMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
@@ -75,7 +107,7 @@ async def test_a2a_streaming():
chunks = []
async for chunk in asend_message_streaming(
request=request,
api_base=api_base,
api_base="http://mock",
):
chunks.append(chunk)
print(f"\nStreaming chunk: {chunk}")
tests/audio_tests/test_audio_speech.py
+64 -44
View File
@@ -35,8 +35,8 @@ import litellm
[
(
"azure/tts",
os.getenv("AZURE_SWEDEN_API_KEY"),
os.getenv("AZURE_SWEDEN_API_BASE"),
os.getenv("AZURE_TTS_API_KEY"),
os.getenv("AZURE_TTS_API_BASE"),
),
("openai/tts-1", os.getenv("OPENAI_API_KEY"), None),
],
@@ -286,9 +286,9 @@ async def test_speech_litellm_vertex_async_with_voice_ssml():
def test_audio_speech_cost_calc():
from litellm.integrations.custom_logger import CustomLogger
model = "azure/azure-tts"
api_base = os.getenv("AZURE_SWEDEN_API_BASE")
api_key = os.getenv("AZURE_SWEDEN_API_KEY")
model = "azure/tts"
api_base = os.getenv("AZURE_TTS_API_BASE")
api_key = os.getenv("AZURE_TTS_API_KEY")
custom_logger = CustomLogger()
litellm.set_verbose = True
@@ -301,7 +301,7 @@ def test_audio_speech_cost_calc():
input="the quick brown fox jumped over the lazy dogs",
api_base=api_base,
api_key=api_key,
base_model="azure/tts-1",
base_model="azure/tts",
)
time.sleep(1)
@@ -337,10 +337,9 @@ async def test_azure_ava_tts_async():
litellm._turn_on_debug()
api_key = os.getenv("AZURE_TTS_API_KEY")
api_base = os.getenv("AZURE_TTS_API_BASE")
speech_file_path = Path(__file__).parent / "azure_speech.mp3"
try:
response = await litellm.aspeech(
model="azure/speech/azure-tts",
@@ -354,30 +353,34 @@ async def test_azure_ava_tts_async():
# Assert the response is HttpxBinaryResponseContent
from litellm.types.llms.openai import HttpxBinaryResponseContent
assert isinstance(response, HttpxBinaryResponseContent)
# Get the binary content
binary_content = response.content
assert len(binary_content) > 0
# MP3 files start with these magic bytes
# ID3 tag or MPEG sync word
assert binary_content[:3] == b"ID3" or binary_content[:2] == b"\xff\xfb" or binary_content[:2] == b"\xff\xf3"
assert (
binary_content[:3] == b"ID3"
or binary_content[:2] == b"\xff\xfb"
or binary_content[:2] == b"\xff\xf3"
)
# Write to file
response.stream_to_file(speech_file_path)
# Verify file was created and has content
assert speech_file_path.exists()
assert speech_file_path.stat().st_size > 0
print(f"Azure TTS audio saved to: {speech_file_path}")
# assert response cost is greater than 0
print("Response cost: ", response._hidden_params["response_cost"])
assert response._hidden_params["response_cost"] > 0
except Exception as e:
pytest.fail(f"Test failed with exception: {str(e)}")
@@ -392,10 +395,9 @@ async def test_runwayml_tts_async():
litellm._turn_on_debug()
api_key = os.getenv("RUNWAYML_API_KEY")
api_base = os.getenv("RUNWAYML_API_BASE")
speech_file_path = Path(__file__).parent / "runwayml_speech.mp3"
try:
response = await litellm.aspeech(
model="runwayml/eleven_multilingual_v2",
@@ -409,30 +411,34 @@ async def test_runwayml_tts_async():
# Assert the response is HttpxBinaryResponseContent
from litellm.types.llms.openai import HttpxBinaryResponseContent
assert isinstance(response, HttpxBinaryResponseContent)
# Get the binary content
binary_content = response.content
assert len(binary_content) > 0
# MP3 files start with these magic bytes
# ID3 tag or MPEG sync word
assert binary_content[:3] == b"ID3" or binary_content[:2] == b"\xff\xfb" or binary_content[:2] == b"\xff\xf3"
assert (
binary_content[:3] == b"ID3"
or binary_content[:2] == b"\xff\xfb"
or binary_content[:2] == b"\xff\xf3"
)
# Write to file
response.stream_to_file(speech_file_path)
# Verify file was created and has content
assert speech_file_path.exists()
assert speech_file_path.stat().st_size > 0
print(f"RunwayML TTS audio saved to: {speech_file_path}")
# assert response cost is greater than 0
print("Response cost: ", response._hidden_params["response_cost"])
assert response._hidden_params["response_cost"] > 0
except Exception as e:
pytest.fail(f"Test failed with exception: {str(e)}")
@@ -445,17 +451,19 @@ async def test_azure_ava_tts_with_custom_voice():
"""
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
# Mock response
mock_response_content = b"fake_audio_data"
mock_httpx_response = MagicMock(spec=httpx.Response)
mock_httpx_response.content = mock_response_content
mock_httpx_response.status_code = 200
mock_httpx_response.headers = {"content-type": "audio/mpeg"}
with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") as mock_post:
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post"
) as mock_post:
mock_post.return_value = mock_httpx_response
response = await litellm.aspeech(
model="azure/speech/azure-tts",
voice="en-US-AndrewNeural",
@@ -464,14 +472,14 @@ async def test_azure_ava_tts_with_custom_voice():
api_key="fake-key",
response_format="mp3",
)
# Verify the mock was called
assert mock_post.called
# Get the call arguments
call_args = mock_post.call_args
ssml_body = call_args.kwargs.get("data")
# Verify the SSML contains the custom voice
assert ssml_body is not None
assert "en-US-AndrewNeural" in ssml_body
@@ -488,17 +496,19 @@ async def test_azure_ava_tts_fable_voice_mapping():
"""
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
# Mock response
mock_response_content = b"fake_audio_data"
mock_httpx_response = MagicMock(spec=httpx.Response)
mock_httpx_response.content = mock_response_content
mock_httpx_response.status_code = 200
mock_httpx_response.headers = {"content-type": "audio/mpeg"}
with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") as mock_post:
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post"
) as mock_post:
mock_post.return_value = mock_httpx_response
response = await litellm.aspeech(
model="azure/speech/azure-tts",
voice="fable",
@@ -507,14 +517,14 @@ async def test_azure_ava_tts_fable_voice_mapping():
api_key="fake-key",
response_format="mp3",
)
# Verify the mock was called
assert mock_post.called
# Get the call arguments
call_args = mock_post.call_args
ssml_body = call_args.kwargs.get("data")
# Verify the SSML contains the mapped voice (en-GB-RyanNeural, not 'fable')
assert ssml_body is not None
assert "en-GB-RyanNeural" in ssml_body
@@ -541,7 +551,9 @@ async def test_aws_polly_tts_with_native_voice():
mock_httpx_response.status_code = 200
mock_httpx_response.headers = {"content-type": "audio/mpeg"}
with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") as mock_post:
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post"
) as mock_post:
mock_post.return_value = mock_httpx_response
response = await litellm.aspeech(
@@ -586,7 +598,9 @@ async def test_aws_polly_tts_with_openai_voice_mapping():
mock_httpx_response.status_code = 200
mock_httpx_response.headers = {"content-type": "audio/mpeg"}
with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") as mock_post:
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post"
) as mock_post:
mock_post.return_value = mock_httpx_response
response = await litellm.aspeech(
@@ -628,7 +642,9 @@ async def test_aws_polly_tts_with_ssml():
ssml_input = '<speak>Hello, <break time="500ms"/> this is SSML.</speak>'
with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") as mock_post:
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post"
) as mock_post:
mock_post.return_value = mock_httpx_response
response = await litellm.aspeech(
@@ -676,7 +692,11 @@ async def test_aws_polly_tts_real_api():
assert len(binary_content) > 0
# MP3 files start with ID3 tag or MPEG sync word
assert binary_content[:3] == b"ID3" or binary_content[:2] == b"\xff\xfb" or binary_content[:2] == b"\xff\xf3"
assert (
binary_content[:3] == b"ID3"
or binary_content[:2] == b"\xff\xfb"
or binary_content[:2] == b"\xff\xf3"
)
response.stream_to_file(speech_file_path)
tests/batches_tests/test_fine_tuning_api.py
+16 -107
View File
@@ -123,81 +123,8 @@ async def test_create_fine_tune_jobs_async():
pass
@pytest.mark.asyncio
async def test_azure_create_fine_tune_jobs_async():
try:
verbose_logger.setLevel(logging.DEBUG)
file_name = "azure_fine_tune.jsonl"
_current_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(_current_dir, file_name)
file_id = "file-5e4b20ecbd724182b9964f3cd2ab7212"
create_fine_tuning_response = await litellm.acreate_fine_tuning_job(
model="gpt-35-turbo-1106",
training_file=file_id,
custom_llm_provider="azure",
api_base="https://exampleopenaiendpoint-production.up.railway.app",
)
print(
"response from litellm.create_fine_tuning_job=", create_fine_tuning_response
)
assert create_fine_tuning_response.id is not None
# response from Example/mocked endpoint
assert create_fine_tuning_response.model == "davinci-002"
# list fine tuning jobs
print("listing ft jobs")
ft_jobs = await litellm.alist_fine_tuning_jobs(
limit=2,
custom_llm_provider="azure",
api_base="https://exampleopenaiendpoint-production.up.railway.app",
)
print("response from litellm.list_fine_tuning_jobs=", ft_jobs)
# cancel ft job
response = await litellm.acancel_fine_tuning_job(
fine_tuning_job_id=create_fine_tuning_response.id,
custom_llm_provider="azure",
api_key=os.getenv("AZURE_SWEDEN_API_KEY"),
api_base="https://exampleopenaiendpoint-production.up.railway.app",
)
print("response from litellm.cancel_fine_tuning_job=", response)
assert response.status == "cancelled"
assert response.id == create_fine_tuning_response.id
except openai.RateLimitError:
pass
except Exception as e:
if "Job has already completed" in str(e):
pass
else:
pytest.fail(f"Error occurred: {e}")
pass
def test_azure_trainingtype_defaults_to_one():
"""
Azure requires trainingType in extra_body. When omitted, AzureOpenAIFineTuningAPI defaults it to 1.
"""
from litellm.llms.azure.fine_tuning.handler import AzureOpenAIFineTuningAPI
handler = AzureOpenAIFineTuningAPI()
create_data = {"model": "gpt-4o-mini", "training_file": "file-test"}
handler._ensure_training_type(create_data)
assert "extra_body" in create_data
assert create_data["extra_body"]["trainingType"] == 1
@pytest.mark.asyncio()
async def test_create_vertex_fine_tune_jobs_mocked():
load_vertex_ai_credentials()
# Define reusable variables for the test
project_id = "633608382793"
location = "us-central1"
@@ -236,7 +163,10 @@ async def test_create_vertex_fine_tune_jobs_mocked():
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=mock_response,
) as mock_post:
) as mock_post, patch(
"litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token",
return_value=("fake-token", project_id),
):
create_fine_tuning_response = await litellm.acreate_fine_tuning_job(
model=base_model,
custom_llm_provider="vertex_ai",
@@ -290,7 +220,6 @@ async def test_create_vertex_fine_tune_jobs_mocked():
@pytest.mark.asyncio()
async def test_create_vertex_fine_tune_jobs_mocked_with_hyperparameters():
load_vertex_ai_credentials()
# Define reusable variables for the test
project_id = "633608382793"
location = "us-central1"
@@ -329,7 +258,10 @@ async def test_create_vertex_fine_tune_jobs_mocked_with_hyperparameters():
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=mock_response,
) as mock_post:
) as mock_post, patch(
"litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token",
return_value=("fake-token", project_id),
):
create_fine_tuning_response = await litellm.acreate_fine_tuning_job(
model=base_model,
custom_llm_provider="vertex_ai",
@@ -478,29 +410,6 @@ def test_convert_basic_openai_request_to_vertex_request():
)
@pytest.mark.asyncio()
@pytest.mark.skip(reason="skipping - we run mock tests for vertex ai")
async def test_create_vertex_fine_tune_jobs():
verbose_logger.setLevel(logging.DEBUG)
# load_vertex_ai_credentials()
vertex_credentials = os.getenv("GCS_PATH_SERVICE_ACCOUNT")
print("creating fine tuning job")
create_fine_tuning_response = await litellm.acreate_fine_tuning_job(
model="gemini-1.0-pro-002",
custom_llm_provider="vertex_ai",
training_file="gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl",
vertex_project="pathrise-convert-1606954137718",
vertex_location="us-central1",
vertex_credentials=vertex_credentials,
)
print("vertex ai create fine tuning response=", create_fine_tuning_response)
assert create_fine_tuning_response.id is not None
assert create_fine_tuning_response.model == "gemini-1.0-pro-002"
assert create_fine_tuning_response.object == "fine_tuning.job"
@pytest.mark.asyncio
async def test_mock_openai_create_fine_tune_job():
"""Test that create_fine_tuning_job sends correct parameters to OpenAI"""
@@ -608,7 +517,6 @@ async def test_mock_openai_retrieve_fine_tune_job():
except Exception as e:
print("error=", e)
# Verify the request
mock_retrieve.assert_called_once_with(fine_tuning_job_id="ft-123")
@@ -616,7 +524,9 @@ async def test_mock_openai_retrieve_fine_tune_job():
@pytest.mark.asyncio
async def test_mock_azure_create_fine_tune_job_with_azure_specific_params():
"""Test that Azure-specific parameters are passed through extra_body"""
from openai.types.fine_tuning.fine_tuning_job import Hyperparameters as OAIHyperparameters
from openai.types.fine_tuning.fine_tuning_job import (
Hyperparameters as OAIHyperparameters,
)
from litellm.types.utils import LiteLLMFineTuningJob
mock_response = LiteLLMFineTuningJob(
@@ -636,7 +546,9 @@ async def test_mock_azure_create_fine_tune_job_with_azure_specific_params():
async def mock_async_create(*args, **kwargs):
return mock_response
with patch("litellm.llms.azure.fine_tuning.handler.AzureOpenAIFineTuningAPI.create_fine_tuning_job") as mock_create:
with patch(
"litellm.llms.azure.fine_tuning.handler.AzureOpenAIFineTuningAPI.create_fine_tuning_job"
) as mock_create:
mock_create.return_value = mock_async_create()
response = await litellm.acreate_fine_tuning_job(
@@ -647,10 +559,7 @@ async def test_mock_azure_create_fine_tune_job_with_azure_specific_params():
api_key="test-key",
api_version="2025-04-01-preview",
trainingType=1,
hyperparameters={
"n_epochs": 3,
"prompt_loss_weight": 0.1
},
hyperparameters={"n_epochs": 3, "prompt_loss_weight": 0.1},
)
# Verify the request
@@ -662,7 +571,7 @@ async def test_mock_azure_create_fine_tune_job_with_azure_specific_params():
assert create_data["model"] == "gpt-4.1-mini-2025-04-14"
assert create_data["training_file"] == "file-123"
assert create_data["hyperparameters"] == {"n_epochs": 3}
# Azure-specific parameters should be in extra_body
assert "extra_body" in create_data
assert create_data["extra_body"]["trainingType"] == 1
tests/batches_tests/test_openai_batches_and_files.py
+28 -24
View File
@@ -54,7 +54,7 @@ def load_vertex_ai_credentials():
print("loading vertex ai credentials")
os.environ["GCS_FLUSH_INTERVAL"] = "1"
filepath = os.path.dirname(os.path.abspath(__file__))
vertex_key_path = filepath + "/pathrise-convert-1606954137718.json"
vertex_key_path = filepath + "/vertex_key.json"
# Read the existing content of the file or create an empty dictionary
try:
@@ -75,8 +75,8 @@ def load_vertex_ai_credentials():
service_account_key_data = {}
# Update the service_account_key_data with environment variables
private_key_id = os.environ.get("GCS_PRIVATE_KEY_ID", "")
private_key = os.environ.get("GCS_PRIVATE_KEY", "")
private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "")
private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "")
private_key = private_key.replace("\\n", "\n")
service_account_key_data["private_key_id"] = private_key_id
service_account_key_data["private_key"] = private_key
@@ -234,9 +234,9 @@ def cleanup_azure_ft_models():
import requests
client = AzureOpenAI(
api_key=os.getenv("AZURE_FT_API_KEY"),
azure_endpoint=os.getenv("AZURE_FT_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
api_key=os.getenv("AZURE_AI_API_KEY"),
azure_endpoint=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_AI_API_VERSION"),
)
_list_ft_jobs = client.fine_tuning.jobs.list()
@@ -577,7 +577,10 @@ async def test_vertex_list_batches(monkeypatch):
monkeypatch.setattr(
"litellm.llms.vertex_ai.batches.handler.VertexAIBatchPrediction._ensure_access_token",
lambda self, credentials, project_id, custom_llm_provider: ("mock-token", "litellm-test-project"),
lambda self, credentials, project_id, custom_llm_provider: (
"mock-token",
"litellm-test-project",
),
)
with patch(
@@ -648,7 +651,7 @@ async def test_vertex_async_create_batch_logs_error_body_on_http_error():
async def test_delete_batch_output_file():
"""
Test that deleting a batch output file works correctly.
This test verifies the fix for:
- When a batch is retrieved and has an output_file_id, the file object is properly stored
- The output file can be deleted without validation errors
@@ -656,11 +659,11 @@ async def test_delete_batch_output_file():
"""
litellm._turn_on_debug()
print("Testing delete batch output file")
file_name = "openai_batch_completions.jsonl"
_current_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(_current_dir, file_name)
# Create file for batch
file_obj = await litellm.acreate_file(
file=open(file_path, "rb"),
@@ -669,7 +672,7 @@ async def test_delete_batch_output_file():
)
print("Response from creating file=", file_obj)
batch_input_file_id = file_obj.id
# Create batch
create_batch_response = await litellm.acreate_batch(
completion_window="24h",
@@ -678,36 +681,37 @@ async def test_delete_batch_output_file():
custom_llm_provider="openai",
)
print("Batch created with ID=", create_batch_response.id)
# Retrieve batch to get output_file_id
retrieved_batch = await litellm.aretrieve_batch(
batch_id=create_batch_response.id,
custom_llm_provider="openai"
batch_id=create_batch_response.id, custom_llm_provider="openai"
)
print("Retrieved batch=", retrieved_batch)
# If batch has completed and has output file, test deleting it
if retrieved_batch.output_file_id:
print(f"Testing deletion of output file: {retrieved_batch.output_file_id}")
# This is the key test - deleting the output file should work
# without validation errors (file_object should not be None)
delete_output_file_response = await litellm.afile_delete(
file_id=retrieved_batch.output_file_id,
custom_llm_provider="openai"
file_id=retrieved_batch.output_file_id, custom_llm_provider="openai"
)
print("Delete output file response=", delete_output_file_response)
assert delete_output_file_response.id == retrieved_batch.output_file_id
assert delete_output_file_response.deleted is True or hasattr(delete_output_file_response, 'id')
assert delete_output_file_response.deleted is True or hasattr(
delete_output_file_response, "id"
)
print("✓ Successfully deleted batch output file")
else:
print("⚠ Batch has not completed yet or no output file available, skipping output file deletion test")
print(
"⚠ Batch has not completed yet or no output file available, skipping output file deletion test"
)
# Clean up - delete the input file
delete_input_file_response = await litellm.afile_delete(
file_id=batch_input_file_id,
custom_llm_provider="openai"
file_id=batch_input_file_id, custom_llm_provider="openai"
)
print("Delete input file response=", delete_input_file_response)
assert delete_input_file_response.id == batch_input_file_id
tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py
+3 -3
View File
@@ -169,9 +169,9 @@ async def test_prometheus_metric_tracking():
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"model_info": {"id": "azure-model-id"},
},
tests/image_gen_tests/test_image_edit.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 2.1 MiB

After

Width:  |  Height:  |  Size: 457 KiB

tests/image_gen_tests/test_image_edits.py
+143 -112
View File
@@ -23,10 +23,11 @@ from litellm.types.utils import StandardLoggingPayload
# Configure pytest marks to avoid warnings
pytestmark = pytest.mark.asyncio
class TestCustomLogger(CustomLogger):
def __init__(self):
self.standard_logging_payload: Optional[StandardLoggingPayload] = None
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.standard_logging_payload = kwargs.get("standard_logging_object", None)
pass
@@ -80,12 +81,12 @@ class BaseLLMImageEditTest(ABC):
result = self.image_edit_function(**call_args)
else:
result = await self.async_image_edit_function(**call_args)
print("result from image edit", result)
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -97,6 +98,7 @@ class BaseLLMImageEditTest(ABC):
except litellm.ContentPolicyViolationError as e:
pass
# Get the current directory of the file being run
pwd = os.path.dirname(os.path.realpath(__file__))
@@ -107,6 +109,7 @@ TEST_IMAGES = [
SINGLE_TEST_IMAGE = open(os.path.join(pwd, "ishaan_github.png"), "rb")
def get_test_images_as_bytesio():
"""Helper function to get test images as BytesIO objects"""
bytesio_images = []
@@ -129,6 +132,7 @@ class TestOpenAIImageEditGPTImage1(BaseLLMImageEditTest):
"image": TEST_IMAGES,
}
class TestOpenAIImageEditDallE2(BaseLLMImageEditTest):
"""
Concrete implementation of BaseLLMImageEditTest for OpenAI DALL-E-2 image edits.
@@ -154,8 +158,8 @@ class TestAzureAIFlux2ImageEdit(BaseLLMImageEditTest):
return {
"model": "azure_ai/flux.2-pro",
"image": SINGLE_TEST_IMAGE,
"api_base": "https://litellm-ci-cd-prod.services.ai.azure.com",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": "preview",
}
@@ -187,7 +191,7 @@ async def test_openai_image_edit_litellm_router():
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -199,17 +203,19 @@ async def test_openai_image_edit_litellm_router():
except litellm.ContentPolicyViolationError as e:
pass
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio
async def test_openai_image_edit_with_bytesio():
"""Test image editing using BytesIO objects instead of file readers"""
from litellm import image_edit, aimage_edit
litellm._turn_on_debug()
try:
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Get images as BytesIO objects
bytesio_images = get_test_images_as_bytesio()
@@ -222,7 +228,7 @@ async def test_openai_image_edit_with_bytesio():
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -239,7 +245,7 @@ async def test_openai_image_edit_with_bytesio():
async def test_azure_image_edit_litellm_sdk():
"""Test Azure image edit with mocked httpx request to validate request body and URL"""
from litellm import image_edit, aimage_edit
# Mock response for Azure image edit
mock_response = {
"created": 1589478378,
@@ -247,7 +253,7 @@ async def test_azure_image_edit_litellm_sdk():
{
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
}
]
],
}
class MockResponse:
@@ -267,16 +273,16 @@ async def test_azure_image_edit_litellm_sdk():
mock_post.return_value = MockResponse(mock_response, 200)
litellm._turn_on_debug()
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Set up test environment variables
test_api_base = "https://ai-api-gw-uae-north.openai.azure.com"
test_api_key = "test-api-key"
test_api_version = "2025-04-01-preview"
result = await aimage_edit(
prompt=prompt,
model="azure/gpt-image-1",
@@ -285,41 +291,54 @@ async def test_azure_image_edit_litellm_sdk():
api_version=test_api_version,
image=TEST_IMAGES,
)
# Verify the request was made correctly
mock_post.assert_called_once()
# Check the URL
call_args = mock_post.call_args
expected_url = f"{test_api_base}/openai/deployments/gpt-image-1/images/edits?api-version={test_api_version}"
actual_url = call_args.args[0] if call_args.args else call_args.kwargs.get('url')
actual_url = (
call_args.args[0] if call_args.args else call_args.kwargs.get("url")
)
print(f"Expected URL: {expected_url}")
print(f"Actual URL: {actual_url}")
assert actual_url == expected_url, f"URL mismatch. Expected: {expected_url}, Got: {actual_url}"
assert (
actual_url == expected_url
), f"URL mismatch. Expected: {expected_url}, Got: {actual_url}"
# Check the request body
if 'data' in call_args.kwargs:
if "data" in call_args.kwargs:
# For multipart form data, check the data parameter
form_data = call_args.kwargs['data']
print("Form data keys:", list(form_data.keys()) if hasattr(form_data, 'keys') else "Not a dict")
form_data = call_args.kwargs["data"]
print(
"Form data keys:",
list(form_data.keys()) if hasattr(form_data, "keys") else "Not a dict",
)
# Validate that model and prompt are in the form data
assert 'model' in form_data, "model should be in form data"
assert 'prompt' in form_data, "prompt should be in form data"
assert form_data['model'] == 'gpt-image-1', f"Expected model 'gpt-image-1', got {form_data['model']}"
assert prompt.strip() in form_data['prompt'], f"Expected prompt to contain '{prompt.strip()}'"
assert "model" in form_data, "model should be in form data"
assert "prompt" in form_data, "prompt should be in form data"
assert (
form_data["model"] == "gpt-image-1"
), f"Expected model 'gpt-image-1', got {form_data['model']}"
assert (
prompt.strip() in form_data["prompt"]
), f"Expected prompt to contain '{prompt.strip()}'"
# Check headers
headers = call_args.kwargs.get('headers', {})
headers = call_args.kwargs.get("headers", {})
print("Request headers:", headers)
assert 'Authorization' in headers, "Authorization header should be present"
assert headers['Authorization'].startswith('Bearer '), "Authorization should be Bearer token"
assert "Authorization" in headers, "Authorization header should be present"
assert headers["Authorization"].startswith(
"Bearer "
), "Authorization should be Bearer token"
print("result from image edit", result)
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -330,15 +349,15 @@ async def test_azure_image_edit_litellm_sdk():
f.write(image_bytes)
@pytest.mark.asyncio
async def test_openai_image_edit_cost_tracking():
"""Test OpenAI image edit cost tracking with custom logger"""
from litellm import image_edit, aimage_edit
test_custom_logger = TestCustomLogger()
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = [test_custom_logger]
# Mock response for Azure image edit with usage data for cost tracking
mock_response = {
"created": 1589478378,
@@ -350,12 +369,9 @@ async def test_openai_image_edit_cost_tracking():
"usage": {
"total_tokens": 1100,
"input_tokens": 100,
"input_tokens_details": {
"image_tokens": 50,
"text_tokens": 50
},
"output_tokens": 1000
}
"input_tokens_details": {"image_tokens": 50, "text_tokens": 50},
"output_tokens": 1000,
},
}
class MockResponse:
@@ -375,26 +391,25 @@ async def test_openai_image_edit_cost_tracking():
mock_post.return_value = MockResponse(mock_response, 200)
litellm._turn_on_debug()
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Set up test environment variables
result = await aimage_edit(
prompt=prompt,
model="openai/gpt-image-1",
image=TEST_IMAGES,
)
# Verify the request was made correctly
mock_post.assert_called_once()
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -403,30 +418,36 @@ async def test_openai_image_edit_cost_tracking():
# Save the image to a file
with open("test_image_edit.png", "wb") as f:
f.write(image_bytes)
await asyncio.sleep(5)
print("standard logging payload", json.dumps(test_custom_logger.standard_logging_payload, indent=4, default=str))
print(
"standard logging payload",
json.dumps(
test_custom_logger.standard_logging_payload, indent=4, default=str
),
)
# check model
assert test_custom_logger.standard_logging_payload["model"] == "gpt-image-1"
assert test_custom_logger.standard_logging_payload["custom_llm_provider"] == "openai"
assert (
test_custom_logger.standard_logging_payload["custom_llm_provider"]
== "openai"
)
# check response_cost
assert test_custom_logger.standard_logging_payload["response_cost"] is not None
assert test_custom_logger.standard_logging_payload["response_cost"] > 0
@pytest.mark.asyncio
async def test_azure_image_edit_cost_tracking():
"""Test Azure image edit cost tracking with custom logger"""
from litellm import image_edit, aimage_edit
test_custom_logger = TestCustomLogger()
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = [test_custom_logger]
# Mock response for Azure image edit with usage data for cost tracking
mock_response = {
"created": 1589478378,
@@ -438,12 +459,9 @@ async def test_azure_image_edit_cost_tracking():
"usage": {
"total_tokens": 1100,
"input_tokens": 100,
"input_tokens_details": {
"image_tokens": 50,
"text_tokens": 50
},
"output_tokens": 1000
}
"input_tokens_details": {"image_tokens": 50, "text_tokens": 50},
"output_tokens": 1000,
},
}
class MockResponse:
@@ -463,27 +481,26 @@ async def test_azure_image_edit_cost_tracking():
mock_post.return_value = MockResponse(mock_response, 200)
litellm._turn_on_debug()
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Set up test environment variables
result = await aimage_edit(
prompt=prompt,
model="azure/CUSTOM_AZURE_DEPLOYMENT_NAME",
base_model="azure/gpt-image-1",
image=TEST_IMAGES,
)
# Verify the request was made correctly
mock_post.assert_called_once()
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
@@ -492,14 +509,24 @@ async def test_azure_image_edit_cost_tracking():
# Save the image to a file
with open("test_image_edit.png", "wb") as f:
f.write(image_bytes)
await asyncio.sleep(5)
print("standard logging payload", json.dumps(test_custom_logger.standard_logging_payload, indent=4, default=str))
print(
"standard logging payload",
json.dumps(
test_custom_logger.standard_logging_payload, indent=4, default=str
),
)
# check model
assert test_custom_logger.standard_logging_payload["model"] == "CUSTOM_AZURE_DEPLOYMENT_NAME"
assert test_custom_logger.standard_logging_payload["custom_llm_provider"] == "azure"
assert (
test_custom_logger.standard_logging_payload["model"]
== "CUSTOM_AZURE_DEPLOYMENT_NAME"
)
assert (
test_custom_logger.standard_logging_payload["custom_llm_provider"]
== "azure"
)
# check response_cost
assert test_custom_logger.standard_logging_payload["response_cost"] is not None
@@ -511,6 +538,7 @@ async def test_azure_image_edit_cost_tracking():
async def test_recraft_image_edit_api():
from litellm import aimage_edit
import requests
litellm._turn_on_debug()
global TEST_IMAGES
try:
@@ -526,10 +554,10 @@ async def test_recraft_image_edit_api():
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_url = result.data[0].url
# download the image
image_bytes = requests.get(image_url).content
with open("test_image_edit.png", "wb") as f:
@@ -545,51 +573,55 @@ def test_recraft_image_edit_config():
from litellm.llms.recraft.image_edit.transformation import RecraftImageEditConfig
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
config = RecraftImageEditConfig()
# Test supported OpenAI params
supported_params = config.get_supported_openai_params("recraftv3")
expected_params = ["n", "response_format", "style"]
assert supported_params == expected_params
# Test parameter mapping (reuses OpenAI logic with filtering)
image_edit_params = ImageEditOptionalRequestParams({
"n": 2,
"response_format": "b64_json",
"style": "realistic_image",
"size": "1024x1024", # Should be dropped
"quality": "high" # Should be dropped
})
mapped_params = config.map_openai_params(image_edit_params, "recraftv3", drop_params=True)
image_edit_params = ImageEditOptionalRequestParams(
{
"n": 2,
"response_format": "b64_json",
"style": "realistic_image",
"size": "1024x1024", # Should be dropped
"quality": "high", # Should be dropped
}
)
mapped_params = config.map_openai_params(
image_edit_params, "recraftv3", drop_params=True
)
# Should only contain supported params
assert mapped_params["n"] == 2
assert mapped_params["response_format"] == "b64_json"
assert mapped_params["style"] == "realistic_image"
assert "size" not in mapped_params # Should be dropped
assert "quality" not in mapped_params # Should be dropped
# Test request transformation (reuses OpenAI file handling)
mock_image = b"fake_image_data"
prompt = "winter landscape"
litellm_params = GenericLiteLLMParams(api_key="test_key")
data, files = config.transform_image_edit_request(
model="recraftv3",
prompt=prompt,
image=mock_image,
image_edit_optional_request_params={"strength": 0.7, "n": 1},
litellm_params=litellm_params,
headers={}
headers={},
)
# Check data structure (like OpenAI but with Recraft additions)
assert data["prompt"] == prompt
assert data["strength"] == 0.7 # Recraft-specific parameter
assert data["model"] == "recraftv3"
# Check file structure (reuses OpenAI logic)
assert len(files) == 1
assert files[0][0] == "image" # Field name (not image[] like OpenAI)
@@ -603,11 +635,12 @@ def test_recraft_image_edit_config():
async def test_multiple_vs_single_image_edit(sync_mode):
"""Test that both single and multiple image editing work correctly"""
from litellm import image_edit, aimage_edit
litellm._turn_on_debug()
try:
prompt = "Add a soft blue tint to the image(s)"
# Test single image
if sync_mode:
single_result = image_edit(
@@ -621,10 +654,10 @@ async def test_multiple_vs_single_image_edit(sync_mode):
model="gpt-image-1",
image=SINGLE_TEST_IMAGE,
)
print("Single image result:", single_result)
ImageResponse.model_validate(single_result)
# Test multiple images
if sync_mode:
multiple_result = image_edit(
@@ -638,10 +671,10 @@ async def test_multiple_vs_single_image_edit(sync_mode):
model="gpt-image-1",
image=TEST_IMAGES,
)
print("Multiple images result:", multiple_result)
ImageResponse.model_validate(multiple_result)
# Both should return valid responses
assert single_result is not None
assert multiple_result is not None
@@ -649,7 +682,7 @@ async def test_multiple_vs_single_image_edit(sync_mode):
assert multiple_result.data is not None
assert len(single_result.data) > 0
assert len(multiple_result.data) > 0
except litellm.ContentPolicyViolationError as e:
pytest.skip(f"Content policy violation: {e}")
@@ -659,36 +692,37 @@ async def test_multiple_vs_single_image_edit(sync_mode):
async def test_multiple_image_edit_with_different_formats():
"""Test multiple images editing with different file formats and types"""
from litellm import aimage_edit
litellm._turn_on_debug()
try:
prompt = "Create a cohesive artistic style across all images"
# Test with mixed BytesIO and file objects
mixed_images = [
SINGLE_TEST_IMAGE, # File object
get_test_images_as_bytesio()[1] # BytesIO object
get_test_images_as_bytesio()[1], # BytesIO object
]
result = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=mixed_images,
)
print("Mixed format images result:", result)
ImageResponse.model_validate(result)
assert result is not None
assert result.data is not None
assert len(result.data) > 0
# Save result if available
if result.data and result.data[0].b64_json:
image_bytes = base64.b64decode(result.data[0].b64_json)
with open("test_multiple_image_edit_mixed.png", "wb") as f:
f.write(image_bytes)
except litellm.ContentPolicyViolationError as e:
pytest.skip(f"Content policy violation: {e}")
@@ -698,7 +732,7 @@ async def test_multiple_image_edit_with_different_formats():
async def test_image_edit_array_handling():
"""Test that the image parameter correctly handles both single items and arrays"""
from litellm import aimage_edit
# Mock response
mock_response = {
"created": 1589478378,
@@ -706,7 +740,7 @@ async def test_image_edit_array_handling():
{
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
}
]
],
}
class MockResponse:
@@ -723,29 +757,26 @@ async def test_image_edit_array_handling():
new_callable=AsyncMock,
) as mock_post:
mock_post.return_value = MockResponse(mock_response, 200)
prompt = "Test prompt"
# Test 1: Single image (should be converted to list internally)
result1 = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=SINGLE_TEST_IMAGE,
)
# Test 2: Multiple images (already a list)
result2 = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=TEST_IMAGES,
)
# Both valid calls should succeed
ImageResponse.model_validate(result1)
ImageResponse.model_validate(result2)
# Verify that both calls were made to the API
assert mock_post.call_count == 2
tests/image_gen_tests/test_image_generation.py
+27 -18
View File
@@ -113,7 +113,7 @@ class TestVertexImageGeneration(BaseImageGenTest):
litellm.in_memory_llm_clients_cache = InMemoryCache()
return {
"model": "vertex_ai/imagen-3.0-fast-generate-001",
"vertex_ai_project": "pathrise-convert-1606954137718",
"vertex_ai_project": "litellm-ci-cd",
"vertex_ai_location": "us-central1",
"n": 1,
}
@@ -121,6 +121,7 @@ class TestVertexImageGeneration(BaseImageGenTest):
class TestVertexAIGeminiImageGeneration(BaseImageGenTest):
"""Test Gemini image generation models (Nano Banana)"""
def get_base_image_generation_call_args(self) -> dict:
# comment this when running locally
load_vertex_ai_credentials()
@@ -128,7 +129,7 @@ class TestVertexAIGeminiImageGeneration(BaseImageGenTest):
litellm.in_memory_llm_clients_cache = InMemoryCache()
return {
"model": "vertex_ai/gemini-2.5-flash-image",
"vertex_ai_project": "pathrise-convert-1606954137718",
"vertex_ai_project": "litellm-ci-cd",
"vertex_ai_location": "us-central1",
"n": 1,
"size": "1024x1024",
@@ -212,7 +213,9 @@ class TestAimlImageGeneration(BaseImageGenTest):
custom_logger = TestCustomLogger()
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = [custom_logger]
base_image_generation_call_args = self.get_base_image_generation_call_args()
base_image_generation_call_args = (
self.get_base_image_generation_call_args()
)
litellm.set_verbose = True
# Pass dummy api_key so validate_environment passes; HTTP is mocked
response = await litellm.aimage_generation(
@@ -229,7 +232,9 @@ class TestAimlImageGeneration(BaseImageGenTest):
# print("response_cost", response._hidden_params["response_cost"])
logged_standard_logging_payload = custom_logger.standard_logging_payload
print("logged_standard_logging_payload", logged_standard_logging_payload)
print(
"logged_standard_logging_payload", logged_standard_logging_payload
)
assert logged_standard_logging_payload is not None
assert logged_standard_logging_payload["response_cost"] is not None
assert logged_standard_logging_payload["response_cost"] > 0
@@ -244,7 +249,9 @@ class TestAimlImageGeneration(BaseImageGenTest):
response_dict["usage"] = dict(response_dict["usage"])
print("response usage=", response_dict.get("usage"))
assert response.data is not None # type guard for iteration (base fails here if None)
assert (
response.data is not None
) # type guard for iteration (base fails here if None)
for d in response.data:
assert isinstance(d, Image)
print("data in response.data", d)
@@ -266,25 +273,27 @@ class TestGoogleImageGen(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
return {"model": "gemini/imagen-4.0-generate-001"}
@pytest.mark.skip(reason="Runwayml image generation API only tested locally")
class TestRunwaymlImageGeneration(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
return {"model": "runwayml/gen4_image"}
class TestAzureOpenAIDalle3(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
return {
"model": "azure/dall-e-3",
"api_version": "2024-02-01",
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_API_KEY"),
"metadata": {
"model_info": {
"base_model": "azure/dall-e-3",
}
},
}
## AZURE AI DALL-E 3 is deprecated and new deployments cannot be made
# class TestAzureOpenAIDalle3(BaseImageGenTest):
# def get_base_image_generation_call_args(self) -> dict:
# return {
# "model": "azure/dall-e-3",
# "api_version": "2024-02-01",
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "metadata": {
# "model_info": {
# "base_model": "azure/dall-e-3",
# }
# },
# }
@pytest.mark.skip(reason="model EOL")
tests/image_gen_tests/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py
+6 -8
View File
@@ -26,22 +26,20 @@ class TestAzureAIAnthropicTokenCounter(BaseTokenCounterTest):
return AzureAIAnthropicTokenCounter()
def get_test_model(self) -> str:
return "claude-3-5-sonnet"
return "claude-sonnet-4-6"
def get_test_messages(self) -> List[Dict[str, Any]]:
return [
{"role": "user", "content": "Hello, how are you today?"}
]
return [{"role": "user", "content": "Hello, how are you today?"}]
def get_deployment_config(self) -> Dict[str, Any]:
api_key = os.getenv("AZURE_AI_API_KEY")
api_base = os.getenv("AZURE_AI_API_BASE")
api_key = os.getenv("AZURE_ANTHROPIC_API_KEY")
api_base = os.getenv("AZURE_AI_SWEDEN_API_BASE")
if not api_key:
pytest.skip("AZURE_AI_API_KEY not set")
if not api_base:
pytest.skip("AZURE_AI_API_BASE not set")
return {
"litellm_params": {
"api_key": api_key,
tests/litellm_utils_tests/test_health_check.py
+16 -40
View File
@@ -21,9 +21,9 @@ async def test_azure_health_check():
model_params={
"model": "azure/gpt-4.1-mini",
"messages": [{"role": "user", "content": "Hey, how's it going?"}],
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
}
)
print(f"response: {response}")
@@ -51,9 +51,9 @@ async def test_azure_embedding_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "azure/text-embedding-ada-002",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
},
input=["test for litellm"],
mode="embedding",
@@ -83,7 +83,9 @@ async def test_openai_img_gen_health_check():
# asyncio.run(test_openai_img_gen_health_check())
@pytest.mark.skip(reason="Azure DALL-E 3 model deployment is deprecated (410 ModelDeprecated)")
@pytest.mark.skip(
reason="Azure DALL-E 3 model deployment is deprecated (410 ModelDeprecated)"
)
@pytest.mark.asyncio
async def test_azure_img_gen_health_check():
"""
@@ -98,8 +100,8 @@ async def test_azure_img_gen_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "azure/dall-e-3",
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
},
mode="image_generation",
prompt="cute baby sea otter",
@@ -244,35 +246,6 @@ async def test_audio_transcription_health_check():
print(response)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model", ["azure/gpt-4o-realtime-preview", "openai/gpt-4o-realtime-preview"]
)
async def test_async_realtime_health_check(model, mocker):
"""
Test Health Check with Valid models passes
"""
mock_websocket = AsyncMock()
mock_connect = AsyncMock().__aenter__.return_value = mock_websocket
mocker.patch("websockets.connect", return_value=mock_connect)
litellm.set_verbose = True
model_params = {
"model": model,
}
if model == "azure/gpt-4o-realtime-preview":
model_params["api_base"] = os.getenv("AZURE_REALTIME_API_BASE")
model_params["api_key"] = os.getenv("AZURE_REALTIME_API_KEY")
model_params["api_version"] = os.getenv("AZURE_REALTIME_API_VERSION")
response = await litellm.ahealth_check(
model_params=model_params,
mode="realtime",
)
print(response)
assert response == {}
def test_update_litellm_params_for_health_check():
"""
Test if _update_litellm_params_for_health_check correctly:
@@ -500,7 +473,9 @@ async def test_perform_health_check_filters_by_model_id():
async def mock_perform_health_check(m_list, details=True, **kwargs):
captured_list.append(m_list)
return [{"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]}], []
return [
{"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]}
], []
with patch(
"litellm.proxy.health_check._perform_health_check",
@@ -657,7 +632,8 @@ async def test_health_check_creates_only_bounded_initial_tasks():
return real_create_task(coro)
with patch("litellm.ahealth_check", side_effect=mock_health_check), patch(
"litellm.proxy.health_check.asyncio.create_task", side_effect=tracked_create_task
"litellm.proxy.health_check.asyncio.create_task",
side_effect=tracked_create_task,
):
perform_task = real_create_task(
_perform_health_check(model_list, max_concurrency=2)
tests/litellm_utils_tests/test_litellm_overhead.py
+1 -1
View File
@@ -101,7 +101,7 @@ async def test_litellm_overhead_non_streaming(model):
kwargs["vertex_project"] = "fake-project"
kwargs["vertex_location"] = "us-central1"
if model == "openai/self_hosted":
kwargs["api_base"] = "https://exampleopenaiendpoint-production.up.railway.app/"
kwargs["api_base"] = os.environ.get("FAKE_OPENAI_API_BASE")
async def _run():
return await litellm.acompletion(**kwargs)
tests/litellm_utils_tests/test_secret_manager.py
+58 -30
View File
@@ -1,3 +1,4 @@
import base64
import os
import sys
import time
@@ -24,7 +25,7 @@ from litellm.secret_managers.main import (
get_secret,
_should_read_secret_from_secret_manager,
)
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, patch, MagicMock
def load_vertex_ai_credentials():
@@ -221,53 +222,79 @@ def test_oidc_env_path():
del os.environ[env_var_name]
@pytest.mark.flaky(retries=6, delay=1)
def test_google_secret_manager():
"""
Test that we can get a secret from Google Secret Manager
"""
os.environ["GOOGLE_SECRET_MANAGER_PROJECT_ID"] = "pathrise-convert-1606954137718"
os.environ["GOOGLE_SECRET_MANAGER_PROJECT_ID"] = "litellm-ci-cd"
from litellm.secret_managers.google_secret_manager import GoogleSecretManager
load_vertex_ai_credentials()
secret_manager = GoogleSecretManager()
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"payload": {
"data": base64.b64encode(b"anything").decode("utf-8"),
}
}
secret_val = secret_manager.get_secret_from_google_secret_manager(
secret_name="OPENAI_API_KEY"
)
print("secret_val: {}".format(secret_val))
with patch(
"litellm.proxy.proxy_server.premium_user", True
), patch.object(
GoogleSecretManager,
"sync_construct_request_headers",
return_value={"Authorization": "Bearer mock_token"},
):
secret_manager = GoogleSecretManager()
secret_manager.sync_httpx_client = MagicMock()
secret_manager.sync_httpx_client.get.return_value = mock_response
assert (
secret_val == "anything"
), "did not get expected secret value. expect 'anything', got '{}'".format(
secret_val
)
secret_val = secret_manager.get_secret_from_google_secret_manager(
secret_name="OPENAI_API_KEY"
)
print("secret_val: {}".format(secret_val))
assert (
secret_val == "anything"
), "did not get expected secret value. expect 'anything', got '{}'".format(
secret_val
)
secret_manager.sync_httpx_client.get.assert_called_once()
call_url = secret_manager.sync_httpx_client.get.call_args[1]["url"]
assert "projects/litellm-ci-cd/secrets/OPENAI_API_KEY" in call_url
def test_google_secret_manager_read_in_memory():
"""
Test that Google Secret manager returs in memory value when it exists
Test that Google Secret manager returns in memory value when it exists
"""
from litellm.secret_managers.google_secret_manager import GoogleSecretManager
load_vertex_ai_credentials()
os.environ["GOOGLE_SECRET_MANAGER_PROJECT_ID"] = "pathrise-convert-1606954137718"
secret_manager = GoogleSecretManager()
secret_manager.cache.cache_dict["UNIQUE_KEY"] = None
secret_manager.cache.cache_dict["UNIQUE_KEY_2"] = "lite-llm"
os.environ["GOOGLE_SECRET_MANAGER_PROJECT_ID"] = "litellm-ci-cd"
secret_val = secret_manager.get_secret_from_google_secret_manager(
secret_name="UNIQUE_KEY"
)
print("secret_val: {}".format(secret_val))
assert secret_val == None
with patch(
"litellm.proxy.proxy_server.premium_user", True
), patch.object(
GoogleSecretManager,
"sync_construct_request_headers",
return_value={"Authorization": "Bearer mock_token"},
):
secret_manager = GoogleSecretManager()
secret_manager.cache.cache_dict["UNIQUE_KEY"] = None
secret_manager.cache.cache_dict["UNIQUE_KEY_2"] = "lite-llm"
secret_val = secret_manager.get_secret_from_google_secret_manager(
secret_name="UNIQUE_KEY_2"
)
print("secret_val: {}".format(secret_val))
assert secret_val == "lite-llm"
secret_val = secret_manager.get_secret_from_google_secret_manager(
secret_name="UNIQUE_KEY"
)
print("secret_val: {}".format(secret_val))
assert secret_val is None
secret_val = secret_manager.get_secret_from_google_secret_manager(
secret_name="UNIQUE_KEY_2"
)
print("secret_val: {}".format(secret_val))
assert secret_val == "lite-llm"
def test_should_read_secret_from_secret_manager():
@@ -337,6 +364,7 @@ def test_get_secret_with_access_mode():
litellm._key_management_settings = KeyManagementSettings()
del os.environ[test_secret_name]
def test_key_management_settings_defaults():
"""
Test that KeyManagementSettings initializes with correct default values.
tests/litellm_utils_tests/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
tests/llm_responses_api_testing/test_azure_responses_api.py
+34 -31
View File
@@ -25,11 +25,11 @@ class TestAzureResponsesAPITest(BaseResponsesAPITest):
return {
"model": "azure/gpt-4.1-mini",
"truncation": "auto",
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": "2025-03-01-preview",
}
def get_advanced_model_for_shell_tool(self) -> Optional[str]:
"""If specified, overrides the model used by test_responses_api_shell_tool_streaming_sees_shell_output (e.g. openai/gpt-5.2 for shell support)."""
return "azure/gpt-5-mini"
@@ -45,8 +45,8 @@ async def test_azure_responses_api_preview_api_version():
model="azure/gpt-5-mini",
truncation="auto",
api_version="preview",
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
input="Hello, can you tell me a short joke?",
)
@@ -108,7 +108,9 @@ async def test_azure_responses_api_status_error():
"role": "assistant",
"type": "message",
"status": "completed",
"content": [{"type": "output_text", "text": "Here's an interesting fact."}],
"content": [
{"type": "output_text", "text": "Here's an interesting fact."}
],
}
],
}
@@ -124,7 +126,7 @@ async def test_azure_responses_api_status_error():
captured_request_body = json.loads(kwargs["data"])
import httpx
# Create a proper httpx Response object
response_content = json.dumps(mock_response_data).encode("utf-8")
response = httpx.Response(
@@ -149,18 +151,17 @@ async def test_azure_responses_api_status_error():
)
# Verify that 'status' field is not present in any of the input messages
print("Final request body:", json.dumps(captured_request_body, indent=4, default=str))
print(
"Final request body:", json.dumps(captured_request_body, indent=4, default=str)
)
assert "input" in captured_request_body, "Request body should contain 'input' field"
expected_input = [
{
"content": "tell me an interesting fact",
"role": "user"
},
{"content": "tell me an interesting fact", "role": "user"},
{
"id": "rs_0ab687487834d9df0068e462a1b2d88197aabbc832c9ba5316",
"summary": [],
"type": "reasoning"
"type": "reasoning",
},
{
"id": "msg_0ab687487834d9df0068e462a1df188197b74b1eef05102c18",
@@ -169,18 +170,15 @@ async def test_azure_responses_api_status_error():
"annotations": [],
"text": "very good morning",
"type": "output_text",
"logprobs": []
"logprobs": [],
}
],
"role": "assistant",
"type": "message"
"type": "message",
},
{
"role": "user",
"content": "tell me another"
}
{"role": "user", "content": "tell me another"},
]
assert captured_request_body["input"] == expected_input, (
f"Request body input should match expected format without 'status' field.\n"
f"Expected: {json.dumps(expected_input, indent=2)}\n"
@@ -193,9 +191,9 @@ async def test_azure_responses_api_headers_with_llm_provider_prefix():
"""
Test that Azure-specific headers like 'x-request-id' and 'apim-request-id'
are properly forwarded with 'llm_provider-' prefix in response._hidden_params["headers"].
Issue: https://github.com/BerriAI/litellm/issues/16538
The fix ensures that processed headers (with llm_provider- prefix) are stored
in response._hidden_params["headers"] instead of additional_headers, making them
accessible via completion.headers in the same way as the completion API.
@@ -253,12 +251,12 @@ async def test_azure_responses_api_headers_with_llm_provider_prefix():
# Check that the response has the expected headers structure
assert hasattr(response, "_hidden_params"), "Response should have _hidden_params"
assert "additional_headers" in response._hidden_params, (
"Response _hidden_params should contain 'additional_headers' with the LLM provider headers"
)
assert (
"additional_headers" in response._hidden_params
), "Response _hidden_params should contain 'additional_headers' with the LLM provider headers"
headers = response._hidden_params["additional_headers"]
# Verify that Azure-specific headers are present with llm_provider- prefix
assert "llm_provider-x-request-id" in headers, (
f"Response should contain 'llm_provider-x-request-id' header. "
@@ -268,12 +266,17 @@ async def test_azure_responses_api_headers_with_llm_provider_prefix():
f"Response should contain 'llm_provider-apim-request-id' header. "
f"Headers: {list(headers.keys())}"
)
# Verify the header values match
assert headers["llm_provider-x-request-id"] == "12086715-aca3-4006-a29f-2f1e1d552043"
assert headers["llm_provider-apim-request-id"] == "25664b0d-cf4b-4e10-8d27-c7272e7efd49"
assert (
headers["llm_provider-x-request-id"] == "12086715-aca3-4006-a29f-2f1e1d552043"
)
assert (
headers["llm_provider-apim-request-id"]
== "25664b0d-cf4b-4e10-8d27-c7272e7efd49"
)
assert headers["llm_provider-x-ms-region"] == "Sweden Central"
# Also verify openai-compatible headers are included
assert "x-ratelimit-limit-tokens" in headers
assert "x-ratelimit-remaining-tokens" in headers
tests/llm_translation/test_anthropic_completion.py
+7
View File
@@ -1350,6 +1350,9 @@ def test_anthropic_text_editor():
@pytest.mark.parametrize("spec", ["anthropic", "openai"])
@pytest.mark.skipif(
os.getenv("ZAPIER_CI_CD_MCP_TOKEN") is None, reason="ZAPIER_CI_CD_MCP_TOKEN not set"
)
def test_anthropic_mcp_server_tool_use(spec: str):
litellm._turn_on_debug()
@@ -1391,6 +1394,9 @@ def test_anthropic_mcp_server_tool_use(spec: str):
@pytest.mark.parametrize(
"model", ["openai/gpt-4.1", "anthropic/claude-sonnet-4-20250514"]
)
@pytest.mark.skipif(
os.getenv("ZAPIER_CI_CD_MCP_TOKEN") is None, reason="ZAPIER_CI_CD_MCP_TOKEN not set"
)
def test_anthropic_mcp_server_responses_api(model: str):
from litellm import responses
@@ -1594,6 +1600,7 @@ def test_anthropic_via_responses_api():
ResponsesAPIStreamEvents.RESPONSE_CREATED,
ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS,
ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, # Can occur multiple times
ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
ResponsesAPIStreamEvents.CONTENT_PART_DONE,
tests/llm_translation/test_azure_ai.py
+75 -71
View File
@@ -34,6 +34,8 @@ from litellm import completion
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import StandardLoggingPayload
AZURE_AI_API_BASE = os.getenv("AZURE_AI_API_BASE")
@pytest.mark.parametrize(
"model_group_header, expected_model",
@@ -188,35 +190,6 @@ def test_azure_ai_services_with_api_version():
)
@pytest.mark.skip(reason="Skipping due to cohere ssl issues")
def test_completion_azure_ai_command_r():
try:
import os
litellm.set_verbose = True
os.environ["AZURE_AI_API_BASE"] = os.getenv("AZURE_COHERE_API_BASE", "")
os.environ["AZURE_AI_API_KEY"] = os.getenv("AZURE_COHERE_API_KEY", "")
response = completion(
model="azure_ai/command-r-plus",
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "What is the meaning of life?"}
],
}
],
) # type: ignore
assert "azure_ai" in response.model
except litellm.Timeout as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_azure_deepseek_reasoning_content():
import json
@@ -283,8 +256,8 @@ async def test_azure_ai_request_format():
litellm._turn_on_debug()
# Set up the test parameters
api_key = os.getenv("AZURE_API_KEY")
api_base = os.getenv("AZURE_API_BASE")
api_key = os.getenv("AZURE_AI_API_KEY")
api_base = os.getenv("AZURE_AI_API_BASE")
model = "azure_ai/gpt-4.1-mini"
messages = [
{"role": "user", "content": "hi"},
@@ -310,17 +283,17 @@ async def test_azure_gpt5_reasoning(model):
messages=[{"role": "user", "content": "What is the capital of France?"}],
reasoning_effort="minimal",
max_tokens=10,
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
)
print("response: ", response)
assert response.choices[0].message.content is not None
def test_completion_azure():
try:
from litellm import completion_cost
litellm.set_verbose = False
## Test azure call
response = completion(
@@ -331,7 +304,7 @@ def test_completion_azure():
"content": "Hello, how are you?",
}
],
api_key="os.environ/AZURE_API_KEY",
api_key="os.environ/AZURE_AI_API_KEY",
)
print(f"response: {response}")
print(f"response hidden params: {response._hidden_params}")
@@ -347,8 +320,8 @@ def test_completion_azure():
@pytest.mark.parametrize(
"api_base",
[
"https://litellm-ci-cd-prod.cognitiveservices.azure.com/",
"https://litellm-ci-cd-prod.cognitiveservices.azure.com/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2023-03-15-preview",
AZURE_AI_API_BASE,
f"{AZURE_AI_API_BASE}/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2023-03-15-preview",
],
)
def test_completion_azure_ai_gpt_4o_with_flexible_api_base(api_base):
@@ -358,7 +331,7 @@ def test_completion_azure_ai_gpt_4o_with_flexible_api_base(api_base):
response = completion(
model="azure_ai/gpt-4.1-mini",
api_base=api_base,
api_key=os.getenv("AZURE_API_KEY"),
api_key=os.getenv("AZURE_AI_API_KEY"),
messages=[{"role": "user", "content": "What is the meaning of life?"}],
)
@@ -374,18 +347,20 @@ async def test_azure_ai_model_router():
"""
Test Azure AI model router non-streaming response cost tracking.
Verifies that the flat cost of $0.14 per M input tokens is applied.
Tests the pattern: azure_ai/model_router/<deployment-name>
Where deployment-name is the Azure deployment (e.g., "azure-model-router").
The model_router prefix is stripped before sending to Azure API.
"""
from litellm.llms.azure_ai.cost_calculator import calculate_azure_model_router_flat_cost
from litellm.llms.azure_ai.cost_calculator import (
calculate_azure_model_router_flat_cost,
)
litellm._turn_on_debug()
response = await litellm.acompletion(
model="azure_ai/model_router/azure-model-router",
messages=[{"role": "user", "content": "hi who is this"}],
api_base="https://ishaa-mh6uutut-swedencentral.cognitiveservices.azure.com/openai/v1/",
api_base=os.getenv("AZURE_MODEL_ROUTER_API_BASE"),
api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"),
)
print("response: ", response)
@@ -394,23 +369,22 @@ async def test_azure_ai_model_router():
tracked_cost = response._hidden_params["response_cost"]
assert tracked_cost > 0
print("Tracked cost: ", tracked_cost)
# Verify flat cost is included using the helper function
usage = response.usage
if usage and usage.prompt_tokens:
expected_flat_cost = calculate_azure_model_router_flat_cost(
model="model_router/azure-model-router",
prompt_tokens=usage.prompt_tokens
model="model_router/azure-model-router", prompt_tokens=usage.prompt_tokens
)
print(f"Prompt tokens: {usage.prompt_tokens}")
print(f"Expected flat cost: ${expected_flat_cost:.9f}")
print(f"Total tracked cost: ${tracked_cost:.9f}")
# Total cost should be at least the flat cost
assert tracked_cost >= expected_flat_cost, (
f"Cost ${tracked_cost:.9f} should be >= flat cost ${expected_flat_cost:.9f}"
)
assert (
tracked_cost >= expected_flat_cost
), f"Cost ${tracked_cost:.9f} should be >= flat cost ${expected_flat_cost:.9f}"
# Verify the flat cost is non-zero
assert expected_flat_cost > 0, "Flat cost should be greater than 0"
@@ -425,7 +399,7 @@ async def test_azure_ai_model_router_streaming_model_in_chunk():
response = await litellm.acompletion(
model="azure_ai/azure-model-router",
messages=[{"role": "user", "content": "hi"}],
api_base="https://ishaa-mh6uutut-swedencentral.cognitiveservices.azure.com/openai/v1/",
api_base=os.getenv("AZURE_MODEL_ROUTER_API_BASE"),
api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"),
stream=True,
)
@@ -445,15 +419,20 @@ async def test_azure_ai_model_router_streaming_model_in_chunk():
# The model should NOT be azure-model-router (the request model)
# It should be the actual model from the response (e.g., gpt-4.1-nano, gpt-5-nano, etc.)
for model in chunks_with_model:
assert model != "azure-model-router", f"Chunk model should be actual model, not request model. Got: {model}"
assert (
model != "azure-model-router"
), f"Chunk model should be actual model, not request model. Got: {model}"
# The actual model should be a real model name like gpt-4.1-nano, gpt-5-nano, etc.
print(f"Verified chunk has actual model: {model}")
class AzureModelRouterStreamingCallback(litellm.integrations.custom_logger.CustomLogger):
class AzureModelRouterStreamingCallback(
litellm.integrations.custom_logger.CustomLogger
):
"""
Custom callback to capture streaming cost tracking for Azure Model Router.
"""
def __init__(self):
self.standard_logging_payload = None
self.response_cost = None
@@ -466,17 +445,21 @@ class AzureModelRouterStreamingCallback(litellm.integrations.custom_logger.Custo
self.async_success_called = True
self.standard_logging_payload = kwargs.get("standard_logging_object")
self.complete_streaming_response = kwargs.get("complete_streaming_response")
if self.standard_logging_payload:
self.response_cost = self.standard_logging_payload.get("response_cost")
print(f"standard_logging_payload model: {self.standard_logging_payload.get('model')}")
print(
f"standard_logging_payload model: {self.standard_logging_payload.get('model')}"
)
print(f"standard_logging_payload response_cost: {self.response_cost}")
if self.complete_streaming_response:
print(f"complete_streaming_response model: {self.complete_streaming_response.model}")
print(f"complete_streaming_response usage: {self.complete_streaming_response.usage}")
print(
f"complete_streaming_response model: {self.complete_streaming_response.model}"
)
print(
f"complete_streaming_response usage: {self.complete_streaming_response.usage}"
)
@pytest.mark.asyncio
@@ -504,10 +487,16 @@ async def test_azure_ai_model_router_streaming_cost_with_stream_options():
full_response = ""
chunks_with_model = []
async for chunk in response:
print(f"Chunk: model={chunk.model}, choices={len(chunk.choices) if chunk.choices else 0}")
print(
f"Chunk: model={chunk.model}, choices={len(chunk.choices) if chunk.choices else 0}"
)
if chunk.model:
chunks_with_model.append(chunk.model)
if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
if (
chunk.choices
and chunk.choices[0].delta
and chunk.choices[0].delta.content
):
full_response += chunk.choices[0].delta.content
print(f"Full streamed response: {full_response}")
@@ -515,27 +504,42 @@ async def test_azure_ai_model_router_streaming_cost_with_stream_options():
# Give async logging time to complete
import asyncio
await asyncio.sleep(1)
# Verify callback was called
assert test_callback.async_success_called is True, "async_log_success_event was not called"
assert test_callback.standard_logging_payload is not None, "standard_logging_payload is None"
assert (
test_callback.async_success_called is True
), "async_log_success_event was not called"
assert (
test_callback.standard_logging_payload is not None
), "standard_logging_payload is None"
# Check response cost
print(f"Final response_cost: {test_callback.response_cost}")
# The first chunk may have the request model (azure-model-router) because it's created
# before the API response is received. Subsequent chunks should have the actual model.
# At least some chunks should have the actual model (not azure-model-router)
actual_model_chunks = [m for m in chunks_with_model if m != "azure-model-router"]
assert len(actual_model_chunks) > 0, "No chunks had the actual model from the API response"
actual_model_chunks = [
m for m in chunks_with_model if m != "azure-model-router"
]
assert (
len(actual_model_chunks) > 0
), "No chunks had the actual model from the API response"
print(f"Chunks with actual model: {actual_model_chunks}")
# Verify response cost is tracked - this is the main goal of this test
assert test_callback.response_cost is not None, "response_cost is None with stream_options"
assert test_callback.response_cost > 0, f"response_cost should be > 0, got {test_callback.response_cost}"
print(f"Streaming cost tracking with stream_options passed. Cost: {test_callback.response_cost}")
assert (
test_callback.response_cost is not None
), "response_cost is None with stream_options"
assert (
test_callback.response_cost > 0
), f"response_cost should be > 0, got {test_callback.response_cost}"
print(
f"Streaming cost tracking with stream_options passed. Cost: {test_callback.response_cost}"
)
finally:
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = []
litellm.callbacks = []
tests/llm_translation/test_azure_o_series.py
+27 -13
View File
@@ -24,9 +24,9 @@ class TestAzureOpenAIO3Mini(BaseOSeriesModelsTest, BaseLLMChatTest):
litellm.in_memory_llm_clients_cache.flush_cache()
return {
"model": "azure/o3-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": "2024-12-01-preview"
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_version": "2024-12-01-preview",
}
def get_client(self):
@@ -187,13 +187,31 @@ async def test_azure_o1_series_response_format_extra_params():
litellm.set_verbose = True
client = AsyncAzureOpenAI(
api_key="fake-api-key",
base_url="https://openai-prod-test.openai.azure.com/openai/deployments/o1/chat/completions?api-version=2025-01-01-preview",
api_version="2025-01-01-preview"
api_key="fake-api-key",
base_url="https://openai-prod-test.openai.azure.com/openai/deployments/o1/chat/completions?api-version=2025-01-01-preview",
api_version="2025-01-01-preview",
)
tools = [{'type': 'function', 'function': {'name': 'get_current_time', 'description': 'Get the current time in a given location.', 'parameters': {'type': 'object', 'properties': {'location': {'type': 'string', 'description': 'The city name, e.g. San Francisco'}}, 'required': ['location']}}}]
response_format = {'type': 'json_object'}
tools = [
{
"type": "function",
"function": {
"name": "get_current_time",
"description": "Get the current time in a given location.",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city name, e.g. San Francisco",
}
},
"required": ["location"],
},
},
}
]
response_format = {"type": "json_object"}
tool_choice = "auto"
with patch.object(
client.chat.completions.with_raw_response, "create"
@@ -208,7 +226,7 @@ async def test_azure_o1_series_response_format_extra_params():
messages=[{"role": "user", "content": "Hello! return a json object"}],
tools=tools,
response_format=response_format,
tool_choice=tool_choice
tool_choice=tool_choice,
)
except Exception as e:
print(f"Error: {e}")
@@ -220,7 +238,3 @@ async def test_azure_o1_series_response_format_extra_params():
assert request_body["tools"] == tools
assert request_body["response_format"] == response_format
assert request_body["tool_choice"] == tool_choice
tests/llm_translation/test_azure_openai.py
+39 -37
View File
@@ -208,8 +208,8 @@ class TestAzureEmbedding(BaseLLMEmbeddingTest):
def get_base_embedding_call_args(self) -> dict:
return {
"model": "azure/text-embedding-ada-002",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
}
def get_custom_llm_provider(self) -> litellm.LlmProviders:
@@ -618,8 +618,8 @@ def test_azure_safety_result():
response = completion(
model="azure/gpt-4.1-mini",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version="2024-12-01-preview",
messages=[{"role": "user", "content": "Hello world"}],
)
@@ -671,6 +671,8 @@ def test_completion_azure_deployment_id():
)
# Add any assertions here to check the response
print(response)
def test_azure_with_content_safety_error():
"""
Verify user can access innererror from the Azure OpenAI exception
@@ -679,55 +681,55 @@ def test_azure_with_content_safety_error():
from litellm.exceptions import ContentPolicyViolationError
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
from unittest.mock import MagicMock
mock_exception = Exception("The response was filtered due to the prompt triggering Azure OpenAI's content management policy")
mock_exception = Exception(
"The response was filtered due to the prompt triggering Azure OpenAI's content management policy"
)
mock_exception.body = {
"innererror": {
"code": "ResponsibleAIPolicyViolation",
"content_filter_result": {
"hate": {
"filtered": False,
"severity": "safe"
},
"jailbreak": {
"filtered": False,
"detected": False
},
"self_harm": {
"filtered": False,
"severity": "safe"
},
"sexual": {
"filtered": False,
"severity": "safe"
},
"violence": {
"filtered": True,
"severity": "high"
}
}
"hate": {"filtered": False, "severity": "safe"},
"jailbreak": {"filtered": False, "detected": False},
"self_harm": {"filtered": False, "severity": "safe"},
"sexual": {"filtered": False, "severity": "safe"},
"violence": {"filtered": True, "severity": "high"},
},
}
}
mock_response = MagicMock()
mock_response.status_code = 400
mock_exception.response = mock_response
with pytest.raises(ContentPolicyViolationError) as exc_info:
exception_type(
model="azure/gpt-4o-new-test",
original_exception=mock_exception,
custom_llm_provider="azure"
custom_llm_provider="azure",
)
e = exc_info.value
print("got exception=", e)
assert e.provider_specific_fields is not None
print("got provider_specific_fields=", e.provider_specific_fields)
assert e.provider_specific_fields.get("innererror") is not None
assert e.provider_specific_fields["innererror"]["code"] == "ResponsibleAIPolicyViolation"
assert e.provider_specific_fields["innererror"]["content_filter_result"]["violence"]["filtered"] is True
assert e.provider_specific_fields["innererror"]["content_filter_result"]["violence"]["severity"] == "high"
assert (
e.provider_specific_fields["innererror"]["code"]
== "ResponsibleAIPolicyViolation"
)
assert (
e.provider_specific_fields["innererror"]["content_filter_result"]["violence"][
"filtered"
]
is True
)
assert (
e.provider_specific_fields["innererror"]["content_filter_result"]["violence"][
"severity"
]
== "high"
)
def test_azure_openai_with_prompt_cache_key():
@@ -737,9 +739,9 @@ def test_azure_openai_with_prompt_cache_key():
litellm._turn_on_debug()
response = litellm.completion(
model="azure/gpt-4.1-mini",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version="2024-12-01-preview",
messages=[{"role": "user", "content": "What is the weather in San Francisco?"}],
prompt_cache_key="test_streaming_azure_openai",
)
)
tests/llm_translation/test_bedrock_completion.py
+15 -60
View File
@@ -217,60 +217,6 @@ def test_completion_bedrock_claude_external_client_auth():
# test_completion_bedrock_claude_external_client_auth()
@pytest.mark.skip(reason="Expired token, need to renew")
def test_completion_bedrock_claude_sts_client_auth():
print("\ncalling bedrock claude external client auth")
import os
aws_access_key_id = os.environ["AWS_TEMP_ACCESS_KEY_ID"]
aws_secret_access_key = os.environ["AWS_TEMP_SECRET_ACCESS_KEY"]
aws_region_name = os.environ["AWS_REGION_NAME"]
aws_role_name = os.environ["AWS_TEMP_ROLE_NAME"]
try:
import boto3
litellm.set_verbose = True
response = completion(
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
messages=messages,
max_tokens=10,
temperature=0.1,
aws_region_name=aws_region_name,
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_role_name=aws_role_name,
aws_session_name="my-test-session",
)
response = embedding(
model="cohere.embed-multilingual-v3",
input=["hello world"],
aws_region_name="us-east-1",
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_role_name=aws_role_name,
aws_session_name="my-test-session",
)
response = completion(
model="gpt-3.5-turbo",
messages=messages,
aws_region_name="us-east-1",
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_role_name=aws_role_name,
aws_session_name="my-test-session",
)
# Add any assertions here to check the response
print(response)
except RateLimitError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.fixture()
def bedrock_session_token_creds():
print("\ncalling oidc auto to get aws_session_token credentials")
@@ -3413,7 +3359,8 @@ def test_bedrock_openai_imported_model():
print(f"URL: {url}")
assert "bedrock-runtime.us-east-1.amazonaws.com" in url
assert (
"arn:aws:bedrock:us-east-1:117159858402:imported-model%2Fm4gc1mrfuddy" in url
"arn:aws:bedrock:us-east-1:117159858402:imported-model%2Fm4gc1mrfuddy"
in url
)
assert "/invoke" in url
@@ -3850,10 +3797,12 @@ def test_bedrock_openai_error_handling():
assert exc_info.value.status_code == 422
print("✓ Error handling works correctly")
# ============================================================================
# Nova Grounding (web_search_options) Unit Tests (Mocked)
# ============================================================================
def test_bedrock_nova_grounding_web_search_options_non_streaming():
"""
Unit test for Nova grounding using web_search_options parameter (non-streaming).
@@ -3907,7 +3856,9 @@ def test_bedrock_nova_grounding_web_search_options_non_streaming():
break
assert system_tool_found, "systemTool with nova_grounding should be present"
print(f"✓ web_search_options correctly transformed to systemTool (non-streaming)")
print(
f"✓ web_search_options correctly transformed to systemTool (non-streaming)"
)
def test_bedrock_nova_grounding_with_function_tools():
@@ -3987,7 +3938,9 @@ def test_bedrock_nova_grounding_with_function_tools():
assert tool["systemTool"]["name"] == "nova_grounding"
system_tool_found = True
assert function_tool_found, "Function tool (get_stock_price) should be present"
assert (
function_tool_found
), "Function tool (get_stock_price) should be present"
assert system_tool_found, "systemTool (nova_grounding) should be present"
print(f"✓ Both function tools and web_search_options correctly combined")
@@ -4092,10 +4045,12 @@ def test_bedrock_nova_grounding_request_transformation():
mock_post.return_value = MagicMock(
status_code=200,
json=lambda: {
"output": {"message": {"role": "assistant", "content": [{"text": "Test"}]}},
"output": {
"message": {"role": "assistant", "content": [{"text": "Test"}]}
},
"stopReason": "end_turn",
"usage": {"inputTokens": 10, "outputTokens": 5}
}
"usage": {"inputTokens": 10, "outputTokens": 5},
},
)
try:
tests/llm_translation/test_clarifai_completion.py
-109
View File
@@ -1,109 +0,0 @@
import sys, os
import traceback
from dotenv import load_dotenv
import asyncio, logging
load_dotenv()
import os, io
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import pytest
import litellm
from litellm import (
embedding,
completion,
acompletion,
acreate,
completion_cost,
Timeout,
ModelResponse,
)
from litellm import RateLimitError
# litellm.num_retries = 3
litellm.cache = None
litellm.success_callback = []
user_message = "Write a short poem about the sky"
messages = [{"content": user_message, "role": "user"}]
@pytest.fixture(autouse=True)
def reset_callbacks():
print("\npytest fixture - resetting callbacks")
litellm.success_callback = []
litellm._async_success_callback = []
litellm.failure_callback = []
litellm.callbacks = []
@pytest.mark.skip(reason="Account rate limited.")
def test_completion_clarifai_claude_2_1():
print("calling clarifai claude completion")
import os
clarifai_pat = os.environ["CLARIFAI_API_KEY"]
try:
response = completion(
model="clarifai/anthropic.completion.claude-2_1",
num_retries=3,
messages=messages,
max_tokens=10,
temperature=0.1,
)
print(response)
except RateLimitError:
pass
except Exception as e:
pytest.fail(f"Error occured: {e}")
@pytest.mark.skip(reason="Account rate limited")
def test_completion_clarifai_mistral_large():
try:
litellm.set_verbose = True
response: ModelResponse = completion(
model="clarifai/mistralai.completion.mistral-small",
messages=messages,
num_retries=3,
max_tokens=10,
temperature=0.78,
)
# Add any assertions here to check the response
assert len(response.choices) > 0
assert len(response.choices[0].message.content) > 0
except RateLimitError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skip(reason="Account rate limited")
@pytest.mark.asyncio
def test_async_completion_clarifai():
import asyncio
litellm.set_verbose = True
async def test_get_response():
user_message = "Hello, how are you?"
messages = [{"content": user_message, "role": "user"}]
try:
response = await acompletion(
model="clarifai/openai.chat-completion.GPT-4",
messages=messages,
num_retries=3,
timeout=10,
api_key=os.getenv("CLARIFAI_API_KEY"),
)
print(f"response: {response}")
except litellm.Timeout as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred: {e}")
asyncio.run(test_get_response())
tests/llm_translation/test_cloudflare.py
+138 -35
View File
@@ -1,42 +1,145 @@
import os
import sys
import traceback
from dotenv import load_dotenv
load_dotenv()
import io
import os
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import asyncio
import json
from typing import Any, Dict
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import litellm
from litellm import RateLimitError, Timeout, completion, completion_cost, embedding
from litellm import acompletion, completion
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
FAKE_API_BASE = "https://fake-cloudflare.example.com/client/v4/accounts/fake-acct/ai/run/"
FAKE_API_KEY = "fake-cf-api-key"
# Cloud flare AI test
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [True, False])
async def test_completion_cloudflare(stream):
try:
litellm.set_verbose = False
response = await litellm.acompletion(
model="cloudflare/@cf/meta/llama-2-7b-chat-int8",
messages=[{"content": "what llm are you", "role": "user"}],
max_tokens=15,
stream=stream,
)
print(response)
if stream is True:
async for chunk in response:
print(chunk)
else:
print(response)
def _make_mock_response(json_data: Dict[str, Any]) -> MagicMock:
mock = MagicMock(spec=httpx.Response)
mock.status_code = 200
mock.headers = {"content-type": "application/json"}
mock.json.return_value = json_data
mock.text = json.dumps(json_data)
return mock
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def _chat_response() -> Dict[str, Any]:
return {
"result": {
"response": "I am a large language model created to assist you.",
},
"success": True,
"errors": [],
"messages": [],
}
def _streaming_chunks() -> list[str]:
return [
json.dumps({"response": "I am"}),
json.dumps({"response": " a language"}),
json.dumps({"response": " model."}),
]
@pytest.mark.parametrize("sync_mode", [True, False])
def test_completion_cloudflare(sync_mode):
messages = [{"role": "user", "content": "what llm are you"}]
mock_resp = _make_mock_response(_chat_response())
if sync_mode:
with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post:
response = completion(
model="cloudflare/@cf/meta/llama-2-7b-chat-int8",
messages=messages,
max_tokens=15,
api_base=FAKE_API_BASE,
api_key=FAKE_API_KEY,
)
mock_post.assert_called_once()
else:
with patch.object(
AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp
) as mock_post:
response = asyncio.run(
acompletion(
model="cloudflare/@cf/meta/llama-2-7b-chat-int8",
messages=messages,
max_tokens=15,
api_base=FAKE_API_BASE,
api_key=FAKE_API_KEY,
)
)
mock_post.assert_called_once()
assert response is not None
assert response.choices[0].message.content is not None
assert "language model" in response.choices[0].message.content.lower()
@pytest.mark.parametrize("sync_mode", [True, False])
def test_completion_cloudflare_stream(sync_mode):
messages = [{"role": "user", "content": "what llm are you"}]
raw_chunks = _streaming_chunks()
if sync_mode:
def _iter_lines():
for chunk in raw_chunks:
yield f"data: {chunk}"
yield "data: [DONE]"
mock_resp = MagicMock()
mock_resp.iter_lines.return_value = _iter_lines()
mock_resp.status_code = 200
mock_resp.headers = {"content-type": "text/event-stream"}
with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post:
response = completion(
model="cloudflare/@cf/meta/llama-2-7b-chat-int8",
messages=messages,
max_tokens=15,
stream=True,
api_base=FAKE_API_BASE,
api_key=FAKE_API_KEY,
)
chunks_received = list(response)
mock_post.assert_called_once()
else:
async def _aiter_lines():
for chunk in raw_chunks:
yield f"data: {chunk}"
yield "data: [DONE]"
mock_resp = MagicMock()
mock_resp.aiter_lines.return_value = _aiter_lines()
mock_resp.status_code = 200
mock_resp.headers = {"content-type": "text/event-stream"}
async def _run():
with patch.object(
AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp
) as mock_post:
resp = await acompletion(
model="cloudflare/@cf/meta/llama-2-7b-chat-int8",
messages=messages,
max_tokens=15,
stream=True,
api_base=FAKE_API_BASE,
api_key=FAKE_API_KEY,
)
received = []
async for chunk in resp:
received.append(chunk)
mock_post.assert_called_once()
return received
chunks_received = asyncio.run(_run())
assert len(chunks_received) > 0
content = "".join(
c.choices[0].delta.content
for c in chunks_received
if c.choices[0].delta.content
)
assert "language" in content.lower()
tests/llm_translation/test_fireworks_ai_translation.py
+1 -15
View File
@@ -77,18 +77,6 @@ def test_map_response_format():
}
@pytest.mark.skip(reason="fireworks is having an active outage")
class TestFireworksAIChatCompletion(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
return {
"model": "fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct"
}
def test_tool_call_no_arguments(self, tool_call_no_arguments):
"""Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833"""
pass
class TestFireworksAIAudioTranscription(BaseLLMAudioTranscriptionTest):
def get_base_audio_transcription_call_args(self) -> dict:
return {
@@ -252,7 +240,5 @@ def test_global_disable_flag_with_transform_messages_helper(monkeypatch):
json_data = json.loads(mock_post.call_args.kwargs["data"])
assert (
"#transform=inline"
not in json_data["messages"][0]["content"][1]["image_url"][
"url"
]
not in json_data["messages"][0]["content"][1]["image_url"]["url"]
)
tests/llm_translation/test_gemini.py
+3 -3
View File
@@ -532,7 +532,7 @@ def test_gemini_with_grounding():
## Check streaming
response = completion(
model="gemini/gemini-2.0-flash",
model="gemini/gemini-2.5-flash",
messages=[{"role": "user", "content": "What is the capital of France?"}],
tools=tools,
stream=True,
@@ -566,7 +566,7 @@ def test_gemini_with_empty_function_call_arguments():
}
]
response = completion(
model="gemini/gemini-2.0-flash",
model="gemini/gemini-2.5-flash",
messages=[{"role": "user", "content": "What is the capital of France?"}],
tools=tools,
)
@@ -775,7 +775,7 @@ def test_gemini_tool_use():
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What's the weather like in Lima, Peru today?"},
],
"model": "gemini/gemini-2.0-flash",
"model": "gemini/gemini-2.5-flash",
"tools": [
{
"type": "function",
tests/llm_translation/test_rerank.py
-42
View File
@@ -148,48 +148,6 @@ async def test_basic_rerank_together_ai(sync_mode):
raise e
@pytest.mark.asyncio()
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.skip(reason="Skipping test due to Cohere RBAC issues")
async def test_basic_rerank_azure_ai(sync_mode):
import os
litellm.set_verbose = True
if sync_mode is True:
response = litellm.rerank(
model="azure_ai/Cohere-rerank-v3-multilingual-ko",
query="hello",
documents=["hello", "world"],
top_n=3,
api_key=os.getenv("AZURE_AI_COHERE_API_KEY"),
api_base=os.getenv("AZURE_AI_COHERE_API_BASE"),
)
print("re rank response: ", response)
assert response.id is not None
assert response.results is not None
assert_response_shape(response, custom_llm_provider="together_ai")
else:
response = await litellm.arerank(
model="azure_ai/Cohere-rerank-v3-multilingual-ko",
query="hello",
documents=["hello", "world"],
top_n=3,
api_key=os.getenv("AZURE_AI_COHERE_API_KEY"),
api_base=os.getenv("AZURE_AI_COHERE_API_BASE"),
)
print("async re rank response: ", response)
assert response.id is not None
assert response.results is not None
assert_response_shape(response, custom_llm_provider="together_ai")
@pytest.mark.asyncio()
@pytest.mark.parametrize("version", ["v1", "v2"])
async def test_rerank_custom_api_base(version):
tests/llm_translation/test_router_llm_translation_tests.py
+8 -8
View File
@@ -66,8 +66,8 @@ def test_router_azure_acompletion():
print("Router Test Azure - Acompletion, Acompletion with stream")
# remove api key from env to repro how proxy passes key to router
old_api_key = os.environ["AZURE_API_KEY"]
os.environ.pop("AZURE_API_KEY", None)
old_api_key = os.environ["AZURE_AI_API_KEY"]
os.environ.pop("AZURE_AI_API_KEY", None)
model_list = [
{
@@ -75,8 +75,8 @@ def test_router_azure_acompletion():
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": old_api_key,
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"rpm": 1800,
},
@@ -85,8 +85,8 @@ def test_router_azure_acompletion():
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": old_api_key,
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"rpm": 1800,
},
@@ -126,9 +126,9 @@ def test_router_azure_acompletion():
asyncio.run(test2())
print("\n Passed Streaming")
os.environ["AZURE_API_KEY"] = old_api_key
os.environ["AZURE_AI_API_KEY"] = old_api_key
router.reset()
except Exception as e:
os.environ["AZURE_API_KEY"] = old_api_key
os.environ["AZURE_AI_API_KEY"] = old_api_key
print(f"FAILED TEST")
pytest.fail(f"Got unexpected exception on router! - {e}")
tests/llm_translation/test_snowflake.py
+3 -5
View File
@@ -1,12 +1,10 @@
import os
import sys
import asyncio
import json
import os
import httpx
from typing import Any, Dict, List
from unittest.mock import Mock, MagicMock, patch
from dotenv import load_dotenv
from unittest.mock import AsyncMock, MagicMock, patch
load_dotenv()
import pytest
from litellm import completion, acompletion, responses
tests/llm_translation/test_watsonx.py
+10
View File
@@ -13,6 +13,16 @@ import pytest
from typing import Optional
@pytest.fixture(autouse=True)
def watsonx_env_vars(monkeypatch):
"""Set required WatsonX env vars so the provider passes validation.
Also clear WATSONX_ZENAPIKEY/WATSONX_TOKEN so they don't bypass the IAM token mock."""
monkeypatch.setenv("WATSONX_URL", "https://us-south.ml.cloud.ibm.com")
monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id")
monkeypatch.delenv("WATSONX_ZENAPIKEY", raising=False)
monkeypatch.delenv("WATSONX_TOKEN", raising=False)
@pytest.fixture
def watsonx_chat_completion_call():
def _call(
tests/load_tests/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
tests/local_testing/adroit-crow-413218-bc47f303efc9.json
-13
View File
@@ -1,13 +0,0 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"private_key_id": "",
"private_key": "",
"client_email": "test-adroit-crow@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "104886546564708740969",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-adroit-crow%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
tests/local_testing/example_config_yaml/azure_config.yaml
+2 -2
View File
@@ -4,12 +4,12 @@ model_list:
model: azure/gpt-4.1-mini
api_base: https://openai-gpt-4-test-v-1.openai.azure.com/
api_version: "2023-05-15"
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
tpm: 20_000
- model_name: gpt-4-team2
litellm_params:
model: azure/gpt-4
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
api_base: https://openai-gpt-4-test-v-2.openai.azure.com/
tpm: 100_000
tests/local_testing/test_acooldowns_router.py
+1 -1
View File
@@ -31,7 +31,7 @@ def _make_model_list():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
tests/local_testing/test_alangfuse.py
-263
View File
@@ -599,233 +599,6 @@ def test_langfuse_logging_function_calling():
# test_langfuse_logging_function_calling()
@pytest.mark.skip(reason="Need to address this on main")
def test_aaalangfuse_existing_trace_id():
"""
When existing trace id is passed, don't set trace params -> prevents overwriting the trace
Pass 1 logging object with a trace
Pass 2nd logging object with the trace id
Assert no changes to the trace
"""
# Test - if the logs were sent to the correct team on langfuse
import datetime
import litellm
from litellm.integrations.langfuse.langfuse import LangFuseLogger
langfuse_Logger = LangFuseLogger(
langfuse_public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"),
langfuse_secret=os.getenv("LANGFUSE_PROJECT2_SECRET"),
)
litellm.success_callback = ["langfuse"]
# langfuse_args = {'kwargs': { 'start_time': 'end_time': datetime.datetime(2024, 5, 1, 7, 31, 29, 903685), 'user_id': None, 'print_verbose': <function print_verbose at 0x109d1f420>, 'level': 'DEFAULT', 'status_message': None}
response_obj = litellm.ModelResponse(
id="chatcmpl-9K5HUAbVRqFrMZKXL0WoC295xhguY",
choices=[
litellm.Choices(
finish_reason="stop",
index=0,
message=litellm.Message(
content="I'm sorry, I am an AI assistant and do not have real-time information. I recommend checking a reliable weather website or app for the most up-to-date weather information in Boston.",
role="assistant",
),
)
],
created=1714573888,
model="gpt-3.5-turbo-0125",
object="chat.completion",
system_fingerprint="fp_3b956da36b",
usage=litellm.Usage(completion_tokens=37, prompt_tokens=14, total_tokens=51),
)
### NEW TRACE ###
message = [{"role": "user", "content": "what's the weather in boston"}]
langfuse_args = {
"response_obj": response_obj,
"kwargs": {
"model": "gpt-3.5-turbo",
"litellm_params": {
"acompletion": False,
"api_key": None,
"force_timeout": 600,
"logger_fn": None,
"verbose": False,
"custom_llm_provider": "openai",
"api_base": "https://api.openai.com/v1/",
"litellm_call_id": None,
"model_alias_map": {},
"completion_call_id": None,
"metadata": None,
"model_info": None,
"proxy_server_request": None,
"preset_cache_key": None,
"no-log": False,
"stream_response": {},
},
"messages": message,
"optional_params": {"temperature": 0.1, "extra_body": {}},
"start_time": "2024-05-01 07:31:27.986164",
"stream": False,
"user": None,
"call_type": "completion",
"litellm_call_id": None,
"completion_start_time": "2024-05-01 07:31:29.903685",
"temperature": 0.1,
"extra_body": {},
"input": [{"role": "user", "content": "what's the weather in boston"}],
"api_key": "my-api-key",
"additional_args": {
"complete_input_dict": {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "what's the weather in boston"}
],
"temperature": 0.1,
"extra_body": {},
}
},
"log_event_type": "successful_api_call",
"end_time": "2024-05-01 07:31:29.903685",
"cache_hit": None,
"response_cost": 6.25e-05,
},
"start_time": datetime.datetime(2024, 5, 1, 7, 31, 27, 986164),
"end_time": datetime.datetime(2024, 5, 1, 7, 31, 29, 903685),
"user_id": None,
"print_verbose": litellm.print_verbose,
"level": "DEFAULT",
"status_message": None,
}
langfuse_response_object = langfuse_Logger.log_event(**langfuse_args)
import langfuse
langfuse_client = langfuse.Langfuse(
public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"),
secret_key=os.getenv("LANGFUSE_PROJECT2_SECRET"),
)
trace_id = langfuse_response_object["trace_id"]
assert trace_id is not None
langfuse_client.flush()
time.sleep(2)
print(langfuse_client.get_trace(id=trace_id))
initial_langfuse_trace = langfuse_client.get_trace(id=trace_id)
### EXISTING TRACE ###
new_metadata = {"existing_trace_id": trace_id}
new_messages = [{"role": "user", "content": "What do you know?"}]
new_response_obj = litellm.ModelResponse(
id="chatcmpl-9K5HUAbVRqFrMZKXL0WoC295xhguY",
choices=[
litellm.Choices(
finish_reason="stop",
index=0,
message=litellm.Message(
content="What do I know?",
role="assistant",
),
)
],
created=1714573888,
model="gpt-3.5-turbo-0125",
object="chat.completion",
system_fingerprint="fp_3b956da36b",
usage=litellm.Usage(completion_tokens=37, prompt_tokens=14, total_tokens=51),
)
langfuse_args = {
"response_obj": new_response_obj,
"kwargs": {
"model": "gpt-3.5-turbo",
"litellm_params": {
"acompletion": False,
"api_key": None,
"force_timeout": 600,
"logger_fn": None,
"verbose": False,
"custom_llm_provider": "openai",
"api_base": "https://api.openai.com/v1/",
"litellm_call_id": "508113a1-c6f1-48ce-a3e1-01c6cce9330e",
"model_alias_map": {},
"completion_call_id": None,
"metadata": new_metadata,
"model_info": None,
"proxy_server_request": None,
"preset_cache_key": None,
"no-log": False,
"stream_response": {},
},
"messages": new_messages,
"optional_params": {"temperature": 0.1, "extra_body": {}},
"start_time": "2024-05-01 07:31:27.986164",
"stream": False,
"user": None,
"call_type": "completion",
"litellm_call_id": "508113a1-c6f1-48ce-a3e1-01c6cce9330e",
"completion_start_time": "2024-05-01 07:31:29.903685",
"temperature": 0.1,
"extra_body": {},
"input": [{"role": "user", "content": "what's the weather in boston"}],
"api_key": "my-api-key",
"additional_args": {
"complete_input_dict": {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "what's the weather in boston"}
],
"temperature": 0.1,
"extra_body": {},
}
},
"log_event_type": "successful_api_call",
"end_time": "2024-05-01 07:31:29.903685",
"cache_hit": None,
"response_cost": 6.25e-05,
},
"start_time": datetime.datetime(2024, 5, 1, 7, 31, 27, 986164),
"end_time": datetime.datetime(2024, 5, 1, 7, 31, 29, 903685),
"user_id": None,
"print_verbose": litellm.print_verbose,
"level": "DEFAULT",
"status_message": None,
}
langfuse_response_object = langfuse_Logger.log_event(**langfuse_args)
new_trace_id = langfuse_response_object["trace_id"]
assert new_trace_id == trace_id
langfuse_client.flush()
time.sleep(2)
print(langfuse_client.get_trace(id=trace_id))
new_langfuse_trace = langfuse_client.get_trace(id=trace_id)
initial_langfuse_trace_dict = dict(initial_langfuse_trace)
initial_langfuse_trace_dict.pop("updatedAt")
initial_langfuse_trace_dict.pop("timestamp")
new_langfuse_trace_dict = dict(new_langfuse_trace)
new_langfuse_trace_dict.pop("updatedAt")
new_langfuse_trace_dict.pop("timestamp")
assert initial_langfuse_trace_dict == new_langfuse_trace_dict
@pytest.mark.skipif(
condition=not os.environ.get("OPENAI_API_KEY", False),
reason="Authentication missing for openai",
@@ -928,42 +701,6 @@ async def test_make_request():
)
@pytest.mark.skip(
reason="local only test, use this to verify if dynamic langfuse logging works as expected"
)
def test_aaalangfuse_dynamic_logging():
"""
pass in langfuse credentials via completion call
assert call is logged.
Covers the team-logging scenario.
"""
from litellm._uuid import uuid
import langfuse
trace_id = str(uuid.uuid4())
_ = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey"}],
mock_response="Hey! how's it going?",
langfuse_public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"),
langfuse_secret_key=os.getenv("LANGFUSE_PROJECT2_SECRET"),
metadata={"trace_id": trace_id},
success_callback=["langfuse"],
)
time.sleep(3)
langfuse_client = langfuse.Langfuse(
public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"),
secret_key=os.getenv("LANGFUSE_PROJECT2_SECRET"),
)
langfuse_client.get_trace(id=trace_id)
import datetime
generation_params = {
tests/local_testing/test_amazing_vertex_completion.py
+95 -90
View File
@@ -167,43 +167,6 @@ async def test_get_response():
pytest.fail(f"An error occurred - {str(e)}")
@pytest.mark.skip(
reason="Local test. Vertex AI Quota is low. Leads to rate limit errors on ci/cd."
)
@pytest.mark.flaky(retries=3, delay=1)
def test_vertex_ai_anthropic_streaming():
try:
load_vertex_ai_credentials()
# litellm.set_verbose = True
model = "claude-3-5-sonnet@20240620"
vertex_ai_project = "pathrise-convert-1606954137718"
vertex_ai_location = "asia-southeast1"
json_obj = get_vertex_ai_creds_json()
vertex_credentials = json.dumps(json_obj)
response = completion(
model="vertex_ai/" + model,
messages=[{"role": "user", "content": "hi"}],
temperature=0.7,
vertex_ai_project=vertex_ai_project,
vertex_ai_location=vertex_ai_location,
stream=True,
)
# print("\nModel Response", response)
for idx, chunk in enumerate(response):
print(f"chunk: {chunk}")
streaming_format_tests(idx=idx, chunk=chunk)
# raise Exception("it worked!")
except litellm.RateLimitError as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_vertex_ai_anthropic_streaming()
@@ -394,10 +357,7 @@ async def test_async_vertexai_response_basic():
user_message = "Hello, how are you?"
messages = [{"content": user_message, "role": "user"}]
response = await acompletion(
model="gemini-2.5-flash",
messages=messages,
temperature=0.7,
timeout=5
model="gemini-2.5-flash", messages=messages, temperature=0.7, timeout=5
)
print(f"response: {response}")
except litellm.NotFoundError as e:
@@ -414,8 +374,6 @@ async def test_async_vertexai_response_basic():
pytest.fail(f"An exception occurred: {e}")
@pytest.mark.flaky(retries=3, delay=1)
@pytest.mark.asyncio
async def test_async_vertexai_streaming_response():
@@ -739,7 +697,9 @@ def test_gemini_pro_grounding(value_in_dict):
# @pytest.mark.skip(reason="exhausted vertex quota. need to refactor to mock the call")
@pytest.mark.parametrize("model", ["vertex_ai_beta/gemini-2.5-flash-lite"]) # "vertex_ai",
@pytest.mark.parametrize(
"model", ["vertex_ai_beta/gemini-2.5-flash-lite"]
) # "vertex_ai",
@pytest.mark.parametrize("sync_mode", [True]) # "vertex_ai",
@pytest.mark.asyncio
@pytest.mark.flaky(retries=6, delay=2)
@@ -808,7 +768,17 @@ async def test_gemini_pro_function_calling_httpx(model, sync_mode):
except Exception as e:
error_msg = str(e)
# Skip test for known transient API issues
if any(x in error_msg for x in ["429 Quota exceeded", "503", "Service unavailable", "timeout", "Timeout", "UNAVAILABLE"]):
if any(
x in error_msg
for x in [
"429 Quota exceeded",
"503",
"Service unavailable",
"timeout",
"Timeout",
"UNAVAILABLE",
]
):
pytest.skip(f"Transient API error: {error_msg}")
else:
pytest.fail(f"An unexpected exception occurred - {error_msg}")
@@ -1396,12 +1366,14 @@ async def test_gemini_pro_json_schema_args_sent_httpx(
# Gemini 2.x+ uses response_json_schema, Gemini 1.x uses response_schema
gen_config = mock_call.call_args.kwargs["json"]["generationConfig"]
assert (
"response_schema" in gen_config or "response_json_schema" in gen_config
"response_schema" in gen_config
or "response_json_schema" in gen_config
), f"Expected response_schema or response_json_schema in {gen_config}"
else:
gen_config = mock_call.call_args.kwargs["json"]["generationConfig"]
assert (
"response_schema" not in gen_config and "response_json_schema" not in gen_config
"response_schema" not in gen_config
and "response_json_schema" not in gen_config
)
assert (
"Use this JSON schema:"
@@ -1577,7 +1549,8 @@ async def test_gemini_pro_json_schema_args_sent_httpx_openai_schema(
# Gemini 2.x+ uses response_json_schema, Gemini 1.x uses response_schema
gen_config = mock_call.call_args.kwargs["json"]["generationConfig"]
assert (
"response_schema" in gen_config or "response_json_schema" in gen_config
"response_schema" in gen_config
or "response_json_schema" in gen_config
), f"Expected response_schema or response_json_schema in {gen_config}"
assert (
"response_mime_type"
@@ -1592,7 +1565,8 @@ async def test_gemini_pro_json_schema_args_sent_httpx_openai_schema(
else:
gen_config = mock_call.call_args.kwargs["json"]["generationConfig"]
assert (
"response_schema" not in gen_config and "response_json_schema" not in gen_config
"response_schema" not in gen_config
and "response_json_schema" not in gen_config
)
assert (
"Use this JSON schema:"
@@ -2313,8 +2287,6 @@ def test_prompt_factory_nested():
), "'text' value not a string."
@pytest.mark.asyncio
async def test_completion_fine_tuned_model():
load_vertex_ai_credentials()
@@ -2380,7 +2352,7 @@ async def test_completion_fine_tuned_model():
# this is the fine-tuned model endpoint
assert (
url[0]
== "https://us-central1-aiplatform.googleapis.com/v1/projects/pathrise-convert-1606954137718/locations/us-central1/endpoints/4965075652664360960:generateContent"
== "https://us-central1-aiplatform.googleapis.com/v1/projects/litellm-ci-cd/locations/us-central1/endpoints/4965075652664360960:generateContent"
)
print("call args = ", kwargs)
@@ -2579,20 +2551,20 @@ async def test_gemini_context_caching_anthropic_format(sync_mode):
async def test_gemini_context_caching_disabled_flag(sync_mode):
"""
Test that disable_anthropic_gemini_context_caching_transform flag properly disables context caching.
When the flag is set to True, messages with cache_control should not trigger caching API calls.
"""
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
litellm.set_verbose = True
# Store original value to restore later
original_flag_value = litellm.disable_anthropic_gemini_context_caching_transform
try:
# Enable the disable flag
litellm.disable_anthropic_gemini_context_caching_transform = True
gemini_context_caching_messages = [
# System Message with cache_control
{
@@ -2633,13 +2605,15 @@ async def test_gemini_context_caching_disabled_flag(sync_mode):
],
},
]
if sync_mode:
client = HTTPHandler(concurrent_limit=1)
else:
client = AsyncHTTPHandler(concurrent_limit=1)
with patch.object(client, "post", side_effect=mock_gemini_request) as mock_client:
with patch.object(
client, "post", side_effect=mock_gemini_request
) as mock_client:
try:
if sync_mode:
response = litellm.completion(
@@ -2662,24 +2636,32 @@ async def test_gemini_context_caching_disabled_flag(sync_mode):
print(e)
# When caching is disabled, should only make 1 call (no separate cache creation call)
assert mock_client.call_count == 1, f"Expected 1 call when caching is disabled, got {mock_client.call_count}"
assert (
mock_client.call_count == 1
), f"Expected 1 call when caching is disabled, got {mock_client.call_count}"
first_call_args = mock_client.call_args_list[0].kwargs
first_call_positional_args = mock_client.call_args_list[0].args
print(f"first_call_args with caching disabled: {first_call_args}")
print(f"first_call_positional_args with caching disabled: {first_call_positional_args}")
print(
f"first_call_positional_args with caching disabled: {first_call_positional_args}"
)
# Assert that cachedContents is NOT in the URL when caching is disabled
url = first_call_args.get("url", first_call_positional_args[0] if first_call_positional_args else "")
assert "cachedContents" not in url, "cachedContents should not be in URL when caching is disabled"
url = first_call_args.get(
"url",
first_call_positional_args[0] if first_call_positional_args else "",
)
assert (
"cachedContents" not in url
), "cachedContents should not be in URL when caching is disabled"
finally:
# Restore original flag value
litellm.disable_anthropic_gemini_context_caching_transform = original_flag_value
@pytest.mark.asyncio
async def test_partner_models_httpx_ai21():
litellm.set_verbose = True
@@ -2776,7 +2758,7 @@ async def test_partner_models_httpx_ai21():
assert (
url[0]
== "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/pathrise-convert-1606954137718/locations/us-central1/publishers/ai21/models/jamba-1.5-mini@001:rawPredict"
== "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/litellm-ci-cd/locations/us-central1/publishers/ai21/models/jamba-1.5-mini@001:rawPredict"
)
# json loads kwargs
@@ -2920,7 +2902,9 @@ def test_gemini_function_call_parameter_in_messages():
"contents": [
{
"role": "user",
"parts": [{"text": "search for weather in boston (use `search`)"}],
"parts": [
{"text": "search for weather in boston (use `search`)"}
],
},
{
"role": "model",
@@ -2947,7 +2931,9 @@ def test_gemini_function_call_parameter_in_messages():
],
},
],
"system_instruction": {"parts": [{"text": "Use search for most queries."}]},
"system_instruction": {
"parts": [{"text": "Use search for most queries."}]
},
"tools": [
{
"function_declarations": [
@@ -3578,8 +3564,9 @@ def test_gemini_tool_calling_working_demo():
},
}
],
"vertex_location": "global",
}
response = completion(model="vertex_ai/gemini-2.0-flash", **args)
response = completion(model="vertex_ai/gemini-3-flash-preview", **args)
print(response)
@@ -3650,8 +3637,9 @@ def test_gemini_tool_calling_not_working():
},
}
],
"vertex_location": "global",
}
response = completion(model="vertex_ai/gemini-2.0-flash", **args)
response = completion(model="vertex_ai/gemini-3-flash-preview", **args)
print(response)
@@ -3746,8 +3734,8 @@ def test_gemini_nullable_object_tool_schema_httpx():
load_vertex_ai_credentials()
litellm._turn_on_debug()
tools = [{
tools = [
{
"type": "function",
"strict": True,
"function": {
@@ -3760,7 +3748,7 @@ def test_gemini_nullable_object_tool_schema_httpx():
"properties": {
"ticket_id": {
"type": "string",
"description": "Unique identifier for the support ticket"
"description": "Unique identifier for the support ticket",
},
"customer_context": {
"type": ["object", "null"],
@@ -3770,18 +3758,19 @@ def test_gemini_nullable_object_tool_schema_httpx():
"properties": {
"user_id": {
"type": "string",
"description": "Internal user identifier"
"description": "Internal user identifier",
},
"plan": {
"type": "string",
"description": "Subscription plan name (e.g. pro, enterprise)"
}
}
}
}
}
}
}]
"description": "Subscription plan name (e.g. pro, enterprise)",
},
},
},
},
},
},
}
]
response = litellm.completion(
model="vertex_ai/gemini-2.5-flash",
@@ -3986,17 +3975,23 @@ def test_vertex_ai_gemini_audio_ogg():
for part in content["parts"]
if "file_data" in part
]
assert len(file_data_parts) == 1, f"Expected 1 file_data part, got: {file_data_parts}"
assert (
len(file_data_parts) == 1
), f"Expected 1 file_data part, got: {file_data_parts}"
file_data = file_data_parts[0]["file_data"]
assert file_data["mime_type"] == "audio/ogg", f"Expected audio/ogg, got: {file_data['mime_type']}"
assert "En-us-public.ogg" in file_data["file_uri"], f"Unexpected file_uri: {file_data['file_uri']}"
assert (
file_data["mime_type"] == "audio/ogg"
), f"Expected audio/ogg, got: {file_data['mime_type']}"
assert (
"En-us-public.ogg" in file_data["file_uri"]
), f"Unexpected file_uri: {file_data['file_uri']}"
print(response)
@pytest.mark.asyncio
async def test_vertex_ai_deepseek():
"""Test that deepseek models use the correct v1 API endpoint instead of v1beta1."""
# load_vertex_ai_credentials()
load_vertex_ai_credentials()
litellm._turn_on_debug()
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
@@ -4041,7 +4036,7 @@ def test_gemini_grounding_on_streaming():
load_vertex_ai_credentials()
# litellm._turn_on_debug()
args = {
"model": "vertex_ai/gemini-2.0-flash",
"model": "vertex_ai/gemini-3-flash-preview",
"messages": [
{
"role": "user",
@@ -4053,6 +4048,7 @@ def test_gemini_grounding_on_streaming():
],
}
],
"vertex_location": "global",
"stream": True,
"tools": [{"googleSearch": {}}],
"fallbacks": [],
@@ -4075,11 +4071,20 @@ def test_gemini_google_maps_tool_simple():
litellm._turn_on_debug()
tools = [{"googleMaps": {"enableWidget": True}}]
tools_with_location = [{"googleMaps": {"enableWidget": True, "latitude": 37.7749, "longitude": -122.4194, "languageCode": "en_US"}}]
tools_with_location = [
{
"googleMaps": {
"enableWidget": True,
"latitude": 37.7749,
"longitude": -122.4194,
"languageCode": "en_US",
}
}
]
try:
for tools in [tools, tools_with_location]:
response = completion(
model="vertex_ai/gemini-2.0-flash",
model="vertex_ai/gemini-3-flash-preview",
messages=[
{
"role": "user",
@@ -4087,6 +4092,7 @@ def test_gemini_google_maps_tool_simple():
}
],
tools=tools,
vertex_location="global",
)
print(f"Response: {response.model_dump_json(indent=4)}")
assert response.choices[0].message.content is not None
@@ -4094,4 +4100,3 @@ def test_gemini_google_maps_tool_simple():
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
tests/local_testing/test_arize_ai.py
+2 -2
View File
@@ -48,8 +48,8 @@ async def test_async_dynamic_arize_config():
messages=[{"role": "user", "content": "hi test from arize dynamic config"}],
temperature=0.1,
user="OTEL_USER",
arize_api_key=os.getenv("ARIZE_SPACE_2_API_KEY"),
arize_space_key=os.getenv("ARIZE_SPACE_2_KEY"),
arize_api_key=os.getenv("ARIZE_SPACE_API_KEY"),
arize_space_key=os.getenv("ARIZE_SPACE_KEY"),
)
await asyncio.sleep(2)
tests/local_testing/test_assistants.py
+3 -4
View File
@@ -37,10 +37,11 @@ V0 Scope:
- Run Thread -> `/v1/threads/{thread_id}/run`
"""
def _add_azure_related_dynamic_params(data: dict) -> dict:
data["api_version"] = "2024-02-15-preview"
data["api_base"] = os.getenv("AZURE_API_BASE")
data["api_key"] = os.getenv("AZURE_API_KEY")
data["api_base"] = os.getenv("AZURE_AI_API_BASE")
data["api_key"] = os.getenv("AZURE_AI_API_KEY")
return data
@@ -236,8 +237,6 @@ async def test_aarun_thread_litellm(sync_mode, provider, is_streaming):
"""
import openai
try:
get_assistants_data = {
"custom_llm_provider": provider,
tests/local_testing/test_azure_openai.py
+6 -4
View File
@@ -40,11 +40,13 @@ async def test_aaaaazure_tenant_id_auth(respx_mock: MockRouter):
PROD Test
"""
litellm.disable_aiohttp_transport = True # since this uses respx, we need to set use_aiohttp_transport to False
litellm.disable_aiohttp_transport = (
True # since this uses respx, we need to set use_aiohttp_transport to False
)
# Clear the HTTP client cache to ensure respx mocking works
# This is critical because respx only intercepts clients created AFTER mocking is active
if hasattr(litellm, 'in_memory_llm_clients_cache'):
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
router = Router(
@@ -53,7 +55,7 @@ async def test_aaaaazure_tenant_id_auth(respx_mock: MockRouter):
"model_name": "gpt-3.5-turbo",
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"tenant_id": os.getenv("AZURE_TENANT_ID"),
"client_id": os.getenv("AZURE_CLIENT_ID"),
"client_secret": os.getenv("AZURE_CLIENT_SECRET"),
tests/local_testing/test_azure_perf.py
+4 -4
View File
@@ -9,8 +9,8 @@
# from openai import AsyncAzureOpenAI
# client = AsyncAzureOpenAI(
# api_key=os.getenv("AZURE_API_KEY"),
# azure_endpoint=os.getenv("AZURE_API_BASE"), # type: ignore
# api_key=os.getenv("AZURE_AI_API_KEY"),
# azure_endpoint=os.getenv("AZURE_AI_API_BASE"), # type: ignore
# api_version=os.getenv("AZURE_API_VERSION"),
# )
@@ -19,8 +19,8 @@
# "model_name": "azure-test",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# },
# }
tests/local_testing/test_caching.py
+116 -37
View File
@@ -147,7 +147,12 @@ def test_caching_dynamic_args(): # test in memory cache
port=_redis_port_env,
password=_redis_password_env,
)
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test")
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test",
)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
print(f"response1: {response1}")
print(f"response2: {response2}")
@@ -173,7 +178,12 @@ def test_caching_v2(): # test in memory cache
try:
litellm.set_verbose = True
litellm.cache = Cache()
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test")
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test",
)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
print(f"response1: {response1}")
print(f"response2: {response2}")
@@ -200,9 +210,18 @@ def test_caching_with_ttl():
litellm.set_verbose = True
litellm.cache = Cache()
response1 = completion(
model="gpt-3.5-turbo", messages=messages, caching=True, ttl=0, mock_response="Hello world from cache test 1"
model="gpt-3.5-turbo",
messages=messages,
caching=True,
ttl=0,
mock_response="Hello world from cache test 1",
)
response2 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test 2",
)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test 2")
print(f"response1: {response1}")
print(f"response2: {response2}")
litellm.cache = None # disable cache
@@ -221,8 +240,18 @@ def test_caching_with_default_ttl():
try:
litellm.set_verbose = True
litellm.cache = Cache(ttl=0)
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test")
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test")
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test",
)
response2 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test",
)
print(f"response1: {response1}")
print(f"response2: {response2}")
litellm.cache = None # disable cache
@@ -247,10 +276,16 @@ async def test_caching_with_cache_controls(sync_flag):
if sync_flag:
## TTL = 0
response1 = completion(
model="gpt-3.5-turbo", messages=messages, cache={"ttl": 0}, mock_response="Hello world"
model="gpt-3.5-turbo",
messages=messages,
cache={"ttl": 0},
mock_response="Hello world",
)
response2 = completion(
model="gpt-3.5-turbo", messages=messages, cache={"s-maxage": 10}, mock_response="Hello world"
model="gpt-3.5-turbo",
messages=messages,
cache={"s-maxage": 10},
mock_response="Hello world",
)
assert response2["id"] != response1["id"]
@@ -322,9 +357,19 @@ def test_caching_with_models_v2():
litellm.cache = Cache()
print("test2 for caching")
litellm.set_verbose = True
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test")
response1 = completion(
model="gpt-3.5-turbo",
messages=messages,
caching=True,
mock_response="Hello world from cache test",
)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
response3 = completion(model="gpt-4.1-nano", messages=messages, caching=True, mock_response="Different model response")
response3 = completion(
model="gpt-4.1-nano",
messages=messages,
caching=True,
mock_response="Different model response",
)
print(f"response1: {response1}")
print(f"response2: {response2}")
print(f"response3: {response3}")
@@ -423,7 +468,10 @@ def test_embedding_caching():
text_to_embed = [embedding_large_text]
start_time = time.time()
embedding1 = embedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True, mock_response="0.1,0.2,0.3,0.4,0.5"
model="text-embedding-ada-002",
input=text_to_embed,
caching=True,
mock_response="0.1,0.2,0.3,0.4,0.5",
)
end_time = time.time()
print(f"Embedding 1 response time: {end_time - start_time} seconds")
@@ -459,12 +507,18 @@ async def test_embedding_caching_individual_items_and_then_list():
"world",
]
embedding1 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed[0], caching=True, mock_response="0.1,0.2,0.3,0.4,0.5"
model="text-embedding-ada-002",
input=text_to_embed[0],
caching=True,
mock_response="0.1,0.2,0.3,0.4,0.5",
)
initial_prompt_tokens = embedding1.usage.prompt_tokens
await asyncio.sleep(1)
embedding2 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed[1], caching=True, mock_response="0.6,0.7,0.8,0.9,1.0"
model="text-embedding-ada-002",
input=text_to_embed[1],
caching=True,
mock_response="0.6,0.7,0.8,0.9,1.0",
)
await asyncio.sleep(1)
embedding3 = await aembedding(
@@ -480,7 +534,10 @@ async def test_embedding_caching_individual_items_and_then_list():
additional_text = "this is a new text"
text_to_embed.append(additional_text)
embedding4 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True, mock_response="0.1,0.2,0.3,0.4,0.5"
model="text-embedding-ada-002",
input=text_to_embed,
caching=True,
mock_response="0.1,0.2,0.3,0.4,0.5",
)
assert embedding4.usage.prompt_tokens > embedding3.usage.prompt_tokens
@@ -490,7 +547,10 @@ async def test_embedding_caching_individual_items():
litellm.cache = Cache()
text_to_embed = "hello"
embedding1 = await aembedding(
model="text-embedding-ada-002", input=text_to_embed, caching=True, mock_response="0.1,0.2,0.3,0.4,0.5"
model="text-embedding-ada-002",
input=text_to_embed,
caching=True,
mock_response="0.1,0.2,0.3,0.4,0.5",
)
await asyncio.sleep(1)
@@ -512,13 +572,13 @@ def test_embedding_caching_azure():
litellm.cache = Cache()
text_to_embed = [embedding_large_text]
api_key = os.environ["AZURE_API_KEY"]
api_base = os.environ["AZURE_API_BASE"]
api_key = os.environ["AZURE_AI_API_KEY"]
api_base = os.environ["AZURE_AI_API_BASE"]
api_version = os.environ["AZURE_API_VERSION"]
os.environ["AZURE_API_VERSION"] = ""
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_AI_API_BASE"] = ""
os.environ["AZURE_AI_API_KEY"] = ""
start_time = time.time()
print("AZURE CONFIGS")
@@ -560,8 +620,8 @@ def test_embedding_caching_azure():
pytest.fail("Error occurred: Embedding caching failed")
os.environ["AZURE_API_VERSION"] = api_version
os.environ["AZURE_API_BASE"] = api_base
os.environ["AZURE_API_KEY"] = api_key
os.environ["AZURE_AI_API_BASE"] = api_base
os.environ["AZURE_AI_API_KEY"] = api_key
# test_embedding_caching_azure()
@@ -851,9 +911,18 @@ def test_redis_cache_completion():
model="gpt-3.5-turbo", messages=messages, caching=True, max_tokens=20
)
response3 = completion(
model="gpt-3.5-turbo", messages=messages, caching=True, temperature=0.5, mock_response="Different params response"
model="gpt-3.5-turbo",
messages=messages,
caching=True,
temperature=0.5,
mock_response="Different params response",
)
response4 = completion(
model="gpt-4o-mini",
messages=messages,
caching=True,
mock_response="Different model response",
)
response4 = completion(model="gpt-4o-mini", messages=messages, caching=True, mock_response="Different model response")
print("\nresponse 1", response1)
print("\nresponse 2", response2)
@@ -1127,7 +1196,11 @@ async def test_redis_cache_atext_completion():
print("test for caching, atext_completion")
response1 = await litellm.atext_completion(
model="gpt-3.5-turbo-instruct", prompt=prompt, max_tokens=40, temperature=1, mock_response="Hello world from cache test"
model="gpt-3.5-turbo-instruct",
prompt=prompt,
max_tokens=40,
temperature=1,
mock_response="Hello world from cache test",
)
await asyncio.sleep(0.5)
@@ -1458,11 +1531,17 @@ def test_cache_override():
# test embedding
response1 = embedding(
model="text-embedding-ada-002", input=["hello who are you"], caching=False, mock_response="0.1,0.2,0.3,0.4,0.5"
model="text-embedding-ada-002",
input=["hello who are you"],
caching=False,
mock_response="0.1,0.2,0.3,0.4,0.5",
)
response2 = embedding(
model="text-embedding-ada-002", input=["hello who are you"], caching=False, mock_response="0.6,0.7,0.8,0.9,1.0"
model="text-embedding-ada-002",
input=["hello who are you"],
caching=False,
mock_response="0.6,0.7,0.8,0.9,1.0",
)
# When caching=False, responses should have different IDs
@@ -2787,7 +2866,7 @@ def test_caching_thinking_args_hit(): # test in memory cache
async def test_cache_key_in_hidden_params_acompletion():
"""
Test that cache_key is present in _hidden_params on cache hits for acompletion.
Validates fix for missing x-litellm-cache-key header on proxy cache hits.
"""
litellm.cache = Cache(
@@ -2796,10 +2875,10 @@ async def test_cache_key_in_hidden_params_acompletion():
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
unique_content = f"test cache key hidden params {uuid.uuid4()}"
messages = [{"role": "user", "content": unique_content}]
# First call - cache miss
response1 = await litellm.acompletion(
model="gpt-3.5-turbo",
@@ -2807,12 +2886,12 @@ async def test_cache_key_in_hidden_params_acompletion():
mock_response="test response",
caching=True,
)
print(f"Response 1 _hidden_params: {response1._hidden_params}")
assert response1._hidden_params.get("cache_hit") is not True
await asyncio.sleep(0.5)
# Second call - cache hit
response2 = await litellm.acompletion(
model="gpt-3.5-turbo",
@@ -2820,17 +2899,17 @@ async def test_cache_key_in_hidden_params_acompletion():
mock_response="test response",
caching=True,
)
print(f"Response 2 _hidden_params: {response2._hidden_params}")
# Verify cache hit occurred
assert response2._hidden_params.get("cache_hit") is True
# Verify cache_key is present in _hidden_params
assert "cache_key" in response2._hidden_params
assert response2._hidden_params["cache_key"] is not None
# Verify both responses have same ID (cache hit)
assert response1.id == response2.id
litellm.cache = None
tests/local_testing/test_caching_ssl.py
+2 -2
View File
@@ -59,9 +59,9 @@ def test_caching_router():
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
tests/local_testing/test_class.py
+4 -4
View File
@@ -56,9 +56,9 @@
# # "model_name": "gpt-3.5-turbo", # openai model name
# # "litellm_params": { # params for litellm completion/embedding call
# # "model": "azure/gpt-4.1-mini",
# # "api_key": os.getenv("AZURE_API_KEY"),
# # "api_key": os.getenv("AZURE_AI_API_KEY"),
# # "api_version": os.getenv("AZURE_API_VERSION"),
# # "api_base": os.getenv("AZURE_API_BASE"),
# # "api_base": os.getenv("AZURE_AI_API_BASE"),
# # },
# # }
# # ]
@@ -94,9 +94,9 @@
# # "model_name": "gpt-3.5-turbo", # openai model name
# # "litellm_params": { # params for litellm completion/embedding call
# # "model": "azure/gpt-4.1-mini",
# # "api_key": os.getenv("AZURE_API_KEY"),
# # "api_key": os.getenv("AZURE_AI_API_KEY"),
# # "api_version": os.getenv("AZURE_API_VERSION"),
# # "api_base": os.getenv("AZURE_API_BASE"),
# # "api_base": os.getenv("AZURE_AI_API_BASE"),
# # },
# # }
# # ],
tests/local_testing/test_completion.py
+143 -225
View File
@@ -132,7 +132,6 @@ def test_null_role_response():
assert response.choices[0].message.role == "assistant"
def predibase_mock_post(url, data=None, json=None, headers=None, timeout=None):
mock_response = MagicMock()
mock_response.status_code = 200
@@ -177,34 +176,6 @@ def predibase_mock_post(url, data=None, json=None, headers=None, timeout=None):
return mock_response
# @pytest.mark.skip(reason="local-only test")
@pytest.mark.asyncio
async def test_completion_predibase():
try:
litellm.set_verbose = True
# with patch("requests.post", side_effect=predibase_mock_post):
response = await litellm.acompletion(
model="predibase/llama-3-8b-instruct",
tenant_id="c4768f95",
api_key=os.getenv("PREDIBASE_API_KEY"),
messages=[{"role": "user", "content": "who are u?"}],
max_tokens=10,
timeout=5,
)
print(response)
except litellm.Timeout as e:
print("got a timeout error from predibase")
pass
except litellm.ServiceUnavailableError as e:
pass
except litellm.InternalServerError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_predibase()
@@ -286,7 +257,9 @@ def test_completion_claude_3_empty_response():
},
]
try:
response = litellm.completion(model="claude-sonnet-4-5-20250929", messages=messages)
response = litellm.completion(
model="claude-sonnet-4-5-20250929", messages=messages
)
print(response)
except litellm.InternalServerError as e:
pytest.skip(f"InternalServerError - {str(e)}")
@@ -849,8 +822,8 @@ def test_completion_mistral_azure():
litellm.set_verbose = True
response = completion(
model="mistral/Mistral-large-nmefg",
api_key=os.environ["MISTRAL_AZURE_API_KEY"],
api_base=os.environ["MISTRAL_AZURE_API_BASE"],
api_key=os.environ["MISTRAL_AZURE_AI_API_KEY"],
api_base=os.environ["MISTRAL_AZURE_AI_API_BASE"],
max_tokens=5,
messages=[
{
@@ -895,78 +868,6 @@ def test_completion_mistral_api_modified_input():
pytest.fail(f"Error occurred: {e}")
# def test_completion_oobabooga():
# try:
# response = completion(
# model="oobabooga/vicuna-1.3b", messages=messages, api_base="http://127.0.0.1:5000"
# )
# # Add any assertions here to check the response
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# test_completion_oobabooga()
# aleph alpha
# def test_completion_aleph_alpha():
# try:
# response = completion(
# model="luminous-base", messages=messages, logger_fn=logger_fn
# )
# # Add any assertions here to check the response
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# test_completion_aleph_alpha()
# def test_completion_aleph_alpha_control_models():
# try:
# response = completion(
# model="luminous-base-control", messages=messages, logger_fn=logger_fn
# )
# # Add any assertions here to check the response
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# test_completion_aleph_alpha_control_models()
import openai
def test_completion_gpt4_turbo():
litellm.set_verbose = True
try:
response = completion(
model="gpt-4-1106-preview",
messages=messages,
max_completion_tokens=10,
)
print(response)
except openai.RateLimitError:
print("got a rate liimt error")
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_gpt4_turbo()
def test_completion_gpt4_turbo_0125():
try:
response = completion(
model="gpt-4-0125-preview",
messages=messages,
max_tokens=10,
)
print(response)
except openai.RateLimitError:
print("got a rate liimt error")
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skip(reason="this test is flaky")
def test_completion_gpt4_vision():
try:
@@ -996,59 +897,6 @@ def test_completion_gpt4_vision():
pytest.fail(f"Error occurred: {e}")
# test_completion_gpt4_vision()
def test_completion_azure_gpt4_vision():
# azure/gpt-4, vision takes 5-seconds to respond
try:
litellm.set_verbose = True
response = completion(
model="azure/gpt-4-vision",
timeout=5,
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "Whats in this image?"},
{
"type": "image_url",
"image_url": {
"url": "https://avatars.githubusercontent.com/u/29436595?v=4"
},
},
],
}
],
base_url="https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions",
api_key=os.getenv("AZURE_VISION_API_KEY"),
enhancements={"ocr": {"enabled": True}, "grounding": {"enabled": True}},
dataSources=[
{
"type": "AzureComputerVision",
"parameters": {
"endpoint": "https://gpt-4-vision-enhancement.cognitiveservices.azure.com/",
"key": os.environ["AZURE_VISION_ENHANCE_KEY"],
},
}
],
)
print(response)
except openai.APIError as e:
pass
except openai.APITimeoutError:
print("got a timeout error")
pass
except openai.RateLimitError as e:
print("got a rate liimt error", e)
pass
except openai.APIStatusError as e:
print("got an api status error", e)
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_azure_gpt4_vision()
@@ -1751,7 +1599,6 @@ def test_completion_openai_pydantic(model, api_version):
pytest.fail(f"Error occurred: {e}")
def test_completion_text_openai():
try:
# litellm.set_verbose =True
@@ -2341,9 +2188,9 @@ def test_completion_azure_extra_headers():
response = completion(
model="azure/gpt-4.1-mini",
messages=messages,
api_base=os.getenv("AZURE_API_BASE"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version="2023-07-01-preview",
api_key=os.getenv("AZURE_API_KEY"),
api_key=os.getenv("AZURE_AI_API_KEY"),
extra_headers={
"Authorization": "my-bad-key",
"Ocp-Apim-Subscription-Key": "hello-world-testing",
@@ -2379,8 +2226,8 @@ def test_completion_azure_ad_token():
litellm.set_verbose = True
old_key = os.environ["AZURE_API_KEY"]
os.environ.pop("AZURE_API_KEY", None)
old_key = os.environ["AZURE_AI_API_KEY"]
os.environ.pop("AZURE_AI_API_KEY", None)
http_client = Client()
@@ -2396,7 +2243,7 @@ def test_completion_azure_ad_token():
except Exception as e:
pass
finally:
os.environ["AZURE_API_KEY"] = old_key
os.environ["AZURE_AI_API_KEY"] = old_key
mock_client.assert_called_once()
request = mock_client.call_args[0][0]
@@ -2412,8 +2259,8 @@ def test_completion_azure_key_completion_arg():
# DO NOT REMOVE THIS TEST. No MATTER WHAT Happens!
# If you want to remove it, speak to Ishaan!
# Ishaan will be very disappointed if this test is removed -> this is a standard way to pass api_key + the router + proxy use this
old_key = os.environ["AZURE_API_KEY"]
os.environ.pop("AZURE_API_KEY", None)
old_key = os.environ["AZURE_AI_API_KEY"]
os.environ.pop("AZURE_AI_API_KEY", None)
try:
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = True
@@ -2430,9 +2277,9 @@ def test_completion_azure_key_completion_arg():
print("Hidden Params", response._hidden_params)
assert response._hidden_params["custom_llm_provider"] == "azure"
os.environ["AZURE_API_KEY"] = old_key
os.environ["AZURE_AI_API_KEY"] = old_key
except Exception as e:
os.environ["AZURE_API_KEY"] = old_key
os.environ["AZURE_AI_API_KEY"] = old_key
pytest.fail(f"Error occurred: {e}")
@@ -2443,8 +2290,8 @@ async def test_re_use_azure_async_client():
import openai
client = openai.AsyncAzureOpenAI(
azure_endpoint=os.environ["AZURE_API_BASE"],
api_key=os.environ["AZURE_API_KEY"],
azure_endpoint=os.environ["AZURE_AI_API_BASE"],
api_key=os.environ["AZURE_AI_API_KEY"],
api_version="2023-07-01-preview",
)
## Test azure call
@@ -2525,13 +2372,13 @@ def test_completion_azure2():
try:
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = False
api_base = os.environ["AZURE_API_BASE"]
api_key = os.environ["AZURE_API_KEY"]
api_base = os.environ["AZURE_AI_API_BASE"]
api_key = os.environ["AZURE_AI_API_KEY"]
api_version = os.environ["AZURE_API_VERSION"]
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_AI_API_BASE"] = ""
os.environ["AZURE_API_VERSION"] = ""
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_AI_API_KEY"] = ""
## Test azure call
response = completion(
@@ -2546,9 +2393,9 @@ def test_completion_azure2():
# Add any assertions here to check the response
print(response)
os.environ["AZURE_API_BASE"] = api_base
os.environ["AZURE_AI_API_BASE"] = api_base
os.environ["AZURE_API_VERSION"] = api_version
os.environ["AZURE_API_KEY"] = api_key
os.environ["AZURE_AI_API_KEY"] = api_key
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@@ -2562,13 +2409,13 @@ def test_completion_azure3():
try:
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = True
litellm.api_base = os.environ["AZURE_API_BASE"]
litellm.api_key = os.environ["AZURE_API_KEY"]
litellm.api_base = os.environ["AZURE_AI_API_BASE"]
litellm.api_key = os.environ["AZURE_AI_API_KEY"]
litellm.api_version = os.environ["AZURE_API_VERSION"]
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_AI_API_BASE"] = ""
os.environ["AZURE_API_VERSION"] = ""
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_AI_API_KEY"] = ""
## Test azure call
response = completion(
@@ -2580,9 +2427,9 @@ def test_completion_azure3():
# Add any assertions here to check the response
print(response)
os.environ["AZURE_API_BASE"] = litellm.api_base
os.environ["AZURE_AI_API_BASE"] = litellm.api_base
os.environ["AZURE_API_VERSION"] = litellm.api_version
os.environ["AZURE_API_KEY"] = litellm.api_key
os.environ["AZURE_AI_API_KEY"] = litellm.api_key
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@@ -2594,7 +2441,7 @@ def test_completion_azure3():
# new azure test for using litellm. vars,
# use the following vars in this test and make an azure_api_call
# litellm.api_type = self.azure_api_type
# litellm.api_base = self.azure_api_base
# litellm.api_base = self.AZURE_AI_API_BASE
# litellm.api_version = self.azure_api_version
# litellm.api_key = self.api_key
def test_completion_azure_with_litellm_key():
@@ -2604,14 +2451,14 @@ def test_completion_azure_with_litellm_key():
#### set litellm vars
litellm.api_type = "azure"
litellm.api_base = os.environ["AZURE_API_BASE"]
litellm.api_base = os.environ["AZURE_AI_API_BASE"]
litellm.api_version = os.environ["AZURE_API_VERSION"]
litellm.api_key = os.environ["AZURE_API_KEY"]
litellm.api_key = os.environ["AZURE_AI_API_KEY"]
######### UNSET ENV VARs for this ################
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_AI_API_BASE"] = ""
os.environ["AZURE_API_VERSION"] = ""
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_AI_API_KEY"] = ""
######### UNSET OpenAI vars for this ##############
openai.api_type = ""
@@ -2627,9 +2474,9 @@ def test_completion_azure_with_litellm_key():
print(response)
######### RESET ENV VARs for this ################
os.environ["AZURE_API_BASE"] = litellm.api_base
os.environ["AZURE_AI_API_BASE"] = litellm.api_base
os.environ["AZURE_API_VERSION"] = litellm.api_version
os.environ["AZURE_API_KEY"] = litellm.api_key
os.environ["AZURE_AI_API_KEY"] = litellm.api_key
######### UNSET litellm vars
litellm.api_type = None
@@ -3081,7 +2928,6 @@ async def test_completion_bedrock_httpx_models(sync_mode, model):
pytest.fail(f"An error occurred - {str(e)}")
# test_completion_bedrock_titan()
@@ -3256,7 +3102,6 @@ def test_completion_anyscale_api():
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skip(reason="anyscale stopped serving public api endpoints")
def test_completion_anyscale_2():
try:
@@ -3293,23 +3138,6 @@ def test_mistral_anyscale_stream():
print(chunk["choices"][0]["delta"].get("content", ""), end="")
# test_completion_anyscale_2()
# def test_completion_with_fallbacks_multiple_keys():
# print(f"backup key 1: {os.getenv('BACKUP_OPENAI_API_KEY_1')}")
# print(f"backup key 2: {os.getenv('BACKUP_OPENAI_API_KEY_2')}")
# backup_keys = [{"api_key": os.getenv("BACKUP_OPENAI_API_KEY_1")}, {"api_key": os.getenv("BACKUP_OPENAI_API_KEY_2")}]
# try:
# api_key = "bad-key"
# response = completion(
# model="gpt-3.5-turbo", messages=messages, force_timeout=120, fallbacks=backup_keys, api_key=api_key
# )
# # Add any assertions here to check the response
# print(response)
# except Exception as e:
# error_str = traceback.format_exc()
# pytest.fail(f"Error occurred: {error_str}")
# test_completion_with_fallbacks_multiple_keys()
def test_petals():
try:
@@ -3392,6 +3220,13 @@ def test_petals():
## test deep infra
@pytest.mark.parametrize("drop_params", [True, False])
def test_completion_deep_infra(drop_params):
"""Test that DeepInfra requests are shaped correctly without making real API calls."""
from unittest.mock import MagicMock, patch
from openai import OpenAI
from openai.types.chat import ChatCompletion, ChatCompletionMessage
from openai.types.chat.chat_completion import Choice
import httpx
litellm.set_verbose = False
model_name = "deepinfra/meta-llama/Llama-2-70b-chat-hf"
tools = [
@@ -3420,7 +3255,51 @@ def test_completion_deep_infra(drop_params):
"content": "What's the weather like in Boston today in Fahrenheit?",
}
]
try:
mock_response = ChatCompletion(
id="chatcmpl-mock",
choices=[
Choice(
finish_reason="stop",
index=0,
message=ChatCompletionMessage(
content="It's sunny.", role="assistant"
),
)
],
created=1234567890,
model="meta-llama/Llama-2-70b-chat-hf",
object="chat.completion",
usage={"completion_tokens": 5, "prompt_tokens": 20, "total_tokens": 25},
)
mock_raw = MagicMock()
mock_raw.parse.return_value = mock_response
mock_raw.headers = httpx.Headers({"content-type": "application/json"})
mock_raw.status_code = 200
with patch(
"litellm.llms.openai.openai.OpenAIChatCompletion.make_sync_openai_chat_completion_request",
return_value=(mock_raw, mock_response),
) as mock_create:
if drop_params is False:
# DeepInfra doesn't support tool_choice, should raise UnsupportedParamsError
with pytest.raises(litellm.exceptions.UnsupportedParamsError):
completion(
model=model_name,
messages=messages,
temperature=0,
max_tokens=10,
tools=tools,
tool_choice={
"type": "function",
"function": {"name": "get_current_weather"},
},
drop_params=drop_params,
api_key="fake-api-key",
)
return
response = completion(
model=model_name,
messages=messages,
@@ -3432,33 +3311,75 @@ def test_completion_deep_infra(drop_params):
"function": {"name": "get_current_weather"},
},
drop_params=drop_params,
api_key="fake-api-key",
)
# Add any assertions here to check the response
print(response)
except Exception as e:
if drop_params is True:
pytest.fail(f"Error occurred: {e}")
# Verify the call was made
mock_create.assert_called_once()
call_kwargs = mock_create.call_args.kwargs
# Verify request shape
data = call_kwargs["data"]
assert data["model"] == "meta-llama/Llama-2-70b-chat-hf"
assert data["messages"] == messages
assert data["temperature"] == 0
assert data["max_tokens"] == 10
# tool_choice should be dropped for unsupported params
assert "tool_choice" not in data
# test_completion_deep_infra()
def test_completion_deep_infra_mistral():
print("deep infra test with temp=0")
"""Test that DeepInfra Mistral requests are shaped correctly without making real API calls."""
from unittest.mock import MagicMock, patch
from openai.types.chat import ChatCompletion, ChatCompletionMessage
from openai.types.chat.chat_completion import Choice
import httpx
model_name = "deepinfra/mistralai/Mistral-7B-Instruct-v0.1"
try:
mock_response = ChatCompletion(
id="chatcmpl-mock",
choices=[
Choice(
finish_reason="stop",
index=0,
message=ChatCompletionMessage(
content="Hello!", role="assistant"
),
)
],
created=1234567890,
model="mistralai/Mistral-7B-Instruct-v0.1",
object="chat.completion",
usage={"completion_tokens": 5, "prompt_tokens": 20, "total_tokens": 25},
)
mock_raw = MagicMock()
mock_raw.parse.return_value = mock_response
mock_raw.headers = httpx.Headers({"content-type": "application/json"})
mock_raw.status_code = 200
with patch(
"litellm.llms.openai.openai.OpenAIChatCompletion.make_sync_openai_chat_completion_request",
return_value=(mock_raw, mock_response),
) as mock_create:
response = completion(
model=model_name,
messages=messages,
temperature=0.01, # mistrail fails with temperature=0
temperature=0.01,
max_tokens=10,
api_key="fake-api-key",
)
# Add any assertions here to check the response
print(response)
except litellm.exceptions.Timeout as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
mock_create.assert_called_once()
call_kwargs = mock_create.call_args.kwargs
data = call_kwargs["data"]
assert data["model"] == "mistralai/Mistral-7B-Instruct-v0.1"
assert data["temperature"] == 0.01
assert data["max_tokens"] == 10
# test_completion_deep_infra_mistral()
@@ -3871,9 +3792,6 @@ async def test_dynamic_azure_params(stream, sync_mode):
raise e
@pytest.mark.parametrize(
"model",
["gpt-4o", "azure/gpt-4.1-mini"],
tests/local_testing/test_config.py
+10 -11
View File
@@ -47,8 +47,8 @@ async def test_delete_deployment():
litellm_params = LiteLLM_Params(
model="azure/gpt-4.1-mini",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
)
encrypted_litellm_params = litellm_params.dict(exclude_none=True)
@@ -131,8 +131,8 @@ async def test_add_existing_deployment():
litellm_params = LiteLLM_Params(
model="gpt-3.5-turbo",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
)
deployment = Deployment(model_name="gpt-3.5-turbo", litellm_params=litellm_params)
@@ -186,8 +186,8 @@ async def test_db_error_new_model_check():
litellm_params = LiteLLM_Params(
model="gpt-3.5-turbo",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
)
deployment = Deployment(model_name="gpt-3.5-turbo", litellm_params=litellm_params)
@@ -233,8 +233,8 @@ async def test_db_error_new_model_check():
litellm_params = LiteLLM_Params(
model="azure/gpt-4.1-mini",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
)
@@ -251,8 +251,8 @@ def _create_model_list(flag_value: Literal[0, 1], master_key: str):
new_litellm_params = LiteLLM_Params(
model="azure/gpt-4.1-mini-3",
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_version=os.getenv("AZURE_API_VERSION"),
)
@@ -421,4 +421,3 @@ def test_litellm_proxy_responses_api_config():
assert (
config.custom_llm_provider == LlmProviders.LITELLM_PROXY
), "custom_llm_provider should be LITELLM_PROXY"
tests/local_testing/test_configs/test_bad_config.yaml
+4 -4
View File
@@ -6,16 +6,16 @@ model_list:
- model_name: working-azure-gpt-3.5-turbo
litellm_params:
model: azure/gpt-4.1-mini
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
- model_name: azure-gpt-3.5-turbo
litellm_params:
model: azure/gpt-4.1-mini
api_base: os.environ/AZURE_API_BASE
api_base: os.environ/AZURE_AI_API_BASE
api_key: bad-key
- model_name: azure-embedding
litellm_params:
model: azure/text-embedding-ada-002
api_base: os.environ/AZURE_API_BASE
api_base: os.environ/AZURE_AI_API_BASE
api_key: bad-key
tests/local_testing/test_configs/test_cloudflare_azure_with_cache_config.yaml
+1 -1
View File
@@ -3,7 +3,7 @@ model_list:
litellm_params:
model: azure/gpt-4.1-mini
api_base: https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-07-01-preview
litellm_settings:
tests/local_testing/test_configs/test_config_no_auth.yaml
+7 -7
View File
@@ -11,7 +11,7 @@ model_list:
model_name: azure-model
- litellm_params:
api_base: https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1
api_key: os.environ/AZURE_API_KEY
api_key: os.environ/AZURE_AI_API_KEY
model: azure/gpt-4.1-mini
model_name: azure-cloudflare-model
- litellm_params:
@@ -49,8 +49,8 @@ model_list:
id: 79fc75bf-8e1b-47d5-8d24-9365a854af03
model_name: test_openai_models
- litellm_params:
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-07-01-preview
model: azure/text-embedding-ada-002
model_info:
@@ -94,16 +94,16 @@ model_list:
mode: image_generation
model_name: dall-e-3
- litellm_params:
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-06-01-preview
model: azure/
model_info:
mode: image_generation
model_name: dall-e-2
- litellm_params:
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: 2023-07-01-preview
model: azure/text-embedding-ada-002
model_info:
tests/local_testing/test_configs/test_custom_logger.yaml
+4 -4
View File
@@ -2,8 +2,8 @@ model_list:
- model_name: Azure OpenAI GPT-4 Canada
litellm_params:
model: azure/gpt-4.1-mini
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: "2023-07-01-preview"
model_info:
mode: chat
@@ -12,8 +12,8 @@ model_list:
- model_name: azure-embedding-model
litellm_params:
model: azure/text-embedding-ada-002
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_key: os.environ/AZURE_AI_API_KEY
api_version: "2023-07-01-preview"
model_info:
mode: embedding
tests/local_testing/test_embedding.py
+79 -67
View File
@@ -108,7 +108,11 @@ def test_openai_embedding_3():
"model, api_base, api_key",
[
# ("azure/text-embedding-ada-002", None, None),
("together_ai/BAAI/bge-base-en-v1.5", None, None), # Updated to current Together AI embedding model
(
"together_ai/BAAI/bge-base-en-v1.5",
None,
None,
), # Updated to current Together AI embedding model
],
)
@pytest.mark.parametrize("sync_mode", [True, False])
@@ -193,9 +197,9 @@ def _azure_ai_image_mock_response(*args, **kwargs):
"model, api_base, api_key",
[
(
"azure_ai/Cohere-embed-v3-multilingual-jzu",
"https://Cohere-embed-v3-multilingual-jzu.eastus2.models.ai.azure.com",
os.getenv("AZURE_AI_COHERE_API_KEY_2"),
"azure_ai/Cohere-embed-v3-multilingual-2",
os.getenv("AZURE_AI_API_BASE"),
os.getenv("AZURE_AI_API_KEY"),
)
],
)
@@ -292,13 +296,13 @@ def test_openai_embedding_timeouts():
def test_openai_azure_embedding():
try:
api_key = os.environ["AZURE_API_KEY"]
api_base = os.environ["AZURE_API_BASE"]
api_key = os.environ["AZURE_AI_API_KEY"]
api_base = os.environ["AZURE_AI_API_BASE"]
api_version = os.environ["AZURE_API_VERSION"]
os.environ["AZURE_API_VERSION"] = ""
os.environ["AZURE_API_BASE"] = ""
os.environ["AZURE_API_KEY"] = ""
os.environ["AZURE_AI_API_BASE"] = ""
os.environ["AZURE_AI_API_KEY"] = ""
response = embedding(
model="azure/text-embedding-ada-002",
@@ -310,8 +314,8 @@ def test_openai_azure_embedding():
print(response)
os.environ["AZURE_API_VERSION"] = api_version
os.environ["AZURE_API_BASE"] = api_base
os.environ["AZURE_API_KEY"] = api_key
os.environ["AZURE_AI_API_BASE"] = api_base
os.environ["AZURE_AI_API_KEY"] = api_key
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@@ -353,11 +357,11 @@ def test_openai_azure_embedding_optional_arg():
)
mock_client.assert_called_once_with(
model="test",
input=["test"],
extra_body={"azure_ad_token": "test"},
timeout=600,
extra_headers={"X-Stainless-Raw-Response": "true"}
model="test",
input=["test"],
extra_body={"azure_ad_token": "test"},
timeout=600,
extra_headers={"X-Stainless-Raw-Response": "true"},
)
# Verify azure_ad_token is passed in extra_body, not as a direct parameter
assert "azure_ad_token" not in mock_client.call_args.kwargs
@@ -545,7 +549,7 @@ def test_bedrock_embedding_cohere():
"good morning from litellm, attempting to embed data",
"lets test a second string for good measure",
],
aws_region_name="os.environ/AWS_REGION_NAME_2",
aws_region_name="os.environ/AWS_REGION_NAME",
)
assert isinstance(
response["data"][0]["embedding"], list
@@ -811,7 +815,7 @@ def test_watsonx_embeddings(monkeypatch):
monkeypatch.setenv("WATSONX_PROJECT_ID", "mock-project-id")
client = HTTPHandler()
# Track the actual request made
captured_request = {}
@@ -820,7 +824,7 @@ def test_watsonx_embeddings(monkeypatch):
captured_request["url"] = url
captured_request["headers"] = kwargs.get("headers", {})
captured_request["data"] = kwargs.get("data")
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/json"}
@@ -843,10 +847,12 @@ def test_watsonx_embeddings(monkeypatch):
print(f"response: {response}")
assert isinstance(response.usage, litellm.Usage)
# Verify the request was made correctly
assert "Authorization" in captured_request["headers"]
assert captured_request["headers"]["Authorization"] == "Bearer mock-watsonx-token"
assert (
captured_request["headers"]["Authorization"] == "Bearer mock-watsonx-token"
)
assert "us-south.ml.cloud.ibm.com" in captured_request["url"]
except litellm.RateLimitError as e:
pass
@@ -1256,9 +1262,7 @@ def test_jina_ai_img_embeddings(input_data, expected_payload_input):
# Call the function we want to test
try:
litellm.embedding(
model="jina_ai/jina-embeddings-v4", input=input_data
)
litellm.embedding(model="jina_ai/jina-embeddings-v4", input=input_data)
except Exception as e:
pytest.fail(
f"litellm.embedding call failed with an unexpected exception: {e}"
@@ -1285,105 +1289,113 @@ def test_jina_ai_img_embeddings(input_data, expected_payload_input):
def test_encoding_format_none_not_omitted_from_openai_sdk():
"""
Test that encoding_format=None is explicitly sent to OpenAI SDK.
This test verifies that when encoding_format is not provided by the user,
liteLLM explicitly sets it to None rather than omitting it. This prevents
the OpenAI SDK from adding its default value of 'base64'.
Without this fix:
- OpenAI SDK adds encoding_format='base64' as default when parameter is missing
- This causes issues with providers that don't support encoding_format (like Gemini)
With this fix:
- encoding_format=None is explicitly passed
- OpenAI SDK respects the explicit None and doesn't add defaults
"""
with patch("litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client") as mock_get_client:
with patch(
"litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client"
) as mock_get_client:
# Create a mock client instance
mock_client_instance = MagicMock()
mock_get_client.return_value = mock_client_instance
# Mock the embeddings.with_raw_response.create method
mock_response = MagicMock()
mock_response.parse.return_value = MagicMock(
model_dump=lambda: {
'data': [{'embedding': [0.1, 0.2, 0.3], 'index': 0}],
'model': 'text-embedding-ada-002',
'object': 'list',
'usage': {'prompt_tokens': 1, 'total_tokens': 1}
"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}],
"model": "text-embedding-ada-002",
"object": "list",
"usage": {"prompt_tokens": 1, "total_tokens": 1},
}
)
mock_response.headers = {}
mock_client_instance.embeddings.with_raw_response.create.return_value = mock_response
mock_client_instance.embeddings.with_raw_response.create.return_value = (
mock_response
)
# Call the embedding function without encoding_format
response = embedding(
model="text-embedding-ada-002",
input="Hello world",
)
# Get the call arguments to verify what was sent to OpenAI SDK
call_args = mock_client_instance.embeddings.with_raw_response.create.call_args
assert call_args is not None, "OpenAI SDK embeddings.create should have been called"
assert (
call_args is not None
), "OpenAI SDK embeddings.create should have been called"
call_kwargs = call_args[1] # Get kwargs
# The key assertion: encoding_format should be in the request with value None
# This prevents OpenAI SDK from adding its default 'base64' value
assert 'encoding_format' in call_kwargs, (
assert "encoding_format" in call_kwargs, (
"encoding_format should be explicitly passed to OpenAI SDK "
"(even if None) to prevent SDK from adding default value"
)
assert call_kwargs['encoding_format'] is None, (
"encoding_format should be None when not provided by user"
)
assert (
call_kwargs["encoding_format"] is None
), "encoding_format should be None when not provided by user"
print("✅ PASS: encoding_format=None is correctly passed to OpenAI SDK")
def test_encoding_format_explicit_value_preserved():
"""
Test that explicitly provided encoding_format values are preserved.
When user provides encoding_format='float' or 'base64', it should be
When user provides encoding_format='float' or 'base64', it should be
sent as-is to the OpenAI SDK.
"""
with patch("litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client") as mock_get_client:
with patch(
"litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client"
) as mock_get_client:
# Create a mock client instance
mock_client_instance = MagicMock()
mock_get_client.return_value = mock_client_instance
# Mock the embeddings.with_raw_response.create method
mock_response = MagicMock()
mock_response.parse.return_value = MagicMock(
model_dump=lambda: {
'data': [{'embedding': [0.1, 0.2, 0.3], 'index': 0}],
'model': 'text-embedding-ada-002',
'object': 'list',
'usage': {'prompt_tokens': 1, 'total_tokens': 1}
"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}],
"model": "text-embedding-ada-002",
"object": "list",
"usage": {"prompt_tokens": 1, "total_tokens": 1},
}
)
mock_response.headers = {}
mock_client_instance.embeddings.with_raw_response.create.return_value = mock_response
mock_client_instance.embeddings.with_raw_response.create.return_value = (
mock_response
)
# Test with explicit encoding_format='float'
response = embedding(
model="text-embedding-ada-002",
input="Hello world",
encoding_format="float"
model="text-embedding-ada-002", input="Hello world", encoding_format="float"
)
# Verify the encoding_format was passed correctly
call_args = mock_client_instance.embeddings.with_raw_response.create.call_args
call_kwargs = call_args[1]
assert 'encoding_format' in call_kwargs, (
"encoding_format should be in the request"
)
assert call_kwargs['encoding_format'] == 'float', (
"encoding_format should be 'float' when explicitly provided"
)
assert (
"encoding_format" in call_kwargs
), "encoding_format should be in the request"
assert (
call_kwargs["encoding_format"] == "float"
), "encoding_format should be 'float' when explicitly provided"
print("✅ PASS: encoding_format='float' is correctly preserved")
tests/local_testing/test_exceptions.py
+60 -47
View File
@@ -24,7 +24,7 @@ from litellm import ( # AuthenticationError,; RateLimitError,; ServiceUnavailab
embedding,
)
litellm.vertex_project = "pathrise-convert-1606954137718"
litellm.vertex_project = "litellm-ci-cd"
litellm.vertex_location = "us-central1"
litellm.num_retries = 0
@@ -162,8 +162,8 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th
temporary_secret_key = os.environ["AWS_SECRET_ACCESS_KEY"]
os.environ["AWS_SECRET_ACCESS_KEY"] = "bad-key"
elif model == "azure/gpt-4.1-mini":
temporary_key = os.environ["AZURE_API_KEY"]
os.environ["AZURE_API_KEY"] = "bad-key"
temporary_key = os.environ["AZURE_AI_API_KEY"]
os.environ["AZURE_AI_API_KEY"] = "bad-key"
elif model == "claude-3-5-haiku-20241022":
temporary_key = os.environ["ANTHROPIC_API_KEY"]
os.environ["ANTHROPIC_API_KEY"] = "bad-key"
@@ -175,9 +175,7 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th
os.environ["AI21_API_KEY"] = "bad-key"
elif "togethercomputer" in model:
temporary_key = os.environ["TOGETHERAI_API_KEY"]
os.environ["TOGETHERAI_API_KEY"] = (
"sk-test-togetherai-key-808"
)
os.environ["TOGETHERAI_API_KEY"] = "sk-test-togetherai-key-808"
elif model in litellm.openrouter_models:
temporary_key = os.environ["OPENROUTER_API_KEY"]
os.environ["OPENROUTER_API_KEY"] = "bad-key"
@@ -185,7 +183,6 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th
temporary_key = os.environ["ALEPH_ALPHA_API_KEY"]
os.environ["ALEPH_ALPHA_API_KEY"] = "bad-key"
elif model in litellm.nlp_cloud_models:
temporary_key = os.environ["NLP_CLOUD_API_KEY"]
os.environ["NLP_CLOUD_API_KEY"] = "bad-key"
elif (
model
@@ -212,7 +209,7 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th
if model == "gpt-3.5-turbo":
os.environ["OPENAI_API_KEY"] = temporary_key
elif model == "chatgpt-test":
os.environ["AZURE_API_KEY"] = temporary_key
os.environ["AZURE_AI_API_KEY"] = temporary_key
azure = True
elif model == "claude-3-5-haiku-20241022":
os.environ["ANTHROPIC_API_KEY"] = temporary_key
@@ -230,7 +227,7 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th
elif model in litellm.aleph_alpha_models:
os.environ["ALEPH_ALPHA_API_KEY"] = temporary_key
elif model in litellm.nlp_cloud_models:
os.environ["NLP_CLOUD_API_KEY"] = temporary_key
os.environ.pop("NLP_CLOUD_API_KEY", None)
elif "bedrock" in model:
os.environ["AWS_ACCESS_KEY_ID"] = temporary_aws_access_key
os.environ["AWS_REGION_NAME"] = temporary_aws_region_name
@@ -259,17 +256,17 @@ def test_completion_azure_exception():
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = True
## Test azure call
old_azure_key = os.environ["AZURE_API_KEY"]
os.environ["AZURE_API_KEY"] = "good morning"
old_azure_key = os.environ["AZURE_AI_API_KEY"]
os.environ["AZURE_AI_API_KEY"] = "good morning"
response = completion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "hello"}],
)
os.environ["AZURE_API_KEY"] = old_azure_key
os.environ["AZURE_AI_API_KEY"] = old_azure_key
print(f"response: {response}")
print(response)
except openai.AuthenticationError as e:
os.environ["AZURE_API_KEY"] = old_azure_key
os.environ["AZURE_AI_API_KEY"] = old_azure_key
print("good job got the correct error for azure when key not set")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@@ -303,8 +300,8 @@ async def asynctest_completion_azure_exception():
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = True
## Test azure call
old_azure_key = os.environ["AZURE_API_KEY"]
os.environ["AZURE_API_KEY"] = "good morning"
old_azure_key = os.environ["AZURE_AI_API_KEY"]
os.environ["AZURE_AI_API_KEY"] = "good morning"
response = await litellm.acompletion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "hello"}],
@@ -312,7 +309,7 @@ async def asynctest_completion_azure_exception():
print(f"response: {response}")
print(response)
except openai.AuthenticationError as e:
os.environ["AZURE_API_KEY"] = old_azure_key
os.environ["AZURE_AI_API_KEY"] = old_azure_key
print("good job got the correct error for azure when key not set")
print(e)
except Exception as e:
@@ -495,6 +492,7 @@ def test_completion_bedrock_invalid_role_exception():
== "litellm.BadRequestError: Invalid Message passed in {'role': 'very-bad-role', 'content': 'hello'}"
)
@pytest.mark.skip(reason="OpenAI exception changed to a generic error")
def test_content_policy_exceptionimage_generation_openai():
try:
@@ -773,7 +771,15 @@ def test_litellm_predibase_exception():
@pytest.mark.parametrize(
"provider", ["predibase", "vertex_ai_beta", "anthropic", "databricks", "watsonx", "fireworks_ai"]
"provider",
[
"predibase",
"vertex_ai_beta",
"anthropic",
"databricks",
"watsonx",
"fireworks_ai",
],
)
def test_exception_mapping(provider):
"""
@@ -826,14 +832,14 @@ def test_fireworks_ai_exception_mapping():
2. Text-based rate limit detection (the main issue fixed)
3. Generic 400 errors that should NOT be rate limits
4. ExceptionCheckers utility function
Related to: https://github.com/BerriAI/litellm/pull/11455
Based on Fireworks AI documentation: https://docs.fireworks.ai/tools-sdks/python-client/api-reference
"""
import litellm
from litellm.llms.fireworks_ai.common_utils import FireworksAIException
from litellm.litellm_core_utils.exception_mapping_utils import ExceptionCheckers
# Test scenarios covering all important cases
test_scenarios = [
{
@@ -855,57 +861,63 @@ def test_fireworks_ai_exception_mapping():
"expected_exception": litellm.BadRequestError,
},
]
# Test each scenario
for scenario in test_scenarios:
mock_exception = FireworksAIException(
status_code=scenario["status_code"],
message=scenario["message"],
headers={}
status_code=scenario["status_code"], message=scenario["message"], headers={}
)
try:
response = litellm.completion(
model="fireworks_ai/llama-v3p1-70b-instruct",
messages=[{"role": "user", "content": "Hello"}],
mock_response=mock_exception,
)
pytest.fail(f"Expected {scenario['expected_exception'].__name__} to be raised")
pytest.fail(
f"Expected {scenario['expected_exception'].__name__} to be raised"
)
except scenario["expected_exception"] as e:
if scenario["expected_exception"] == litellm.RateLimitError:
assert "rate limit" in str(e).lower() or "429" in str(e)
except Exception as e:
pytest.fail(f"Expected {scenario['expected_exception'].__name__} but got {type(e).__name__}: {e}")
pytest.fail(
f"Expected {scenario['expected_exception'].__name__} but got {type(e).__name__}: {e}"
)
# Test ExceptionCheckers.is_error_str_rate_limit() method directly
# Test cases that should return True (rate limit detected)
rate_limit_strings = [
"429 rate limit exceeded",
"Rate limit exceeded, please try again later",
"Rate limit exceeded, please try again later",
"RATE LIMIT ERROR",
"Error 429: rate limit",
'{"error":{"type":"invalid_request_error","message":"rate limit exceeded, please try again later"}}',
"HTTP 429 Too Many Requests",
]
for error_str in rate_limit_strings:
assert ExceptionCheckers.is_error_str_rate_limit(error_str), f"Should detect rate limit in: {error_str}"
assert ExceptionCheckers.is_error_str_rate_limit(
error_str
), f"Should detect rate limit in: {error_str}"
# Test cases that should return False (not rate limit)
non_rate_limit_strings = [
"400 Bad Request",
"Authentication failed",
"Authentication failed",
"Invalid model specified",
"Context window exceeded",
"Internal server error",
"",
"Some other error message",
]
for error_str in non_rate_limit_strings:
assert not ExceptionCheckers.is_error_str_rate_limit(error_str), f"Should NOT detect rate limit in: {error_str}"
assert not ExceptionCheckers.is_error_str_rate_limit(
error_str
), f"Should NOT detect rate limit in: {error_str}"
# Test edge cases
assert not ExceptionCheckers.is_error_str_rate_limit(None) # type: ignore
assert not ExceptionCheckers.is_error_str_rate_limit(42) # type: ignore
@@ -1142,6 +1154,7 @@ def test_openai_gateway_timeout_error():
"""
openai_client = OpenAI()
mapped_target = openai_client.chat.completions.with_raw_response # type: ignore
def _return_exception(*args, **kwargs):
import datetime
@@ -1175,13 +1188,17 @@ def test_openai_gateway_timeout_error():
setattr(exception, k, v)
raise exception
try:
try:
with patch.object(
mapped_target,
"create",
side_effect=_return_exception,
):
litellm.completion(model="openai/gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello world"}], client=openai_client)
litellm.completion(
model="openai/gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello world"}],
client=openai_client,
)
pytest.fail("Expected to raise Timeout")
except litellm.Timeout as e:
assert e.status_code == 504
@@ -1350,7 +1367,7 @@ def test_context_window_exceeded_error_from_litellm_proxy():
def test_bad_request_error_with_response_without_request():
"""
Test that BadRequestError handles Response objects without a request attribute.
This simulates a real scenario where a Response is created without a request
(e.g., in tests or when manually creating error responses), and we need to
ensure it doesn't raise RuntimeError when the exception is created.
@@ -1362,8 +1379,7 @@ def test_bad_request_error_with_response_without_request():
# Create a Response without a request (simulates the scenario that was failing)
response_without_request = Response(status_code=400, text="Bad Request")
# Test that extract_and_raise_litellm_exception can handle this
args = {
"response": response_without_request,
@@ -1371,17 +1387,17 @@ def test_bad_request_error_with_response_without_request():
"model": "gpt-3.5-turbo",
"custom_llm_provider": "openai",
}
# This should raise BadRequestError without RuntimeError
with pytest.raises(litellm.BadRequestError) as exc_info:
extract_and_raise_litellm_exception(**args)
# Verify the exception was created successfully
error = exc_info.value
assert error is not None
assert error.model == "gpt-3.5-turbo"
assert error.llm_provider == "openai"
# Verify the exception has a response (should be minimal error response)
assert error.response is not None
# The response should have a request (minimal error response has one)
@@ -1420,6 +1436,3 @@ async def test_exception_bubbling_up(sync_mode, stream_mode, model):
assert exc_info.value.code == "invalid_value"
assert exc_info.value.param is not None
assert exc_info.value.type == "invalid_request_error"
tests/local_testing/test_gcs_bucket.py
+172 -649
View File
@@ -21,176 +21,68 @@ from litellm.integrations.gcs_bucket.gcs_bucket import (
StandardLoggingPayload,
)
from litellm.types.utils import StandardCallbackDynamicParams
from unittest.mock import patch
from litellm.types.integrations.gcs_bucket import GCSLoggingConfig
from unittest.mock import patch, AsyncMock, MagicMock
verbose_logger.setLevel(logging.DEBUG)
def load_vertex_ai_credentials():
# Define the path to the vertex_key.json file
print("loading vertex ai credentials")
os.environ["GCS_FLUSH_INTERVAL"] = "1"
os.environ["GCS_USE_BATCHED_LOGGING"] = "false"
filepath = os.path.dirname(os.path.abspath(__file__))
vertex_key_path = filepath + "/vertex_key.json"
# Read the existing content of the file or create an empty dictionary
try:
with open(vertex_key_path, "r") as file:
# Read the file content
print("Read vertexai file path")
content = file.read()
# If the file is empty or not valid JSON, create an empty dictionary
if not content or not content.strip():
service_account_key_data = {}
else:
# Attempt to load the existing JSON content
file.seek(0)
service_account_key_data = json.load(file)
except FileNotFoundError:
# If the file doesn't exist, create an empty dictionary
service_account_key_data = {}
# Update the service_account_key_data with environment variables
private_key_id = os.environ.get("GCS_PRIVATE_KEY_ID", "")
private_key = os.environ.get("GCS_PRIVATE_KEY", "")
private_key = private_key.replace("\\n", "\n")
service_account_key_data["private_key_id"] = private_key_id
service_account_key_data["private_key"] = private_key
# Create a temporary file
with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file:
# Write the updated content to the temporary files
json.dump(service_account_key_data, temp_file, indent=2)
# Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS
os.environ["GCS_PATH_SERVICE_ACCOUNT"] = os.path.abspath(temp_file.name)
print("created gcs path service account=", os.environ["GCS_PATH_SERVICE_ACCOUNT"])
def _make_mock_gcs_logging_config():
return GCSLoggingConfig(
bucket_name="test-bucket",
vertex_instance=MagicMock(),
path_service_account=None,
)
@pytest.mark.asyncio
async def test_aaabasic_gcs_logger():
load_vertex_ai_credentials()
gcs_logger = GCSBucketLogger()
print("GCSBucketLogger", gcs_logger)
os.environ["GCS_FLUSH_INTERVAL"] = "1"
os.environ["GCS_USE_BATCHED_LOGGING"] = "false"
os.environ["GCS_BUCKET_NAME"] = "test-bucket"
litellm.callbacks = [gcs_logger]
response = await litellm.acompletion(
model="gpt-3.5-turbo",
temperature=0.7,
messages=[{"role": "user", "content": "This is a test"}],
max_tokens=10,
user="ishaan-2",
mock_response="Hi!",
metadata={
"tags": ["model-anthropic-claude-v2.1", "app-ishaan-prod"],
"user_api_key": "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456",
"user_api_key_alias": None,
"user_api_end_user_max_budget": None,
"litellm_api_version": "0.0.0",
"global_max_parallel_requests": None,
"user_api_key_user_id": "116544810872468347480",
"user_api_key_org_id": None,
"user_api_key_team_id": None,
"user_api_key_team_alias": None,
"user_api_key_metadata": {},
"requester_ip_address": "127.0.0.1",
"requester_metadata": {"foo": "bar"},
"spend_logs_metadata": {"hello": "world"},
"headers": {
"content-type": "application/json",
"user-agent": "PostmanRuntime/7.32.3",
"accept": "*/*",
"postman-token": "92300061-eeaa-423b-a420-0b44896ecdc4",
"host": "localhost:4000",
"accept-encoding": "gzip, deflate, br",
"connection": "keep-alive",
"content-length": "163",
},
"endpoint": "http://localhost:4000/chat/completions",
"model_group": "gpt-3.5-turbo",
"model_info": {
"id": "4bad40a1eb6bebd1682800f16f44b9f06c52a6703444c99c7f9f32e9de3693b4",
"db_model": False,
},
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"caching_groups": None,
"raw_request": "\n\nPOST Request Sent from LiteLLM:\ncurl -X POST \\\nhttps://openai-gpt-4-test-v-1.openai.azure.com//openai/ \\\n-H 'Authorization: *****' \\\n-d '{'model': 'chatgpt-v-3', 'messages': [{'role': 'system', 'content': 'you are a helpful assistant.\\n'}, {'role': 'user', 'content': 'bom dia'}], 'stream': False, 'max_tokens': 10, 'user': '116544810872468347480', 'extra_body': {}}'\n",
},
)
captured_payloads = []
print("response", response)
async def mock_log_json_data_on_gcs(
self, headers, bucket_name, object_name, logging_payload
):
captured_payloads.append(
{
"bucket_name": bucket_name,
"object_name": object_name,
"logging_payload": logging_payload,
}
)
return {"kind": "storage#object", "name": object_name}
await asyncio.sleep(5)
with patch(
"litellm.proxy.proxy_server.premium_user", True
), patch.object(
GCSBucketLogger,
"construct_request_headers",
new_callable=AsyncMock,
return_value={"Authorization": "Bearer mock_token"},
), patch.object(
GCSBucketLogger,
"get_gcs_logging_config",
new_callable=AsyncMock,
return_value=_make_mock_gcs_logging_config(),
), patch.object(
GCSBucketLogger,
"_log_json_data_on_gcs",
mock_log_json_data_on_gcs,
):
gcs_logger = GCSBucketLogger()
# Get the current date
# Get the current date
current_date = datetime.now().strftime("%Y-%m-%d")
# Modify the object_name to include the date-based folder
object_name = f"{current_date}%2F{response.id}"
print("object_name", object_name)
# Check if object landed on GCS
object_from_gcs = await gcs_logger.download_gcs_object(object_name=object_name)
print("object from gcs=", object_from_gcs)
# convert object_from_gcs from bytes to DICT
parsed_data = json.loads(object_from_gcs)
print("object_from_gcs as dict", parsed_data)
print("type of object_from_gcs", type(parsed_data))
gcs_payload = StandardLoggingPayload(**parsed_data)
print("gcs_payload", gcs_payload)
assert gcs_payload["model"] == "gpt-3.5-turbo"
assert gcs_payload["messages"] == [{"role": "user", "content": "This is a test"}]
assert gcs_payload["response"]["choices"][0]["message"]["content"] == "Hi!"
assert gcs_payload["response_cost"] > 0.0
assert gcs_payload["status"] == "success"
assert (
gcs_payload["metadata"]["user_api_key_hash"]
== "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456"
)
assert gcs_payload["metadata"]["user_api_key_user_id"] == "116544810872468347480"
assert gcs_payload["metadata"]["requester_metadata"] == {"foo": "bar"}
# Delete Object from GCS
print("deleting object from GCS")
await gcs_logger.delete_gcs_object(object_name=object_name)
@pytest.mark.asyncio
async def test_basic_gcs_logger_failure():
load_vertex_ai_credentials()
gcs_logger = GCSBucketLogger()
print("GCSBucketLogger", gcs_logger)
gcs_log_id = f"failure-test-{uuid.uuid4().hex}"
litellm.callbacks = [gcs_logger]
try:
litellm.callbacks = [gcs_logger]
response = await litellm.acompletion(
model="gpt-3.5-turbo",
temperature=0.7,
messages=[{"role": "user", "content": "This is a test"}],
max_tokens=10,
user="ishaan-2",
mock_response=litellm.BadRequestError(
model="gpt-3.5-turbo",
message="Error: 400: Bad Request: Invalid API key, please check your API key and try again.",
llm_provider="openai",
),
mock_response="Hi!",
metadata={
"gcs_log_id": gcs_log_id,
"tags": ["model-anthropic-claude-v2.1", "app-ishaan-prod"],
"user_api_key": "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456",
"user_api_key_alias": None,
@@ -203,6 +95,7 @@ async def test_basic_gcs_logger_failure():
"user_api_key_team_alias": None,
"user_api_key_metadata": {},
"requester_ip_address": "127.0.0.1",
"requester_metadata": {"foo": "bar"},
"spend_logs_metadata": {"hello": "world"},
"headers": {
"content-type": "application/json",
@@ -225,523 +118,153 @@ async def test_basic_gcs_logger_failure():
"raw_request": "\n\nPOST Request Sent from LiteLLM:\ncurl -X POST \\\nhttps://openai-gpt-4-test-v-1.openai.azure.com//openai/ \\\n-H 'Authorization: *****' \\\n-d '{'model': 'chatgpt-v-3', 'messages': [{'role': 'system', 'content': 'you are a helpful assistant.\\n'}, {'role': 'user', 'content': 'bom dia'}], 'stream': False, 'max_tokens': 10, 'user': '116544810872468347480', 'extra_body': {}}'\n",
},
)
except Exception:
pass
await asyncio.sleep(5)
print("response", response)
# Get the current date
# Get the current date
current_date = datetime.now().strftime("%Y-%m-%d")
await asyncio.sleep(3)
# Modify the object_name to include the date-based folder
object_name = gcs_log_id
assert len(captured_payloads) == 1, (
f"Expected 1 GCS upload, got {len(captured_payloads)}"
)
print("object_name", object_name)
gcs_payload = captured_payloads[0]["logging_payload"]
# Check if object landed on GCS
object_from_gcs = await gcs_logger.download_gcs_object(object_name=object_name)
print("object from gcs=", object_from_gcs)
# convert object_from_gcs from bytes to DICT
parsed_data = json.loads(object_from_gcs)
print("object_from_gcs as dict", parsed_data)
assert gcs_payload["model"] == "gpt-3.5-turbo"
assert gcs_payload["messages"] == [
{"role": "user", "content": "This is a test"}
]
print("type of object_from_gcs", type(parsed_data))
assert gcs_payload["response"]["choices"][0]["message"]["content"] == "Hi!"
gcs_payload = StandardLoggingPayload(**parsed_data)
assert gcs_payload["response_cost"] > 0.0
print("gcs_payload", gcs_payload)
assert gcs_payload["status"] == "success"
assert gcs_payload["model"] == "gpt-3.5-turbo"
assert gcs_payload["messages"] == [{"role": "user", "content": "This is a test"}]
assert (
gcs_payload["metadata"]["user_api_key_hash"]
== "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456"
)
assert (
gcs_payload["metadata"]["user_api_key_user_id"] == "116544810872468347480"
)
assert gcs_payload["response_cost"] == 0
assert gcs_payload["status"] == "failure"
assert (
gcs_payload["metadata"]["user_api_key_hash"]
== "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456"
)
assert gcs_payload["metadata"]["user_api_key_user_id"] == "116544810872468347480"
# Delete Object from GCS
print("deleting object from GCS")
await gcs_logger.delete_gcs_object(object_name=object_name)
assert gcs_payload["metadata"]["requester_metadata"] == {"foo": "bar"}
@pytest.mark.skip(reason="This test is flaky")
@pytest.mark.asyncio
async def test_basic_gcs_logging_per_request_with_callback_set():
"""
Test GCS Bucket logging per request
async def test_basic_gcs_logger_failure():
os.environ["GCS_FLUSH_INTERVAL"] = "1"
os.environ["GCS_USE_BATCHED_LOGGING"] = "false"
os.environ["GCS_BUCKET_NAME"] = "test-bucket"
Request 1 - pass gcs_bucket_name in kwargs
Request 2 - don't pass gcs_bucket_name in kwargs - ensure 'litellm-testing-bucket'
"""
import logging
from litellm._logging import verbose_logger
captured_payloads = []
verbose_logger.setLevel(logging.DEBUG)
load_vertex_ai_credentials()
gcs_logger = GCSBucketLogger()
print("GCSBucketLogger", gcs_logger)
litellm.callbacks = [gcs_logger]
GCS_BUCKET_NAME = "example-bucket-1-litellm"
standard_callback_dynamic_params: StandardCallbackDynamicParams = (
StandardCallbackDynamicParams(gcs_bucket_name=GCS_BUCKET_NAME)
)
try:
response = await litellm.acompletion(
model="gpt-4o-mini",
temperature=0.7,
messages=[{"role": "user", "content": "This is a test"}],
max_tokens=10,
user="ishaan-2",
gcs_bucket_name=GCS_BUCKET_NAME,
async def mock_log_json_data_on_gcs(
self, headers, bucket_name, object_name, logging_payload
):
captured_payloads.append(
{
"bucket_name": bucket_name,
"object_name": object_name,
"logging_payload": logging_payload,
}
)
except:
pass
return {"kind": "storage#object", "name": object_name}
await asyncio.sleep(5)
# Get the current date
# Get the current date
current_date = datetime.now().strftime("%Y-%m-%d")
# Modify the object_name to include the date-based folder
object_name = f"{current_date}%2F{response.id}"
print("object_name", object_name)
# Check if object landed on GCS
object_from_gcs = await gcs_logger.download_gcs_object(
object_name=object_name,
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
print("object from gcs=", object_from_gcs)
# convert object_from_gcs from bytes to DICT
parsed_data = json.loads(object_from_gcs)
print("object_from_gcs as dict", parsed_data)
print("type of object_from_gcs", type(parsed_data))
gcs_payload = StandardLoggingPayload(**parsed_data)
assert gcs_payload["model"] == "gpt-4o-mini"
assert gcs_payload["messages"] == [{"role": "user", "content": "This is a test"}]
assert gcs_payload["response_cost"] > 0.0
assert gcs_payload["status"] == "success"
# clean up the object from GCS
await gcs_logger.delete_gcs_object(
object_name=object_name,
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
# Request 2 - don't pass gcs_bucket_name in kwargs - ensure 'litellm-testing-bucket'
try:
response = await litellm.acompletion(
model="gpt-4o-mini",
temperature=0.7,
messages=[{"role": "user", "content": "This is a test"}],
max_tokens=10,
user="ishaan-2",
mock_response="Hi!",
)
except:
pass
await asyncio.sleep(5)
# Get the current date
# Get the current date
current_date = datetime.now().strftime("%Y-%m-%d")
standard_callback_dynamic_params = StandardCallbackDynamicParams(
gcs_bucket_name="litellm-testing-bucket"
)
# Modify the object_name to include the date-based folder
object_name = f"{current_date}%2F{response.id}"
print("object_name", object_name)
# Check if object landed on GCS
object_from_gcs = await gcs_logger.download_gcs_object(
object_name=object_name,
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
print("object from gcs=", object_from_gcs)
# convert object_from_gcs from bytes to DICT
parsed_data = json.loads(object_from_gcs)
print("object_from_gcs as dict", parsed_data)
print("type of object_from_gcs", type(parsed_data))
gcs_payload = StandardLoggingPayload(**parsed_data)
assert gcs_payload["model"] == "gpt-4o-mini"
assert gcs_payload["messages"] == [{"role": "user", "content": "This is a test"}]
assert gcs_payload["response_cost"] > 0.0
assert gcs_payload["status"] == "success"
# clean up the object from GCS
await gcs_logger.delete_gcs_object(
object_name=object_name,
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
@pytest.mark.skip(reason="This test is flaky")
@pytest.mark.asyncio
async def test_basic_gcs_logging_per_request_with_no_litellm_callback_set():
"""
Test GCS Bucket logging per request
key difference: no litellm.callbacks set
Request 1 - pass gcs_bucket_name in kwargs
Request 2 - don't pass gcs_bucket_name in kwargs - ensure 'litellm-testing-bucket'
"""
import logging
from litellm._logging import verbose_logger
verbose_logger.setLevel(logging.DEBUG)
load_vertex_ai_credentials()
gcs_logger = GCSBucketLogger()
GCS_BUCKET_NAME = "example-bucket-1-litellm"
standard_callback_dynamic_params: StandardCallbackDynamicParams = (
StandardCallbackDynamicParams(gcs_bucket_name=GCS_BUCKET_NAME)
)
try:
response = await litellm.acompletion(
model="gpt-4o-mini",
temperature=0.7,
messages=[{"role": "user", "content": "This is a test"}],
max_tokens=10,
user="ishaan-2",
gcs_bucket_name=GCS_BUCKET_NAME,
success_callback=["gcs_bucket"],
failure_callback=["gcs_bucket"],
)
except:
pass
await asyncio.sleep(5)
# Get the current date
# Get the current date
current_date = datetime.now().strftime("%Y-%m-%d")
# Modify the object_name to include the date-based folder
object_name = f"{current_date}%2F{response.id}"
print("object_name", object_name)
# Check if object landed on GCS
object_from_gcs = await gcs_logger.download_gcs_object(
object_name=object_name,
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
print("object from gcs=", object_from_gcs)
# convert object_from_gcs from bytes to DICT
parsed_data = json.loads(object_from_gcs)
print("object_from_gcs as dict", parsed_data)
print("type of object_from_gcs", type(parsed_data))
gcs_payload = StandardLoggingPayload(**parsed_data)
assert gcs_payload["model"] == "gpt-4o-mini"
assert gcs_payload["messages"] == [{"role": "user", "content": "This is a test"}]
assert gcs_payload["response_cost"] > 0.0
assert gcs_payload["status"] == "success"
# clean up the object from GCS
await gcs_logger.delete_gcs_object(
object_name=object_name,
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
# make a failure request - assert that failure callback is hit
gcs_log_id = f"failure-test-{uuid.uuid4().hex}"
try:
response = await litellm.acompletion(
model="gpt-4o-mini",
temperature=0.7,
messages=[{"role": "user", "content": "This is a test"}],
max_tokens=10,
user="ishaan-2",
mock_response=litellm.BadRequestError(
with patch(
"litellm.proxy.proxy_server.premium_user", True
), patch.object(
GCSBucketLogger,
"construct_request_headers",
new_callable=AsyncMock,
return_value={"Authorization": "Bearer mock_token"},
), patch.object(
GCSBucketLogger,
"get_gcs_logging_config",
new_callable=AsyncMock,
return_value=_make_mock_gcs_logging_config(),
), patch.object(
GCSBucketLogger,
"_log_json_data_on_gcs",
mock_log_json_data_on_gcs,
):
gcs_logger = GCSBucketLogger()
litellm.callbacks = [gcs_logger]
try:
response = await litellm.acompletion(
model="gpt-3.5-turbo",
message="Error: 400: Bad Request: Invalid API key, please check your API key and try again.",
llm_provider="openai",
),
success_callback=["gcs_bucket"],
failure_callback=["gcs_bucket"],
gcs_bucket_name=GCS_BUCKET_NAME,
metadata={
"gcs_log_id": gcs_log_id,
},
)
except:
pass
await asyncio.sleep(5)
# check if the failure object is logged in GCS
object_from_gcs = await gcs_logger.download_gcs_object(
object_name=gcs_log_id,
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
print("object from gcs=", object_from_gcs)
# convert object_from_gcs from bytes to DICT
parsed_data = json.loads(object_from_gcs)
print("object_from_gcs as dict", parsed_data)
gcs_payload = StandardLoggingPayload(**parsed_data)
assert gcs_payload["model"] == "gpt-4o-mini"
assert gcs_payload["messages"] == [{"role": "user", "content": "This is a test"}]
assert gcs_payload["response_cost"] == 0
assert gcs_payload["status"] == "failure"
# clean up the object from GCS
await gcs_logger.delete_gcs_object(
object_name=gcs_log_id,
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
@pytest.mark.skip(reason="This test is flaky")
@pytest.mark.asyncio
async def test_aaaget_gcs_logging_config_without_service_account():
"""
Test the get_gcs_logging_config works for IAM auth on GCS
1. Key based logging without a service account
2. Default Callback without a service account
"""
load_vertex_ai_credentials()
_old_gcs_bucket_name = os.environ.get("GCS_BUCKET_NAME")
os.environ.pop("GCS_BUCKET_NAME", None)
_old_gcs_service_acct = os.environ.get("GCS_PATH_SERVICE_ACCOUNT")
os.environ.pop("GCS_PATH_SERVICE_ACCOUNT", None)
# Mock the load_auth function to avoid credential loading issues
# Test 1: With standard_callback_dynamic_params (with service account)
gcs_logger = GCSBucketLogger()
dynamic_params = StandardCallbackDynamicParams(
gcs_bucket_name="dynamic-bucket",
)
config = await gcs_logger.get_gcs_logging_config(
{"standard_callback_dynamic_params": dynamic_params}
)
assert config["bucket_name"] == "dynamic-bucket"
assert config["path_service_account"] is None
assert config["vertex_instance"] is not None
# Test 2: With standard_callback_dynamic_params (without service account - this is IAM auth)
dynamic_params = StandardCallbackDynamicParams(
gcs_bucket_name="dynamic-bucket", gcs_path_service_account=None
)
config = await gcs_logger.get_gcs_logging_config(
{"standard_callback_dynamic_params": dynamic_params}
)
assert config["bucket_name"] == "dynamic-bucket"
assert config["path_service_account"] is None
assert config["vertex_instance"] is not None
# Test 5: With missing bucket name
with pytest.raises(ValueError, match="GCS_BUCKET_NAME is not set"):
gcs_logger = GCSBucketLogger(bucket_name=None)
await gcs_logger.get_gcs_logging_config({})
if _old_gcs_bucket_name is not None:
os.environ["GCS_BUCKET_NAME"] = _old_gcs_bucket_name
if _old_gcs_service_acct is not None:
os.environ["GCS_PATH_SERVICE_ACCOUNT"] = _old_gcs_service_acct
@pytest.mark.skip(reason="This test is flaky")
@pytest.mark.asyncio
async def test_basic_gcs_logger_with_folder_in_bucket_name():
load_vertex_ai_credentials()
gcs_logger = GCSBucketLogger()
bucket_name = "litellm-testing-bucket/test-folder-logs"
old_bucket_name = os.environ.get("GCS_BUCKET_NAME")
os.environ["GCS_BUCKET_NAME"] = bucket_name
print("GCSBucketLogger", gcs_logger)
litellm.callbacks = [gcs_logger]
response = await litellm.acompletion(
model="gpt-3.5-turbo",
temperature=0.7,
messages=[{"role": "user", "content": "This is a test"}],
max_tokens=10,
user="ishaan-2",
mock_response="Hi!",
metadata={
"tags": ["model-anthropic-claude-v2.1", "app-ishaan-prod"],
"user_api_key": "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456",
"user_api_key_alias": None,
"user_api_end_user_max_budget": None,
"litellm_api_version": "0.0.0",
"global_max_parallel_requests": None,
"user_api_key_user_id": "116544810872468347480",
"user_api_key_org_id": None,
"user_api_key_team_id": None,
"user_api_key_team_alias": None,
"user_api_key_metadata": {},
"requester_ip_address": "127.0.0.1",
"requester_metadata": {"foo": "bar"},
"spend_logs_metadata": {"hello": "world"},
"headers": {
"content-type": "application/json",
"user-agent": "PostmanRuntime/7.32.3",
"accept": "*/*",
"postman-token": "92300061-eeaa-423b-a420-0b44896ecdc4",
"host": "localhost:4000",
"accept-encoding": "gzip, deflate, br",
"connection": "keep-alive",
"content-length": "163",
},
"endpoint": "http://localhost:4000/chat/completions",
"model_group": "gpt-3.5-turbo",
"model_info": {
"id": "4bad40a1eb6bebd1682800f16f44b9f06c52a6703444c99c7f9f32e9de3693b4",
"db_model": False,
},
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"caching_groups": None,
"raw_request": "\n\nPOST Request Sent from LiteLLM:\ncurl -X POST \\\nhttps://openai-gpt-4-test-v-1.openai.azure.com//openai/ \\\n-H 'Authorization: *****' \\\n-d '{'model': 'chatgpt-v-3', 'messages': [{'role': 'system', 'content': 'you are a helpful assistant.\\n'}, {'role': 'user', 'content': 'bom dia'}], 'stream': False, 'max_tokens': 10, 'user': '116544810872468347480', 'extra_body': {}}'\n",
},
)
print("response", response)
await asyncio.sleep(5)
# Get the current date
# Get the current date
current_date = datetime.now().strftime("%Y-%m-%d")
# Modify the object_name to include the date-based folder
object_name = f"{current_date}%2F{response.id}"
print("object_name", object_name)
# Check if object landed on GCS
object_from_gcs = await gcs_logger.download_gcs_object(object_name=object_name)
print("object from gcs=", object_from_gcs)
# convert object_from_gcs from bytes to DICT
parsed_data = json.loads(object_from_gcs)
print("object_from_gcs as dict", parsed_data)
print("type of object_from_gcs", type(parsed_data))
gcs_payload = StandardLoggingPayload(**parsed_data)
print("gcs_payload", gcs_payload)
assert gcs_payload["model"] == "gpt-3.5-turbo"
assert gcs_payload["messages"] == [{"role": "user", "content": "This is a test"}]
assert gcs_payload["response"]["choices"][0]["message"]["content"] == "Hi!"
assert gcs_payload["response_cost"] > 0.0
assert gcs_payload["status"] == "success"
assert (
gcs_payload["metadata"]["user_api_key_hash"]
== "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456"
)
assert gcs_payload["metadata"]["user_api_key_user_id"] == "116544810872468347480"
assert gcs_payload["metadata"]["requester_metadata"] == {"foo": "bar"}
# Delete Object from GCS
print("deleting object from GCS")
await gcs_logger.delete_gcs_object(object_name=object_name)
# clean up
if old_bucket_name is not None:
os.environ["GCS_BUCKET_NAME"] = old_bucket_name
@pytest.mark.skip(reason="This test is flaky on ci/cd")
def test_create_file_e2e():
"""
Asserts 'create_file' is called with the correct arguments
"""
load_vertex_ai_credentials()
test_file_content = b"test audio content"
test_file = ("test.wav", test_file_content, "audio/wav")
from litellm import create_file
response = create_file(
file=test_file,
purpose="user_data",
custom_llm_provider="vertex_ai",
)
print("response", response)
assert response is not None
@pytest.mark.skip(reason="This test is flaky on ci/cd")
def test_create_file_e2e_jsonl():
"""
Asserts 'create_file' is called with the correct arguments
"""
load_vertex_ai_credentials()
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
example_jsonl = [{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-1.5-flash-001", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}},{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-1.5-flash-001", "messages": [{"role": "system", "content": "You are an unhelpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}}]
# Create and write to the file
file_path = "example.jsonl"
with open(file_path, "w") as f:
for item in example_jsonl:
f.write(json.dumps(item) + "\n")
# Verify file content
with open(file_path, "r") as f:
content = f.read()
print("File content:", content)
assert len(content) > 0, "File is empty"
from litellm import create_file
with patch.object(client, "post") as mock_create_file:
try:
response = create_file(
file=open(file_path, "rb"),
purpose="user_data",
custom_llm_provider="vertex_ai",
client=client,
temperature=0.7,
messages=[{"role": "user", "content": "This is a test"}],
max_tokens=10,
user="ishaan-2",
mock_response=litellm.BadRequestError(
model="gpt-3.5-turbo",
message="Error: 400: Bad Request: Invalid API key, please check your API key and try again.",
llm_provider="openai",
),
metadata={
"gcs_log_id": gcs_log_id,
"tags": ["model-anthropic-claude-v2.1", "app-ishaan-prod"],
"user_api_key": "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456",
"user_api_key_alias": None,
"user_api_end_user_max_budget": None,
"litellm_api_version": "0.0.0",
"global_max_parallel_requests": None,
"user_api_key_user_id": "116544810872468347480",
"user_api_key_org_id": None,
"user_api_key_team_id": None,
"user_api_key_team_alias": None,
"user_api_key_metadata": {},
"requester_ip_address": "127.0.0.1",
"spend_logs_metadata": {"hello": "world"},
"headers": {
"content-type": "application/json",
"user-agent": "PostmanRuntime/7.32.3",
"accept": "*/*",
"postman-token": "92300061-eeaa-423b-a420-0b44896ecdc4",
"host": "localhost:4000",
"accept-encoding": "gzip, deflate, br",
"connection": "keep-alive",
"content-length": "163",
},
"endpoint": "http://localhost:4000/chat/completions",
"model_group": "gpt-3.5-turbo",
"model_info": {
"id": "4bad40a1eb6bebd1682800f16f44b9f06c52a6703444c99c7f9f32e9de3693b4",
"db_model": False,
},
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"caching_groups": None,
"raw_request": "\n\nPOST Request Sent from LiteLLM:\ncurl -X POST \\\nhttps://openai-gpt-4-test-v-1.openai.azure.com//openai/ \\\n-H 'Authorization: *****' \\\n-d '{'model': 'chatgpt-v-3', 'messages': [{'role': 'system', 'content': 'you are a helpful assistant.\\n'}, {'role': 'user', 'content': 'bom dia'}], 'stream': False, 'max_tokens': 10, 'user': '116544810872468347480', 'extra_body': {}}'\n",
},
)
except Exception as e:
print("error", e)
except Exception:
pass
mock_create_file.assert_called_once()
await asyncio.sleep(3)
print(f"kwargs: {mock_create_file.call_args.kwargs}")
assert len(captured_payloads) == 1, (
f"Expected 1 GCS upload, got {len(captured_payloads)}"
)
assert mock_create_file.call_args.kwargs["data"] is not None and len(mock_create_file.call_args.kwargs["data"]) > 0
gcs_payload = captured_payloads[0]["logging_payload"]
assert gcs_payload["model"] == "gpt-3.5-turbo"
assert gcs_payload["messages"] == [
{"role": "user", "content": "This is a test"}
]
assert gcs_payload["response_cost"] == 0
assert gcs_payload["status"] == "failure"
assert (
gcs_payload["metadata"]["user_api_key_hash"]
== "a1b2c3d4e5f6789012345678901234567890abcdef1234567890abcdef123456"
)
assert (
gcs_payload["metadata"]["user_api_key_user_id"] == "116544810872468347480"
)
tests/local_testing/test_loadtest_router.py
+2 -2
View File
@@ -39,8 +39,8 @@
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# },
# },
tests/local_testing/test_prompt_injection_detection.py
+2 -2
View File
@@ -108,9 +108,9 @@ async def test_prompt_injection_llm_eval():
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
tests/local_testing/test_router.py
+50 -92
View File
@@ -126,7 +126,9 @@ async def test_router_provider_wildcard_routing():
print("response 3 = ", response3)
response4 = await router.acompletion(
model=os.environ.get("CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"),
model=os.environ.get(
"CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"
),
messages=[{"role": "user", "content": "hello"}],
)
@@ -356,51 +358,6 @@ async def test_router_retries(sync_mode):
print(response.choices[0].message)
@pytest.mark.parametrize(
"mistral_api_base",
[
"os.environ/AZURE_MISTRAL_API_BASE",
"https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com/v1/",
"https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com/v1",
"https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com/",
"https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com",
],
)
@pytest.mark.skip(
reason="Router no longer creates clients, this is delegated to the provider integration."
)
def test_router_azure_ai_studio_init(mistral_api_base):
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "azure/mistral-large-latest",
"api_key": "os.environ/AZURE_MISTRAL_API_KEY",
"api_base": mistral_api_base,
},
"model_info": {"id": 1234},
}
]
)
# model_client = router._get_client(
# deployment={"model_info": {"id": 1234}}, client_type="sync_client", kwargs={}
# )
# url = getattr(model_client, "_base_url")
# uri_reference = str(getattr(url, "_uri_reference"))
# print(f"uri_reference: {uri_reference}")
# assert "/v1/" in uri_reference
# assert uri_reference.count("v1") == 1
response = router.completion(
model="azure/mistral-large-latest",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
)
assert response is not None
def test_exception_raising():
# this tests if the router raises an exception when invalid params are set
# in this test both deployments have bad keys - Keep this test. It validates if the router raises the most recent exception
@@ -409,8 +366,8 @@ def test_exception_raising():
try:
print("testing if router raises an exception")
old_api_key = os.environ["AZURE_API_KEY"]
os.environ["AZURE_API_KEY"] = ""
old_api_key = os.environ["AZURE_AI_API_KEY"]
os.environ["AZURE_AI_API_KEY"] = ""
model_list = [
{
"model_name": "gpt-3.5-turbo", # openai model name
@@ -418,7 +375,7 @@ def test_exception_raising():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -446,16 +403,16 @@ def test_exception_raising():
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hello this request will fail"}],
)
os.environ["AZURE_API_KEY"] = old_api_key
os.environ["AZURE_AI_API_KEY"] = old_api_key
pytest.fail(f"Should have raised an Auth Error")
except openai.AuthenticationError:
print(
"Test Passed: Caught an OPENAI AUTH Error, Good job. This is what we needed!"
)
os.environ["AZURE_API_KEY"] = old_api_key
os.environ["AZURE_AI_API_KEY"] = old_api_key
router.reset()
except Exception as e:
os.environ["AZURE_API_KEY"] = old_api_key
os.environ["AZURE_AI_API_KEY"] = old_api_key
print("Got unexpected exception on router!", e)
@@ -530,7 +487,7 @@ def test_call_one_endpoint():
# this test makes a completion calls azure/gpt-4.1-mini, it should work
try:
print("Testing calling a specific deployment")
old_api_key = os.environ["AZURE_API_KEY"]
old_api_key = os.environ["AZURE_AI_API_KEY"]
model_list = [
{
@@ -539,7 +496,7 @@ def test_call_one_endpoint():
"model": "azure/gpt-4.1-mini",
"api_key": old_api_key,
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -548,8 +505,8 @@ def test_call_one_endpoint():
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "azure/text-embedding-ada-002",
"api_key": os.environ["AZURE_API_KEY"],
"api_base": os.environ["AZURE_API_BASE"],
"api_key": os.environ["AZURE_AI_API_KEY"],
"api_base": os.environ["AZURE_AI_API_BASE"],
},
"tpm": 100000,
"rpm": 10000,
@@ -562,7 +519,7 @@ def test_call_one_endpoint():
set_verbose=True,
num_retries=1,
) # type: ignore
old_api_base = os.environ.pop("AZURE_API_BASE", None)
old_api_base = os.environ.pop("AZURE_AI_API_BASE", None)
async def call_azure_completion():
response = await router.acompletion(
@@ -584,8 +541,8 @@ def test_call_one_endpoint():
asyncio.run(call_azure_completion())
asyncio.run(call_azure_embedding())
os.environ["AZURE_API_BASE"] = old_api_base
os.environ["AZURE_API_KEY"] = old_api_key
os.environ["AZURE_AI_API_BASE"] = old_api_base
os.environ["AZURE_AI_API_KEY"] = old_api_key
except Exception as e:
print(f"FAILED TEST")
pytest.fail(f"Got unexpected exception on router! - {e}")
@@ -594,7 +551,6 @@ def test_call_one_endpoint():
# test_call_one_endpoint()
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_async_router_context_window_fallback(sync_mode):
@@ -708,9 +664,9 @@ def test_router_context_window_check_pre_call_check_in_group_custom_model_info()
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"base_model": "azure/gpt-35-turbo",
"mock_response": "Hello world 1!",
},
@@ -762,9 +718,9 @@ def test_router_context_window_check_pre_call_check():
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"base_model": "azure/gpt-35-turbo",
"mock_response": "Hello world 1!",
},
@@ -816,9 +772,9 @@ def test_router_context_window_check_pre_call_check_out_group():
"model_name": "gpt-3.5-turbo-small", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"base_model": "azure/gpt-35-turbo",
},
},
@@ -896,9 +852,9 @@ def test_router_region_pre_call_check(allowed_model_region):
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"base_model": "azure/gpt-35-turbo",
"region_name": allowed_model_region,
},
@@ -1173,8 +1129,8 @@ def test_azure_embedding_on_router():
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "azure/text-embedding-ada-002",
"api_key": os.environ["AZURE_API_KEY"],
"api_base": os.environ["AZURE_API_BASE"],
"api_key": os.environ["AZURE_AI_API_KEY"],
"api_base": os.environ["AZURE_AI_API_BASE"],
},
"tpm": 100000,
"rpm": 10000,
@@ -1381,8 +1337,8 @@ def test_reading_keys_os_environ():
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": "os.environ/AZURE_API_KEY",
"api_base": "os.environ/AZURE_API_BASE",
"api_key": "os.environ/AZURE_AI_API_KEY",
"api_base": "os.environ/AZURE_AI_API_BASE",
"api_version": "os.environ/AZURE_API_VERSION",
"timeout": "os.environ/AZURE_TIMEOUT",
"stream_timeout": "os.environ/AZURE_STREAM_TIMEOUT",
@@ -1394,11 +1350,11 @@ def test_reading_keys_os_environ():
router = Router(model_list=model_list)
for model in router.model_list:
assert (
model["litellm_params"]["api_key"] == os.environ["AZURE_API_KEY"]
), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_API_KEY']}"
model["litellm_params"]["api_key"] == os.environ["AZURE_AI_API_KEY"]
), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_AI_API_KEY']}"
assert (
model["litellm_params"]["api_base"] == os.environ["AZURE_API_BASE"]
), f"{model['litellm_params']['api_base']} vs {os.environ['AZURE_API_BASE']}"
model["litellm_params"]["api_base"] == os.environ["AZURE_AI_API_BASE"]
), f"{model['litellm_params']['api_base']} vs {os.environ['AZURE_AI_API_BASE']}"
assert (
model["litellm_params"]["api_version"]
== os.environ["AZURE_API_VERSION"]
@@ -1415,8 +1371,8 @@ def test_reading_keys_os_environ():
print("passed testing of reading keys from os.environ")
model_id = model["model_info"]["id"]
async_client: openai.AsyncAzureOpenAI = router.cache.get_cache(f"{model_id}_async_client") # type: ignore
assert async_client.api_key == os.environ["AZURE_API_KEY"]
assert async_client.base_url == os.environ["AZURE_API_BASE"]
assert async_client.api_key == os.environ["AZURE_AI_API_KEY"]
assert async_client.base_url == os.environ["AZURE_AI_API_BASE"]
assert async_client.max_retries == int(
os.environ["AZURE_MAX_RETRIES"]
), f"{async_client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}"
@@ -1428,8 +1384,8 @@ def test_reading_keys_os_environ():
print("\n Testing async streaming client")
stream_async_client: openai.AsyncAzureOpenAI = router.cache.get_cache(f"{model_id}_stream_async_client") # type: ignore
assert stream_async_client.api_key == os.environ["AZURE_API_KEY"]
assert stream_async_client.base_url == os.environ["AZURE_API_BASE"]
assert stream_async_client.api_key == os.environ["AZURE_AI_API_KEY"]
assert stream_async_client.base_url == os.environ["AZURE_AI_API_BASE"]
assert stream_async_client.max_retries == int(
os.environ["AZURE_MAX_RETRIES"]
), f"{stream_async_client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}"
@@ -1440,8 +1396,8 @@ def test_reading_keys_os_environ():
print("\n Testing sync client")
client: openai.AzureOpenAI = router.cache.get_cache(f"{model_id}_client") # type: ignore
assert client.api_key == os.environ["AZURE_API_KEY"]
assert client.base_url == os.environ["AZURE_API_BASE"]
assert client.api_key == os.environ["AZURE_AI_API_KEY"]
assert client.base_url == os.environ["AZURE_AI_API_BASE"]
assert client.max_retries == int(
os.environ["AZURE_MAX_RETRIES"]
), f"{client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}"
@@ -1452,8 +1408,8 @@ def test_reading_keys_os_environ():
print("\n Testing sync stream client")
stream_client: openai.AzureOpenAI = router.cache.get_cache(f"{model_id}_stream_client") # type: ignore
assert stream_client.api_key == os.environ["AZURE_API_KEY"]
assert stream_client.base_url == os.environ["AZURE_API_BASE"]
assert stream_client.api_key == os.environ["AZURE_AI_API_KEY"]
assert stream_client.base_url == os.environ["AZURE_AI_API_BASE"]
assert stream_client.max_retries == int(
os.environ["AZURE_MAX_RETRIES"]
), f"{stream_client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}"
@@ -1503,7 +1459,7 @@ def test_reading_openai_keys_os_environ():
for model in router.model_list:
assert (
model["litellm_params"]["api_key"] == os.environ["OPENAI_API_KEY"]
), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_API_KEY']}"
), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_AI_API_KEY']}"
assert float(model["litellm_params"]["timeout"]) == float(
os.environ["AZURE_TIMEOUT"]
), f"{model['litellm_params']['timeout']} vs {os.environ['AZURE_TIMEOUT']}"
@@ -1574,7 +1530,9 @@ def test_router_anthropic_key_dynamic():
{
"model_name": "anthropic-claude",
"litellm_params": {
"model": os.environ.get("CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"),
"model": os.environ.get(
"CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"
),
"api_key": anthropic_api_key,
},
}
@@ -2273,8 +2231,8 @@ async def test_router_batch_endpoints(provider):
"model_name": "my-custom-name",
"litellm_params": {
"model": "azure/gpt-4o-mini",
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
},
},
]
@@ -2452,8 +2410,8 @@ def test_is_team_specific_model():
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "tpm": 100000,
# "rpm": 100000,
# },
@@ -2462,8 +2420,8 @@ def test_is_team_specific_model():
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "tpm": 500,
# "rpm": 500,
# },
tests/local_testing/test_router_budget_limiter.py
+3 -2
View File
@@ -75,9 +75,9 @@ async def test_provider_budgets_e2e_test():
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"model_info": {"id": "azure-model-id"},
},
@@ -609,6 +609,7 @@ async def test_deployment_budgets_e2e_test_expect_to_fail():
assert "Exceeded budget for deployment" in str(exc_info.value)
@pytest.mark.flaky(retries=6, delay=2)
@pytest.mark.asyncio
async def test_tag_budgets_e2e_test_expect_to_fail():
tests/local_testing/test_router_caching.py
+2 -2
View File
@@ -268,8 +268,8 @@ async def test_acompletion_caching_on_router_caching_groups():
"model_name": "azure-gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
},
"tpm": 100000,
tests/local_testing/test_router_client_init.py
+4 -4
View File
@@ -85,7 +85,7 @@ def test_router_init_azure_service_principal_with_secret_with_environment_variab
To allow for local testing without real credentials, first must mock Azure SDK authentication functions
and environment variables.
"""
monkeypatch.delenv("AZURE_API_KEY", raising=False)
monkeypatch.delenv("AZURE_AI_API_KEY", raising=False)
litellm.enable_azure_ad_token_refresh = True
# mock the token provider function
mocked_func_generating_token = MagicMock(return_value="test_token")
@@ -174,9 +174,9 @@ async def test_audio_speech_router():
{
"model_name": "tts",
"litellm_params": {
"model": "azure/azure-tts",
"api_base": os.getenv("AZURE_SWEDEN_API_BASE"),
"api_key": os.getenv("AZURE_SWEDEN_API_KEY"),
"model": "azure/tts",
"api_base": os.getenv("AZURE_TTS_API_BASE"),
"api_key": os.getenv("AZURE_TTS_API_KEY"),
},
},
]
tests/local_testing/test_router_cooldown_handlers.py
+2 -2
View File
@@ -45,9 +45,9 @@ async def test_cooldown_badrequest_error():
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
}
],
tests/local_testing/test_router_debug_logs.py
+2 -2
View File
@@ -34,9 +34,9 @@ def test_async_fallbacks(caplog):
"model_name": "azure/gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"mock_response": "Hello world",
},
"tpm": 240000,
tests/local_testing/test_router_fallbacks.py
+43 -35
View File
@@ -70,7 +70,7 @@ def test_sync_fallbacks():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -79,9 +79,9 @@ def test_sync_fallbacks():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -92,7 +92,7 @@ def test_sync_fallbacks():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -132,7 +132,9 @@ def test_sync_fallbacks():
response = router.completion(**kwargs)
print(f"response: {response}")
time.sleep(0.05) # allow a delay as success_callbacks are on a separate thread
assert customHandler.previous_models == 3 # 1 init call + 2 retries (fallback not counted as previous)
assert (
customHandler.previous_models == 3
) # 1 init call + 2 retries (fallback not counted as previous)
print("Passed ! Test router_fallbacks: test_sync_fallbacks()")
router.reset()
@@ -153,7 +155,7 @@ async def test_async_fallbacks():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -162,9 +164,9 @@ async def test_async_fallbacks():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -175,7 +177,7 @@ async def test_async_fallbacks():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -220,7 +222,9 @@ async def test_async_fallbacks():
await asyncio.sleep(
0.05
) # allow a delay as success_callbacks are on a separate thread
assert customHandler.previous_models == 3 # 1 init call + 2 retries (fallback not counted as previous)
assert (
customHandler.previous_models == 3
) # 1 init call + 2 retries (fallback not counted as previous)
router.reset()
except litellm.Timeout as e:
pass
@@ -242,7 +246,7 @@ def test_sync_fallbacks_embeddings():
"model": "azure/text-embedding-ada-002",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -292,7 +296,7 @@ async def test_async_fallbacks_embeddings():
"model": "azure/text-embedding-ada-002",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -348,7 +352,7 @@ def test_dynamic_fallbacks_sync():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -357,9 +361,9 @@ def test_dynamic_fallbacks_sync():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -370,7 +374,7 @@ def test_dynamic_fallbacks_sync():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -403,7 +407,9 @@ def test_dynamic_fallbacks_sync():
response = router.completion(**kwargs)
print(f"response: {response}")
time.sleep(0.05) # allow a delay as success_callbacks are on a separate thread
assert customHandler.previous_models >= 3 # 1 init call, retries, 1 fallback (count varies with cooldown timing)
assert (
customHandler.previous_models >= 3
) # 1 init call, retries, 1 fallback (count varies with cooldown timing)
router.reset()
except Exception as e:
pytest.fail(f"An exception occurred - {e}")
@@ -425,7 +431,7 @@ async def test_dynamic_fallbacks_async():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -434,9 +440,9 @@ async def test_dynamic_fallbacks_async():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -447,7 +453,7 @@ async def test_dynamic_fallbacks_async():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -489,7 +495,9 @@ async def test_dynamic_fallbacks_async():
await asyncio.sleep(
0.05
) # allow a delay as success_callbacks are on a separate thread
assert customHandler.previous_models >= 3 # 1 init call, retries, 1 fallback (count varies with cooldown timing)
assert (
customHandler.previous_models >= 3
) # 1 init call, retries, 1 fallback (count varies with cooldown timing)
router.reset()
except Exception as e:
pytest.fail(f"An exception occurred - {e}")
@@ -562,7 +570,7 @@ def test_sync_fallbacks_streaming():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -571,9 +579,9 @@ def test_sync_fallbacks_streaming():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -584,7 +592,7 @@ def test_sync_fallbacks_streaming():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -643,7 +651,7 @@ async def test_async_fallbacks_max_retries_per_request():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -652,9 +660,9 @@ async def test_async_fallbacks_max_retries_per_request():
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -665,7 +673,7 @@ async def test_async_fallbacks_max_retries_per_request():
"model": "azure/chatgpt-functioncalling",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -750,9 +758,9 @@ def test_ausage_based_routing_fallbacks():
def get_azure_params(deployment_name: str):
params = {
"model": f"azure/{deployment_name}",
"api_key": os.environ["AZURE_API_KEY"],
"api_key": os.environ["AZURE_AI_API_KEY"],
"api_version": os.environ["AZURE_API_VERSION"],
"api_base": os.environ["AZURE_API_BASE"],
"api_base": os.environ["AZURE_AI_API_BASE"],
}
return params
@@ -855,7 +863,7 @@ def test_custom_cooldown_times():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 24000000,
},
@@ -863,9 +871,9 @@ def test_custom_cooldown_times():
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 1,
},
tests/local_testing/test_router_init.py
-704
View File
@@ -1,704 +0,0 @@
# # this tests if the router is initialized correctly
# import asyncio
# import os
# import sys
# import time
# import traceback
# import pytest
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# from collections import defaultdict
# from concurrent.futures import ThreadPoolExecutor
# from dotenv import load_dotenv
# import litellm
# from litellm import Router
# load_dotenv()
# # every time we load the router we should have 4 clients:
# # Async
# # Sync
# # Async + Stream
# # Sync + Stream
# def test_init_clients():
# litellm.set_verbose = True
# import logging
# from litellm._logging import verbose_router_logger
# verbose_router_logger.setLevel(logging.DEBUG)
# try:
# print("testing init 4 clients with diff timeouts")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# },
# ]
# router = Router(model_list=model_list, set_verbose=True)
# for elem in router.model_list:
# model_id = elem["model_info"]["id"]
# assert router.cache.get_cache(f"{model_id}_client") is not None
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
# # check if timeout for stream/non stream clients is set correctly
# async_client = router.cache.get_cache(f"{model_id}_async_client")
# stream_async_client = router.cache.get_cache(
# f"{model_id}_stream_async_client"
# )
# assert async_client.timeout == 0.01
# assert stream_async_client.timeout == 0.000_001
# print(vars(async_client))
# print()
# print(async_client._base_url)
# assert (
# async_client._base_url
# == "https://openai-gpt-4-test-v-1.openai.azure.com/openai/"
# )
# assert (
# stream_async_client._base_url
# == "https://openai-gpt-4-test-v-1.openai.azure.com/openai/"
# )
# print("PASSED !")
# except Exception as e:
# traceback.print_exc()
# pytest.fail(f"Error occurred: {e}")
# # test_init_clients()
# def test_init_clients_basic():
# litellm.set_verbose = True
# try:
# print("Test basic client init")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# },
# },
# ]
# router = Router(model_list=model_list)
# for elem in router.model_list:
# model_id = elem["model_info"]["id"]
# assert router.cache.get_cache(f"{model_id}_client") is not None
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
# print("PASSED !")
# # see if we can init clients without timeout or max retries set
# except Exception as e:
# traceback.print_exc()
# pytest.fail(f"Error occurred: {e}")
# # test_init_clients_basic()
# def test_init_clients_basic_azure_cloudflare():
# # init azure + cloudflare
# # init OpenAI gpt-3.5
# # init OpenAI text-embedding
# # init OpenAI comptaible - Mistral/mistral-medium
# # init OpenAI compatible - xinference/bge
# litellm.set_verbose = True
# try:
# print("Test basic client init")
# model_list = [
# {
# "model_name": "azure-cloudflare",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": "https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1",
# },
# },
# {
# "model_name": "gpt-openai",
# "litellm_params": {
# "model": "gpt-3.5-turbo",
# "api_key": os.getenv("OPENAI_API_KEY"),
# },
# },
# {
# "model_name": "text-embedding-ada-002",
# "litellm_params": {
# "model": "text-embedding-ada-002",
# "api_key": os.getenv("OPENAI_API_KEY"),
# },
# },
# {
# "model_name": "mistral",
# "litellm_params": {
# "model": "mistral/mistral-tiny",
# "api_key": os.getenv("MISTRAL_API_KEY"),
# },
# },
# {
# "model_name": "bge-base-en",
# "litellm_params": {
# "model": "xinference/bge-base-en",
# "api_base": "http://127.0.0.1:9997/v1",
# "api_key": os.getenv("OPENAI_API_KEY"),
# },
# },
# ]
# router = Router(model_list=model_list)
# for elem in router.model_list:
# model_id = elem["model_info"]["id"]
# assert router.cache.get_cache(f"{model_id}_client") is not None
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
# print("PASSED !")
# # see if we can init clients without timeout or max retries set
# except Exception as e:
# traceback.print_exc()
# pytest.fail(f"Error occurred: {e}")
# # test_init_clients_basic_azure_cloudflare()
# def test_timeouts_router():
# """
# Test the timeouts of the router with multiple clients. This HASas to raise a timeout error
# """
# import openai
# litellm.set_verbose = True
# try:
# print("testing init 4 clients with diff timeouts")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "timeout": 0.000001,
# "stream_timeout": 0.000_001,
# },
# },
# ]
# router = Router(model_list=model_list, num_retries=0)
# print("PASSED !")
# async def test():
# try:
# await router.acompletion(
# model="gpt-3.5-turbo",
# messages=[
# {"role": "user", "content": "hello, write a 20 pg essay"}
# ],
# )
# except Exception as e:
# raise e
# asyncio.run(test())
# except openai.APITimeoutError as e:
# print(
# "Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e
# )
# print(type(e))
# pass
# except Exception as e:
# pytest.fail(
# f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}"
# )
# # test_timeouts_router()
# def test_stream_timeouts_router():
# """
# Test the stream timeouts router. See if it selected the correct client with stream timeout
# """
# import openai
# litellm.set_verbose = True
# try:
# print("testing init 4 clients with diff timeouts")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "timeout": 200, # regular calls will not timeout, stream calls will
# "stream_timeout": 10,
# },
# },
# ]
# router = Router(model_list=model_list)
# print("PASSED !")
# data = {
# "model": "gpt-3.5-turbo",
# "messages": [{"role": "user", "content": "hello, write a 20 pg essay"}],
# "stream": True,
# }
# selected_client = router._get_client(
# deployment=router.model_list[0],
# kwargs=data,
# client_type=None,
# )
# print("Select client timeout", selected_client.timeout)
# assert selected_client.timeout == 10
# # make actual call
# response = router.completion(**data)
# for chunk in response:
# print(f"chunk: {chunk}")
# except openai.APITimeoutError as e:
# print(
# "Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e
# )
# print(type(e))
# pass
# except Exception as e:
# pytest.fail(
# f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}"
# )
# # test_stream_timeouts_router()
# def test_xinference_embedding():
# # [Test Init Xinference] this tests if we init xinference on the router correctly
# # [Test Exception Mapping] tests that xinference is an openai comptiable provider
# print("Testing init xinference")
# print(
# "this tests if we create an OpenAI client for Xinference, with the correct API BASE"
# )
# model_list = [
# {
# "model_name": "xinference",
# "litellm_params": {
# "model": "xinference/bge-base-en",
# "api_base": "os.environ/XINFERENCE_API_BASE",
# },
# }
# ]
# router = Router(model_list=model_list)
# print(router.model_list)
# print(router.model_list[0])
# assert (
# router.model_list[0]["litellm_params"]["api_base"] == "http://0.0.0.0:9997"
# ) # set in env
# openai_client = router._get_client(
# deployment=router.model_list[0],
# kwargs={"input": ["hello"], "model": "xinference"},
# )
# assert openai_client._base_url == "http://0.0.0.0:9997"
# assert "xinference" in litellm.openai_compatible_providers
# print("passed")
# # test_xinference_embedding()
# def test_router_init_gpt_4_vision_enhancements():
# try:
# # tests base_url set when any base_url with /openai/deployments passed to router
# print("Testing Azure GPT_Vision enhancements")
# model_list = [
# {
# "model_name": "gpt-4-vision-enhancements",
# "litellm_params": {
# "model": "azure/gpt-4-vision",
# "api_key": os.getenv("AZURE_API_KEY"),
# "base_url": "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/",
# "dataSources": [
# {
# "type": "AzureComputerVision",
# "parameters": {
# "endpoint": "os.environ/AZURE_VISION_ENHANCE_ENDPOINT",
# "key": "os.environ/AZURE_VISION_ENHANCE_KEY",
# },
# }
# ],
# },
# }
# ]
# router = Router(model_list=model_list)
# print(router.model_list)
# print(router.model_list[0])
# assert (
# router.model_list[0]["litellm_params"]["base_url"]
# == "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/"
# ) # set in env
# assert (
# router.model_list[0]["litellm_params"]["dataSources"][0]["parameters"][
# "endpoint"
# ]
# == os.environ["AZURE_VISION_ENHANCE_ENDPOINT"]
# )
# assert (
# router.model_list[0]["litellm_params"]["dataSources"][0]["parameters"][
# "key"
# ]
# == os.environ["AZURE_VISION_ENHANCE_KEY"]
# )
# azure_client = router._get_client(
# deployment=router.model_list[0],
# kwargs={"stream": True, "model": "gpt-4-vision-enhancements"},
# client_type="async",
# )
# assert (
# azure_client._base_url
# == "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/"
# )
# print("passed")
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# @pytest.mark.parametrize("sync_mode", [True, False])
# @pytest.mark.asyncio
# async def test_openai_with_organization(sync_mode):
# try:
# print("Testing OpenAI with organization")
# model_list = [
# {
# "model_name": "openai-bad-org",
# "litellm_params": {
# "model": "gpt-3.5-turbo",
# "organization": "org-ikDc4ex8NB",
# },
# },
# {
# "model_name": "openai-good-org",
# "litellm_params": {"model": "gpt-3.5-turbo"},
# },
# ]
# router = Router(model_list=model_list)
# print(router.model_list)
# print(router.model_list[0])
# if sync_mode:
# openai_client = router._get_client(
# deployment=router.model_list[0],
# kwargs={"input": ["hello"], "model": "openai-bad-org"},
# )
# print(vars(openai_client))
# assert openai_client.organization == "org-ikDc4ex8NB"
# # bad org raises error
# try:
# response = router.completion(
# model="openai-bad-org",
# messages=[{"role": "user", "content": "this is a test"}],
# )
# pytest.fail(
# "Request should have failed - This organization does not exist"
# )
# except Exception as e:
# print("Got exception: " + str(e))
# assert "header should match organization for API key" in str(
# e
# ) or "No such organization" in str(e)
# # good org works
# response = router.completion(
# model="openai-good-org",
# messages=[{"role": "user", "content": "this is a test"}],
# max_tokens=5,
# )
# else:
# openai_client = router._get_client(
# deployment=router.model_list[0],
# kwargs={"input": ["hello"], "model": "openai-bad-org"},
# client_type="async",
# )
# print(vars(openai_client))
# assert openai_client.organization == "org-ikDc4ex8NB"
# # bad org raises error
# try:
# response = await router.acompletion(
# model="openai-bad-org",
# messages=[{"role": "user", "content": "this is a test"}],
# )
# pytest.fail(
# "Request should have failed - This organization does not exist"
# )
# except Exception as e:
# print("Got exception: " + str(e))
# assert "header should match organization for API key" in str(
# e
# ) or "No such organization" in str(e)
# # good org works
# response = await router.acompletion(
# model="openai-good-org",
# messages=[{"role": "user", "content": "this is a test"}],
# max_tokens=5,
# )
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# def test_init_clients_azure_command_r_plus():
# # This tests that the router uses the OpenAI client for Azure/Command-R+
# # For azure/command-r-plus we need to use openai.OpenAI because of how the Azure provider requires requests being sent
# litellm.set_verbose = True
# import logging
# from litellm._logging import verbose_router_logger
# verbose_router_logger.setLevel(logging.DEBUG)
# try:
# print("testing init 4 clients with diff timeouts")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/command-r-plus",
# "api_key": os.getenv("AZURE_COHERE_API_KEY"),
# "api_base": os.getenv("AZURE_COHERE_API_BASE"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# },
# ]
# router = Router(model_list=model_list, set_verbose=True)
# for elem in router.model_list:
# model_id = elem["model_info"]["id"]
# async_client = router.cache.get_cache(f"{model_id}_async_client")
# stream_async_client = router.cache.get_cache(
# f"{model_id}_stream_async_client"
# )
# # Assert the Async Clients used are OpenAI clients and not Azure
# # For using Azure/Command-R-Plus and Azure/Mistral the clients NEED to be OpenAI clients used
# # this is weirdness introduced on Azure's side
# assert "openai.AsyncOpenAI" in str(async_client)
# assert "openai.AsyncOpenAI" in str(stream_async_client)
# print("PASSED !")
# except Exception as e:
# traceback.print_exc()
# pytest.fail(f"Error occurred: {e}")
# @pytest.mark.asyncio
# async def test_aaaaatext_completion_with_organization():
# try:
# print("Testing Text OpenAI with organization")
# model_list = [
# {
# "model_name": "openai-bad-org",
# "litellm_params": {
# "model": "text-completion-openai/gpt-3.5-turbo-instruct",
# "api_key": os.getenv("OPENAI_API_KEY", None),
# "organization": "org-ikDc4ex8NB",
# },
# },
# {
# "model_name": "openai-good-org",
# "litellm_params": {
# "model": "text-completion-openai/gpt-3.5-turbo-instruct",
# "api_key": os.getenv("OPENAI_API_KEY", None),
# "organization": os.getenv("OPENAI_ORGANIZATION", None),
# },
# },
# ]
# router = Router(model_list=model_list)
# print(router.model_list)
# print(router.model_list[0])
# openai_client = router._get_client(
# deployment=router.model_list[0],
# kwargs={"input": ["hello"], "model": "openai-bad-org"},
# )
# print(vars(openai_client))
# assert openai_client.organization == "org-ikDc4ex8NB"
# # bad org raises error
# try:
# response = await router.atext_completion(
# model="openai-bad-org",
# prompt="this is a test",
# )
# pytest.fail("Request should have failed - This organization does not exist")
# except Exception as e:
# print("Got exception: " + str(e))
# assert "header should match organization for API key" in str(
# e
# ) or "No such organization" in str(e)
# # good org works
# response = await router.atext_completion(
# model="openai-good-org",
# prompt="this is a test",
# max_tokens=5,
# )
# print("working response: ", response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# def test_init_clients_async_mode():
# litellm.set_verbose = True
# import logging
# from litellm._logging import verbose_router_logger
# from litellm.types.router import RouterGeneralSettings
# verbose_router_logger.setLevel(logging.DEBUG)
# try:
# print("testing init 4 clients with diff timeouts")
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# },
# ]
# router = Router(
# model_list=model_list,
# set_verbose=True,
# router_general_settings=RouterGeneralSettings(async_only_mode=True),
# )
# for elem in router.model_list:
# model_id = elem["model_info"]["id"]
# # sync clients not initialized in async_only_mode=True
# assert router.cache.get_cache(f"{model_id}_client") is None
# assert router.cache.get_cache(f"{model_id}_stream_client") is None
# # only async clients initialized in async_only_mode=True
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# @pytest.mark.parametrize(
# "environment,expected_models",
# [
# ("development", ["gpt-3.5-turbo"]),
# ("production", ["gpt-4", "gpt-3.5-turbo", "gpt-4o"]),
# ],
# )
# def test_init_router_with_supported_environments(environment, expected_models):
# """
# Tests that the correct models are setup on router when LITELLM_ENVIRONMENT is set
# """
# os.environ["LITELLM_ENVIRONMENT"] = environment
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# "model_info": {"supported_environments": ["development", "production"]},
# },
# {
# "model_name": "gpt-4",
# "litellm_params": {
# "model": "openai/gpt-4",
# "api_key": os.getenv("OPENAI_API_KEY"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# "model_info": {"supported_environments": ["production"]},
# },
# {
# "model_name": "gpt-4o",
# "litellm_params": {
# "model": "openai/gpt-4o",
# "api_key": os.getenv("OPENAI_API_KEY"),
# "timeout": 0.01,
# "stream_timeout": 0.000_001,
# "max_retries": 7,
# },
# "model_info": {"supported_environments": ["production"]},
# },
# ]
# router = Router(model_list=model_list, set_verbose=True)
# _model_list = router.get_model_names()
# print("model_list: ", _model_list)
# print("expected_models: ", expected_models)
# assert set(_model_list) == set(expected_models)
# os.environ.pop("LITELLM_ENVIRONMENT")
tests/local_testing/test_router_timeout.py
+2 -2
View File
@@ -31,8 +31,8 @@ def test_router_timeouts():
"model_name": "openai-gpt-4",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": "os.environ/AZURE_API_KEY",
"api_base": "os.environ/AZURE_API_BASE",
"api_key": "os.environ/AZURE_AI_API_KEY",
"api_base": "os.environ/AZURE_AI_API_BASE",
"api_version": "os.environ/AZURE_API_VERSION",
},
"tpm": 80000,
tests/local_testing/test_router_utils.py
+20 -9
View File
@@ -35,7 +35,7 @@ def test_returned_settings():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -99,7 +99,7 @@ def test_update_kwargs_before_fallbacks_unit_test():
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
}
],
@@ -136,7 +136,7 @@ async def test_update_kwargs_before_fallbacks(call_type):
"model": "azure/gpt-4.1-mini",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
}
],
@@ -266,6 +266,7 @@ async def test_call_router_callbacks_on_success():
)
assert increment["increment_value"] == 1
@pytest.mark.serial
@pytest.mark.asyncio
async def test_call_router_callbacks_on_failure():
@@ -486,7 +487,9 @@ def test_router_get_deployment_credentials_with_provider():
)
# Test getting credentials by model_id
credentials = router.get_deployment_credentials_with_provider(model_id="openai-deployment-1")
credentials = router.get_deployment_credentials_with_provider(
model_id="openai-deployment-1"
)
assert credentials is not None
assert credentials["api_key"] == "sk-test-123"
assert credentials["custom_llm_provider"] == "openai"
@@ -499,14 +502,16 @@ def test_router_get_deployment_credentials_with_provider():
assert credentials2["custom_llm_provider"] == "anthropic"
# Test with non-existent model
credentials3 = router.get_deployment_credentials_with_provider(model_id="non-existent")
credentials3 = router.get_deployment_credentials_with_provider(
model_id="non-existent"
)
assert credentials3 is None
def test_router_get_deployment_credentials_with_provider_wildcard():
"""
Test that get_deployment_credentials_with_provider handles wildcard patterns.
When a model like openai/gpt-4o is requested and the config has openai/*,
the method should resolve the wildcard pattern and return credentials.
"""
@@ -533,20 +538,26 @@ def test_router_get_deployment_credentials_with_provider_wildcard():
)
# Test wildcard pattern matching for OpenAI
credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-4o")
credentials = router.get_deployment_credentials_with_provider(
model_id="openai/gpt-4o"
)
assert credentials is not None
assert credentials["api_key"] == "sk-wildcard-123"
assert credentials["custom_llm_provider"] == "openai"
assert credentials["api_base"] == "https://api.openai.com/v1"
# Test wildcard pattern matching for Anthropic
credentials2 = router.get_deployment_credentials_with_provider(model_id="anthropic/claude-3-opus")
credentials2 = router.get_deployment_credentials_with_provider(
model_id="anthropic/claude-3-opus"
)
assert credentials2 is not None
assert credentials2["api_key"] == "sk-ant-wildcard-456"
assert credentials2["custom_llm_provider"] == "anthropic"
# Test with non-matching model
credentials3 = router.get_deployment_credentials_with_provider(model_id="vertex_ai/gemini-pro")
credentials3 = router.get_deployment_credentials_with_provider(
model_id="vertex_ai/gemini-pro"
)
assert credentials3 is None
tests/local_testing/test_simple_shuffle.py
-53
View File
@@ -1,53 +0,0 @@
# What is this?
## unit tests for 'simple-shuffle'
import sys, os, asyncio, time, random
from datetime import datetime
import traceback
from dotenv import load_dotenv
load_dotenv()
import os
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import pytest
from litellm import Router
"""
Test random shuffle
- async
- sync
"""
async def test_simple_shuffle():
model_list = [
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-turbo",
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
"api_base": "https://openai-france-1234.openai.azure.com",
"rpm": 1440,
},
"model_info": {"id": 1},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
"rpm": 6,
},
"model_info": {"id": 2},
},
]
router = Router(
model_list=model_list,
routing_strategy="usage-based-routing-v2",
set_verbose=False,
num_retries=3,
) # type: ignore
tests/local_testing/test_streaming.py
+10 -202
View File
@@ -471,86 +471,6 @@ def test_completion_azure_stream():
# test_completion_azure_stream()
@pytest.mark.skip("Skipping predibase streaming test - ran out of credits")
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_completion_predibase_streaming(sync_mode):
try:
litellm.set_verbose = True
litellm._turn_on_debug()
if sync_mode:
response = completion(
model="predibase/llama-3-8b-instruct",
timeout=5,
tenant_id="c4768f95",
max_tokens=10,
api_base="https://serving.app.predibase.com",
api_key=os.getenv("PREDIBASE_API_KEY"),
messages=[{"role": "user", "content": "What is the meaning of life?"}],
stream=True,
)
complete_response = ""
for idx, init_chunk in enumerate(response):
chunk, finished = streaming_format_tests(idx, init_chunk)
complete_response += chunk
custom_llm_provider = init_chunk._hidden_params["custom_llm_provider"]
print(f"custom_llm_provider: {custom_llm_provider}")
assert custom_llm_provider == "predibase"
if finished:
assert isinstance(
init_chunk.choices[0], litellm.utils.StreamingChoices
)
break
if complete_response.strip() == "":
raise Exception("Empty response received")
else:
response = await litellm.acompletion(
model="predibase/llama-3-8b-instruct",
tenant_id="c4768f95",
timeout=5,
max_tokens=10,
api_base="https://serving.app.predibase.com",
api_key=os.getenv("PREDIBASE_API_KEY"),
messages=[{"role": "user", "content": "What is the meaning of life?"}],
stream=True,
)
# await response
complete_response = ""
idx = 0
async for init_chunk in response:
chunk, finished = streaming_format_tests(idx, init_chunk)
complete_response += chunk
custom_llm_provider = init_chunk._hidden_params["custom_llm_provider"]
print(f"custom_llm_provider: {custom_llm_provider}")
assert custom_llm_provider == "predibase"
idx += 1
if finished:
assert isinstance(
init_chunk.choices[0], litellm.utils.StreamingChoices
)
break
if complete_response.strip() == "":
raise Exception("Empty response received")
print(f"complete_response: {complete_response}")
except litellm.Timeout:
pass
except litellm.InternalServerError:
pass
except litellm.ServiceUnavailableError:
pass
except litellm.APIConnectionError:
pass
except Exception as e:
print("ERROR class", e.__class__)
print("ERROR message", e)
print("ERROR traceback", traceback.format_exc())
pytest.fail(f"Error occurred: {e}")
def test_completion_azure_function_calling_stream():
@@ -937,49 +857,6 @@ def test_completion_mistral_api_mistral_large_function_call_with_streaming():
# test_completion_mistral_api_stream()
def test_completion_deep_infra_stream():
# deep infra,currently includes role in the 2nd chunk
# waiting for them to make a fix on this
litellm.set_verbose = True
try:
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{
"role": "user",
"content": "how does a court case get to the Supreme Court?",
},
]
print("testing deep infra streaming")
response = completion(
model="deepinfra/meta-llama/Llama-2-70b-chat-hf",
messages=messages,
stream=True,
max_tokens=80,
)
complete_response = ""
# Add any assertions here to check the response
has_finish_reason = False
for idx, chunk in enumerate(response):
chunk, finished = streaming_format_tests(idx, chunk)
if finished:
has_finish_reason = True
break
complete_response += chunk
if has_finish_reason == False:
raise Exception("finish reason not set")
if complete_response.strip() == "":
raise Exception("Empty response received")
print(f"completion_response: {complete_response}")
except Exception as e:
if "Model busy, retry later" in str(e):
pass
pytest.fail(f"Error occurred: {e}")
# test_completion_deep_infra_stream()
@pytest.mark.skip()
def test_completion_nlp_cloud_stream():
try:
@@ -1068,7 +945,6 @@ def test_vertex_ai_stream(provider):
load_vertex_ai_credentials()
litellm.set_verbose = True
litellm.vertex_project = "pathrise-convert-1606954137718"
import random
test_models = ["gemini-2.5-flash-lite"]
@@ -1187,6 +1063,7 @@ def test_vertex_ai_stream(provider):
# test_completion_vertexai_stream_bad_key()
@pytest.mark.skip(reason="Replicate extremely flaky.")
@pytest.mark.parametrize("sync_mode", [False, True])
@pytest.mark.asyncio
async def test_completion_replicate_llama3_streaming(sync_mode):
@@ -1655,80 +1532,9 @@ def test_sagemaker_weird_response():
# test_sagemaker_weird_response()
@pytest.mark.skip(reason="Move to being a mock endpoint")
@pytest.mark.asyncio
async def test_sagemaker_streaming_async():
try:
messages = [{"role": "user", "content": "Hey, how's it going?"}]
litellm.set_verbose = True
response = await litellm.acompletion(
model="sagemaker/jumpstart-dft-hf-llm-mistral-7b-ins-20240329-150233",
model_id="huggingface-llm-mistral-7b-instruct-20240329-150233",
messages=messages,
temperature=0.2,
max_tokens=80,
aws_region_name=os.getenv("AWS_REGION_NAME_2"),
aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID_2"),
aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY_2"),
stream=True,
)
# Add any assertions here to check the response
print(response)
complete_response = ""
has_finish_reason = False
# Add any assertions here to check the response
idx = 0
async for chunk in response:
# print
chunk, finished = streaming_format_tests(idx, chunk)
has_finish_reason = finished
complete_response += chunk
if finished:
break
idx += 1
if has_finish_reason is False:
raise Exception("finish reason not set for last chunk")
if complete_response.strip() == "":
raise Exception("Empty response received")
print(f"completion_response: {complete_response}")
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# asyncio.run(test_sagemaker_streaming_async())
@pytest.mark.skip(reason="costly sagemaker deployment. Move to mock implementation")
def test_completion_sagemaker_stream():
try:
response = completion(
model="sagemaker/jumpstart-dft-hf-llm-mistral-7b-ins-20240329-150233",
model_id="huggingface-llm-mistral-7b-instruct-20240329-150233",
messages=messages,
temperature=0.2,
max_tokens=80,
aws_region_name=os.getenv("AWS_REGION_NAME_2"),
aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID_2"),
aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY_2"),
stream=True,
)
complete_response = ""
has_finish_reason = False
# Add any assertions here to check the response
for idx, chunk in enumerate(response):
chunk, finished = streaming_format_tests(idx, chunk)
has_finish_reason = finished
if finished:
break
complete_response += chunk
if has_finish_reason is False:
raise Exception("finish reason not set for last chunk")
if complete_response.strip() == "":
raise Exception("Empty response received")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skip(reason="Account deleted by IBM.")
@pytest.mark.asyncio
async def test_completion_watsonx_stream():
@@ -2725,8 +2531,8 @@ def test_azure_streaming_and_function_calling():
tool_choice="auto",
messages=messages,
stream=True,
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_version="2024-02-15-preview",
)
# Add any assertions here to check the response
@@ -2796,8 +2602,8 @@ async def test_azure_astreaming_and_function_calling():
tool_choice="auto",
messages=messages,
stream=True,
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_version="2024-02-15-preview",
caching=True,
)
@@ -2827,8 +2633,8 @@ async def test_azure_astreaming_and_function_calling():
tool_choice="auto",
messages=messages,
stream=True,
api_base=os.getenv("AZURE_API_BASE"),
api_key=os.getenv("AZURE_API_KEY"),
api_base=os.getenv("AZURE_AI_API_BASE"),
api_key=os.getenv("AZURE_AI_API_KEY"),
api_version="2024-02-15-preview",
caching=True,
)
@@ -3109,7 +2915,9 @@ def test_unit_test_custom_stream_wrapper_repeating_chunk(
print(f"expected_chunk_fail: {expected_chunk_fail}")
if (loop_amount > litellm.REPEATED_STREAMING_CHUNK_LIMIT) and expected_chunk_fail:
with pytest.raises((litellm.InternalServerError, litellm.exceptions.MidStreamFallbackError)):
with pytest.raises(
(litellm.InternalServerError, litellm.exceptions.MidStreamFallbackError)
):
for chunk in response:
continue
else:
tests/local_testing/test_timeout.py
+4 -4
View File
@@ -111,8 +111,8 @@ def test_hanging_request_azure():
"model_name": "azure-gpt",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_base": os.environ["AZURE_API_BASE"],
"api_key": os.environ["AZURE_API_KEY"],
"api_base": os.environ["AZURE_AI_API_BASE"],
"api_key": os.environ["AZURE_AI_API_KEY"],
},
},
{
@@ -175,8 +175,8 @@ def test_hanging_request_openai():
"model_name": "azure-gpt",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_base": os.environ["AZURE_API_BASE"],
"api_key": os.environ["AZURE_API_KEY"],
"api_base": os.environ["AZURE_AI_API_BASE"],
"api_key": os.environ["AZURE_AI_API_KEY"],
},
},
{
tests/local_testing/test_tpm_rpm_routing_v2.py
+5 -11
View File
@@ -39,9 +39,7 @@ from create_mock_standard_logging_payload import create_standard_logging_payload
def test_tpm_rpm_updated():
test_cache = DualCache()
lowest_tpm_logger = LowestTPMLoggingHandler(
router_cache=test_cache
)
lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache)
model_group = "gpt-3.5-turbo"
deployment_id = "1234"
deployment = "azure/gpt-4.1-mini"
@@ -108,9 +106,7 @@ def test_get_available_deployments():
"model_info": {"id": "5678"},
},
]
lowest_tpm_logger = LowestTPMLoggingHandler(
router_cache=test_cache
)
lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache)
model_group = "gpt-3.5-turbo"
## DEPLOYMENT 1 ##
total_tokens = 50
@@ -669,9 +665,7 @@ def test_return_potential_deployments():
"""
test_cache = DualCache()
lowest_tpm_logger = LowestTPMLoggingHandler(
router_cache=test_cache
)
lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache)
args: Dict = {
"healthy_deployments": [
@@ -731,8 +725,8 @@ async def test_tpm_rpm_routing_model_name_checks():
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"mock_response": "Hey, how's it going?",
},
}
tests/local_testing/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json
+1 -1
View File
@@ -11,7 +11,7 @@
"user": "",
"team_id": "",
"organization_id": "",
"metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"guardrail_information\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}",
"metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"guardrail_information\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}",
"cache_key": "Cache OFF",
"spend": 0.00022500000000000002,
"total_tokens": 30,
tests/logging_callback_tests/test_alerting.py
+10 -49
View File
@@ -641,7 +641,7 @@ async def test_outage_alerting_called(
"model_name": model,
"litellm_params": {
"model": model,
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": api_base,
"vertex_location": vertex_location,
"vertex_project": vertex_project,
@@ -749,7 +749,7 @@ async def test_region_outage_alerting_called(
"model_name": model,
"litellm_params": {
"model": model,
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": api_base,
"vertex_location": vertex_location,
"vertex_project": vertex_project,
@@ -760,7 +760,7 @@ async def test_region_outage_alerting_called(
"model_name": model,
"litellm_params": {
"model": model,
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": api_base,
"vertex_location": vertex_location,
"vertex_project": "vertex_project-2",
@@ -788,40 +788,6 @@ async def test_region_outage_alerting_called(
mock_send_alert.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.skip(reason="test only needs to run locally ")
async def test_alerting():
router = litellm.Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": "bad_key",
},
}
],
debug_level="DEBUG",
set_verbose=True,
alerting_config=AlertingConfig(
alerting_threshold=10, # threshold for slow / hanging llm responses (in seconds). Defaults to 300 seconds
webhook_url=os.getenv(
"SLACK_WEBHOOK_URL"
), # webhook you want to send alerts to
),
)
try:
await router.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
)
except Exception:
pass
finally:
await asyncio.sleep(3)
@pytest.mark.asyncio
async def test_langfuse_trace_id():
"""
@@ -868,7 +834,9 @@ async def test_langfuse_trace_id():
returned_trace_id = trace_url.split("/")[-1]
assert returned_trace_id == litellm_logging_obj._get_trace_id(service_name="langfuse")
assert returned_trace_id == litellm_logging_obj._get_trace_id(
service_name="langfuse"
)
@pytest.mark.asyncio
@@ -1007,7 +975,7 @@ async def test_soft_budget_alerts():
# Verify alert message contains correct percentage
alert_message = mock_send_alert.call_args[1]["message"]
print("GOT MESSAGE\n\n", alert_message)
expected_message = (
@@ -1077,10 +1045,10 @@ key_no_max_budget_info = CallInfo(
async def test_soft_budget_alerts_webhook(entity_info):
"""
Tests that soft budget alerts are triggered for different entity types.
Tests:
- Key with max budget
- Team
- Team
- User
- Key without max budget
"""
@@ -1097,7 +1065,7 @@ async def test_soft_budget_alerts_webhook(entity_info):
# Verify the webhook event
call_args = mock_send_alert.call_args[1]
logged_webhook_event: WebhookEvent = call_args["user_info"]
# Validate the webhook event has all expected fields
assert logged_webhook_event.spend == entity_info.spend
assert logged_webhook_event.soft_budget == entity_info.soft_budget
@@ -1106,10 +1074,3 @@ async def test_soft_budget_alerts_webhook(entity_info):
assert logged_webhook_event.user_email == entity_info.user_email
assert logged_webhook_event.key_alias == entity_info.key_alias
assert logged_webhook_event.event_group == entity_info.event_group
tests/logging_callback_tests/test_amazing_s3_logs.py
+28 -94
View File
@@ -74,15 +74,13 @@ async def test_basic_s3_logging(sync_mode, streaming):
s3.delete_object(Bucket="load-testing-oct", Key=key)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"streaming", [(True)]
)
@pytest.mark.parametrize("streaming", [True])
@pytest.mark.flaky(retries=3, delay=1)
async def test_basic_s3_v2_logging(streaming):
from blockbuster import BlockBuster
from litellm.integrations.s3_v2 import S3Logger
s3_v2_logger = S3Logger(s3_flush_interval=1)
litellm.callbacks = [s3_v2_logger]
blockbuster = BlockBuster()
@@ -120,7 +118,7 @@ async def test_basic_s3_v2_logging(streaming):
print(f"all_s3_keys: {all_s3_keys}")
#assert that atlest one key has response.id in it
# assert that atlest one key has response.id in it
assert any(response_id in key for key in all_s3_keys)
s3 = boto3.client("s3")
# delete all objects
@@ -134,22 +132,22 @@ async def test_basic_s3_v2_logging_failure():
"""Test that S3 v2 logger makes httpx PUT request when logging failures"""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.integrations.s3_v2 import S3Logger
# Create S3 logger with short flush interval
s3_v2_logger = S3Logger(s3_flush_interval=1)
# Mock the httpx client to capture the PUT request
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.raise_for_status = MagicMock()
s3_v2_logger.async_httpx_client = AsyncMock()
s3_v2_logger.async_httpx_client.put.return_value = mock_response
# Track the upload method calls
original_upload = s3_v2_logger.async_upload_data_to_s3
upload_called = False
async def mock_upload(batch_logging_element):
nonlocal upload_called
upload_called = True
@@ -157,12 +155,12 @@ async def test_basic_s3_v2_logging_failure():
url = f"https://test-bucket.s3.us-west-2.amazonaws.com/{batch_logging_element.s3_object_key}"
headers = {"Content-Type": "application/json"}
data = '{"model": "gpt-4o-mini"}'
# Make the actual httpx call we want to test
await s3_v2_logger.async_httpx_client.put(url=url, headers=headers, data=data)
s3_v2_logger.async_upload_data_to_s3 = mock_upload
# Configure S3 callback params
litellm.callbacks = [s3_v2_logger]
litellm.s3_callback_params = {
@@ -172,7 +170,7 @@ async def test_basic_s3_v2_logging_failure():
"s3_region_name": "us-west-2",
}
litellm.set_verbose = True
# Trigger a failure by using invalid API key
try:
response = await litellm.acompletion(
@@ -182,33 +180,33 @@ async def test_basic_s3_v2_logging_failure():
)
except Exception as e:
print(f"Expected error: {e}")
# Wait for logger to process the failure
await asyncio.sleep(5)
# Verify that our mock upload was called
assert upload_called, "S3 upload method was not called"
print("✓ S3 upload method was called")
# Verify that httpx PUT was called
s3_v2_logger.async_httpx_client.put.assert_called()
# Get the call arguments to verify the S3 URL
call_args = s3_v2_logger.async_httpx_client.put.call_args
assert call_args is not None
url = call_args[1]['url'] if 'url' in call_args[1] else call_args[0][0]
url = call_args[1]["url"] if "url" in call_args[1] else call_args[0][0]
# Verify the URL contains expected S3 endpoint
assert "test-bucket.s3.us-west-2.amazonaws.com" in url
print(f"✓ S3 PUT request made to: {url}")
# Verify headers include expected content type
headers = call_args[1]['headers']
assert headers['Content-Type'] == 'application/json'
headers = call_args[1]["headers"]
assert headers["Content-Type"] == "application/json"
print("✓ S3 request headers are correct")
# Verify JSON data was included
data = call_args[1]['data']
data = call_args[1]["data"]
assert data is not None
assert '"model": "gpt-4o-mini"' in data
print("✓ S3 request data contains expected log payload")
@@ -411,83 +409,19 @@ async def make_async_calls():
return total_time
@pytest.mark.skip(reason="flaky test on ci/cd")
def test_s3_logging_r2():
# all s3 requests need to be in one test function
# since we are modifying stdout, and pytests runs tests in parallel
# on circle ci - we only test litellm.acompletion()
try:
# redirect stdout to log_file
# litellm.cache = litellm.Cache(
# type="s3", s3_bucket_name="litellm-r2-bucket", s3_region_name="us-west-2"
# )
litellm.set_verbose = True
from litellm._logging import verbose_logger
import logging
verbose_logger.setLevel(level=logging.DEBUG)
litellm.success_callback = ["s3"]
litellm.s3_callback_params = {
"s3_bucket_name": "litellm-r2-bucket",
"s3_aws_secret_access_key": "os.environ/R2_S3_ACCESS_KEY",
"s3_aws_access_key_id": "os.environ/R2_S3_ACCESS_ID",
"s3_endpoint_url": "os.environ/R2_S3_URL",
"s3_region_name": "os.environ/R2_S3_REGION_NAME",
}
print("Testing async s3 logging")
expected_keys = []
import time
curr_time = str(time.time())
async def _test():
return await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": f"This is a test {curr_time}"}],
max_tokens=10,
temperature=0.7,
user="ishaan-2",
)
response = asyncio.run(_test())
print(f"response: {response}")
expected_keys.append(response.id)
import boto3
s3 = boto3.client(
"s3",
endpoint_url=os.getenv("R2_S3_URL"),
region_name=os.getenv("R2_S3_REGION_NAME"),
aws_access_key_id=os.getenv("R2_S3_ACCESS_ID"),
aws_secret_access_key=os.getenv("R2_S3_ACCESS_KEY"),
)
bucket_name = "litellm-r2-bucket"
# List objects in the bucket
response = s3.list_objects(Bucket=bucket_name)
except Exception as e:
pytest.fail(f"An exception occurred - {e}")
finally:
# post, close log file and verify
# Reset stdout to the original value
print("Passed! Testing async s3 logging")
from litellm.integrations.s3_v2 import S3Logger
class TestS3Logger(S3Logger):
def __init__(self, *args, **kwargs):
self.recorded_requests = {}
self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None
super().__init__(*args, **kwargs)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.recorded_requests[response_obj["id"]] = start_time
print("recorded request", self.recorded_requests)
self.logged_standard_logging_payload = kwargs["standard_logging_object"]
return await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
return await super().async_log_success_event(
kwargs, response_obj, start_time, end_time
)
tests/logging_callback_tests/test_azure_blob_storage.py
-45
View File
@@ -1,45 +0,0 @@
import io
import os
import sys
sys.path.insert(0, os.path.abspath("../.."))
import asyncio
import gzip
import json
import logging
import time
from unittest.mock import AsyncMock, patch
import pytest
import litellm
from litellm import completion
from litellm._logging import verbose_logger
from litellm.integrations.datadog.datadog import *
from datetime import datetime, timedelta
from litellm.types.utils import (
StandardLoggingPayload,
StandardLoggingModelInformation,
StandardLoggingMetadata,
StandardLoggingHiddenParams,
)
from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger
verbose_logger.setLevel(logging.DEBUG)
@pytest.mark.asyncio
async def test_azure_blob_storage():
azure_storage_logger = AzureBlobStorageLogger(flush_interval=1)
litellm.callbacks = [azure_storage_logger]
response = await litellm.acompletion(
model="gpt-4o",
messages=[{"role": "user", "content": "Hello, world!"}],
)
print(response)
await asyncio.sleep(3)
pass
tests/logging_callback_tests/test_custom_callback_router.py
+27 -21
View File
@@ -267,7 +267,10 @@ class CompletionCustomHandler(
try:
print("CompletionCustomHandler.async_log_success_event, kwargs: ", kwargs)
self.states.append("async_success")
print("############### CompletionCustomHandler async success, kwargs: ", kwargs)
print(
"############### CompletionCustomHandler async success, kwargs: ",
kwargs,
)
## START TIME
assert isinstance(start_time, datetime)
## END TIME
@@ -396,9 +399,9 @@ async def test_async_chat_azure():
"model_name": "gpt-4.1-nano", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"model_info": {"base_model": "azure/gpt-4.1-mini"},
"tpm": 240000,
@@ -443,7 +446,7 @@ async def test_async_chat_azure():
"model": "azure/gpt-4o-new-test",
"api_key": "my-bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -483,9 +486,9 @@ async def test_async_embedding_azure():
"model_name": "azure-embedding-model", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/text-embedding-ada-002",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -506,7 +509,7 @@ async def test_async_embedding_azure():
"model": "azure/text-embedding-ada-002",
"api_key": "my-bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -549,7 +552,7 @@ async def test_async_chat_azure_with_fallbacks():
"model": "azure/gpt-4.1-mini",
"api_key": "my-bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -608,9 +611,9 @@ async def test_async_completion_azure_caching():
"model_name": "gpt-4.1-nano", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
@@ -664,23 +667,23 @@ async def test_async_completion_azure_caching_streaming():
)
litellm.callbacks = [customHandler_caching]
unique_time = uuid.uuid4()
# Use Router instead of direct litellm.acompletion to get router-specific metadata
model_list = [
{
"model_name": "gpt-4.1-nano",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
},
]
router = Router(model_list=model_list)
response1 = await router.acompletion(
model="gpt-4.1-nano",
messages=[
@@ -725,12 +728,16 @@ async def test_async_embedding_azure_caching():
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
router = Router(model_list=[{
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "openai/text-embedding-ada-002",
},
}])
router = Router(
model_list=[
{
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "openai/text-embedding-ada-002",
},
}
]
)
litellm.callbacks = [customHandler_caching]
unique_time = time.time()
response1 = await router.aembedding(
@@ -818,4 +825,3 @@ async def test_rate_limit_error_callback():
assert "original_model_group" in mock_client.call_args.kwargs
assert mock_client.call_args.kwargs["original_model_group"] == "my-test-gpt"
tests/logging_callback_tests/test_gcs_pub_sub.py
+8 -4
View File
@@ -190,7 +190,9 @@ async def test_async_gcs_pub_sub():
mock_post.return_value.text = "Accepted"
# Initialize the GcsPubSubLogger and set the mock
gcs_pub_sub_logger = GcsPubSubLogger(flush_interval=1)
gcs_pub_sub_logger = GcsPubSubLogger(
project_id="STUBBED_PROJECT_ID", topic_id="STUBBED_TOPIC_ID", flush_interval=1
)
gcs_pub_sub_logger.async_httpx_client.post = mock_post
mock_construct_request_headers = AsyncMock()
@@ -215,7 +217,7 @@ async def test_async_gcs_pub_sub():
print("sent to url", actual_url)
assert (
actual_url
== "https://pubsub.googleapis.com/v1/projects/reliableKeys/topics/litellmDB:publish"
== "https://pubsub.googleapis.com/v1/projects/STUBBED_PROJECT_ID/topics/STUBBED_TOPIC_ID:publish"
)
actual_request = mock_post.call_args[1]["json"]
@@ -245,7 +247,9 @@ async def test_async_gcs_pub_sub_v1():
mock_post.return_value.text = "Accepted"
# Initialize the GcsPubSubLogger and set the mock
gcs_pub_sub_logger = GcsPubSubLogger(flush_interval=1)
gcs_pub_sub_logger = GcsPubSubLogger(
project_id="STUBBED_PROJECT_ID", topic_id="STUBBED_TOPIC_ID", flush_interval=1
)
gcs_pub_sub_logger.async_httpx_client.post = mock_post
mock_construct_request_headers = AsyncMock()
@@ -270,7 +274,7 @@ async def test_async_gcs_pub_sub_v1():
print("sent to url", actual_url)
assert (
actual_url
== "https://pubsub.googleapis.com/v1/projects/reliableKeys/topics/litellmDB:publish"
== "https://pubsub.googleapis.com/v1/projects/STUBBED_PROJECT_ID/topics/STUBBED_TOPIC_ID:publish"
)
actual_request = mock_post.call_args[1]["json"]
tests/ocr_tests/vertex_key.json
+4 -4
View File
@@ -1,13 +1,13 @@
{
"type": "service_account",
"project_id": "pathrise-convert-1606954137718",
"project_id": "litellm-ci-cd",
"private_key_id": "",
"private_key": "",
"client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_id": "109577393201924326488",
"client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com",
"client_id": "116563532503305622785",
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com",
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com",
"universe_domain": "googleapis.com"
}
tests/old_proxy_tests/tests/load_test_q.py
+2 -2
View File
@@ -26,7 +26,7 @@ config = {
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.environ["AZURE_API_KEY"],
"api_key": os.environ["AZURE_AI_API_KEY"],
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"api_version": "2023-07-01-preview",
},
@@ -34,7 +34,7 @@ config = {
]
}
print("STARTING LOAD TEST Q")
print(os.environ["AZURE_API_KEY"])
print(os.environ["AZURE_AI_API_KEY"])
response = requests.post(
url=f"{base_url}/key/generate",

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