diff --git a/.circleci/config.yml b/.circleci/config.yml index d998ee467b..16019fe047 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -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 diff --git a/.github/workflows/scorecard.yml b/.github/workflows/scorecard.yml new file mode 100644 index 0000000000..7cd12bb219 --- /dev/null +++ b/.github/workflows/scorecard.yml @@ -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 diff --git a/docs/my-website/blog/security_townhall_updates/index.md b/docs/my-website/blog/security_townhall_updates/index.md index 6b2260156d..b997de9c18 100644 --- a/docs/my-website/blog/security_townhall_updates/index.md +++ b/docs/my-website/blog/security_townhall_updates/index.md @@ -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: diff --git a/docs/my-website/blog/vanta_compliance_recertification/index.md b/docs/my-website/blog/vanta_compliance_recertification/index.md new file mode 100644 index 0000000000..d05c113967 --- /dev/null +++ b/docs/my-website/blog/vanta_compliance_recertification/index.md @@ -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. diff --git a/docs/my-website/docs/learn/sdk_quickstart.md b/docs/my-website/docs/learn/sdk_quickstart.md index bdf7b63eb5..0fb8c3f02a 100644 --- a/docs/my-website/docs/learn/sdk_quickstart.md +++ b/docs/my-website/docs/learn/sdk_quickstart.md @@ -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 diff --git a/docs/my-website/static/img/blog/vanta_soc2_recertification.png b/docs/my-website/static/img/blog/vanta_soc2_recertification.png new file mode 100644 index 0000000000..c1c5644fbe Binary files /dev/null and b/docs/my-website/static/img/blog/vanta_soc2_recertification.png differ diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py index b6c9104b23..de02a0c4da 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py @@ -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, diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 6e6070f7f2..a387418ef5 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -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( { diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 8de01aac60..294b9d464f 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -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 diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index 91e7e9d555..fe05bf90f8 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -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], diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b6139a2f84..ad888f49e3 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -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. diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index dd8b1b0a69..90b501a9cb 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -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, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2010a8053b..d4f986edd9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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, diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 880ce3fb32..fef7235585 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -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." ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9dd2ab18c8..53ae08aefb 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d1aebe4dce..9e8a1585c4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 853ceb63c9..5cc1e0b4f3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) ) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 3eacc19a6d..8ea93453ed 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -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, diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 0672b03bcd..8e75ffdff6 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -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__( diff --git a/litellm/setup_wizard.py b/litellm/setup_wizard.py index ee5918e127..666915669f 100644 --- a/litellm/setup_wizard.py +++ b/litellm/setup_wizard.py @@ -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))}"' diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 05913e609f..0d1501664b 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -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 diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 9e66965576..80b6190db8 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -747,6 +747,7 @@ class ChatCompletionAssistantMessage(OpenAIChatCompletionAssistantMessage, total thinking_blocks: Optional[ List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] ] + reasoning_items: Optional[List[ChatCompletionReasoningItem]] class ChatCompletionToolMessage(TypedDict): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8b94fbdad6..82557513a8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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] diff --git a/litellm/utils.py b/litellm/utils.py index 0e3792773a..83d1242f3f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2010a8053b..d4f986edd9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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, diff --git a/tests/agent_tests/test_a2a_agent.py b/tests/agent_tests/test_a2a_agent.py index ace1f8c54c..ed7a5ab982 100644 --- a/tests/agent_tests/test_a2a_agent.py +++ b/tests/agent_tests/test_a2a_agent.py @@ -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}") diff --git a/tests/audio_tests/test_audio_speech.py b/tests/audio_tests/test_audio_speech.py index 67e0dbffa6..abf0df22a6 100644 --- a/tests/audio_tests/test_audio_speech.py +++ b/tests/audio_tests/test_audio_speech.py @@ -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 = 'Hello, this is SSML.' - 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) diff --git a/tests/batches_tests/test_fine_tuning_api.py b/tests/batches_tests/test_fine_tuning_api.py index 7e220612ac..cb570ff3c3 100644 --- a/tests/batches_tests/test_fine_tuning_api.py +++ b/tests/batches_tests/test_fine_tuning_api.py @@ -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 diff --git a/tests/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index e1165812e2..0aed224c25 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -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 diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py index ab37a0a84c..76a5778347 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py @@ -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"}, }, diff --git a/tests/image_gen_tests/test_image_edit.png b/tests/image_gen_tests/test_image_edit.png index 3b9d865ce6..6ccc3026c6 100644 Binary files a/tests/image_gen_tests/test_image_edit.png and b/tests/image_gen_tests/test_image_edit.png differ diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py index 393b4cb67a..3504065a10 100644 --- a/tests/image_gen_tests/test_image_edits.py +++ b/tests/image_gen_tests/test_image_edits.py @@ -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 - - diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 3b4abeeb82..b22a18b49b 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -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") diff --git a/tests/image_gen_tests/vertex_key.json b/tests/image_gen_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/image_gen_tests/vertex_key.json +++ b/tests/image_gen_tests/vertex_key.json @@ -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" } diff --git a/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py b/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py index 031502cbec..2686c28cb1 100644 --- a/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py +++ b/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py @@ -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, diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index b048590d51..d014b3f72a 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -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) diff --git a/tests/litellm_utils_tests/test_litellm_overhead.py b/tests/litellm_utils_tests/test_litellm_overhead.py index 006fbea8d4..76a1894327 100644 --- a/tests/litellm_utils_tests/test_litellm_overhead.py +++ b/tests/litellm_utils_tests/test_litellm_overhead.py @@ -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) diff --git a/tests/litellm_utils_tests/test_secret_manager.py b/tests/litellm_utils_tests/test_secret_manager.py index 7569c673ec..a1190193ea 100644 --- a/tests/litellm_utils_tests/test_secret_manager.py +++ b/tests/litellm_utils_tests/test_secret_manager.py @@ -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. diff --git a/tests/litellm_utils_tests/vertex_key.json b/tests/litellm_utils_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/litellm_utils_tests/vertex_key.json +++ b/tests/litellm_utils_tests/vertex_key.json @@ -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" } diff --git a/tests/llm_responses_api_testing/test_azure_responses_api.py b/tests/llm_responses_api_testing/test_azure_responses_api.py index 5e876ce084..fed9e9e11f 100644 --- a/tests/llm_responses_api_testing/test_azure_responses_api.py +++ b/tests/llm_responses_api_testing/test_azure_responses_api.py @@ -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 diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py index 3c1162cd01..fdf8c24ac9 100644 --- a/tests/llm_translation/test_anthropic_completion.py +++ b/tests/llm_translation/test_anthropic_completion.py @@ -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, diff --git a/tests/llm_translation/test_azure_ai.py b/tests/llm_translation/test_azure_ai.py index 972ba34a17..d5abca1a49 100644 --- a/tests/llm_translation/test_azure_ai.py +++ b/tests/llm_translation/test_azure_ai.py @@ -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/ 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 = [] \ No newline at end of file + litellm.callbacks = [] diff --git a/tests/llm_translation/test_azure_o_series.py b/tests/llm_translation/test_azure_o_series.py index 925453e68c..181d3b677d 100644 --- a/tests/llm_translation/test_azure_o_series.py +++ b/tests/llm_translation/test_azure_o_series.py @@ -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 - - - - diff --git a/tests/llm_translation/test_azure_openai.py b/tests/llm_translation/test_azure_openai.py index 1da380b57a..f3dc954020 100644 --- a/tests/llm_translation/test_azure_openai.py +++ b/tests/llm_translation/test_azure_openai.py @@ -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", - ) \ No newline at end of file + ) diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index b71e4e5187..f57c77db54 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -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: diff --git a/tests/llm_translation/test_clarifai_completion.py b/tests/llm_translation/test_clarifai_completion.py deleted file mode 100644 index 5080413f2e..0000000000 --- a/tests/llm_translation/test_clarifai_completion.py +++ /dev/null @@ -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()) diff --git a/tests/llm_translation/test_cloudflare.py b/tests/llm_translation/test_cloudflare.py index 109e5a8632..5d8e3e5990 100644 --- a/tests/llm_translation/test_cloudflare.py +++ b/tests/llm_translation/test_cloudflare.py @@ -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() diff --git a/tests/llm_translation/test_fireworks_ai_translation.py b/tests/llm_translation/test_fireworks_ai_translation.py index 24c0d546e2..1cc6aabdca 100644 --- a/tests/llm_translation/test_fireworks_ai_translation.py +++ b/tests/llm_translation/test_fireworks_ai_translation.py @@ -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"] ) diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index b10a7d699c..1ad71d25a0 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -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", diff --git a/tests/llm_translation/test_rerank.py b/tests/llm_translation/test_rerank.py index edd5981c16..d784677060 100644 --- a/tests/llm_translation/test_rerank.py +++ b/tests/llm_translation/test_rerank.py @@ -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): diff --git a/tests/llm_translation/test_router_llm_translation_tests.py b/tests/llm_translation/test_router_llm_translation_tests.py index 61446ce613..26456ab0a3 100644 --- a/tests/llm_translation/test_router_llm_translation_tests.py +++ b/tests/llm_translation/test_router_llm_translation_tests.py @@ -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}") diff --git a/tests/llm_translation/test_snowflake.py b/tests/llm_translation/test_snowflake.py index 44b51c0e83..6861c2c7ec 100644 --- a/tests/llm_translation/test_snowflake.py +++ b/tests/llm_translation/test_snowflake.py @@ -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 diff --git a/tests/llm_translation/test_watsonx.py b/tests/llm_translation/test_watsonx.py index d6a82b4496..ce02d3aac6 100644 --- a/tests/llm_translation/test_watsonx.py +++ b/tests/llm_translation/test_watsonx.py @@ -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( diff --git a/tests/load_tests/vertex_key.json b/tests/load_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/load_tests/vertex_key.json +++ b/tests/load_tests/vertex_key.json @@ -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" } diff --git a/tests/local_testing/adroit-crow-413218-bc47f303efc9.json b/tests/local_testing/adroit-crow-413218-bc47f303efc9.json deleted file mode 100644 index 7e02c82136..0000000000 --- a/tests/local_testing/adroit-crow-413218-bc47f303efc9.json +++ /dev/null @@ -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" -} diff --git a/tests/local_testing/example_config_yaml/azure_config.yaml b/tests/local_testing/example_config_yaml/azure_config.yaml index 0a015aefde..05ba0c9bf5 100644 --- a/tests/local_testing/example_config_yaml/azure_config.yaml +++ b/tests/local_testing/example_config_yaml/azure_config.yaml @@ -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 diff --git a/tests/local_testing/test_acooldowns_router.py b/tests/local_testing/test_acooldowns_router.py index ff99210298..18dc26bda9 100644 --- a/tests/local_testing/test_acooldowns_router.py +++ b/tests/local_testing/test_acooldowns_router.py @@ -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, diff --git a/tests/local_testing/test_alangfuse.py b/tests/local_testing/test_alangfuse.py index 306c7749f1..7c2ec7e9f6 100644 --- a/tests/local_testing/test_alangfuse.py +++ b/tests/local_testing/test_alangfuse.py @@ -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': , '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 = { diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 6f7c371bdb..a1684e2376 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -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}") - diff --git a/tests/local_testing/test_arize_ai.py b/tests/local_testing/test_arize_ai.py index 3b497d638a..138858cee0 100644 --- a/tests/local_testing/test_arize_ai.py +++ b/tests/local_testing/test_arize_ai.py @@ -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) diff --git a/tests/local_testing/test_assistants.py b/tests/local_testing/test_assistants.py index df0df6b79c..93ebba3c34 100644 --- a/tests/local_testing/test_assistants.py +++ b/tests/local_testing/test_assistants.py @@ -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, diff --git a/tests/local_testing/test_azure_openai.py b/tests/local_testing/test_azure_openai.py index e95c1b6fcc..1b99140b6e 100644 --- a/tests/local_testing/test_azure_openai.py +++ b/tests/local_testing/test_azure_openai.py @@ -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"), diff --git a/tests/local_testing/test_azure_perf.py b/tests/local_testing/test_azure_perf.py index 1e2d5cc4f7..57d56a24a1 100644 --- a/tests/local_testing/test_azure_perf.py +++ b/tests/local_testing/test_azure_perf.py @@ -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"), # }, # } diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py index 01004e4bfa..0b06b7195b 100644 --- a/tests/local_testing/test_caching.py +++ b/tests/local_testing/test_caching.py @@ -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 diff --git a/tests/local_testing/test_caching_ssl.py b/tests/local_testing/test_caching_ssl.py index 523976f123..a8a38e747a 100644 --- a/tests/local_testing/test_caching_ssl.py +++ b/tests/local_testing/test_caching_ssl.py @@ -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, diff --git a/tests/local_testing/test_class.py b/tests/local_testing/test_class.py index e02f59a294..b4b4f85a9d 100644 --- a/tests/local_testing/test_class.py +++ b/tests/local_testing/test_class.py @@ -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"), # # }, # # } # # ], diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index e6f5cd8651..4f56916057 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -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"], diff --git a/tests/local_testing/test_config.py b/tests/local_testing/test_config.py index e74f92ca7a..2c5d04d381 100644 --- a/tests/local_testing/test_config.py +++ b/tests/local_testing/test_config.py @@ -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" - diff --git a/tests/local_testing/test_configs/test_bad_config.yaml b/tests/local_testing/test_configs/test_bad_config.yaml index 4a70886a93..4bc4c7cc54 100644 --- a/tests/local_testing/test_configs/test_bad_config.yaml +++ b/tests/local_testing/test_configs/test_bad_config.yaml @@ -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 \ No newline at end of file diff --git a/tests/local_testing/test_configs/test_cloudflare_azure_with_cache_config.yaml b/tests/local_testing/test_configs/test_cloudflare_azure_with_cache_config.yaml index 9902835618..24240008fe 100644 --- a/tests/local_testing/test_configs/test_cloudflare_azure_with_cache_config.yaml +++ b/tests/local_testing/test_configs/test_cloudflare_azure_with_cache_config.yaml @@ -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: diff --git a/tests/local_testing/test_configs/test_config_no_auth.yaml b/tests/local_testing/test_configs/test_config_no_auth.yaml index cdc447a5ee..f489621704 100644 --- a/tests/local_testing/test_configs/test_config_no_auth.yaml +++ b/tests/local_testing/test_configs/test_config_no_auth.yaml @@ -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: diff --git a/tests/local_testing/test_configs/test_custom_logger.yaml b/tests/local_testing/test_configs/test_custom_logger.yaml index 22bbfe42be..464d66bd78 100644 --- a/tests/local_testing/test_configs/test_custom_logger.yaml +++ b/tests/local_testing/test_configs/test_custom_logger.yaml @@ -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 diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index 48d535fe45..e6776e3329 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -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") diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index 2c950d7906..e02d9e2117 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -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" - - - diff --git a/tests/local_testing/test_gcs_bucket.py b/tests/local_testing/test_gcs_bucket.py index e2f51eb76c..9d72ff873a 100644 --- a/tests/local_testing/test_gcs_bucket.py +++ b/tests/local_testing/test_gcs_bucket.py @@ -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 \ No newline at end of file + 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" + ) diff --git a/tests/local_testing/test_loadtest_router.py b/tests/local_testing/test_loadtest_router.py index 3f6e4af4fb..3d1062f0d2 100644 --- a/tests/local_testing/test_loadtest_router.py +++ b/tests/local_testing/test_loadtest_router.py @@ -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"), # }, # }, diff --git a/tests/local_testing/test_prompt_injection_detection.py b/tests/local_testing/test_prompt_injection_detection.py index c4cc4cde32..b1a9aff158 100644 --- a/tests/local_testing/test_prompt_injection_detection.py +++ b/tests/local_testing/test_prompt_injection_detection.py @@ -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, diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index e68d271a11..f7885fb8a0 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -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, # }, diff --git a/tests/local_testing/test_router_budget_limiter.py b/tests/local_testing/test_router_budget_limiter.py index 05ce7c1f53..1a36e9de8f 100644 --- a/tests/local_testing/test_router_budget_limiter.py +++ b/tests/local_testing/test_router_budget_limiter.py @@ -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(): diff --git a/tests/local_testing/test_router_caching.py b/tests/local_testing/test_router_caching.py index 6fc220bf72..cb223b661b 100644 --- a/tests/local_testing/test_router_caching.py +++ b/tests/local_testing/test_router_caching.py @@ -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, diff --git a/tests/local_testing/test_router_client_init.py b/tests/local_testing/test_router_client_init.py index f2541601c3..6cf8b4e59b 100644 --- a/tests/local_testing/test_router_client_init.py +++ b/tests/local_testing/test_router_client_init.py @@ -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"), }, }, ] diff --git a/tests/local_testing/test_router_cooldown_handlers.py b/tests/local_testing/test_router_cooldown_handlers.py index 7be8289abf..fdc89fc04e 100644 --- a/tests/local_testing/test_router_cooldown_handlers.py +++ b/tests/local_testing/test_router_cooldown_handlers.py @@ -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"), }, } ], diff --git a/tests/local_testing/test_router_debug_logs.py b/tests/local_testing/test_router_debug_logs.py index 1004e7747e..0b4771b526 100644 --- a/tests/local_testing/test_router_debug_logs.py +++ b/tests/local_testing/test_router_debug_logs.py @@ -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, diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index c586fa8c93..383ad10457 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -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, }, diff --git a/tests/local_testing/test_router_init.py b/tests/local_testing/test_router_init.py deleted file mode 100644 index e232ff105d..0000000000 --- a/tests/local_testing/test_router_init.py +++ /dev/null @@ -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") diff --git a/tests/local_testing/test_router_timeout.py b/tests/local_testing/test_router_timeout.py index 1d09f1f1e0..943a5413d6 100644 --- a/tests/local_testing/test_router_timeout.py +++ b/tests/local_testing/test_router_timeout.py @@ -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, diff --git a/tests/local_testing/test_router_utils.py b/tests/local_testing/test_router_utils.py index 4f7f53cef0..f2fd2fdf55 100644 --- a/tests/local_testing/test_router_utils.py +++ b/tests/local_testing/test_router_utils.py @@ -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 diff --git a/tests/local_testing/test_simple_shuffle.py b/tests/local_testing/test_simple_shuffle.py deleted file mode 100644 index 8837e91126..0000000000 --- a/tests/local_testing/test_simple_shuffle.py +++ /dev/null @@ -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 diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index 56f0e5fe82..3aed069960 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -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: diff --git a/tests/local_testing/test_timeout.py b/tests/local_testing/test_timeout.py index 4128a595d7..e2ec9ba9a0 100644 --- a/tests/local_testing/test_timeout.py +++ b/tests/local_testing/test_timeout.py @@ -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"], }, }, { diff --git a/tests/local_testing/test_tpm_rpm_routing_v2.py b/tests/local_testing/test_tpm_rpm_routing_v2.py index a3218bc987..c7449ef6e2 100644 --- a/tests/local_testing/test_tpm_rpm_routing_v2.py +++ b/tests/local_testing/test_tpm_rpm_routing_v2.py @@ -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?", }, } diff --git a/tests/local_testing/vertex_key.json b/tests/local_testing/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/local_testing/vertex_key.json +++ b/tests/local_testing/vertex_key.json @@ -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" } diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index 7dcbd2467f..a4c50d3c57 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -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, diff --git a/tests/logging_callback_tests/test_alerting.py b/tests/logging_callback_tests/test_alerting.py index f77cbcb4bb..86588bbd14 100644 --- a/tests/logging_callback_tests/test_alerting.py +++ b/tests/logging_callback_tests/test_alerting.py @@ -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 - - - - - - - \ No newline at end of file diff --git a/tests/logging_callback_tests/test_amazing_s3_logs.py b/tests/logging_callback_tests/test_amazing_s3_logs.py index f074026926..c1e6884278 100644 --- a/tests/logging_callback_tests/test_amazing_s3_logs.py +++ b/tests/logging_callback_tests/test_amazing_s3_logs.py @@ -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 + ) diff --git a/tests/logging_callback_tests/test_azure_blob_storage.py b/tests/logging_callback_tests/test_azure_blob_storage.py deleted file mode 100644 index a90f253cc9..0000000000 --- a/tests/logging_callback_tests/test_azure_blob_storage.py +++ /dev/null @@ -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 diff --git a/tests/logging_callback_tests/test_custom_callback_router.py b/tests/logging_callback_tests/test_custom_callback_router.py index f6c7f2fa02..63d8b14f48 100644 --- a/tests/logging_callback_tests/test_custom_callback_router.py +++ b/tests/logging_callback_tests/test_custom_callback_router.py @@ -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" - diff --git a/tests/logging_callback_tests/test_gcs_pub_sub.py b/tests/logging_callback_tests/test_gcs_pub_sub.py index 540fb59ab0..aa846e34f6 100644 --- a/tests/logging_callback_tests/test_gcs_pub_sub.py +++ b/tests/logging_callback_tests/test_gcs_pub_sub.py @@ -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"] diff --git a/tests/ocr_tests/vertex_key.json b/tests/ocr_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/ocr_tests/vertex_key.json +++ b/tests/ocr_tests/vertex_key.json @@ -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" } diff --git a/tests/old_proxy_tests/tests/load_test_q.py b/tests/old_proxy_tests/tests/load_test_q.py index a8f2c0a322..89137c306a 100644 --- a/tests/old_proxy_tests/tests/load_test_q.py +++ b/tests/old_proxy_tests/tests/load_test_q.py @@ -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", diff --git a/tests/old_proxy_tests/tests/test_vtx_sdk_embedding.py b/tests/old_proxy_tests/tests/test_vtx_sdk_embedding.py index ff7f835a8d..a71718a204 100644 --- a/tests/old_proxy_tests/tests/test_vtx_sdk_embedding.py +++ b/tests/old_proxy_tests/tests/test_vtx_sdk_embedding.py @@ -37,7 +37,7 @@ class CredentialsWrapper(Credentials): credentials = CredentialsWrapper(token=LITELLM_PROXY_API_KEY) vertexai.init( - project="pathrise-convert-1606954137718", + project="litellm-ci-cd", location="us-central1", api_endpoint=LITELLM_PROXY_BASE, credentials=credentials, diff --git a/tests/pass_through_tests/vertex_key.json b/tests/pass_through_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/pass_through_tests/vertex_key.json +++ b/tests/pass_through_tests/vertex_key.json @@ -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" } diff --git a/tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py b/tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py index da46016b35..b2470bf6b6 100644 --- a/tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py +++ b/tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py @@ -30,7 +30,7 @@ class TestAzureAnthropicStructuredOutput(BaseAnthropicMessagesStructuredOutputTe return "azure_ai/claude-opus-4-5" def get_api_base(self) -> Optional[str]: - return "https://krish-mh44t553-eastus2.services.ai.azure.com/" + return "https://krris-mnb3t0vd-swedencentral.services.ai.azure.com" def get_api_key(self) -> Optional[str]: - return os.environ.get("AZURE_ANTHROPIC_API_KEY") \ No newline at end of file + return os.environ.get("AZURE_ANTHROPIC_API_KEY") diff --git a/tests/proxy_unit_tests/adroit-crow-413218-bc47f303efc9.json b/tests/proxy_unit_tests/adroit-crow-413218-bc47f303efc9.json deleted file mode 100644 index 7e02c82136..0000000000 --- a/tests/proxy_unit_tests/adroit-crow-413218-bc47f303efc9.json +++ /dev/null @@ -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" -} diff --git a/tests/proxy_unit_tests/example_config_yaml/azure_config.yaml b/tests/proxy_unit_tests/example_config_yaml/azure_config.yaml index 0a015aefde..05ba0c9bf5 100644 --- a/tests/proxy_unit_tests/example_config_yaml/azure_config.yaml +++ b/tests/proxy_unit_tests/example_config_yaml/azure_config.yaml @@ -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 diff --git a/tests/proxy_unit_tests/test_configs/test_bad_config.yaml b/tests/proxy_unit_tests/test_configs/test_bad_config.yaml index 4a70886a93..4bc4c7cc54 100644 --- a/tests/proxy_unit_tests/test_configs/test_bad_config.yaml +++ b/tests/proxy_unit_tests/test_configs/test_bad_config.yaml @@ -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 \ No newline at end of file diff --git a/tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml b/tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml index 9902835618..24240008fe 100644 --- a/tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml +++ b/tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml @@ -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: diff --git a/tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml b/tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml index cdc447a5ee..f489621704 100644 --- a/tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml +++ b/tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml @@ -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: diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index b67dd2792f..66e5b3839b 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -135,6 +135,81 @@ async def test_jwt_to_virtual_key_mapping_no_mapping(): prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called() +# ────────────────────────────────────────────── +# Tests: OIDC / JWT routing in user_api_key_auth +# ────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_virtual_key_mapping_oidc_enabled_jwt_token_uses_auth_jwt(): + """ + Regression test for the is_jwt routing fix in user_api_key_auth.py. + + When oidc_userinfo_enabled=True and virtual_key_claim_field is set, but + the token is a well-formed JWT (3-part header.payload.sig), the virtual-key + claim lookup must call auth_jwt — not get_oidc_userinfo. + """ + # Three-part token: is_jwt() returns True + api_key = "eyJhbGciOiJSUzI1NiJ9.eyJlbWFpbCI6InVzZXJAZXhhbXBsZS5jb20ifQ.sig" + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + oidc_userinfo_enabled=True, + virtual_key_claim_field="email", + ) + + # Confirm our fixture token is treated as a JWT + assert jwt_handler.is_jwt(token=api_key) is True + + auth_jwt_mock = AsyncMock(return_value={"email": "user@example.com", "sub": "123"}) + oidc_userinfo_mock = AsyncMock(return_value={"email": "user@example.com"}) + + # Simulate the routing condition from user_api_key_auth.py + if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt( + token=api_key + ): + jwt_claims = await oidc_userinfo_mock(token=api_key) + else: + jwt_claims = await auth_jwt_mock(token=api_key) + + auth_jwt_mock.assert_called_once_with(token=api_key) + oidc_userinfo_mock.assert_not_called() + assert jwt_claims["email"] == "user@example.com" + + +@pytest.mark.asyncio +async def test_virtual_key_mapping_oidc_enabled_opaque_token_uses_oidc_userinfo(): + """ + Complement of the test above: when oidc_userinfo_enabled=True and the token + is an opaque access token (not a JWT), the virtual-key claim lookup must + call get_oidc_userinfo — not auth_jwt. + """ + # Opaque token: no dots → is_jwt() returns False + api_key = "some_opaque_access_token_with_no_dots" + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + oidc_userinfo_enabled=True, + virtual_key_claim_field="email", + ) + + assert jwt_handler.is_jwt(token=api_key) is False + + auth_jwt_mock = AsyncMock(return_value={"email": "user@example.com"}) + oidc_userinfo_mock = AsyncMock(return_value={"email": "user@example.com", "sub": "123"}) + + if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt( + token=api_key + ): + jwt_claims = await oidc_userinfo_mock(token=api_key) + else: + jwt_claims = await auth_jwt_mock(token=api_key) + + oidc_userinfo_mock.assert_called_once_with(token=api_key) + auth_jwt_mock.assert_not_called() + assert jwt_claims["sub"] == "123" + + # ────────────────────────────────────────────── # Tests: _to_response redacts hashed token # ────────────────────────────────────────────── diff --git a/tests/proxy_unit_tests/test_proxy_pass_user_config.py b/tests/proxy_unit_tests/test_proxy_pass_user_config.py index 4828014e33..6beb86eca7 100644 --- a/tests/proxy_unit_tests/test_proxy_pass_user_config.py +++ b/tests/proxy_unit_tests/test_proxy_pass_user_config.py @@ -54,8 +54,9 @@ def client_no_auth(): @pytest.mark.skipif( - os.environ.get("AZURE_API_KEY") is None or os.environ.get("OPENAI_API_KEY") is None, - reason="AZURE_API_KEY or OPENAI_API_KEY not set - skipping integration test" + os.environ.get("AZURE_AI_API_KEY") is None + or os.environ.get("OPENAI_API_KEY") is None, + reason="AZURE_AI_API_KEY or OPENAI_API_KEY not set - skipping integration test", ) def test_chat_completion(client_no_auth): global headers @@ -69,9 +70,9 @@ def test_chat_completion(client_no_auth): model_name="user-azure-instance", litellm_params=CompletionRequest( 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"), timeout=10, ), tpm=240000, diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 61a2f3055a..59b9297d1b 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -119,7 +119,7 @@ def fake_env_vars(monkeypatch): # Set some fake environment variables monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key") monkeypatch.setenv("OPENAI_API_BASE", "http://fake-openai-api-base") - monkeypatch.setenv("AZURE_API_BASE", "http://fake-azure-api-base") + monkeypatch.setenv("AZURE_AI_API_BASE", "http://fake-azure-api-base") monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake_azure_openai_api_key") monkeypatch.setenv("AZURE_SWEDEN_API_BASE", "http://fake-azure-sweden-api-base") monkeypatch.setenv("REDIS_HOST", "localhost") @@ -178,7 +178,7 @@ def test_chat_completion(mock_acompletion, client_no_auth): def test_chat_completion_malformed_messages_returns_400(client_no_auth): """ Test that malformed messages (strings instead of dicts) return 400 instead of 500. - + This test verifies that when a client sends messages as raw strings instead of {role, content} objects, LiteLLM returns a 400 invalid_request_error instead of a 500 Internal Server Error. @@ -188,33 +188,41 @@ def test_chat_completion_malformed_messages_returns_400(client_no_auth): # Test data with malformed messages (string instead of dict) test_data = { "model": "gpt-3.5-turbo", - "messages": ["hi how are you"], # Invalid: should be [{"role": "user", "content": "hi how are you"}] + "messages": [ + "hi how are you" + ], # Invalid: should be [{"role": "user", "content": "hi how are you"}] } print("testing proxy server with malformed messages") - response = client_no_auth.post("/v1/chat/completions", json=test_data, headers=headers) - + response = client_no_auth.post( + "/v1/chat/completions", json=test_data, headers=headers + ) + print(f"response status: {response.status_code}") print(f"response text: {response.text}") - + # Should return 400, not 500 - assert response.status_code == 400, f"Expected 400, got {response.status_code}. Response: {response.text}" - + assert ( + response.status_code == 400 + ), f"Expected 400, got {response.status_code}. Response: {response.text}" + # Verify error format result = response.json() assert "error" in result, "Response should contain 'error' key" error = result["error"] - + # Verify error type and message - assert error.get("type") == "invalid_request_error" or error.get("type") is None, \ - f"Expected invalid_request_error or None, got {error.get('type')}" - assert error.get("code") == "400" or error.get("code") == 400, \ - f"Expected code 400, got {error.get('code')}" - + assert ( + error.get("type") == "invalid_request_error" or error.get("type") is None + ), f"Expected invalid_request_error or None, got {error.get('type')}" + assert ( + error.get("code") == "400" or error.get("code") == 400 + ), f"Expected code 400, got {error.get('code')}" + # Error message should indicate invalid request format error_message = error.get("message", "") assert len(error_message) > 0, "Error message should not be empty" - + except Exception as e: pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") @@ -342,7 +350,7 @@ def test_chat_completion_forward_llm_provider_auth_headers( """ Test that LLM provider auth headers (x-api-key, x-goog-api-key) are forwarded when forward_llm_provider_auth_headers=True. - + This allows clients to send their own LLM provider API keys through the proxy. """ try: @@ -351,7 +359,7 @@ def test_chat_completion_forward_llm_provider_auth_headers( gs["forward_client_headers_to_llm_api"] = True gs["forward_llm_provider_auth_headers"] = forward_llm_auth_headers setattr(litellm.proxy.proxy_server, "general_settings", gs) - + # Test data test_data = { "model": "gpt-3.5-turbo", @@ -360,7 +368,7 @@ def test_chat_completion_forward_llm_provider_auth_headers( ], "max_tokens": 10, } - + # Headers including LLM provider auth request_headers = { "Authorization": "Bearer sk-proxy-auth-123", # Proxy auth (should be stripped) @@ -368,17 +376,17 @@ def test_chat_completion_forward_llm_provider_auth_headers( "x-goog-api-key": "google-api-key-123", # Google API key "X-Custom-Header": "custom-value", # Custom header (should be forwarded) } - + # Make request response = client_no_auth.post( "/v1/chat/completions", json=test_data, headers=request_headers ) - + assert response.status_code == 200 - + # Check forwarded headers forwarded_headers = mock_acompletion.call_args.kwargs.get("headers", {}) - + if forward_llm_auth_headers: # LLM provider auth headers should be forwarded assert "x-api-key" in forwarded_headers @@ -389,19 +397,23 @@ def test_chat_completion_forward_llm_provider_auth_headers( # LLM provider auth headers should be stripped assert "x-api-key" not in forwarded_headers assert "x-goog-api-key" not in forwarded_headers - + # Custom headers should always be forwarded (when forward_client_headers_to_llm_api=True) assert "x-custom-header" in forwarded_headers assert forwarded_headers["x-custom-header"] == "custom-value" - + # Proxy Authorization should never be forwarded assert "authorization" not in forwarded_headers - - print(f"✓ Test passed with forward_llm_provider_auth_headers={forward_llm_auth_headers}") + + print( + f"✓ Test passed with forward_llm_provider_auth_headers={forward_llm_auth_headers}" + ) print(f" Forwarded headers: {list(forwarded_headers.keys())}") - + except Exception as e: - pytest.fail(f"Test failed with forward_llm_auth_headers={forward_llm_auth_headers}: {str(e)}") + pytest.fail( + f"Test failed with forward_llm_auth_headers={forward_llm_auth_headers}: {str(e)}" + ) finally: # Clean up gs = getattr(litellm.proxy.proxy_server, "general_settings") @@ -2406,9 +2418,7 @@ async def test_run_background_health_check_reflects_llm_model_list(monkeypatch): test_model_list_2 = [{"model_name": "model-b"}] called_model_lists = [] - async def fake_perform_health_check( - model_list, details, max_concurrency=None - ): + async def fake_perform_health_check(model_list, details, max_concurrency=None): called_model_lists.append(copy.deepcopy(model_list)) return (["healthy"], ["unhealthy"]) @@ -2452,13 +2462,14 @@ async def test_background_health_check_skip_disabled_models(monkeypatch): test_model_list = [ {"model_name": "model-a"}, - {"model_name": "model-b", "model_info": {"disable_background_health_check": True}}, + { + "model_name": "model-b", + "model_info": {"disable_background_health_check": True}, + }, ] called_model_lists = [] - async def fake_perform_health_check( - model_list, details, max_concurrency=None - ): + async def fake_perform_health_check(model_list, details, max_concurrency=None): called_model_lists.append(copy.deepcopy(model_list)) return (["healthy"], []) @@ -2500,15 +2511,15 @@ def test_get_timeout_from_request(): @pytest.mark.parametrize( "ui_exists, ui_has_content", [ - (True, True), # UI path exists and has content + (True, True), # UI path exists and has content (True, False), # UI path exists but is empty - (False, False), # UI path doesn't exist + (False, False), # UI path doesn't exist ], ) def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content): """ Test the non-root Docker UI path detection logic. - + Tests that when LITELLM_NON_ROOT is set to "true": - If UI path exists and has content, it should be used - If UI path doesn't exist or is empty, proper error logging occurs @@ -2516,44 +2527,54 @@ def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content import tempfile import shutil from unittest.mock import MagicMock - + # Create a temporary directory to act as /tmp/litellm_ui test_ui_path = tmp_path / "litellm_ui" - + if ui_exists: test_ui_path.mkdir(parents=True, exist_ok=True) if ui_has_content: # Create some dummy files to simulate built UI (test_ui_path / "index.html").write_text("") (test_ui_path / "app.js").write_text("console.log('test');") - + # Mock the environment variable and os.path operations monkeypatch.setenv("LITELLM_NON_ROOT", "true") - + # Create a mock logger to capture log messages mock_logger = MagicMock() - + # We need to reimport or reload the relevant code section # Since this is module-level code, we'll test the logic directly ui_path = None non_root_ui_path = str(test_ui_path) - + # Simulate the logic from proxy_server.py lines 909-920 if os.getenv("LITELLM_NON_ROOT", "").lower() == "true": if os.path.exists(non_root_ui_path) and os.listdir(non_root_ui_path): - mock_logger.info(f"Using pre-built UI for non-root Docker: {non_root_ui_path}") - mock_logger.info(f"UI files found: {len(os.listdir(non_root_ui_path))} items") + mock_logger.info( + f"Using pre-built UI for non-root Docker: {non_root_ui_path}" + ) + mock_logger.info( + f"UI files found: {len(os.listdir(non_root_ui_path))} items" + ) ui_path = non_root_ui_path else: - mock_logger.error(f"UI not found at {non_root_ui_path}. UI will not be available.") - mock_logger.error(f"Path exists: {os.path.exists(non_root_ui_path)}, Has content: {os.path.exists(non_root_ui_path) and bool(os.listdir(non_root_ui_path))}") - + mock_logger.error( + f"UI not found at {non_root_ui_path}. UI will not be available." + ) + mock_logger.error( + f"Path exists: {os.path.exists(non_root_ui_path)}, Has content: {os.path.exists(non_root_ui_path) and bool(os.listdir(non_root_ui_path))}" + ) + # Verify behavior based on test parameters if ui_exists and ui_has_content: # UI should be found and used assert ui_path == non_root_ui_path assert mock_logger.info.call_count == 2 - mock_logger.info.assert_any_call(f"Using pre-built UI for non-root Docker: {non_root_ui_path}") + mock_logger.info.assert_any_call( + f"Using pre-built UI for non-root Docker: {non_root_ui_path}" + ) # Verify the second info call mentions the number of items info_calls = [call[0][0] for call in mock_logger.info.call_args_list] assert any("UI files found:" in call and "items" in call for call in info_calls) @@ -2562,7 +2583,9 @@ def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content # UI should not be found, error should be logged assert ui_path is None assert mock_logger.error.call_count == 2 - mock_logger.error.assert_any_call(f"UI not found at {non_root_ui_path}. UI will not be available.") + mock_logger.error.assert_any_call( + f"UI not found at {non_root_ui_path}. UI will not be available." + ) # Verify the second error call has path existence info error_calls = [call[0][0] for call in mock_logger.error.call_args_list] assert any("Path exists:" in call for call in error_calls) @@ -2574,17 +2597,17 @@ async def test_get_config_callbacks_with_all_types(client_no_auth): """ Test that /get/config/callbacks returns all three callback types: - success_callback with type="success" - - failure_callback with type="failure" + - failure_callback with type="failure" - callbacks (success_and_failure) with type="success_and_failure" """ from litellm.proxy.proxy_server import ProxyConfig - + # Create a mock config with all three callback types mock_config_data = { "litellm_settings": { "success_callback": ["langfuse", "braintrust"], "failure_callback": ["sentry"], - "callbacks": ["otel", "langsmith"] + "callbacks": ["otel", "langsmith"], }, "environment_variables": { "LANGFUSE_PUBLIC_KEY": "test-public-key", @@ -2595,51 +2618,53 @@ async def test_get_config_callbacks_with_all_types(client_no_auth): "OTEL_ENDPOINT": "http://localhost:4317", "LANGSMITH_API_KEY": "test-langsmith-key", }, - "general_settings": {} + "general_settings": {}, } - + proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config") - + with patch.object( proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data) ): response = client_no_auth.get("/get/config/callbacks") - + assert response.status_code == 200 result = response.json() - + # Verify response structure assert "status" in result assert result["status"] == "success" assert "callbacks" in result - + callbacks = result["callbacks"] - + # Verify we have all 5 callbacks (2 success + 1 failure + 2 success_and_failure) assert len(callbacks) == 5 - + # Group callbacks by type success_callbacks = [cb for cb in callbacks if cb.get("type") == "success"] failure_callbacks = [cb for cb in callbacks if cb.get("type") == "failure"] - success_and_failure_callbacks = [cb for cb in callbacks if cb.get("type") == "success_and_failure"] - + success_and_failure_callbacks = [ + cb for cb in callbacks if cb.get("type") == "success_and_failure" + ] + # Verify all callbacks have required fields for callback in callbacks: assert "name" in callback assert "variables" in callback assert "type" in callback assert callback["type"] in ["success", "failure", "success_and_failure"] - + # Verify success callbacks assert len(success_callbacks) == 2 success_names = [cb["name"] for cb in success_callbacks] assert "langfuse" in success_names assert "braintrust" in success_names - + # Verify failure callbacks assert len(failure_callbacks) == 1 assert failure_callbacks[0]["name"] == "sentry" - + # Verify success_and_failure callbacks assert len(success_and_failure_callbacks) == 2 success_and_failure_names = [cb["name"] for cb in success_and_failure_callbacks] @@ -2654,13 +2679,13 @@ async def test_get_config_callbacks_environment_variables(client_no_auth): for each callback type. Values are returned as-is from the config (no decryption). """ from litellm.proxy.proxy_server import ProxyConfig - + # Create a mock config with callbacks and their env vars mock_config_data = { "litellm_settings": { "success_callback": ["langfuse"], "failure_callback": [], - "callbacks": ["otel"] + "callbacks": ["otel"], }, "environment_variables": { "LANGFUSE_PUBLIC_KEY": "test-public-key", @@ -2670,21 +2695,21 @@ async def test_get_config_callbacks_environment_variables(client_no_auth): "OTEL_ENDPOINT": "http://localhost:4317", "OTEL_HEADERS": "key=value", }, - "general_settings": {} + "general_settings": {}, } - + proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config") - + with patch.object( proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data) ): response = client_no_auth.get("/get/config/callbacks") - + assert response.status_code == 200 result = response.json() - + callbacks = result["callbacks"] - + # Find langfuse callback (success type) langfuse_callback = next( (cb for cb in callbacks if cb["name"] == "langfuse"), None @@ -2692,7 +2717,7 @@ async def test_get_config_callbacks_environment_variables(client_no_auth): assert langfuse_callback is not None assert langfuse_callback["type"] == "success" assert "variables" in langfuse_callback - + # Verify langfuse env vars are present (values returned as-is, no decryption) langfuse_vars = langfuse_callback["variables"] assert "LANGFUSE_PUBLIC_KEY" in langfuse_vars @@ -2701,15 +2726,13 @@ async def test_get_config_callbacks_environment_variables(client_no_auth): assert langfuse_vars["LANGFUSE_SECRET_KEY"] == "test-secret-key" assert "LANGFUSE_HOST" in langfuse_vars assert langfuse_vars["LANGFUSE_HOST"] == "https://cloud.langfuse.com" - + # Find otel callback (success_and_failure type) - otel_callback = next( - (cb for cb in callbacks if cb["name"] == "otel"), None - ) + otel_callback = next((cb for cb in callbacks if cb["name"] == "otel"), None) assert otel_callback is not None assert otel_callback["type"] == "success_and_failure" assert "variables" in otel_callback - + # Verify otel env vars are present otel_vars = otel_callback["variables"] assert "OTEL_EXPORTER" in otel_vars diff --git a/tests/proxy_unit_tests/vertex_key.json b/tests/proxy_unit_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/proxy_unit_tests/vertex_key.json +++ b/tests/proxy_unit_tests/vertex_key.json @@ -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" } diff --git a/tests/store_model_in_db_tests/test_adding_passthrough_model.py b/tests/store_model_in_db_tests/test_adding_passthrough_model.py index e901be5bd7..80cd7e1aab 100644 --- a/tests/store_model_in_db_tests/test_adding_passthrough_model.py +++ b/tests/store_model_in_db_tests/test_adding_passthrough_model.py @@ -1,12 +1,12 @@ """ Test adding a pass through assemblyai model + api key + api base to the db -wait 20 seconds -make request +wait 20 seconds +make request -Cases to cover -1. user points api base to /assemblyai +Cases to cover +1. user points api base to /assemblyai 2. user points api base to /asssemblyai/us -3. user points api base to /assemblyai/eu +3. user points api base to /assemblyai/eu 4. Bad API Key / credential - 401 """ @@ -21,7 +21,7 @@ TEST_MASTER_KEY = "sk-1234" PROXY_BASE_URL = "http://0.0.0.0:4000" US_BASE_URL = f"{PROXY_BASE_URL}/assemblyai" EU_BASE_URL = f"{PROXY_BASE_URL}/eu.assemblyai" -ASSEMBLYAI_API_KEY_ENV_VAR = "TEST_SPECIAL_ASSEMBLYAI_API_KEY" +ASSEMBLYAI_API_KEY_ENV_VAR = "ASSEMBLYAI_API_KEY" def _delete_all_assemblyai_models_from_db(): diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json new file mode 100644 index 0000000000..d0a519fcd1 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json @@ -0,0 +1,20 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "cancelled", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3, + "batch_size": null, + "learning_rate_multiplier": null + }, + "organization_id": "", + "result_files": [], + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json new file mode 100644 index 0000000000..0093a04f70 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json @@ -0,0 +1,18 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "canceled", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3 + }, + "organization_id": null, + "result_files": null, + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json new file mode 100644 index 0000000000..bdbbdbad07 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json @@ -0,0 +1,3 @@ +{ + "fine_tuning_job_id": "ftjob-azure-create-123" +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json b/tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json new file mode 100644 index 0000000000..201ae3047d --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json @@ -0,0 +1,20 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "running", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3, + "batch_size": null, + "learning_rate_multiplier": null + }, + "organization_id": "", + "result_files": [], + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json b/tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json new file mode 100644 index 0000000000..857b249938 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json @@ -0,0 +1,18 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "running", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3 + }, + "organization_id": null, + "result_files": null, + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_request.json b/tests/test_litellm/expected_fine_tuning_api/azure_create_request.json new file mode 100644 index 0000000000..57319b2698 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_create_request.json @@ -0,0 +1,8 @@ +{ + "model": "gpt-35-turbo-1106", + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": {}, + "extra_body": { + "trainingType": 1 + } +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json b/tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json new file mode 100644 index 0000000000..9986e7d4c5 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json @@ -0,0 +1,20 @@ +{ + "object": "list", + "data": [ + { + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "running" + }, + { + "id": "ftjob-azure-prev-000", + "object": "fine_tuning.job", + "created_at": 1735603200, + "model": "davinci-002", + "status": "succeeded" + } + ], + "has_more": false +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_list_request.json b/tests/test_litellm/expected_fine_tuning_api/azure_list_request.json new file mode 100644 index 0000000000..6bfbb2ffec --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_list_request.json @@ -0,0 +1,4 @@ +{ + "after": "ftjob-azure-prev-000", + "limit": 2 +} diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py index cd76ba1e86..9bcf08fdd7 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -2,7 +2,7 @@ Unit tests for Prometheus user and team count metrics """ from datetime import datetime, timezone -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from prometheus_client import REGISTRY @@ -523,3 +523,193 @@ async def test_set_user_budget_metrics_after_api_request_inf_when_genuinely_no_b assert actual_value == float("inf"), ( "remaining_user_budget_metric should be +Inf when user truly has no budget" ) + + +# --------------------------------------------------------------------------- +# Org budget metric tests +# --------------------------------------------------------------------------- + + +def test_org_budget_metrics_initialized(prometheus_logger): + """Test that the 3 org budget gauge metrics are initialized.""" + assert hasattr(prometheus_logger, "litellm_remaining_org_budget_metric") + assert hasattr(prometheus_logger, "litellm_org_max_budget_metric") + assert hasattr(prometheus_logger, "litellm_org_budget_remaining_hours_metric") + assert prometheus_logger.litellm_remaining_org_budget_metric is not None + assert prometheus_logger.litellm_org_max_budget_metric is not None + assert prometheus_logger.litellm_org_budget_remaining_hours_metric is not None + + +def test_set_org_budget_metrics_remaining_budget(prometheus_logger): + """_set_org_budget_metrics sets remaining budget gauge correctly.""" + prometheus_logger.litellm_remaining_org_budget_metric = MagicMock() + prometheus_logger.litellm_org_max_budget_metric = MagicMock() + prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock() + + prometheus_logger._set_org_budget_metrics( + org_id="org-abc", + org_alias="my-org", + spend=200.0, + max_budget=500.0, + budget_reset_at=None, + ) + + set_call = prometheus_logger.litellm_remaining_org_budget_metric.labels().set + set_call.assert_called_once() + actual = set_call.call_args[0][0] + assert abs(actual - 300.0) < 0.01, f"Expected 300.0, got {actual}" + + +def test_set_org_budget_metrics_max_budget(prometheus_logger): + """_set_org_budget_metrics sets max budget gauge when max_budget is not None.""" + prometheus_logger.litellm_remaining_org_budget_metric = MagicMock() + prometheus_logger.litellm_org_max_budget_metric = MagicMock() + prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock() + + prometheus_logger._set_org_budget_metrics( + org_id="org-abc", + org_alias="my-org", + spend=100.0, + max_budget=1000.0, + budget_reset_at=None, + ) + + prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_called_once_with( + 1000.0 + ) + + +def test_set_org_budget_metrics_no_max_budget(prometheus_logger): + """_set_org_budget_metrics does not set max budget gauge when max_budget is None.""" + prometheus_logger.litellm_remaining_org_budget_metric = MagicMock() + prometheus_logger.litellm_org_max_budget_metric = MagicMock() + prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock() + + prometheus_logger._set_org_budget_metrics( + org_id="org-abc", + org_alias="my-org", + spend=50.0, + max_budget=None, + budget_reset_at=None, + ) + + prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_not_called() + + +def test_set_org_budget_metrics_remaining_hours(prometheus_logger): + """_set_org_budget_metrics sets remaining hours gauge when budget_reset_at is set.""" + prometheus_logger.litellm_remaining_org_budget_metric = MagicMock() + prometheus_logger.litellm_org_max_budget_metric = MagicMock() + prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock() + + future_reset = datetime(2099, 1, 1, tzinfo=timezone.utc) + prometheus_logger._set_org_budget_metrics( + org_id="org-abc", + org_alias="my-org", + spend=10.0, + max_budget=500.0, + budget_reset_at=future_reset, + ) + + prometheus_logger.litellm_org_budget_remaining_hours_metric.labels().set.assert_called_once() + + +@pytest.mark.asyncio +async def test_set_org_budget_metrics_after_api_request(prometheus_logger): + """_set_org_budget_metrics_after_api_request uses cache helper and accounts for response_cost.""" + import sys + + prometheus_logger.litellm_remaining_org_budget_metric = MagicMock() + prometheus_logger.litellm_org_max_budget_metric = MagicMock() + prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock() + + budget_mock = MagicMock() + budget_mock.max_budget = 1000.0 + budget_mock.budget_reset_at = datetime(2099, 1, 1, tzinfo=timezone.utc) + + org_mock = MagicMock() + org_mock.organization_id = "org-xyz" + org_mock.organization_alias = "test-org" + org_mock.spend = 300.0 + org_mock.litellm_budget_table = budget_mock + + mock_prisma = MagicMock() + mock_proxy_server = MagicMock() + mock_proxy_server.prisma_client = mock_prisma + mock_proxy_server.user_api_key_cache = MagicMock() + + with ( + patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}), + patch( + "litellm.proxy.auth.auth_checks.get_org_object", + AsyncMock(return_value=org_mock), + ), + ): + await prometheus_logger._set_org_budget_metrics_after_api_request( + org_id="org-xyz", + response_cost=50.0, + ) + + # remaining budget should reflect spend + response_cost (300 + 50 = 350, remaining = 1000 - 350 = 650) + remaining_call = prometheus_logger.litellm_remaining_org_budget_metric.labels().set.call_args + assert remaining_call is not None + assert remaining_call[0][0] == pytest.approx(650.0) + + prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_called_once_with( + 1000.0 + ) + prometheus_logger.litellm_org_budget_remaining_hours_metric.labels().set.assert_called_once() + + +@pytest.mark.asyncio +async def test_set_org_budget_metrics_after_api_request_no_org_id(prometheus_logger): + """_set_org_budget_metrics_after_api_request is a no-op when org_id is None.""" + prometheus_logger.litellm_remaining_org_budget_metric = MagicMock() + prometheus_logger.litellm_org_max_budget_metric = MagicMock() + prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock() + + await prometheus_logger._set_org_budget_metrics_after_api_request( + org_id=None, + response_cost=1.0, + ) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_not_called() + prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_not_called() + prometheus_logger.litellm_org_budget_remaining_hours_metric.labels().set.assert_not_called() + + +@pytest.mark.asyncio +async def test_initialize_org_budget_metrics(prometheus_logger): + """_initialize_org_budget_metrics fetches all orgs and sets gauges for each.""" + import sys + + prometheus_logger.litellm_remaining_org_budget_metric = MagicMock() + prometheus_logger.litellm_org_max_budget_metric = MagicMock() + prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock() + + budget_mock = MagicMock() + budget_mock.max_budget = 500.0 + budget_mock.budget_reset_at = None + + org_mock = MagicMock() + org_mock.organization_id = "org-init" + org_mock.organization_alias = "init-org" + org_mock.spend = 100.0 + org_mock.litellm_budget_table = budget_mock + + mock_prisma = MagicMock() + mock_prisma.db.litellm_organizationtable.find_many = AsyncMock( + return_value=[org_mock] + ) + mock_prisma.db.litellm_organizationtable.count = AsyncMock(return_value=1) + + mock_proxy_server = MagicMock() + mock_proxy_server.prisma_client = mock_prisma + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._initialize_org_budget_metrics() + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once() + prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_called_once_with( + 500.0 + ) diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index d689c67658..9ed801b360 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -460,7 +460,7 @@ async def test_ensure_initialize_azure_sdk_client_always_used(call_type): "api_key": "test-api-key", "api_version": os.getenv("AZURE_API_VERSION", "2023-05-15"), "api_base": os.getenv( - "AZURE_API_BASE", "https://test.openai.azure.com" + "AZURE_AI_API_BASE", "https://test.openai.azure.com" ), }, } @@ -539,7 +539,11 @@ async def test_ensure_initialize_azure_sdk_client_always_used(call_type): patch_target = ( "litellm.rerank_api.main.azure_rerank.initialize_azure_sdk_client" ) - elif call_type == CallTypes.acreate_batch or call_type == CallTypes.aretrieve_batch or call_type == CallTypes.acancel_batch: + elif ( + call_type == CallTypes.acreate_batch + or call_type == CallTypes.aretrieve_batch + or call_type == CallTypes.acancel_batch + ): patch_target = ( "litellm.batches.main.azure_batches_instance.initialize_azure_sdk_client" ) @@ -570,7 +574,9 @@ async def test_ensure_initialize_azure_sdk_client_always_used(call_type): or call_type == CallTypes.avideo_extension ): # Skip video call types as they don't use Azure SDK client initialization - pytest.skip(f"Skipping {call_type.value} because Azure video calls don't use initialize_azure_sdk_client") + pytest.skip( + f"Skipping {call_type.value} because Azure video calls don't use initialize_azure_sdk_client" + ) elif ( call_type == CallTypes.alist_containers or call_type == CallTypes.aretrieve_container @@ -580,13 +586,26 @@ async def test_ensure_initialize_azure_sdk_client_always_used(call_type): or call_type == CallTypes.aupload_container_file ): # Skip container call types as they're not supported for Azure (only OpenAI) - pytest.skip(f"Skipping {call_type.value} because Azure doesn't support container operations") - elif call_type == CallTypes.avector_store_file_create or call_type == CallTypes.avector_store_file_list or call_type == CallTypes.avector_store_file_retrieve or call_type == CallTypes.avector_store_file_content or call_type == CallTypes.avector_store_file_update or call_type == CallTypes.avector_store_file_delete: + pytest.skip( + f"Skipping {call_type.value} because Azure doesn't support container operations" + ) + elif ( + call_type == CallTypes.avector_store_file_create + or call_type == CallTypes.avector_store_file_list + or call_type == CallTypes.avector_store_file_retrieve + or call_type == CallTypes.avector_store_file_content + or call_type == CallTypes.avector_store_file_update + or call_type == CallTypes.avector_store_file_delete + ): # Skip vector store file call types as they're not supported for Azure (only OpenAI) - pytest.skip(f"Skipping {call_type.value} because Azure doesn't support vector store file operations") + pytest.skip( + f"Skipping {call_type.value} because Azure doesn't support vector store file operations" + ) elif call_type == CallTypes.aocr or call_type == CallTypes.ocr: # Skip OCR call types as they don't use Azure SDK client initialization - pytest.skip(f"Skipping {call_type.value} because OCR calls don't use initialize_azure_sdk_client") + pytest.skip( + f"Skipping {call_type.value} because OCR calls don't use initialize_azure_sdk_client" + ) # Mock the initialize_azure_sdk_client function with patch(patch_target) as mock_init_azure: # Also mock async_function_with_fallbacks to prevent actual API calls @@ -651,7 +670,7 @@ async def test_ensure_initialize_azure_sdk_client_always_used_azure_text(call_ty "api_key": "test-api-key", "api_version": os.getenv("AZURE_API_VERSION", "2023-05-15"), "api_base": os.getenv( - "AZURE_API_BASE", "https://test.openai.azure.com" + "AZURE_AI_API_BASE", "https://test.openai.azure.com" ), }, } @@ -767,7 +786,7 @@ AZURE_API_FUNCTION_PARAMS = [ "speech", False, { - "model": "azure/tts-1", + "model": "azure/tts", "input": "Hello, this is a test of text to speech", "voice": "alloy", "api_key": "test-api-key", @@ -1434,43 +1453,44 @@ def test_token_provider_raises_exception(setup_mocks): def test_get_azure_ad_token_provider_with_default_azure_credential(): """ - Test that get_azure_ad_token_provider correctly uses DefaultAzureCredential + Test that get_azure_ad_token_provider correctly uses DefaultAzureCredential when explicitly specified as the credential type. This verifies that the function can dynamically instantiate DefaultAzureCredential and return a working token provider. """ # Mock Azure identity classes - with patch('azure.identity.DefaultAzureCredential') as mock_default_cred, \ - patch('azure.identity.get_bearer_token_provider') as mock_token_provider: - + with patch("azure.identity.DefaultAzureCredential") as mock_default_cred, patch( + "azure.identity.get_bearer_token_provider" + ) as mock_token_provider: # Configure mocks mock_credential_instance = MagicMock() mock_default_cred.return_value = mock_credential_instance mock_token_provider.return_value = lambda: "test-default-azure-token" - + # Test with DefaultAzureCredential specified explicitly token_provider = get_azure_ad_token_provider( azure_scope="https://cognitiveservices.azure.com/.default", - azure_credential=AzureCredentialType.DefaultAzureCredential + azure_credential=AzureCredentialType.DefaultAzureCredential, ) - + # Verify DefaultAzureCredential was instantiated mock_default_cred.assert_called_once_with() - + # Verify get_bearer_token_provider was called with the right parameters mock_token_provider.assert_called_once_with( - mock_credential_instance, - "https://cognitiveservices.azure.com/.default" + mock_credential_instance, "https://cognitiveservices.azure.com/.default" ) - + # Verify the returned token provider works token = token_provider() assert token == "test-default-azure-token" -def test_get_azure_ad_token_fallback_to_default_azure_credential(setup_mocks, monkeypatch): +def test_get_azure_ad_token_fallback_to_default_azure_credential( + setup_mocks, monkeypatch +): """ - Test that get_azure_ad_token falls back to DefaultAzureCredential when the - service principal method fails but token refresh is enabled. This tests the + Test that get_azure_ad_token falls back to DefaultAzureCredential when the + service principal method fails but token refresh is enabled. This tests the complete fallback flow from service principal to DefaultAzureCredential. """ # Clear environment variables that might interfere @@ -1486,7 +1506,7 @@ def test_get_azure_ad_token_fallback_to_default_azure_credential(setup_mocks, mo # Enable token refresh setup_mocks["litellm"].enable_azure_ad_token_refresh = True - # Configure get_azure_ad_token_provider to fail first (service principal) + # Configure get_azure_ad_token_provider to fail first (service principal) # but succeed on second call (DefaultAzureCredential) def mock_token_provider_side_effect(*args, **kwargs): # If called with azure_credential=DefaultAzureCredential, return a working provider @@ -1512,19 +1532,22 @@ def test_get_azure_ad_token_fallback_to_default_azure_credential(setup_mocks, mo # 1. First with just azure_scope (service principal attempt) # 2. Second with azure_credential=DefaultAzureCredential (fallback) assert setup_mocks["token_provider"].call_count == 2 - + # Verify the calls were made with expected parameters calls = setup_mocks["token_provider"].call_args_list - + # First call should be service principal attempt (no azure_credential) first_call_kwargs = calls[0][1] assert "azure_scope" in first_call_kwargs assert first_call_kwargs.get("azure_credential") is None - + # Second call should be DefaultAzureCredential attempt second_call_kwargs = calls[1][1] assert "azure_scope" in second_call_kwargs - assert second_call_kwargs.get("azure_credential") == AzureCredentialType.DefaultAzureCredential + assert ( + second_call_kwargs.get("azure_credential") + == AzureCredentialType.DefaultAzureCredential + ) # Verify the token is what we expect from our DefaultAzureCredential mock assert token == "mock-default-azure-credential-token" @@ -1584,9 +1607,13 @@ def test_azure_v1_api_uses_openai_client(api_version): ) # Should be OpenAI client, not AzureOpenAI - assert isinstance(client, OpenAI), f"Expected OpenAI client for api_version={api_version}" + assert isinstance( + client, OpenAI + ), f"Expected OpenAI client for api_version={api_version}" # base_url should be /openai/v1/ (not /deployments/) - assert "/openai/v1/" in str(client.base_url), f"base_url should contain /openai/v1/, got {client.base_url}" + assert "/openai/v1/" in str( + client.base_url + ), f"base_url should contain /openai/v1/, got {client.base_url}" # Test async client with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: @@ -1606,9 +1633,13 @@ def test_azure_v1_api_uses_openai_client(api_version): ) # Should be AsyncOpenAI client, not AsyncAzureOpenAI - assert isinstance(async_client, AsyncOpenAI), f"Expected AsyncOpenAI client for api_version={api_version}" + assert isinstance( + async_client, AsyncOpenAI + ), f"Expected AsyncOpenAI client for api_version={api_version}" # base_url should be /openai/v1/ - assert "/openai/v1/" in str(async_client.base_url), f"base_url should contain /openai/v1/, got {async_client.base_url}" + assert "/openai/v1/" in str( + async_client.base_url + ), f"base_url should contain /openai/v1/, got {async_client.base_url}" def test_azure_traditional_api_uses_azure_openai_client(): @@ -1643,7 +1674,9 @@ def test_azure_traditional_api_uses_azure_openai_client(): ) # Should be AzureOpenAI client - assert isinstance(client, AzureOpenAI), f"Expected AzureOpenAI client for api_version={api_version}" + assert isinstance( + client, AzureOpenAI + ), f"Expected AzureOpenAI client for api_version={api_version}" # Test async client with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: @@ -1663,4 +1696,6 @@ def test_azure_traditional_api_uses_azure_openai_client(): ) # Should be AsyncAzureOpenAI client - assert isinstance(async_client, AsyncAzureOpenAI), f"Expected AsyncAzureOpenAI client for api_version={api_version}" + assert isinstance( + async_client, AsyncAzureOpenAI + ), f"Expected AsyncAzureOpenAI client for api_version={api_version}" diff --git a/tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py b/tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py new file mode 100644 index 0000000000..8d008d6c07 --- /dev/null +++ b/tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py @@ -0,0 +1,150 @@ +import json +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest +from openai import AsyncAzureOpenAI + +import litellm +from litellm.llms.azure.fine_tuning.handler import AzureOpenAIFineTuningAPI + + +def _expected_dir() -> Path: + return Path(__file__).resolve().parent.parent.parent / "expected_fine_tuning_api" + + +def _load_json(file_name: str) -> dict: + path = _expected_dir() / file_name + assert path.exists(), f"Expected fixture file not found: {path}" + with open(path) as f: + return json.load(f) + + +class _MockSDKResponse: + def __init__(self, payload: dict): + self._payload = payload + + def model_dump(self) -> dict: + return self._payload + + +def _mock_azure_client( + create_payload: dict | None = None, + list_payload: dict | None = None, + cancel_payload: dict | None = None, +): + client = AsyncAzureOpenAI( + api_key="test-key", + api_version="2024-10-21", + azure_endpoint="https://exampleopenaiendpoint-production.up.railway.app", + ) + client.fine_tuning.jobs.create = AsyncMock( + return_value=( + _MockSDKResponse(create_payload) if create_payload is not None else None + ) + ) # type: ignore[method-assign] + client.fine_tuning.jobs.list = AsyncMock( + return_value=list_payload + ) # type: ignore[method-assign] + client.fine_tuning.jobs.cancel = AsyncMock( + return_value=( + _MockSDKResponse(cancel_payload) if cancel_payload is not None else None + ) + ) # type: ignore[method-assign] + return client + + +@pytest.mark.asyncio +async def test_azure_acreate_fine_tuning_job_request_and_output_match_expected_json(): + expected_request = _load_json("azure_create_request.json") + raw_response = _load_json("azure_create_raw_response.json") + expected_output = _load_json("azure_create_expected_output.json") + + mock_client = _mock_azure_client(create_payload=raw_response) + + with patch.object( + AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client + ): + response = await litellm.acreate_fine_tuning_job( + model="gpt-35-turbo-1106", + training_file="file-5e4b20ecbd724182b9964f3cd2ab7212", + custom_llm_provider="azure", + api_base="https://exampleopenaiendpoint-production.up.railway.app", + api_key="test-key", + api_version="2024-10-21", + ) + + request_kwargs = mock_client.fine_tuning.jobs.create.call_args.kwargs + assert request_kwargs == expected_request + + response_dict = response.model_dump(exclude={"_hidden_params"}) + for key, expected_value in expected_output.items(): + assert key in response_dict, f"Missing key in response: {key}" + assert response_dict[key] == expected_value + + assert response.id is not None + assert response.model == "davinci-002" + + +@pytest.mark.asyncio +async def test_azure_alist_fine_tuning_jobs_request_matches_expected_json(): + expected_request = _load_json("azure_list_request.json") + raw_list_response = _load_json("azure_list_raw_response.json") + + mock_client = _mock_azure_client(list_payload=raw_list_response) + + with patch.object( + AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client + ): + response = await litellm.alist_fine_tuning_jobs( + after=expected_request["after"], + limit=expected_request["limit"], + custom_llm_provider="azure", + api_base="https://exampleopenaiendpoint-production.up.railway.app", + api_key="test-key", + api_version="2024-10-21", + ) + + request_kwargs = mock_client.fine_tuning.jobs.list.call_args.kwargs + assert request_kwargs == expected_request + assert response == raw_list_response + + +@pytest.mark.asyncio +async def test_azure_acancel_fine_tuning_job_request_and_output_match_expected_json(): + expected_request = _load_json("azure_cancel_request.json") + raw_response = _load_json("azure_cancel_raw_response.json") + expected_output = _load_json("azure_cancel_expected_output.json") + + mock_client = _mock_azure_client(cancel_payload=raw_response) + + with patch.object( + AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client + ): + response = await litellm.acancel_fine_tuning_job( + fine_tuning_job_id=expected_request["fine_tuning_job_id"], + custom_llm_provider="azure", + api_base="https://exampleopenaiendpoint-production.up.railway.app", + api_key="test-key", + api_version="2024-10-21", + ) + + request_kwargs = mock_client.fine_tuning.jobs.cancel.call_args.kwargs + assert request_kwargs == expected_request + + response_dict = response.model_dump(exclude={"_hidden_params"}) + for key, expected_value in expected_output.items(): + assert key in response_dict, f"Missing key in response: {key}" + assert response_dict[key] == expected_value + + assert response.status == "cancelled" + + +def test_azure_trainingtype_defaults_to_one(): + 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 diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py index 7880683168..d66798a572 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py @@ -3,6 +3,7 @@ Tests for Azure AI Anthropic CountTokens transformation. Verifies that the CountTokens API uses the correct authentication headers. """ + import os import sys @@ -40,7 +41,7 @@ class TestAzureAIAnthropicCountTokensConfig: assert headers["anthropic-version"] == "2023-06-01" assert "anthropic-beta" in headers - def test_get_required_headers_includes_azure_api_key(self): + def test_get_required_headers_includes_AZURE_AI_API_KEY(self): """ Test that get_required_headers includes Azure api-key header. diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index e9aaa97a42..867a3e61bb 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -6,9 +6,7 @@ import sys import pytest from fastapi.testclient import TestClient -sys.path.insert( - 0, os.path.abspath("../../../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path from unittest.mock import MagicMock, patch import litellm @@ -37,10 +35,7 @@ def test_transform_usage(): ) assert openai_usage.completion_tokens == usage["outputTokens"] assert openai_usage.total_tokens == usage["totalTokens"] - assert ( - openai_usage.prompt_tokens_details.cached_tokens - == usage["cacheReadInputTokens"] - ) + assert openai_usage.prompt_tokens_details.cached_tokens == usage["cacheReadInputTokens"] assert openai_usage._cache_creation_input_tokens == usage["cacheWriteInputTokens"] assert openai_usage._cache_read_input_tokens == usage["cacheReadInputTokens"] # completion_tokens_details should always be populated @@ -194,14 +189,10 @@ def test_apply_tool_call_transformation_if_needed(): role="user", content=json.dumps(tool_response), ) - transformed_message, _ = config.apply_tool_call_transformation_if_needed( - message, tool_calls - ) + transformed_message, _ = config.apply_tool_call_transformation_if_needed(message, tool_calls) assert len(transformed_message.tool_calls) == 1 assert transformed_message.tool_calls[0].function.name == "test_function" - assert transformed_message.tool_calls[0].function.arguments == json.dumps( - tool_response["parameters"] - ) + assert transformed_message.tool_calls[0].function.arguments == json.dumps(tool_response["parameters"]) def test_transform_tool_call_with_cache_control(): @@ -250,12 +241,7 @@ def test_transform_tool_call_with_cache_control(): print(function_out_msg) assert function_out_msg["toolSpec"]["name"] == "get_location" assert function_out_msg["toolSpec"]["description"] == "Get the user's location" - assert ( - function_out_msg["toolSpec"]["inputSchema"]["json"]["properties"]["location"][ - "type" - ] - == "string" - ) + assert function_out_msg["toolSpec"]["inputSchema"]["json"]["properties"]["location"]["type"] == "string" transformed_cache_msg = result["toolConfig"]["tools"][1] assert "cachePoint" in transformed_cache_msg @@ -285,6 +271,7 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto(): assert optional_params["tool_choice"] == {"auto": {}} + def test_get_supported_openai_params(): config = AmazonConverseConfig() supported_params = config.get_supported_openai_params( @@ -307,15 +294,13 @@ def test_get_supported_openai_params_bedrock_converse(): for model in litellm.BEDROCK_CONVERSE_MODELS: print(f"Testing model: {model}") config = AmazonConverseConfig() - supported_params_without_prefix = config.get_supported_openai_params( - model=model - ) + supported_params_without_prefix = config.get_supported_openai_params(model=model) - supported_params_with_prefix = config.get_supported_openai_params( - model=f"bedrock/converse/{model}" - ) + supported_params_with_prefix = config.get_supported_openai_params(model=f"bedrock/converse/{model}") - assert set(supported_params_without_prefix) == set(supported_params_with_prefix), f"Supported params mismatch for model: {model}. Without prefix: {supported_params_without_prefix}, With prefix: {supported_params_with_prefix}" + assert set(supported_params_without_prefix) == set(supported_params_with_prefix), ( + f"Supported params mismatch for model: {model}. Without prefix: {supported_params_without_prefix}, With prefix: {supported_params_with_prefix}" + ) print(f"✅ Passed for model: {model}") @@ -382,7 +367,7 @@ def test_transform_response_with_computer_use_tool(): }, } } - ] + ], } }, "stopReason": "tool_use", @@ -396,10 +381,12 @@ def test_transform_response_with_computer_use_tool(): "cacheWriteInputTokens": 0, }, } + # Mock httpx.Response class MockResponse: def json(self): return response_json + @property def text(self): return json.dumps(response_json) @@ -468,12 +455,10 @@ def test_transform_response_with_bash_tool(): "toolUse": { "toolUseId": "tooluse_456", "name": "bash", - "input": { - "command": "ls -la *.py" - }, + "input": {"command": "ls -la *.py"}, } } - ] + ], } }, "stopReason": "tool_use", @@ -487,10 +472,12 @@ def test_transform_response_with_bash_tool(): "cacheWriteInputTokens": 0, }, } + # Mock httpx.Response class MockResponse: def json(self): return response_json + @property def text(self): return json.dumps(response_json) @@ -549,10 +536,11 @@ def test_transform_response_with_structured_response_being_called(): "name": "json_tool_call", "input": { "Current_Temperature": 62, - "Weather_Explanation": "San Francisco typically has mild, cool weather year-round due to its coastal location and marine influence. The city is known for its fog, moderate temperatures, and relatively stable climate with little seasonal variation."}, + "Weather_Explanation": "San Francisco typically has mild, cool weather year-round due to its coastal location and marine influence. The city is known for its fog, moderate temperatures, and relatively stable climate with little seasonal variation.", + }, } } - ] + ], } }, "stopReason": "tool_use", @@ -566,10 +554,12 @@ def test_transform_response_with_structured_response_being_called(): "cacheWriteInputTokens": 0, }, } + # Mock httpx.Response class MockResponse: def json(self): return response_json + @property def text(self): return json.dumps(response_json) @@ -580,49 +570,43 @@ def test_transform_response_with_structured_response_being_called(): "json_mode": True, "tools": [ { - 'type': 'function', - 'function': { - 'name': 'get_weather', - 'description': 'Get the current weather in a given location', - 'parameters': { - 'type': 'object', - 'properties': { - 'location': { - 'type': 'string', - 'description': 'The city and state, e.g. San Francisco, CA' - }, - 'unit': { - 'type': 'string', - 'enum': ['celsius', 'fahrenheit'] - } + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city and state, e.g. San Francisco, CA"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, }, - 'required': ['location'] - } - } + "required": ["location"], + }, + }, }, { - 'type': 'function', - 'function': { - 'name': 'json_tool_call', - 'parameters': { - '$schema': 'http://json-schema.org/draft-07/schema#', - 'type': 'object', - 'required': ['Weather_Explanation', 'Current_Temperature'], - 'properties': { - 'Weather_Explanation': { - 'type': ['string', 'null'], - 'description': '1-2 sentences explaining the weather in the location' + "type": "function", + "function": { + "name": "json_tool_call", + "parameters": { + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "required": ["Weather_Explanation", "Current_Temperature"], + "properties": { + "Weather_Explanation": { + "type": ["string", "null"], + "description": "1-2 sentences explaining the weather in the location", + }, + "Current_Temperature": { + "type": ["number", "null"], + "description": "Current temperature in the location", }, - 'Current_Temperature': { - 'type': ['number', 'null'], - 'description': 'Current temperature in the location' - } }, - 'additionalProperties': False - } - } - } - ] + "additionalProperties": False, + }, + }, + }, + ], } # Call the transformation logic result = config._transform_response( @@ -641,7 +625,11 @@ def test_transform_response_with_structured_response_being_called(): assert result.choices[0].message.tool_calls is None assert result.choices[0].message.content is not None - assert result.choices[0].message.content == '{"Current_Temperature": 62, "Weather_Explanation": "San Francisco typically has mild, cool weather year-round due to its coastal location and marine influence. The city is known for its fog, moderate temperatures, and relatively stable climate with little seasonal variation."}' + assert ( + result.choices[0].message.content + == '{"Current_Temperature": 62, "Weather_Explanation": "San Francisco typically has mild, cool weather year-round due to its coastal location and marine influence. The city is known for its fog, moderate temperatures, and relatively stable climate with little seasonal variation."}' + ) + def test_transform_response_with_structured_response_calling_tool(): """Test response transformation with structured response.""" @@ -650,28 +638,20 @@ def test_transform_response_with_structured_response_calling_tool(): # Simulate a Bedrock Converse response with a bash tool call response_json = { - "metrics": { - "latencyMs": 1148 - }, + "metrics": {"latencyMs": 1148}, "output": { - "message": - { + "message": { "content": [ - { - "text": "I\'ll check the current weather in San Francisco for you." - }, + {"text": "I'll check the current weather in San Francisco for you."}, { "toolUse": { - "input": { - "location": "San Francisco, CA", - "unit": "celsius" - }, + "input": {"location": "San Francisco, CA", "unit": "celsius"}, "name": "get_weather", - "toolUseId": "tooluse_oKk__QrqSUmufMw3Q7vGaQ" + "toolUseId": "tooluse_oKk__QrqSUmufMw3Q7vGaQ", } - } + }, ], - "role": "assistant" + "role": "assistant", } }, "stopReason": "tool_use", @@ -682,13 +662,15 @@ def test_transform_response_with_structured_response_calling_tool(): "cacheWriteInputTokens": 0, "inputTokens": 534, "outputTokens": 69, - "totalTokens": 603 - } + "totalTokens": 603, + }, } + # Mock httpx.Response class MockResponse: def json(self): return response_json + @property def text(self): return json.dumps(response_json) @@ -699,49 +681,43 @@ def test_transform_response_with_structured_response_calling_tool(): "json_mode": True, "tools": [ { - 'type': 'function', - 'function': { - 'name': 'get_weather', - 'description': 'Get the current weather in a given location', - 'parameters': { - 'type': 'object', - 'properties': { - 'location': { - 'type': 'string', - 'description': 'The city and state, e.g. San Francisco, CA' - }, - 'unit': { - 'type': 'string', - 'enum': ['celsius', 'fahrenheit'] - } + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city and state, e.g. San Francisco, CA"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, }, - 'required': ['location'] - } - } + "required": ["location"], + }, + }, }, { - 'type': 'function', - 'function': { - 'name': 'json_tool_call', - 'parameters': { - '$schema': 'http://json-schema.org/draft-07/schema#', - 'type': 'object', - 'required': ['Weather_Explanation', 'Current_Temperature'], - 'properties': { - 'Weather_Explanation': { - 'type': ['string', 'null'], - 'description': '1-2 sentences explaining the weather in the location' + "type": "function", + "function": { + "name": "json_tool_call", + "parameters": { + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "required": ["Weather_Explanation", "Current_Temperature"], + "properties": { + "Weather_Explanation": { + "type": ["string", "null"], + "description": "1-2 sentences explaining the weather in the location", + }, + "Current_Temperature": { + "type": ["number", "null"], + "description": "Current temperature in the location", }, - 'Current_Temperature': { - 'type': ['number', 'null'], - 'description': 'Current temperature in the location' - } }, - 'additionalProperties': False - } - } - } - ] + "additionalProperties": False, + }, + }, + }, + ], } # Call the transformation logic result = config._transform_response( @@ -760,7 +736,10 @@ def test_transform_response_with_structured_response_calling_tool(): assert result.choices[0].message.tool_calls is not None assert len(result.choices[0].message.tool_calls) == 1 assert result.choices[0].message.tool_calls[0].function.name == "get_weather" - assert result.choices[0].message.tool_calls[0].function.arguments == '{"location": "San Francisco, CA", "unit": "celsius"}' + assert ( + result.choices[0].message.tool_calls[0].function.arguments + == '{"location": "San Francisco, CA", "unit": "celsius"}' + ) @pytest.mark.asyncio @@ -775,12 +754,7 @@ async def test_bedrock_bash_tool_acompletion(): } ] - messages = [ - { - "role": "user", - "content": "run ls command and find all python files" - } - ] + messages = [{"role": "user", "content": "run ls command and find all python files"}] try: response = await litellm.acompletion( @@ -788,7 +762,7 @@ async def test_bedrock_bash_tool_acompletion(): messages=messages, tools=tools, # Using dummy API key - test should fail with auth error, proving request formatting works - api_key="dummy-key-for-testing" + api_key="dummy-key-for-testing", ) # If we get here, something's wrong - we expect an auth error assert False, "Expected authentication error but got successful response" @@ -797,8 +771,16 @@ async def test_bedrock_bash_tool_acompletion(): # Check if it's an expected authentication/credentials error auth_error_indicators = [ - "credentials", "authentication", "unauthorized", "access denied", - "aws", "region", "profile", "token", "invalid", "signature" + "credentials", + "authentication", + "unauthorized", + "access denied", + "aws", + "region", + "profile", + "token", + "invalid", + "signature", ] if any(auth_error in error_str for auth_error in auth_error_indicators): @@ -828,17 +810,14 @@ async def test_bedrock_computer_use_acompletion(): { "role": "user", "content": [ - { - "type": "text", - "text": "Go to the bedrock console" - }, + {"type": "text", "text": "Go to the bedrock console"}, { "type": "image_url", "image_url": { "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" - } - } - ] + }, + }, + ], } ] @@ -848,7 +827,7 @@ async def test_bedrock_computer_use_acompletion(): messages=messages, tools=tools, # Using dummy API key - test should fail with auth error, proving request formatting works - api_key="dummy-key-for-testing" + api_key="dummy-key-for-testing", ) # If we get here, something's wrong - we expect an auth error assert False, "Expected authentication error but got successful response" @@ -857,8 +836,16 @@ async def test_bedrock_computer_use_acompletion(): # Check if it's an expected authentication/credentials error auth_error_indicators = [ - "credentials", "authentication", "unauthorized", "access denied", - "aws", "region", "profile", "token", "invalid", "signature" + "credentials", + "authentication", + "unauthorized", + "access denied", + "aws", + "region", + "profile", + "token", + "invalid", + "signature", ] if any(auth_error in error_str for auth_error in auth_error_indicators): @@ -886,15 +873,10 @@ async def test_transformation_directly(): { "type": "bash_20241022", "name": "bash", - } + }, ] - messages = [ - { - "role": "user", - "content": "run ls command and find all python files" - } - ] + messages = [{"role": "user", "content": "run ls command and find all python files"}] # Transform request request_data = config.transform_request( @@ -902,7 +884,7 @@ async def test_transformation_directly(): messages=messages, optional_params={"tools": tools}, litellm_params={}, - headers={} + headers={}, ) # Verify the structure @@ -994,16 +976,11 @@ def test_transform_request_with_multiple_tools(): }, "required": ["location"], }, - } - } + }, + }, ] - messages = [ - { - "role": "user", - "content": "run ls command and find all python files" - } - ] + messages = [{"role": "user", "content": "run ls command and find all python files"}] # Transform request request_data = config.transform_request( @@ -1011,7 +988,7 @@ def test_transform_request_with_multiple_tools(): messages=messages, optional_params={"tools": tools}, litellm_params={}, - headers={} + headers={}, ) # Verify the structure @@ -1054,17 +1031,14 @@ def test_transform_request_with_computer_tool_only(): { "role": "user", "content": [ - { - "type": "text", - "text": "Go to the bedrock console" - }, + {"type": "text", "text": "Go to the bedrock console"}, { "type": "image_url", "image_url": { "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" - } - } - ] + }, + }, + ], } ] @@ -1074,7 +1048,7 @@ def test_transform_request_with_computer_tool_only(): messages=messages, optional_params={"tools": tools}, litellm_params={}, - headers={} + headers={}, ) # Verify the structure @@ -1102,12 +1076,7 @@ def test_transform_request_with_bash_tool_only(): } ] - messages = [ - { - "role": "user", - "content": "run ls command and find all python files" - } - ] + messages = [{"role": "user", "content": "run ls command and find all python files"}] # Transform request request_data = config.transform_request( @@ -1115,7 +1084,7 @@ def test_transform_request_with_bash_tool_only(): messages=messages, optional_params={"tools": tools}, litellm_params={}, - headers={} + headers={}, ) # Verify the structure @@ -1143,12 +1112,7 @@ def test_transform_request_with_text_editor_tool(): } ] - messages = [ - { - "role": "user", - "content": "Edit this text file" - } - ] + messages = [{"role": "user", "content": "Edit this text file"}] # Transform request request_data = config.transform_request( @@ -1156,7 +1120,7 @@ def test_transform_request_with_text_editor_tool(): messages=messages, optional_params={"tools": tools}, litellm_params={}, - headers={} + headers={}, ) # Verify the structure @@ -1194,16 +1158,11 @@ def test_transform_request_with_function_tool(): }, "required": ["location"], }, - } + }, } ] - messages = [ - { - "role": "user", - "content": "What's the weather like in San Francisco?" - } - ] + messages = [{"role": "user", "content": "What's the weather like in San Francisco?"}] # Transform request request_data = config.transform_request( @@ -1211,7 +1170,7 @@ def test_transform_request_with_function_tool(): messages=messages, optional_params={"tools": tools}, litellm_params={}, - headers={} + headers={}, ) # Verify the structure @@ -1247,7 +1206,7 @@ def test_map_openai_params_with_response_format(): }, "required": ["location"], }, - } + }, } ] @@ -1279,7 +1238,7 @@ def test_map_openai_params_with_response_format(): non_default_params={"response_format": json_schema}, optional_params={"tools": tools}, model="eu.anthropic.claude-sonnet-4-20250514-v1:0", - drop_params=False + drop_params=False, ) assert "tools" in optional_params @@ -1299,31 +1258,21 @@ async def test_assistant_message_cache_control(): # Test assistant message with string content and cache_control messages = [ {"role": "user", "content": "Hello"}, - { - "role": "assistant", - "content": "Hi there!", - "cache_control": {"type": "ephemeral"} - } + {"role": "assistant", "content": "Hi there!", "cache_control": {"type": "ephemeral"}}, ] result = _bedrock_converse_messages_pt( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) assert result == async_result async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) assert result == async_result @@ -1353,26 +1302,16 @@ async def test_assistant_message_list_content_cache_control(): {"role": "user", "content": "Hello"}, { "role": "assistant", - "content": [ - { - "type": "text", - "text": "This should be cached", - "cache_control": {"type": "ephemeral"} - } - ] - } + "content": [{"type": "text", "text": "This should be cached", "cache_control": {"type": "ephemeral"}}], + }, ] result = _bedrock_converse_messages_pt( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) assert result == async_result @@ -1399,36 +1338,22 @@ async def test_tool_message_cache_control(): "role": "assistant", "content": None, "tool_calls": [ - { - "id": "call_123", - "type": "function", - "function": {"name": "get_weather", "arguments": "{}"} - } - ] + {"id": "call_123", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], }, { "role": "tool", "tool_call_id": "call_123", - "content": [ - { - "type": "text", - "text": "Weather data: sunny, 25°C", - "cache_control": {"type": "ephemeral"} - } - ] - } + "content": [{"type": "text", "text": "Weather data: sunny, 25°C", "cache_control": {"type": "ephemeral"}}], + }, ] result = _bedrock_converse_messages_pt( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) assert result == async_result @@ -1463,31 +1388,23 @@ async def test_tool_message_string_content_cache_control(): "role": "assistant", "content": None, "tool_calls": [ - { - "id": "call_123", - "type": "function", - "function": {"name": "get_weather", "arguments": "{}"} - } - ] + {"id": "call_123", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], }, { "role": "tool", "tool_call_id": "call_123", "content": "Weather: sunny, 25°C", - "cache_control": {"type": "ephemeral"} - } + "cache_control": {"type": "ephemeral"}, + }, ] result = _bedrock_converse_messages_pt( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) assert result == async_result @@ -1523,22 +1440,18 @@ async def test_assistant_tool_calls_cache_control(): "id": "call_proxy_123", "type": "function", "function": {"name": "calc", "arguments": "{}"}, - "cache_control": {"type": "ephemeral"} + "cache_control": {"type": "ephemeral"}, } - ] - } + ], + }, ] result = _bedrock_converse_messages_pt( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) assert result == async_result @@ -1575,28 +1488,24 @@ async def test_multiple_tool_calls_with_mixed_cache_control(): "id": "call_1", "type": "function", "function": {"name": "calc", "arguments": '{"expr": "2+2"}'}, - "cache_control": {"type": "ephemeral"} + "cache_control": {"type": "ephemeral"}, }, { "id": "call_2", "type": "function", - "function": {"name": "calc", "arguments": '{"expr": "3+3"}'} + "function": {"name": "calc", "arguments": '{"expr": "3+3"}'}, # No cache_control - } - ] - } + }, + ], + }, ] result = _bedrock_converse_messages_pt( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) assert result == async_result @@ -1632,20 +1541,16 @@ async def test_no_cache_control_no_cache_point(): { "role": "tool", "tool_call_id": "call_123", - "content": "Tool result" # No cache_control - } + "content": "Tool result", # No cache_control + }, ] result = _bedrock_converse_messages_pt( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( - messages=messages, - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) assert result == async_result @@ -1665,6 +1570,7 @@ async def test_no_cache_control_no_cache_point(): # Guarded Text Feature Tests # ============================================================================ + def test_guarded_text_wraps_in_guardrail_converse_content(): """Test that guarded_text content type gets wrapped in guardContent blocks.""" from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -1677,15 +1583,13 @@ def test_guarded_text_wraps_in_guardrail_converse_content(): "content": [ {"type": "text", "text": "Regular text content"}, {"type": "guarded_text", "text": "This should be guarded"}, - {"type": "text", "text": "More regular text"} - ] + {"type": "text", "text": "More regular text"}, + ], } ] result = _bedrock_converse_messages_pt( - messages=messages, - model="us.amazon.nova-pro-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="us.amazon.nova-pro-v1:0", llm_provider="bedrock_converse" ) # Should have 1 message @@ -1705,6 +1609,7 @@ def test_guarded_text_wraps_in_guardrail_converse_content(): assert "guardContent" in content[1] assert content[1]["guardContent"]["text"]["text"] == "This should be guarded" + def test_guarded_text_with_system_messages(): """Test guarded_text with system messages using the full transformation.""" config = AmazonConverseConfig() @@ -1715,24 +1620,22 @@ def test_guarded_text_with_system_messages(): "role": "user", "content": [ {"type": "text", "text": "What is the main topic of this legal document?"}, - {"type": "guarded_text", "text": "This is a set of very long instructions that you will follow. Here is a legal document that you will use to answer the user's question."} - ] - } + { + "type": "guarded_text", + "text": "This is a set of very long instructions that you will follow. Here is a legal document that you will use to answer the user's question.", + }, + ], + }, ] - optional_params = { - "guardrailConfig": { - "guardrailIdentifier": "gr-abc123", - "guardrailVersion": "DRAFT" - } - } + optional_params = {"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "DRAFT"}} result = config._transform_request( model="us.amazon.nova-pro-v1:0", messages=messages, optional_params=optional_params, litellm_params={}, - headers={} + headers={}, ) # Should have system content blocks @@ -1755,7 +1658,10 @@ def test_guarded_text_with_system_messages(): assert content[0]["text"] == "What is the main topic of this legal document?" # Second should be guardContent assert "guardContent" in content[1] - assert content[1]["guardContent"]["text"]["text"] == "This is a set of very long instructions that you will follow. Here is a legal document that you will use to answer the user's question." + assert ( + content[1]["guardContent"]["text"]["text"] + == "This is a set of very long instructions that you will follow. Here is a legal document that you will use to answer the user's question." + ) def test_guarded_text_with_mixed_content_types(): @@ -1770,15 +1676,13 @@ def test_guarded_text_with_mixed_content_types(): "content": [ {"type": "text", "text": "Look at this image"}, {"type": "image_url", "image_url": {"url": "data:image/png;base64,test"}}, - {"type": "guarded_text", "text": "This sensitive content should be guarded"} - ] + {"type": "guarded_text", "text": "This sensitive content should be guarded"}, + ], } ] result = _bedrock_converse_messages_pt( - messages=messages, - model="us.amazon.nova-pro-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="us.amazon.nova-pro-v1:0", llm_provider="bedrock_converse" ) # Should have 1 message @@ -1800,6 +1704,7 @@ def test_guarded_text_with_mixed_content_types(): assert "guardContent" in content[2] assert content[2]["guardContent"]["text"]["text"] == "This sensitive content should be guarded" + @pytest.mark.asyncio async def test_async_guarded_text(): """Test async version of guarded_text processing.""" @@ -1810,17 +1715,12 @@ async def test_async_guarded_text(): messages = [ { "role": "user", - "content": [ - {"type": "text", "text": "Hello"}, - {"type": "guarded_text", "text": "This should be guarded"} - ] + "content": [{"type": "text", "text": "Hello"}, {"type": "guarded_text", "text": "This should be guarded"}], } ] result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( - messages=messages, - model="us.amazon.nova-pro-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="us.amazon.nova-pro-v1:0", llm_provider="bedrock_converse" ) # Should have 1 message @@ -1851,31 +1751,21 @@ def test_guarded_text_with_tool_calls(): "role": "user", "content": [ {"type": "text", "text": "What's the weather?"}, - {"type": "guarded_text", "text": "Please be careful with sensitive information"} - ] + {"type": "guarded_text", "text": "Please be careful with sensitive information"}, + ], }, { "role": "assistant", "content": None, "tool_calls": [ - { - "id": "call_123", - "type": "function", - "function": {"name": "get_weather", "arguments": "{}"} - } - ] + {"id": "call_123", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], }, - { - "role": "tool", - "tool_call_id": "call_123", - "content": "It's sunny and 25°C" - } + {"role": "tool", "tool_call_id": "call_123", "content": "It's sunny and 25°C"}, ] result = _bedrock_converse_messages_pt( - messages=messages, - model="us.amazon.nova-pro-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="us.amazon.nova-pro-v1:0", llm_provider="bedrock_converse" ) # Should have 3 messages @@ -1909,26 +1799,18 @@ def test_guarded_text_guardrail_config_preserved(): messages = [ { "role": "user", - "content": [ - {"type": "text", "text": "Hello"}, - {"type": "guarded_text", "text": "This should be guarded"} - ] + "content": [{"type": "text", "text": "Hello"}, {"type": "guarded_text", "text": "This should be guarded"}], } ] - optional_params = { - "guardrailConfig": { - "guardrailIdentifier": "gr-abc123", - "guardrailVersion": "DRAFT" - } - } + optional_params = {"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "DRAFT"}} result = config._transform_request( model="us.amazon.nova-pro-v1:0", messages=messages, optional_params=optional_params, litellm_params={}, - headers={} + headers={}, ) # GuardrailConfig should be present at top level @@ -1946,23 +1828,10 @@ def test_auto_convert_last_user_message_to_guarded_text(): config = AmazonConverseConfig() messages = [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "What is the main topic of this legal document?" - } - ] - } + {"role": "user", "content": [{"type": "text", "text": "What is the main topic of this legal document?"}]} ] - optional_params = { - "guardrailConfig": { - "guardrailIdentifier": "gr-abc123", - "guardrailVersion": "1" - } - } + optional_params = {"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "1"}} # Test the helper method directly converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) @@ -1979,19 +1848,9 @@ def test_auto_convert_last_user_message_string_content(): """Test that last user message with string content is automatically converted to guarded_text when guardrailConfig is present.""" config = AmazonConverseConfig() - messages = [ - { - "role": "user", - "content": "What is the main topic of this legal document?" - } - ] + messages = [{"role": "user", "content": "What is the main topic of this legal document?"}] - optional_params = { - "guardrailConfig": { - "guardrailIdentifier": "gr-abc123", - "guardrailVersion": "1" - } - } + optional_params = {"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "1"}} # Test the helper method directly converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) @@ -2009,15 +1868,7 @@ def test_no_conversion_when_no_guardrail_config(): config = AmazonConverseConfig() messages = [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "What is the main topic of this legal document?" - } - ] - } + {"role": "user", "content": [{"type": "text", "text": "What is the main topic of this legal document?"}]} ] optional_params = {} @@ -2033,24 +1884,9 @@ def test_no_conversion_when_guarded_text_already_present(): """Test that no conversion happens when guarded_text is already present in the last user message.""" config = AmazonConverseConfig() - messages = [ - { - "role": "user", - "content": [ - { - "type": "guarded_text", - "text": "This is already guarded" - } - ] - } - ] + messages = [{"role": "user", "content": [{"type": "guarded_text", "text": "This is already guarded"}]}] - optional_params = { - "guardrailConfig": { - "guardrailIdentifier": "gr-abc123", - "guardrailVersion": "1" - } - } + optional_params = {"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "1"}} # Test the helper method directly converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) @@ -2067,24 +1903,13 @@ def test_auto_convert_with_mixed_content(): { "role": "user", "content": [ - { - "type": "text", - "text": "What is the main topic of this legal document?" - }, - { - "type": "image_url", - "image_url": {"url": "https://example.com/image.jpg"} - } - ] + {"type": "text", "text": "What is the main topic of this legal document?"}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.jpg"}}, + ], } ] - optional_params = { - "guardrailConfig": { - "guardrailIdentifier": "gr-abc123", - "guardrailVersion": "1" - } - } + optional_params = {"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "1"}} # Test the helper method directly converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) @@ -2108,23 +1933,10 @@ def test_auto_convert_in_full_transformation(): config = AmazonConverseConfig() messages = [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "What is the main topic of this legal document?" - } - ] - } + {"role": "user", "content": [{"type": "text", "text": "What is the main topic of this legal document?"}]} ] - optional_params = { - "guardrailConfig": { - "guardrailIdentifier": "gr-abc123", - "guardrailVersion": "1" - } - } + optional_params = {"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "1"}} # Test the full transformation result = config._transform_request( @@ -2132,7 +1944,7 @@ def test_auto_convert_in_full_transformation(): messages=messages, optional_params=optional_params, litellm_params={}, - headers={} + headers={}, ) # Verify the transformation worked @@ -2152,45 +1964,13 @@ def test_convert_consecutive_user_messages_to_guarded_text(): config = AmazonConverseConfig() messages = [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "First user message" - } - ] - }, - { - "role": "assistant", - "content": "Assistant response" - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Second user message" - } - ] - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Third user message" - } - ] - } + {"role": "user", "content": [{"type": "text", "text": "First user message"}]}, + {"role": "assistant", "content": "Assistant response"}, + {"role": "user", "content": [{"type": "text", "text": "Second user message"}]}, + {"role": "user", "content": [{"type": "text", "text": "Third user message"}]}, ] - optional_params = { - "guardrailConfig": { - "guardrailIdentifier": "gr-abc123", - "guardrailVersion": "1" - } - } + optional_params = {"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "1"}} # Test the helper method directly converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) @@ -2223,41 +2003,12 @@ def test_convert_all_user_messages_when_all_consecutive(): config = AmazonConverseConfig() messages = [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "First user message" - } - ] - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Second user message" - } - ] - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Third user message" - } - ] - } + {"role": "user", "content": [{"type": "text", "text": "First user message"}]}, + {"role": "user", "content": [{"type": "text", "text": "Second user message"}]}, + {"role": "user", "content": [{"type": "text", "text": "Third user message"}]}, ] - optional_params = { - "guardrailConfig": { - "guardrailIdentifier": "gr-abc123", - "guardrailVersion": "1" - } - } + optional_params = {"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "1"}} # Test the helper method directly converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) @@ -2279,26 +2030,12 @@ def test_convert_consecutive_user_messages_with_string_content(): config = AmazonConverseConfig() messages = [ - { - "role": "assistant", - "content": "Assistant response" - }, - { - "role": "user", - "content": "First user message" - }, - { - "role": "user", - "content": "Second user message" - } + {"role": "assistant", "content": "Assistant response"}, + {"role": "user", "content": "First user message"}, + {"role": "user", "content": "Second user message"}, ] - optional_params = { - "guardrailConfig": { - "guardrailIdentifier": "gr-abc123", - "guardrailVersion": "1" - } - } + optional_params = {"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "1"}} # Test the helper method directly converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) @@ -2327,32 +2064,11 @@ def test_skip_consecutive_user_messages_with_existing_guarded_text(): config = AmazonConverseConfig() messages = [ - { - "role": "user", - "content": [ - { - "type": "guarded_text", - "text": "Already guarded" - } - ] - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Should be converted" - } - ] - } + {"role": "user", "content": [{"type": "guarded_text", "text": "Already guarded"}]}, + {"role": "user", "content": [{"type": "text", "text": "Should be converted"}]}, ] - optional_params = { - "guardrailConfig": { - "guardrailIdentifier": "gr-abc123", - "guardrailVersion": "1" - } - } + optional_params = {"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "1"}} # Test the helper method directly converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) @@ -2384,11 +2100,7 @@ def test_request_metadata_transformation(): """Test that requestMetadata is properly transformed to top-level field.""" config = AmazonConverseConfig() - request_metadata = { - "cost_center": "engineering", - "user_id": "user123", - "session_id": "sess_abc123" - } + request_metadata = {"cost_center": "engineering", "user_id": "user123", "session_id": "sess_abc123"} messages = [ {"role": "user", "content": "Hello!"}, @@ -2400,7 +2112,7 @@ def test_request_metadata_transformation(): messages=messages, optional_params={"requestMetadata": request_metadata}, litellm_params={}, - headers={} + headers={}, ) # Verify that requestMetadata appears as top-level field @@ -2426,7 +2138,7 @@ def test_request_metadata_validation(): messages=messages, optional_params={"requestMetadata": valid_metadata}, litellm_params={}, - headers={} + headers={}, ) # Test too many items (max 16) @@ -2438,7 +2150,7 @@ def test_request_metadata_validation(): messages=messages, optional_params={"requestMetadata": too_many_items}, litellm_params={}, - headers={} + headers={}, ) assert False, "Should have raised validation error for too many items" except Exception as e: @@ -2461,7 +2173,7 @@ def test_request_metadata_key_constraints(): messages=messages, optional_params={"requestMetadata": invalid_metadata}, litellm_params={}, - headers={} + headers={}, ) assert False, "Should have raised validation error for key too long" except Exception as e: @@ -2476,7 +2188,7 @@ def test_request_metadata_key_constraints(): messages=messages, optional_params={"requestMetadata": invalid_metadata}, litellm_params={}, - headers={} + headers={}, ) assert False, "Should have raised validation error for empty key" except Exception as e: @@ -2499,7 +2211,7 @@ def test_request_metadata_value_constraints(): messages=messages, optional_params={"requestMetadata": invalid_metadata}, litellm_params={}, - headers={} + headers={}, ) assert False, "Should have raised validation error for value too long" except Exception as e: @@ -2514,7 +2226,7 @@ def test_request_metadata_value_constraints(): messages=messages, optional_params={"requestMetadata": valid_metadata}, litellm_params={}, - headers={} + headers={}, ) @@ -2537,7 +2249,7 @@ def test_request_metadata_character_pattern(): messages=messages, optional_params={"requestMetadata": valid_metadata}, litellm_params={}, - headers={} + headers={}, ) @@ -2545,10 +2257,7 @@ def test_request_metadata_with_other_params(): """Test that requestMetadata works alongside other parameters.""" config = AmazonConverseConfig() - request_metadata = { - "experiment": "test_A", - "user_type": "premium" - } + request_metadata = {"experiment": "test_A", "user_type": "premium"} messages = [ {"role": "user", "content": "What's the weather?"}, @@ -2562,12 +2271,10 @@ def test_request_metadata_with_other_params(): "description": "Get the current weather", "parameters": { "type": "object", - "properties": { - "location": {"type": "string"} - }, - "required": ["location"] - } - } + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, } ] @@ -2575,14 +2282,9 @@ def test_request_metadata_with_other_params(): request_data = config.transform_request( model="anthropic.claude-3-5-sonnet-20240620-v1:0", messages=messages, - optional_params={ - "requestMetadata": request_metadata, - "tools": tools, - "max_tokens": 100, - "temperature": 0.7 - }, + optional_params={"requestMetadata": request_metadata, "tools": tools, "max_tokens": 100, "temperature": 0.7}, litellm_params={}, - headers={} + headers={}, ) # Verify requestMetadata is at top level @@ -2607,7 +2309,7 @@ def test_request_metadata_empty(): messages=messages, optional_params={"requestMetadata": {}}, litellm_params={}, - headers={} + headers={}, ) assert "requestMetadata" in request_data @@ -2626,7 +2328,7 @@ def test_request_metadata_not_provided(): messages=messages, optional_params={}, litellm_params={}, - headers={} + headers={}, ) # requestMetadata should not be in the request @@ -2649,16 +2351,14 @@ def test_empty_assistant_message_handling(): messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": ""}, # Empty content - {"role": "user", "content": "How are you?"} + {"role": "user", "content": "How are you?"}, ] # Use patch to ensure we modify the litellm reference that factory.py actually uses # This avoids issues with module reloading during parallel test execution with patch.object(factory_module.litellm, "modify_params", True): result = _bedrock_converse_messages_pt( - messages=messages, - model="anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) # Should have 3 messages: user, assistant (with placeholder), user @@ -2676,13 +2376,11 @@ def test_empty_assistant_message_handling(): messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": " "}, # Whitespace-only content - {"role": "user", "content": "How are you?"} + {"role": "user", "content": "How are you?"}, ] result = _bedrock_converse_messages_pt( - messages=messages, - model="anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) # Assistant message should have placeholder text instead of whitespace @@ -2693,13 +2391,11 @@ def test_empty_assistant_message_handling(): messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": [{"type": "text", "text": ""}]}, # Empty text in list - {"role": "user", "content": "How are you?"} + {"role": "user", "content": "How are you?"}, ] result = _bedrock_converse_messages_pt( - messages=messages, - model="anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) # Assistant message should have placeholder text instead of empty text @@ -2710,13 +2406,11 @@ def test_empty_assistant_message_handling(): messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "I'm doing well, thank you!"}, # Normal content - {"role": "user", "content": "How are you?"} + {"role": "user", "content": "How are you?"}, ] result = _bedrock_converse_messages_pt( - messages=messages, - model="anthropic.claude-3-5-sonnet-20240620-v1:0", - llm_provider="bedrock_converse" + messages=messages, model="anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) # Assistant message should keep original content @@ -2828,6 +2522,7 @@ def test_thinking_with_max_completion_tokens(): assert result["thinking"]["type"] == "enabled" assert result["thinking"]["budget_tokens"] == 5000 + def test_drop_thinking_param_when_thinking_blocks_missing(): """ Test that thinking param is dropped when modify_params=True and @@ -2868,24 +2563,21 @@ def test_drop_thinking_param_when_thinking_blocks_missing(): optional_params = {"thinking": {"type": "enabled", "budget_tokens": 1000}} # Verify the condition is detected - assert last_assistant_with_tool_calls_has_no_thinking_blocks( - messages_without_thinking_blocks - ), "Should detect missing thinking_blocks" + assert last_assistant_with_tool_calls_has_no_thinking_blocks(messages_without_thinking_blocks), ( + "Should detect missing thinking_blocks" + ) # Simulate what _transform_request_helper does if ( optional_params.get("thinking") is not None and messages_without_thinking_blocks is not None - and last_assistant_with_tool_calls_has_no_thinking_blocks( - messages_without_thinking_blocks - ) + and last_assistant_with_tool_calls_has_no_thinking_blocks(messages_without_thinking_blocks) ): if litellm.modify_params: optional_params.pop("thinking", None) assert "thinking" not in optional_params, ( - "thinking param should be dropped when modify_params=True " - "and thinking_blocks are missing" + "thinking param should be dropped when modify_params=True and thinking_blocks are missing" ) # Test case 2: thinking should NOT be dropped when thinking_blocks are present @@ -2901,29 +2593,23 @@ def test_drop_thinking_param_when_thinking_blocks_missing(): "function": {"name": "search", "arguments": "{}"}, } ], - "thinking_blocks": [ - {"type": "thinking", "thinking": "Let me search for weather..."} - ], + "thinking_blocks": [{"type": "thinking", "thinking": "Let me search for weather..."}], }, {"role": "tool", "content": "Weather is sunny", "tool_call_id": "call_123"}, ] - optional_params_with_thinking = { - "thinking": {"type": "enabled", "budget_tokens": 1000} - } + optional_params_with_thinking = {"thinking": {"type": "enabled", "budget_tokens": 1000}} # Verify the condition is NOT detected when thinking_blocks are present - assert not last_assistant_with_tool_calls_has_no_thinking_blocks( - messages_with_thinking_blocks - ), "Should NOT detect missing thinking_blocks when they are present" + assert not last_assistant_with_tool_calls_has_no_thinking_blocks(messages_with_thinking_blocks), ( + "Should NOT detect missing thinking_blocks when they are present" + ) # Simulate what _transform_request_helper does if ( optional_params_with_thinking.get("thinking") is not None and messages_with_thinking_blocks is not None - and last_assistant_with_tool_calls_has_no_thinking_blocks( - messages_with_thinking_blocks - ) + and last_assistant_with_tool_calls_has_no_thinking_blocks(messages_with_thinking_blocks) ): if litellm.modify_params: optional_params_with_thinking.pop("thinking", None) @@ -2935,24 +2621,18 @@ def test_drop_thinking_param_when_thinking_blocks_missing(): # Test case 3: thinking should NOT be dropped when modify_params=False litellm.modify_params = False - optional_params_no_modify = { - "thinking": {"type": "enabled", "budget_tokens": 1000} - } + optional_params_no_modify = {"thinking": {"type": "enabled", "budget_tokens": 1000}} # Simulate what _transform_request_helper does if ( optional_params_no_modify.get("thinking") is not None and messages_without_thinking_blocks is not None - and last_assistant_with_tool_calls_has_no_thinking_blocks( - messages_without_thinking_blocks - ) + and last_assistant_with_tool_calls_has_no_thinking_blocks(messages_without_thinking_blocks) ): if litellm.modify_params: optional_params_no_modify.pop("thinking", None) - assert "thinking" in optional_params_no_modify, ( - "thinking param should NOT be dropped when modify_params=False" - ) + assert "thinking" in optional_params_no_modify, "thinking param should NOT be dropped when modify_params=False" finally: # Restore original modify_params setting @@ -2960,46 +2640,53 @@ def test_drop_thinking_param_when_thinking_blocks_missing(): def test_supports_native_structured_outputs(): - """Test model detection for native structured outputs support.""" - config = AmazonConverseConfig() + """Test model detection for native structured outputs support. - # Supported models - assert config._supports_native_structured_outputs( - "anthropic.claude-sonnet-4-5-20250929-v1:0" - ) - assert config._supports_native_structured_outputs( - "anthropic.claude-haiku-4-5-20251001-v1:0" - ) - assert config._supports_native_structured_outputs( - "anthropic.claude-opus-4-6-v1:0" - ) - assert config._supports_native_structured_outputs( - "eu.anthropic.claude-opus-4-5-20260101-v1:0" - ) - assert config._supports_native_structured_outputs("qwen.qwen3-235b-instruct-v1:0") - assert config._supports_native_structured_outputs("mistral.mistral-large-3-v1:0") - assert config._supports_native_structured_outputs("deepseek.deepseek-v3.1-v1:0") + Support is driven by the ``supports_native_structured_output`` flag in the + cost JSON (litellm.model_cost), not a hardcoded model set. + """ + old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") + old_cost = litellm.model_cost + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + try: + config = AmazonConverseConfig() - # Unsupported models — should fall back to tool-call approach - assert not config._supports_native_structured_outputs( - "anthropic.claude-3-5-sonnet-20241022-v2:0" - ) - assert not config._supports_native_structured_outputs( - "anthropic.claude-sonnet-4-20250514-v1:0" - ) - assert not config._supports_native_structured_outputs( - "meta.llama3-3-70b-instruct-v1:0" - ) - assert not config._supports_native_structured_outputs( - "amazon.nova-pro-v1:0" - ) - # Excluded despite AWS listing them: broken constrained decoding on Bedrock - assert not config._supports_native_structured_outputs( - "openai.gpt-oss-120b-1:0" - ) - assert not config._supports_native_structured_outputs( - "mistral.magistral-small-2509" - ) + # Supported models (have supports_native_structured_output=true in cost JSON) + assert config._supports_native_structured_outputs("anthropic.claude-sonnet-4-5-20250929-v1:0") + assert config._supports_native_structured_outputs("anthropic.claude-haiku-4-5-20251001-v1:0") + assert config._supports_native_structured_outputs("anthropic.claude-opus-4-6-v1") + # Regional prefix is stripped by get_bedrock_base_model + assert config._supports_native_structured_outputs("eu.anthropic.claude-opus-4-5-20251101-v1:0") + # Claude 4.6 Sonnet + assert config._supports_native_structured_outputs("anthropic.claude-sonnet-4-6") + assert config._supports_native_structured_outputs("us.anthropic.claude-sonnet-4-6") + # Non-Anthropic models + assert config._supports_native_structured_outputs("qwen.qwen3-235b-a22b-2507-v1:0") + assert config._supports_native_structured_outputs("mistral.mistral-large-3-675b-instruct") + assert config._supports_native_structured_outputs("minimax.minimax-m2") + assert config._supports_native_structured_outputs("moonshot.kimi-k2-thinking") + assert config._supports_native_structured_outputs("nvidia.nemotron-nano-3-30b") + # DeepSeek: old substring "deepseek-v3.1" didn't match real ID + assert config._supports_native_structured_outputs("deepseek.v3-v1:0") + + # Unsupported models -- should fall back to tool-call approach + assert not config._supports_native_structured_outputs("anthropic.claude-3-5-sonnet-20241022-v2:0") + assert not config._supports_native_structured_outputs("anthropic.claude-sonnet-4-20250514-v1:0") + assert not config._supports_native_structured_outputs("meta.llama3-3-70b-instruct-v1:0") + assert not config._supports_native_structured_outputs("amazon.nova-pro-v1:0") + # Excluded: broken constrained decoding on Bedrock + assert not config._supports_native_structured_outputs("openai.gpt-oss-120b-1:0") + assert not config._supports_native_structured_outputs("mistral.magistral-small-2509") + # Excluded: ignores schema or broken on Bedrock + assert not config._supports_native_structured_outputs("google.gemma-3-27b-it") + assert not config._supports_native_structured_outputs("nvidia.nemotron-nano-12b-v2") + finally: + litellm.model_cost = old_cost + if old_env is None: + os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) + else: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env def test_create_output_config_for_response_format(): @@ -3039,49 +2726,57 @@ def test_create_output_config_for_response_format(): def test_translate_response_format_native_output_config(): """For supported models, _translate_response_format_param should produce outputConfig.""" - config = AmazonConverseConfig() + old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") + old_cost = litellm.model_cost + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + try: + config = AmazonConverseConfig() - response_format = { - "type": "json_schema", - "json_schema": { - "name": "WeatherResult", - "description": "Weather info", - "schema": { - "type": "object", - "properties": { - "temp": {"type": "number"}, + response_format = { + "type": "json_schema", + "json_schema": { + "name": "WeatherResult", + "description": "Weather info", + "schema": { + "type": "object", + "properties": { + "temp": {"type": "number"}, + }, + "required": ["temp"], }, - "required": ["temp"], }, - }, - } + } - optional_params: dict = {} - result = config._translate_response_format_param( - value=response_format, - model="anthropic.claude-sonnet-4-5-20250929-v1:0", - optional_params=optional_params, - non_default_params={"response_format": response_format}, - is_thinking_enabled=False, - ) + optional_params: dict = {} + result = config._translate_response_format_param( + value=response_format, + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + optional_params=optional_params, + non_default_params={"response_format": response_format}, + is_thinking_enabled=False, + ) - # Should have outputConfig, NOT tools - assert "outputConfig" in result - assert "tools" not in result - assert "tool_choice" not in result - assert result["json_mode"] is True - # No fake_stream for native approach - assert "fake_stream" not in result + # Should have outputConfig, NOT tools + assert "outputConfig" in result + assert "tools" not in result + assert "tool_choice" not in result + assert result["json_mode"] is True + # No fake_stream for native approach + assert "fake_stream" not in result - # Verify the schema content (additionalProperties: false is added by normalization) - schema_str = result["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"] - parsed_schema = json.loads(schema_str) - expected_schema = {**response_format["json_schema"]["schema"], "additionalProperties": False} - assert parsed_schema == expected_schema - assert ( - result["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["name"] - == "WeatherResult" - ) + # Verify the schema content (additionalProperties: false is added by normalization) + schema_str = result["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"] + parsed_schema = json.loads(schema_str) + expected_schema = {**response_format["json_schema"]["schema"], "additionalProperties": False} + assert parsed_schema == expected_schema + assert result["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["name"] == "WeatherResult" + finally: + litellm.model_cost = old_cost + if old_env is None: + os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) + else: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env def test_translate_response_format_fallback_tool_call(): @@ -3118,42 +2813,53 @@ def test_translate_response_format_fallback_tool_call(): def test_native_structured_output_no_fake_stream(): """When using native structured outputs with streaming, fake_stream should NOT be set.""" - config = AmazonConverseConfig() + old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") + old_cost = litellm.model_cost + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + try: + config = AmazonConverseConfig() - response_format = { - "type": "json_schema", - "json_schema": { - "name": "Result", - "schema": { - "type": "object", - "properties": { - "answer": {"type": "string"}, + response_format = { + "type": "json_schema", + "json_schema": { + "name": "Result", + "schema": { + "type": "object", + "properties": { + "answer": {"type": "string"}, + }, }, }, - }, - } + } - optional_params: dict = {} - result = config._translate_response_format_param( - value=response_format, - model="anthropic.claude-sonnet-4-5-20250929-v1:0", - optional_params=optional_params, - non_default_params={"response_format": response_format, "stream": True}, - is_thinking_enabled=False, - ) + optional_params: dict = {} + result = config._translate_response_format_param( + value=response_format, + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + optional_params=optional_params, + non_default_params={"response_format": response_format, "stream": True}, + is_thinking_enabled=False, + ) - assert "outputConfig" in result - assert result["json_mode"] is True - # No fake_stream for native approach - assert "fake_stream" not in result + assert "outputConfig" in result + assert result["json_mode"] is True + # No fake_stream for native approach + assert "fake_stream" not in result - # Verify the schema content - schema_str = result["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"] - assert json.loads(schema_str) == { - "type": "object", - "properties": {"answer": {"type": "string"}}, - "additionalProperties": False, - } + # Verify the schema content + schema_str = result["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"] + assert json.loads(schema_str) == { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "additionalProperties": False, + } + finally: + litellm.model_cost = old_cost + if old_env is None: + os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) + else: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env def test_transform_request_with_output_config(): @@ -3227,11 +2933,7 @@ def test_transform_response_native_structured_output(): "output": { "message": { "role": "assistant", - "content": [ - { - "text": '{"temp": 62, "description": "Mild and foggy"}' - } - ], + "content": [{"text": '{"temp": 62, "description": "Mild and foggy"}'}], } }, "stopReason": "end_turn", @@ -3388,23 +3090,34 @@ def test_add_additional_properties_definitions(): def test_json_object_no_schema_falls_back_to_tool_call(): """response_format: {type: json_object} with no schema should use tool-call fallback, even for models that support native structured outputs.""" - config = AmazonConverseConfig() - optional_params: dict = {} - non_default_params = {"response_format": {"type": "json_object"}} + old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") + old_cost = litellm.model_cost + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + try: + config = AmazonConverseConfig() + optional_params: dict = {} + non_default_params = {"response_format": {"type": "json_object"}} - result = config._translate_response_format_param( - value=non_default_params["response_format"], - model="anthropic.claude-sonnet-4-5-20250929-v1:0", - optional_params=optional_params, - non_default_params=non_default_params, - is_thinking_enabled=False, - ) + result = config._translate_response_format_param( + value=non_default_params["response_format"], + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + optional_params=optional_params, + non_default_params=non_default_params, + is_thinking_enabled=False, + ) - # Should NOT use native outputConfig (no schema provided) - assert "outputConfig" not in result - # Should use tool-call fallback - assert "tools" in result - assert result["json_mode"] is True + # Should NOT use native outputConfig (no schema provided) + assert "outputConfig" not in result + # Should use tool-call fallback + assert "tools" in result + assert result["json_mode"] is True + finally: + litellm.model_cost = old_cost + if old_env is None: + os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) + else: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env def test_output_config_applies_additional_properties(): @@ -3427,7 +3140,6 @@ def test_output_config_applies_additional_properties(): assert parsed["properties"]["nested"]["additionalProperties"] is False - _TOOL_PARAM = [ { "type": "function", @@ -3505,9 +3217,7 @@ def test_parallel_tool_calls_older_model_drops_disable_flag(): class TestBedrockMinThinkingBudgetTokens: """Test that thinking.budget_tokens is clamped to the Bedrock minimum (1024).""" - def _map_params( - self, thinking_value, model="anthropic.claude-3-7-sonnet-20250219-v1:0" - ): + def _map_params(self, thinking_value, model="anthropic.claude-3-7-sonnet-20250219-v1:0"): """Helper to call map_openai_params with the given thinking value.""" config = AmazonConverseConfig() non_default_params = {"thinking": thinking_value} @@ -3545,6 +3255,7 @@ class TestBedrockMinThinkingBudgetTokens: ) assert "thinking" not in result or result.get("thinking") is None + def test_transform_response_with_both_json_tool_call_and_real_tool(): """ When Bedrock returns BOTH json_tool_call AND a real tool (get_weather), @@ -3733,9 +3444,7 @@ def test_streaming_filters_json_tool_call_with_real_tools(): # Chunk 2: json_tool_call delta — should become text, not tool_use json_delta = ContentBlockDeltaEvent(toolUse={"input": '{"temp": 62}'}) - text_2, tool_use_2, _, _, _ = decoder._handle_converse_delta_event( - json_delta, index=0 - ) + text_2, tool_use_2, _, _, _ = decoder._handle_converse_delta_event(json_delta, index=0) assert text_2 == '{"temp": 62}' assert tool_use_2 is None @@ -3758,12 +3467,8 @@ def test_streaming_filters_json_tool_call_with_real_tools(): assert decoder.tool_calls_index == 0 # Chunk 5: real tool delta - real_delta = ContentBlockDeltaEvent( - toolUse={"input": '{"location": "SF"}'} - ) - text_5, tool_use_5, _, _, _ = decoder._handle_converse_delta_event( - real_delta, index=1 - ) + real_delta = ContentBlockDeltaEvent(toolUse={"input": '{"location": "SF"}'}) + text_5, tool_use_5, _, _, _ = decoder._handle_converse_delta_event(real_delta, index=1) assert text_5 == "" assert tool_use_5 is not None assert tool_use_5["function"]["arguments"] == '{"location": "SF"}' @@ -3796,9 +3501,7 @@ def test_streaming_without_json_mode_passes_all_tools(): # json_tool_call delta — should be a tool_use, not text json_delta = ContentBlockDeltaEvent(toolUse={"input": '{"data": 1}'}) - text, tool_use_delta, _, _, _ = decoder._handle_converse_delta_event( - json_delta, index=0 - ) + text, tool_use_delta, _, _, _ = decoder._handle_converse_delta_event(json_delta, index=0) assert text == "" assert tool_use_delta is not None assert tool_use_delta["function"]["arguments"] == '{"data": 1}' @@ -3899,4 +3602,3 @@ def test_cache_control_injection_tool_config_not_added_without_injection_point() tools = result["toolConfig"]["tools"] # No cachePoint should be appended assert all("cachePoint" not in tool for tool in tools) - diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index d3214a8801..802868aa7a 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -336,7 +336,9 @@ class TestOpenAIResponsesAPIConfig: ) assert isinstance(result, ImageGenerationPartialImageEvent) - assert result.type == ResponsesAPIStreamEvents.IMAGE_GENERATION_PARTIAL_IMAGE + assert ( + result.type == ResponsesAPIStreamEvents.IMAGE_GENERATION_PARTIAL_IMAGE + ) assert result.partial_image_index == idx assert result.b64_json == chunk["b64_json"] @@ -689,9 +691,7 @@ class TestTransformListInputItemsRequest: def test_openai_transform_compact_response_api_request_query_params_preserved(self): """Test compact URL construction preserves query params and appends path.""" # Setup - azure_style_api_base = ( - "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" - ) + azure_style_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" # Execute url, data = self.openai_config.transform_compact_response_api_request( @@ -731,12 +731,12 @@ class TestTransformListInputItemsRequest: def test_azure_transform_list_input_items_request_minimal(self): """Test Azure implementation with minimal parameters""" # Setup - azure_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" + AZURE_AI_API_BASE = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" # Execute url, params = self.azure_config.transform_list_input_items_request( response_id=self.response_id, - api_base=azure_api_base, + api_base=AZURE_AI_API_BASE, litellm_params=self.litellm_params, headers=self.headers, ) @@ -749,12 +749,12 @@ class TestTransformListInputItemsRequest: def test_azure_transform_list_input_items_request_url_construction(self): """Test Azure implementation URL construction with response_id in path""" # Setup - azure_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" + AZURE_AI_API_BASE = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" # Execute url, params = self.azure_config.transform_list_input_items_request( response_id=self.response_id, - api_base=azure_api_base, + api_base=AZURE_AI_API_BASE, litellm_params=self.litellm_params, headers=self.headers, ) @@ -768,12 +768,12 @@ class TestTransformListInputItemsRequest: def test_azure_transform_list_input_items_request_with_all_params(self): """Test Azure implementation with all optional parameters""" # Setup - azure_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" + AZURE_AI_API_BASE = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" # Execute url, params = self.azure_config.transform_list_input_items_request( response_id=self.response_id, - api_base=azure_api_base, + api_base=AZURE_AI_API_BASE, litellm_params=self.litellm_params, headers=self.headers, after="cursor_after_123", @@ -1128,9 +1128,9 @@ class TestPhaseParameter: phase = getattr(output_item, "phase", None) expected = "commentary" if idx == 0 else "final_answer" - assert phase == expected, ( - f"output[{idx}] phase={phase!r}, expected {expected!r}" - ) + assert ( + phase == expected + ), f"output[{idx}] phase={phase!r}, expected {expected!r}" def test_streaming_output_item_done_preserves_phase(self): """OutputItemDoneEvent must preserve phase on its item.""" diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py index 67863be6be..5b18618fdf 100644 --- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -9,7 +9,7 @@ import copy import json from typing import Any, Dict, List -from unittest.mock import patch, Mock, MagicMock +from unittest.mock import AsyncMock, patch, Mock, MagicMock import httpx import pytest @@ -444,7 +444,19 @@ class TestSnowFlakeCompletion: os.environ.pop("SNOWFLAKE_JWT", None) -def _mock_snowflake_chat_response() -> Dict[str, Any]: +FAKE_API_BASE = "https://fake-snowflake.example.com/api/v2/cortex/inference:chat" + + +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 + + +def _chat_response() -> Dict[str, Any]: return { "id": "chatcmpl-snowflake-123", "object": "chat.completion", @@ -468,7 +480,7 @@ def _mock_snowflake_chat_response() -> Dict[str, Any]: } -def _mock_snowflake_streaming_chunks() -> List[str]: +def _streaming_chunks() -> List[str]: base = { "id": "chatcmpl-snowflake-stream-123", "object": "chat.completion.chunk", @@ -480,13 +492,20 @@ def _mock_snowflake_streaming_chunks() -> List[str]: {"content": " sky"}, {"content": " is blue"}, ] - finish_reasons = [None, None, "stop"] - return [ - json.dumps( - {**base, "choices": [{"index": 0, "delta": d, "finish_reason": fr}]} + chunks = [] + for i, delta in enumerate(deltas): + finish = "stop" if i == len(deltas) - 1 else None + chunks.append( + json.dumps( + { + **base, + "choices": [ + {"index": 0, "delta": delta, "finish_reason": finish} + ], + } + ) ) - for d, fr in zip(deltas, finish_reasons) - ] + return chunks class TestSnowflakeChatCompletion: @@ -496,54 +515,56 @@ class TestSnowflakeChatCompletion: @pytest.mark.parametrize("sync_mode", [True, False]) def test_chat_completion_snowflake(self, sync_mode): - mock_response = Mock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = _mock_snowflake_chat_response() + mock_resp = _make_mock_response(_chat_response()) if sync_mode: - handler = HTTPHandler() - with patch.object(HTTPHandler, "post", return_value=mock_response): + with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: response = completion( model="snowflake/mistral-7b", messages=self.messages, api_key="fake-jwt", account_id="FAKE-ACCOUNT", - client=handler, + api_base=FAKE_API_BASE, ) + mock_post.assert_called_once() else: - handler = AsyncHTTPHandler() - with patch.object(AsyncHTTPHandler, "post", return_value=mock_response): + with patch.object( + AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp + ) as mock_post: response = asyncio.run( acompletion( model="snowflake/mistral-7b", messages=self.messages, api_key="fake-jwt", account_id="FAKE-ACCOUNT", - client=handler, + api_base=FAKE_API_BASE, ) ) + mock_post.assert_called_once() assert response is not None assert response.choices[0].message.content is not None assert "sky" in response.choices[0].message.content.lower() + assert response.usage.prompt_tokens == 10 + assert response.usage.completion_tokens == 30 @pytest.mark.parametrize("sync_mode", [True, False]) def test_chat_completion_snowflake_stream(self, sync_mode): - mock_chunks = _mock_snowflake_streaming_chunks() + raw_chunks = _streaming_chunks() if sync_mode: - handler = HTTPHandler() - def mock_iter_lines(): - for chunk in mock_chunks: + def _iter_lines(): + for chunk in raw_chunks: yield f"data: {chunk}" yield "data: [DONE]" mock_resp = MagicMock() - mock_resp.iter_lines.side_effect = mock_iter_lines + 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): + with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: response = completion( model="snowflake/mistral-7b", messages=self.messages, @@ -551,25 +572,26 @@ class TestSnowflakeChatCompletion: stream=True, api_key="fake-jwt", account_id="FAKE-ACCOUNT", - client=handler, + api_base=FAKE_API_BASE, ) chunks_received = list(response) - assert len(chunks_received) > 0 + mock_post.assert_called_once() else: - handler = AsyncHTTPHandler() - async def mock_iter_lines(): - for chunk in mock_chunks: + async def _aiter_lines(): + for chunk in raw_chunks: yield f"data: {chunk}" yield "data: [DONE]" mock_resp = MagicMock() - mock_resp.iter_lines.side_effect = mock_iter_lines + mock_resp.aiter_lines.return_value = _aiter_lines() mock_resp.status_code = 200 + mock_resp.headers = {"content-type": "text/event-stream"} - with patch.object(AsyncHTTPHandler, "post", return_value=mock_resp): - - async def run(): + async def _run(): + with patch.object( + AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp + ) as mock_post: resp = await acompletion( model="snowflake/mistral-7b", messages=self.messages, @@ -577,9 +599,17 @@ class TestSnowflakeChatCompletion: stream=True, api_key="fake-jwt", account_id="FAKE-ACCOUNT", - client=handler, + api_base=FAKE_API_BASE, ) - return [chunk async for chunk in resp] + 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 + 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 + ) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index ada67fbba8..5303da6fbc 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -305,24 +305,28 @@ async def test_sync_user_role_and_teams(): # Create mock objects for required types mock_user_api_key_cache = MagicMock() mock_proxy_logging_obj = MagicMock() - + jwt_handler = JWTHandler() jwt_handler.update_environment( prisma_client=None, user_api_key_cache=mock_user_api_key_cache, litellm_jwtauth=LiteLLM_JWTAuth( jwt_litellm_role_map=[ - JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN) + JWTLiteLLMRoleMap( + jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN + ) ], roles_jwt_field="roles", team_ids_jwt_field="my_id_teams", - sync_user_role_and_teams=True + sync_user_role_and_teams=True, ), ) token = {"roles": ["ADMIN"], "my_id_teams": ["team1", "team2"]} - user = LiteLLM_UserTable(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER.value, teams=["team2"]) + user = LiteLLM_UserTable( + user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER.value, teams=["team2"] + ) prisma = AsyncMock() prisma.db.litellm_usertable.update = AsyncMock() @@ -350,7 +354,9 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_role_change(): user_api_key_cache=AsyncMock(), litellm_jwtauth=LiteLLM_JWTAuth( jwt_litellm_role_map=[ - JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN) + JWTLiteLLMRoleMap( + jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN + ) ], roles_jwt_field="roles", team_ids_jwt_field="my_id_teams", @@ -375,7 +381,9 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_role_change(): mock_cache.async_set_cache.assert_called_once() call_kwargs = mock_cache.async_set_cache.call_args assert call_kwargs.kwargs["key"] == "u1" - assert call_kwargs.kwargs["value"]["user_role"] == LitellmUserRoles.PROXY_ADMIN.value + assert ( + call_kwargs.kwargs["value"]["user_role"] == LitellmUserRoles.PROXY_ADMIN.value + ) @pytest.mark.asyncio @@ -389,7 +397,9 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_team_change(): user_api_key_cache=AsyncMock(), litellm_jwtauth=LiteLLM_JWTAuth( jwt_litellm_role_map=[ - JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN) + JWTLiteLLMRoleMap( + jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN + ) ], roles_jwt_field="roles", team_ids_jwt_field="my_id_teams", @@ -432,7 +442,9 @@ async def test_sync_user_role_and_teams_no_cache_write_when_nothing_changes(): user_api_key_cache=AsyncMock(), litellm_jwtauth=LiteLLM_JWTAuth( jwt_litellm_role_map=[ - JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN) + JWTLiteLLMRoleMap( + jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN + ) ], roles_jwt_field="roles", team_ids_jwt_field="my_id_teams", @@ -463,7 +475,7 @@ async def test_map_jwt_role_to_litellm_role(): # Create mock objects for required types mock_user_api_key_cache = MagicMock() - + jwt_handler = JWTHandler() jwt_handler.update_environment( prisma_client=None, @@ -471,13 +483,21 @@ async def test_map_jwt_role_to_litellm_role(): litellm_jwtauth=LiteLLM_JWTAuth( jwt_litellm_role_map=[ # Exact match - JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN), + JWTLiteLLMRoleMap( + jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN + ), # Wildcard patterns - JWTLiteLLMRoleMap(jwt_role="user_*", litellm_role=LitellmUserRoles.INTERNAL_USER), - JWTLiteLLMRoleMap(jwt_role="team_?", litellm_role=LitellmUserRoles.TEAM), - JWTLiteLLMRoleMap(jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER), + JWTLiteLLMRoleMap( + jwt_role="user_*", litellm_role=LitellmUserRoles.INTERNAL_USER + ), + JWTLiteLLMRoleMap( + jwt_role="team_?", litellm_role=LitellmUserRoles.TEAM + ), + JWTLiteLLMRoleMap( + jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER + ), ], - roles_jwt_field="roles" + roles_jwt_field="roles", ), ) @@ -547,7 +567,9 @@ async def test_map_jwt_role_to_litellm_role(): # Test patterns that don't match character classes jwt_handler.litellm_jwtauth.jwt_litellm_role_map = [ - JWTLiteLLMRoleMap(jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER), + JWTLiteLLMRoleMap( + jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER + ), ] token = {"roles": ["dev_4"]} # 4 is not in [123] result = jwt_handler.map_jwt_role_to_litellm_role(token) @@ -570,7 +592,7 @@ async def test_map_jwt_role_to_litellm_role(): async def test_nested_jwt_field_access(): """ Test that all JWT fields support dot notation for nested access - + This test verifies that: 1. All JWT field methods can access nested values using dot notation 2. Backward compatibility is maintained for flat field names @@ -581,33 +603,18 @@ async def test_nested_jwt_field_access(): # Create JWT handler jwt_handler = JWTHandler() - + # Test token with nested claims nested_token = { - "user": { - "sub": "u123", - "email": "user@example.com" - }, - "resource_access": { - "my-client": { - "roles": ["admin", "user"] - } - }, + "user": {"sub": "u123", "email": "user@example.com"}, + "resource_access": {"my-client": {"roles": ["admin", "user"]}}, "groups": ["team1", "team2"], - "organization": { - "id": "org456" - }, - "profile": { - "object_id": "obj789" - }, - "customer": { - "end_user_id": "customer123" - }, - "tenant": { - "team_id": "team456" - } + "organization": {"id": "org456"}, + "profile": {"object_id": "obj789"}, + "customer": {"end_user_id": "customer123"}, + "tenant": {"team_id": "team456"}, } - + # Test flat token for backward compatibility flat_token = { "sub": "u123", @@ -617,13 +624,13 @@ async def test_nested_jwt_field_access(): "org_id": "org456", "object_id": "obj789", "end_user_id": "customer123", - "team_id": "team456" + "team_id": "team456", } # Test 1: user_id_jwt_field with nested access jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_id_jwt_field="user.sub") assert jwt_handler.get_user_id(nested_token, None) == "u123" - + # Test 1b: user_id_jwt_field with flat access (backward compatibility) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_id_jwt_field="sub") assert jwt_handler.get_user_id(flat_token, None) == "u123" @@ -631,7 +638,7 @@ async def test_nested_jwt_field_access(): # Test 2: user_email_jwt_field with nested access jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_email_jwt_field="user.email") assert jwt_handler.get_user_email(nested_token, None) == "user@example.com" - + # Test 2b: user_email_jwt_field with flat access (backward compatibility) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_email_jwt_field="email") assert jwt_handler.get_user_email(flat_token, None) == "user@example.com" @@ -639,7 +646,7 @@ async def test_nested_jwt_field_access(): # Test 3: team_ids_jwt_field with nested access jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups") assert jwt_handler.get_team_ids_from_jwt(nested_token) == ["team1", "team2"] - + # Test 3b: team_ids_jwt_field with flat access (backward compatibility) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups") assert jwt_handler.get_team_ids_from_jwt(flat_token) == ["team1", "team2"] @@ -647,30 +654,37 @@ async def test_nested_jwt_field_access(): # Test 4: org_id_jwt_field with nested access jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_id_jwt_field="organization.id") assert jwt_handler.get_org_id(nested_token, None) == "org456" - + # Test 4b: org_id_jwt_field with flat access (backward compatibility) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_id_jwt_field="org_id") assert jwt_handler.get_org_id(flat_token, None) == "org456" # Test 5: object_id_jwt_field with nested access (requires role_mappings) from litellm.proxy._types import LitellmUserRoles, RoleMapping + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( object_id_jwt_field="profile.object_id", - role_mappings=[RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER)] + role_mappings=[ + RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER) + ], ) assert jwt_handler.get_object_id(nested_token, None) == "obj789" - + # Test 5b: object_id_jwt_field with flat access (backward compatibility) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( object_id_jwt_field="object_id", - role_mappings=[RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER)] + role_mappings=[ + RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER) + ], ) assert jwt_handler.get_object_id(flat_token, None) == "obj789" # Test 6: end_user_id_jwt_field with nested access - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(end_user_id_jwt_field="customer.end_user_id") + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + end_user_id_jwt_field="customer.end_user_id" + ) assert jwt_handler.get_end_user_id(nested_token, None) == "customer123" - + # Test 6b: end_user_id_jwt_field with flat access (backward compatibility) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(end_user_id_jwt_field="end_user_id") assert jwt_handler.get_end_user_id(flat_token, None) == "customer123" @@ -678,19 +692,21 @@ async def test_nested_jwt_field_access(): # Test 7: team_id_jwt_field with nested access jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="tenant.team_id") assert jwt_handler.get_team_id(nested_token, None) == "team456" - + # Test 7b: team_id_jwt_field with flat access (backward compatibility) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id") assert jwt_handler.get_team_id(flat_token, None) == "team456" # Test 8: roles_jwt_field with deeply nested access (already supported, but testing) - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(roles_jwt_field="resource_access.my-client.roles") + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + roles_jwt_field="resource_access.my-client.roles" + ) assert jwt_handler.get_jwt_role(nested_token, []) == ["admin", "user"] # Test 9: user_roles_jwt_field with nested access (already supported, but testing) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( user_roles_jwt_field="resource_access.my-client.roles", - user_allowed_roles=["admin", "user"] + user_allowed_roles=["admin", "user"], ) assert jwt_handler.get_user_roles(nested_token, []) == ["admin", "user"] @@ -699,7 +715,7 @@ async def test_nested_jwt_field_access(): async def test_nested_jwt_field_missing_paths(): """ Test handling of missing nested paths in JWT tokens - + This test verifies that: 1. Missing nested paths return appropriate defaults 2. Partial paths that exist but don't have the final key return defaults @@ -710,7 +726,7 @@ async def test_nested_jwt_field_missing_paths(): # Create JWT handler jwt_handler = JWTHandler() - + # Test token with missing nested paths incomplete_token = { "user": { @@ -718,9 +734,7 @@ async def test_nested_jwt_field_missing_paths(): # missing "sub" and "email" }, "resource_access": { - "other-client": { - "roles": ["viewer"] - } + "other-client": {"roles": ["viewer"]} # missing "my-client" } # missing "organization", "profile", "customer", "tenant", "groups" @@ -732,7 +746,10 @@ async def test_nested_jwt_field_missing_paths(): # Test 2: Missing user.email should return default jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_email_jwt_field="user.email") - assert jwt_handler.get_user_email(incomplete_token, "default@example.com") == "default@example.com" + assert ( + jwt_handler.get_user_email(incomplete_token, "default@example.com") + == "default@example.com" + ) # Test 3: Missing groups should return empty list jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups") @@ -744,40 +761,53 @@ async def test_nested_jwt_field_missing_paths(): # Test 5: Missing profile.object_id should return default (requires role_mappings) from litellm.proxy._types import LitellmUserRoles, RoleMapping + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( object_id_jwt_field="profile.object_id", - role_mappings=[RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER)] + role_mappings=[ + RoleMapping(role="admin", internal_role=LitellmUserRoles.INTERNAL_USER) + ], ) assert jwt_handler.get_object_id(incomplete_token, "default_obj") == "default_obj" # Test 6: Missing customer.end_user_id should return default - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(end_user_id_jwt_field="customer.end_user_id") - assert jwt_handler.get_end_user_id(incomplete_token, "default_customer") == "default_customer" + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + end_user_id_jwt_field="customer.end_user_id" + ) + assert ( + jwt_handler.get_end_user_id(incomplete_token, "default_customer") + == "default_customer" + ) # Test 7: Missing tenant.team_id should use team_id_default fallback jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - team_id_jwt_field="tenant.team_id", - team_id_default="fallback_team" + team_id_jwt_field="tenant.team_id", team_id_default="fallback_team" ) assert jwt_handler.get_team_id(incomplete_token, "default_team") == "fallback_team" # Test 8: Missing resource_access.my-client.roles should return default - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(roles_jwt_field="resource_access.my-client.roles") - assert jwt_handler.get_jwt_role(incomplete_token, ["default_role"]) == ["default_role"] + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + roles_jwt_field="resource_access.my-client.roles" + ) + assert jwt_handler.get_jwt_role(incomplete_token, ["default_role"]) == [ + "default_role" + ] # Test 9: Missing nested user roles should return default jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( user_roles_jwt_field="resource_access.my-client.roles", - user_allowed_roles=["admin", "user"] + user_allowed_roles=["admin", "user"], ) - assert jwt_handler.get_user_roles(incomplete_token, ["default_user_role"]) == ["default_user_role"] + assert jwt_handler.get_user_roles(incomplete_token, ["default_user_role"]) == [ + "default_user_role" + ] -@pytest.mark.asyncio +@pytest.mark.asyncio async def test_metadata_prefix_handling_in_nested_fields(): """ Test that metadata. prefix is properly handled in nested JWT field access - + The get_nested_value function should remove metadata. prefix before traversing """ from litellm.proxy._types import LiteLLM_JWTAuth @@ -785,17 +815,19 @@ async def test_metadata_prefix_handling_in_nested_fields(): # Create JWT handler jwt_handler = JWTHandler() - + # Test token with proper structure for metadata prefix removal token = { "user": { "email": "user@example.com" # This will be accessed when metadata.user.email is used }, - "sub": "u123" + "sub": "u123", } # Test 1: metadata.user.email should access user.email after prefix removal - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(user_email_jwt_field="metadata.user.email") + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + user_email_jwt_field="metadata.user.email" + ) # The get_nested_value function removes "metadata." prefix, so "metadata.user.email" becomes "user.email" assert jwt_handler.get_user_email(token, None) == "user@example.com" @@ -871,24 +903,21 @@ async def test_auth_builder_returns_team_membership_object(): # Create mock objects from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamMembership - + mock_team_membership = LiteLLM_TeamMembership( user_id=_user_id, team_id=_team_id, budget_id="budget_123", spend=10.5, litellm_budget_table=LiteLLM_BudgetTable( - budget_id="budget_123", - rpm_limit=100, - tpm_limit=5000 - ) + budget_id="budget_123", rpm_limit=100, tpm_limit=5000 + ), ) - + user_object = LiteLLM_UserTable( - user_id=_user_id, - user_role=LitellmUserRoles.INTERNAL_USER + user_id=_user_id, user_role=LitellmUserRoles.INTERNAL_USER ) - + team_object = LiteLLM_TeamTable(team_id=_team_id) # Create mock JWT handler @@ -958,12 +987,24 @@ async def test_auth_builder_returns_team_membership_object(): ) # Verify that team_membership_object is returned - assert result["team_membership"] is not None, "team_membership should be present" - assert result["team_membership"] == mock_team_membership, "team_membership should match the mock object" - assert result["team_membership"].user_id == _user_id, "team_membership user_id should match" - assert result["team_membership"].team_id == _team_id, "team_membership team_id should match" - assert result["team_membership"].budget_id == "budget_123", "team_membership budget_id should match" - assert result["team_membership"].spend == 10.5, "team_membership spend should match" + assert ( + result["team_membership"] is not None + ), "team_membership should be present" + assert ( + result["team_membership"] == mock_team_membership + ), "team_membership should match the mock object" + assert ( + result["team_membership"].user_id == _user_id + ), "team_membership user_id should match" + assert ( + result["team_membership"].team_id == _team_id + ), "team_membership team_id should match" + assert ( + result["team_membership"].budget_id == "budget_123" + ), "team_membership budget_id should match" + assert ( + result["team_membership"].spend == 10.5 + ), "team_membership spend should match" @pytest.mark.asyncio @@ -979,16 +1020,16 @@ async def test_auth_builder_with_oidc_userinfo_enabled(): request_data = {"model": "gpt-4"} general_settings = {"enforce_rbac": False} route = "/chat/completions" - + user_object = LiteLLM_UserTable( user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER ) - + # Create JWT handler with OIDC UserInfo enabled jwt_handler = JWTHandler() user_api_key_cache = DualCache() proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) - + jwt_handler.update_environment( prisma_client=None, user_api_key_cache=user_api_key_cache, @@ -999,14 +1040,14 @@ async def test_auth_builder_with_oidc_userinfo_enabled(): user_email_jwt_field="email", ), ) - + # Mock OIDC UserInfo response userinfo_response = { "sub": "test_user_1", "email": "test@example.com", "scope": "", } - + # Mock all the dependencies with patch.object( jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock @@ -1057,7 +1098,7 @@ async def test_auth_builder_with_oidc_userinfo_enabled(): ) as mock_sync_user: # Set up mock return values mock_get_userinfo.return_value = userinfo_response - + # Call auth_builder result = await JWTAuthManager.auth_builder( api_key=api_key, @@ -1070,11 +1111,11 @@ async def test_auth_builder_with_oidc_userinfo_enabled(): parent_otel_span=None, proxy_logging_obj=proxy_logging_obj, ) - + # Verify that get_oidc_userinfo was called instead of auth_jwt mock_get_userinfo.assert_called_once_with(token=api_key) mock_auth_jwt.assert_not_called() # Should not be called when OIDC is enabled - + # Verify the result assert result["user_id"] == "test_user_1" assert result["user_object"] == user_object @@ -1093,16 +1134,16 @@ async def test_auth_builder_with_oidc_userinfo_disabled(): request_data = {"model": "gpt-4"} general_settings = {"enforce_rbac": False} route = "/chat/completions" - + user_object = LiteLLM_UserTable( user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER ) - + # Create JWT handler with OIDC UserInfo disabled jwt_handler = JWTHandler() user_api_key_cache = DualCache() proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) - + jwt_handler.update_environment( prisma_client=None, user_api_key_cache=user_api_key_cache, @@ -1111,13 +1152,13 @@ async def test_auth_builder_with_oidc_userinfo_disabled(): user_id_jwt_field="sub", ), ) - + # Mock JWT validation response jwt_response = { "sub": "test_user_1", "scope": "", } - + # Mock all the dependencies with patch.object( jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock @@ -1168,7 +1209,7 @@ async def test_auth_builder_with_oidc_userinfo_disabled(): ) as mock_sync_user: # Set up mock return values mock_auth_jwt.return_value = jwt_response - + # Call auth_builder result = await JWTAuthManager.auth_builder( api_key=api_key, @@ -1181,16 +1222,125 @@ async def test_auth_builder_with_oidc_userinfo_disabled(): parent_otel_span=None, proxy_logging_obj=proxy_logging_obj, ) - + # Verify that auth_jwt was called instead of get_oidc_userinfo mock_auth_jwt.assert_called_once_with(token=api_key) mock_get_userinfo.assert_not_called() # Should not be called when OIDC is disabled - + # Verify the result assert result["user_id"] == "test_user_1" assert result["user_object"] == user_object +@pytest.mark.asyncio +async def test_auth_builder_oidc_enabled_falls_back_to_jwt_auth_for_jwt_tokens(): + """ + Regression test for the is_jwt routing fix. + + When oidc_userinfo_enabled=True but the supplied token is a well-formed + JWT (three dot-separated parts), auth_builder must call auth_jwt and skip + get_oidc_userinfo. Sending a standard JWT to the OIDC UserInfo endpoint + is incorrect — the endpoint expects an opaque access token. + """ + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + # Three-part token: recognised as a JWT by is_jwt() + api_key = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ0ZXN0X3VzZXIifQ.some_signature" + request_data = {"model": "gpt-4"} + general_settings = {"enforce_rbac": False} + route = "/chat/completions" + + user_object = LiteLLM_UserTable( + user_id="test_user_1", user_role=LitellmUserRoles.INTERNAL_USER + ) + + jwt_handler = JWTHandler() + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=user_api_key_cache, + litellm_jwtauth=LiteLLM_JWTAuth( + oidc_userinfo_enabled=True, + oidc_userinfo_endpoint="https://example.com/oauth2/userinfo", + user_id_jwt_field="sub", + ), + ) + + jwt_response = {"sub": "test_user_1", "scope": ""} + + with patch.object( + jwt_handler, "get_oidc_userinfo", new_callable=AsyncMock + ) as mock_get_userinfo, patch.object( + jwt_handler, "auth_jwt", new_callable=AsyncMock + ) as mock_auth_jwt, patch.object( + JWTAuthManager, "check_rbac_role", new_callable=AsyncMock + ), patch.object( + jwt_handler, "get_rbac_role", return_value=None + ), patch.object( + jwt_handler, "get_scopes", return_value=[] + ), patch.object( + jwt_handler, "get_object_id", return_value=None + ), patch.object( + JWTAuthManager, + "get_user_info", + new_callable=AsyncMock, + return_value=("test_user_1", None, None), + ), patch.object( + jwt_handler, "get_org_id", return_value=None + ), patch.object( + jwt_handler, "get_end_user_id", return_value=None + ), patch.object( + JWTAuthManager, "check_admin_access", new_callable=AsyncMock, return_value=None + ), patch.object( + JWTAuthManager, + "find_and_validate_specific_team_id", + new_callable=AsyncMock, + return_value=(None, None), + ), patch.object( + JWTAuthManager, "get_all_team_ids", return_value=set() + ), patch.object( + JWTAuthManager, + "find_team_with_model_access", + new_callable=AsyncMock, + return_value=(None, None), + ), patch.object( + JWTAuthManager, + "get_objects", + new_callable=AsyncMock, + return_value=(user_object, None, None, None), + ), patch.object( + JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock + ), patch.object( + JWTAuthManager, "validate_object_id", return_value=True + ), patch.object( + JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock + ): + mock_auth_jwt.return_value = jwt_response + + result = await JWTAuthManager.auth_builder( + api_key=api_key, + jwt_handler=jwt_handler, + request_data=request_data, + general_settings=general_settings, + route=route, + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + # Token is a JWT, so standard JWT auth must be used even when + # oidc_userinfo_enabled is True. + mock_auth_jwt.assert_called_once_with(token=api_key) + mock_get_userinfo.assert_not_called() + + assert result["user_id"] == "test_user_1" + assert result["user_object"] == user_object + + def test_get_team_id_from_header(): """Test get_team_id_from_header returns team when valid, None when missing, raises on invalid.""" from fastapi import HTTPException @@ -1236,17 +1386,33 @@ async def test_auth_builder_uses_team_from_header_e2e(): ) team_object = LiteLLM_TeamTable(team_id="team-2") - user_object = LiteLLM_UserTable(user_id="user-1", user_role=LitellmUserRoles.INTERNAL_USER) + user_object = LiteLLM_UserTable( + user_id="user-1", user_role=LitellmUserRoles.INTERNAL_USER + ) - with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, \ - patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), \ - patch.object(JWTAuthManager, "check_admin_access", new_callable=AsyncMock, return_value=None), \ - patch("litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock) as mock_get_team, \ - patch.object(JWTAuthManager, "get_objects", new_callable=AsyncMock, return_value=(user_object, None, None, None)), \ - patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), \ - patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock): - - mock_auth_jwt.return_value = {"sub": "user-1", "scope": "", "groups": ["team-1", "team-2"]} + with patch.object( + jwt_handler, "auth_jwt", new_callable=AsyncMock + ) as mock_auth_jwt, patch.object( + JWTAuthManager, "check_rbac_role", new_callable=AsyncMock + ), patch.object( + JWTAuthManager, "check_admin_access", new_callable=AsyncMock, return_value=None + ), patch( + "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock + ) as mock_get_team, patch.object( + JWTAuthManager, + "get_objects", + new_callable=AsyncMock, + return_value=(user_object, None, None, None), + ), patch.object( + JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock + ), patch.object( + JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock + ): + mock_auth_jwt.return_value = { + "sub": "user-1", + "scope": "", + "groups": ["team-1", "team-2"], + } mock_get_team.return_value = team_object result = await JWTAuthManager.auth_builder( @@ -1275,29 +1441,29 @@ async def test_get_team_alias_with_nested_fields(): from litellm.proxy.auth.handle_jwt import JWTHandler jwt_handler = JWTHandler() - + # Test token with nested team name nested_token = { - "organization": { - "team": { - "name": "engineering-team" - } - }, - "team_name": "flat-team" + "organization": {"team": {"name": "engineering-team"}}, + "team_name": "flat-team", } - + # Test nested access - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="organization.team.name") + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + team_alias_jwt_field="organization.team.name" + ) assert jwt_handler.get_team_alias(nested_token, None) == "engineering-team" - + # Test flat access (backward compatibility) jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="team_name") assert jwt_handler.get_team_alias(nested_token, None) == "flat-team" - + # Test missing field returns default - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="nonexistent.field") + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + team_alias_jwt_field="nonexistent.field" + ) assert jwt_handler.get_team_alias(nested_token, "default-team") == "default-team" - + # Test with team_alias_jwt_field not configured jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() # team_alias_jwt_field is None assert jwt_handler.get_team_alias(nested_token, "default") is None @@ -1312,23 +1478,22 @@ async def test_is_required_team_id_with_team_alias_field(): from litellm.proxy.auth.handle_jwt import JWTHandler jwt_handler = JWTHandler() - + # Neither field set - should return False jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() assert jwt_handler.is_required_team_id() is False - + # Only team_id_jwt_field set - should return True jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id") assert jwt_handler.is_required_team_id() is True - + # Only team_alias_jwt_field set - should return True jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="team_name") assert jwt_handler.is_required_team_id() is True - + # Both fields set - should return True jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - team_id_jwt_field="team_id", - team_alias_jwt_field="team_name" + team_id_jwt_field="team_id", team_alias_jwt_field="team_name" ) assert jwt_handler.is_required_team_id() is True @@ -1348,30 +1513,24 @@ async def test_find_and_validate_specific_team_id_with_team_alias(): jwt_handler = JWTHandler() user_api_key_cache = DualCache() proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) - + jwt_handler.update_environment( prisma_client=None, user_api_key_cache=user_api_key_cache, - litellm_jwtauth=LiteLLM_JWTAuth( - team_alias_jwt_field="team_alias" - ), + litellm_jwtauth=LiteLLM_JWTAuth(team_alias_jwt_field="team_alias"), ) - + # Token with team name (no team_id) - jwt_token = { - "sub": "user-1", - "team_alias": "my-team" - } - + jwt_token = {"sub": "user-1", "team_alias": "my-team"} + # Mock team object returned by get_team_object_by_alias team_object = LiteLLM_TeamTable(team_id="resolved-team-id", team_alias="my-team") - + with patch( - "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", - new_callable=AsyncMock + "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", new_callable=AsyncMock ) as mock_get_by_alias: mock_get_by_alias.return_value = team_object - + team_id, result_team = await JWTAuthManager.find_and_validate_specific_team_id( jwt_handler=jwt_handler, jwt_valid_token=jwt_token, @@ -1380,7 +1539,7 @@ async def test_find_and_validate_specific_team_id_with_team_alias(): parent_otel_span=None, proxy_logging_obj=proxy_logging_obj, ) - + # Should have resolved team_id from team name assert team_id == "resolved-team-id" assert result_team == team_object @@ -1408,35 +1567,28 @@ async def test_find_and_validate_team_id_takes_precedence_over_name(): jwt_handler = JWTHandler() user_api_key_cache = DualCache() proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) - + jwt_handler.update_environment( prisma_client=None, user_api_key_cache=user_api_key_cache, litellm_jwtauth=LiteLLM_JWTAuth( - team_id_jwt_field="team_id", - team_alias_jwt_field="team_alias" + team_id_jwt_field="team_id", team_alias_jwt_field="team_alias" ), ) - + # Token with both team_id and team name - jwt_token = { - "sub": "user-1", - "team_id": "direct-team-id", - "team_alias": "my-team" - } - + jwt_token = {"sub": "user-1", "team_id": "direct-team-id", "team_alias": "my-team"} + # Mock team object returned by get_team_object (by ID) team_object = LiteLLM_TeamTable(team_id="direct-team-id") - + with patch( - "litellm.proxy.auth.handle_jwt.get_team_object", - new_callable=AsyncMock + "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock ) as mock_get_by_id, patch( - "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", - new_callable=AsyncMock + "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", new_callable=AsyncMock ) as mock_get_by_alias: mock_get_by_id.return_value = team_object - + team_id, result_team = await JWTAuthManager.find_and_validate_specific_team_id( jwt_handler=jwt_handler, jwt_valid_token=jwt_token, @@ -1445,7 +1597,7 @@ async def test_find_and_validate_team_id_takes_precedence_over_name(): parent_otel_span=None, proxy_logging_obj=proxy_logging_obj, ) - + # Should use team_id directly, not resolve by name assert team_id == "direct-team-id" assert result_team == team_object @@ -1466,7 +1618,7 @@ async def test_find_and_validate_raises_when_required_team_not_found(): jwt_handler = JWTHandler() user_api_key_cache = DualCache() proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) - + jwt_handler.update_environment( prisma_client=None, user_api_key_cache=user_api_key_cache, @@ -1474,12 +1626,10 @@ async def test_find_and_validate_raises_when_required_team_not_found(): team_alias_jwt_field="team_alias" # Required, but not in token ), ) - + # Token without team info - jwt_token = { - "sub": "user-1" - } - + jwt_token = {"sub": "user-1"} + with pytest.raises(Exception) as exc_info: await JWTAuthManager.find_and_validate_specific_team_id( jwt_handler=jwt_handler, @@ -1489,7 +1639,7 @@ async def test_find_and_validate_raises_when_required_team_not_found(): parent_otel_span=None, proxy_logging_obj=proxy_logging_obj, ) - + assert "No team found in token" in str(exc_info.value) assert "team_alias field 'team_alias'" in str(exc_info.value) @@ -1503,29 +1653,29 @@ async def test_get_org_alias_with_nested_fields(): from litellm.proxy.auth.handle_jwt import JWTHandler jwt_handler = JWTHandler() - + # Test token with nested org name nested_token = { - "company": { - "organization": { - "name": "acme-corp" - } - }, - "org_name": "flat-org" + "company": {"organization": {"name": "acme-corp"}}, + "org_name": "flat-org", } - + # Test nested access - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="company.organization.name") + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + org_alias_jwt_field="company.organization.name" + ) assert jwt_handler.get_org_alias(nested_token, None) == "acme-corp" - + # Test flat access jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="org_name") assert jwt_handler.get_org_alias(nested_token, None) == "flat-org" - + # Test missing field returns default - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="nonexistent.field") + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + org_alias_jwt_field="nonexistent.field" + ) assert jwt_handler.get_org_alias(nested_token, "default-org") == "default-org" - + # Test with org_alias_jwt_field not configured jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() assert jwt_handler.get_org_alias(nested_token, "default") is None @@ -1544,15 +1694,13 @@ async def test_get_objects_resolves_org_by_name(): jwt_handler = JWTHandler() user_api_key_cache = DualCache() proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) - + jwt_handler.update_environment( prisma_client=None, user_api_key_cache=user_api_key_cache, - litellm_jwtauth=LiteLLM_JWTAuth( - org_alias_jwt_field="org_alias" - ), + litellm_jwtauth=LiteLLM_JWTAuth(org_alias_jwt_field="org_alias"), ) - + # Mock org object returned by get_org_object_by_alias org_object = LiteLLM_OrganizationTable( organization_id="resolved-org-id", @@ -1560,15 +1708,14 @@ async def test_get_objects_resolves_org_by_name(): budget_id="budget-1", created_by="admin", updated_by="admin", - models=[] + models=[], ) - + with patch( - "litellm.proxy.auth.handle_jwt.get_org_object_by_alias", - new_callable=AsyncMock + "litellm.proxy.auth.handle_jwt.get_org_object_by_alias", new_callable=AsyncMock ) as mock_get_by_alias: mock_get_by_alias.return_value = org_object - + ( result_user_obj, result_org_obj, @@ -1589,7 +1736,7 @@ async def test_get_objects_resolves_org_by_name(): route="/chat/completions", org_alias="my-org", ) - + # Should resolve org by alias - org_id can be derived from org_object.organization_id assert result_org_obj == org_object assert result_org_obj.organization_id == "resolved-org-id" @@ -1643,7 +1790,9 @@ async def test_resolve_jwks_url_resolves_oidc_discovery_document(): litellm_jwtauth=LiteLLM_JWTAuth(), ) - discovery_url = "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" + discovery_url = ( + "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" + ) jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys" mock_response = MagicMock() @@ -1674,7 +1823,9 @@ async def test_resolve_jwks_url_caches_resolved_jwks_uri(): litellm_jwtauth=LiteLLM_JWTAuth(), ) - discovery_url = "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" + discovery_url = ( + "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" + ) jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys" mock_response = MagicMock() @@ -1818,9 +1969,9 @@ async def test_find_and_validate_specific_team_id_hints_bracket_notation(): error_msg = str(exc_info.value) # Should mention the bad field name and suggest the fix assert "roles.0" in error_msg, f"Expected field name in: {error_msg}" - assert "roles" in error_msg and "list" in error_msg, ( - f"Expected hint about using 'roles' instead: {error_msg}" - ) + assert ( + "roles" in error_msg and "list" in error_msg + ), f"Expected hint about using 'roles' instead: {error_msg}" @pytest.mark.asyncio @@ -1848,9 +1999,9 @@ async def test_find_and_validate_specific_team_id_hints_bracket_index_notation() error_msg = str(exc_info.value) assert "roles[0]" in error_msg, f"Expected field name in: {error_msg}" - assert "roles" in error_msg and "list" in error_msg, ( - f"Expected hint about using 'roles' instead: {error_msg}" - ) + assert ( + "roles" in error_msg and "list" in error_msg + ), f"Expected hint about using 'roles' instead: {error_msg}" @pytest.mark.asyncio @@ -1878,4 +2029,3 @@ async def test_find_and_validate_specific_team_id_no_hint_for_valid_field(): error_msg = str(exc_info.value) assert "Hint" not in error_msg - diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 6e000f9c0e..74c7f9bca5 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -334,7 +334,6 @@ async def test_proxy_admin_expired_key_from_cache(): "litellm.proxy.auth.user_api_key_auth._delete_cache_key_object", new_callable=AsyncMock, ) as mock_delete_cache: - mock_get_key_object.return_value = expired_token # Set attributes on proxy_server module (these are imported inside _user_api_key_auth_builder) @@ -609,7 +608,6 @@ class TestJWTOAuth2Coexistence: "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", new_callable=AsyncMock, ) as mock_jwt_auth: - litellm.proxy.proxy_server.jwt_handler.update_environment( prisma_client=None, user_api_key_cache=DualCache(), @@ -674,7 +672,6 @@ class TestJWTOAuth2Coexistence: new_callable=AsyncMock, return_value=mock_jwt_result, ) as mock_jwt_auth: - litellm.proxy.proxy_server.jwt_handler.update_environment( prisma_client=None, user_api_key_cache=DualCache(), @@ -726,7 +723,6 @@ class TestJWTOAuth2Coexistence: new_callable=AsyncMock, return_value=mock_oauth2_response, ) as mock_oauth2: - result = await user_api_key_auth( request=mock_request, api_key=f"Bearer {jwt_like_token}", diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index ed5b8cbd81..c6a03cf4ec 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -34,8 +34,8 @@ def llm_router() -> Router: "model_name": "azure-gpt-3-5-turbo", "litellm_params": { "model": "azure/chatgpt-v-2", - "api_key": "azure_api_key", - "api_base": "azure_api_base", + "api_key": "AZURE_AI_API_KEY", + "api_base": "AZURE_AI_API_BASE", "api_version": "azure_api_version", }, "model_info": { @@ -106,16 +106,18 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: Asserts 'create_file' is called with the correct arguments """ import litellm + import litellm.proxy.proxy_server as ps from litellm import Router from litellm.proxy._types import LitellmUserRoles - import litellm.proxy.proxy_server as ps from litellm.proxy.utils import ProxyLogging monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) # Mock create_file as an async function - mock_create_file = mocker.patch("litellm.files.main.create_file", new=mocker.AsyncMock()) + mock_create_file = mocker.patch( + "litellm.files.main.create_file", new=mocker.AsyncMock() + ) proxy_logging_obj = ProxyLogging( user_api_key_cache=DualCache(default_in_memory_ttl=1) @@ -127,7 +129,14 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: from litellm.llms.base_llm.files.transformation import BaseFileEndpoints class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Handle both dict and object forms of create_file_request if isinstance(create_file_request, dict): file_data = create_file_request.get("file") @@ -135,12 +144,12 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: else: file_data = create_file_request.file purpose_data = create_file_request.purpose - + # Call the mocked litellm.files.main.create_file to ensure asserts work await litellm.files.main.create_file( custom_llm_provider="azure", model="azure/chatgpt-v-2", - api_key="azure_api_key", + api_key="AZURE_AI_API_KEY", file=file_data[1], purpose=purpose_data, ) @@ -153,6 +162,7 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: ) # Return a dummy response object as needed by the test from litellm.types.llms.openai import OpenAIFileObject + return OpenAIFileObject( id="dummy-id", object="file", @@ -162,17 +172,21 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: purpose=purpose_data, status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") # Manually add the hook to the proxy_hook_mapping @@ -214,7 +228,7 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: if ( kwargs.get("custom_llm_provider") == "azure" and kwargs.get("model") == "azure/chatgpt-v-2" - and kwargs.get("api_key") == "azure_api_key" + and kwargs.get("api_key") == "AZURE_AI_API_KEY" ): azure_call_found = True break @@ -245,8 +259,8 @@ def test_target_storage_invokes_storage_backend( """ Ensure target_storage is parsed and invokes the storage backend service. """ - from litellm.proxy._types import LitellmUserRoles import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) @@ -304,8 +318,8 @@ def test_target_storage_with_target_models( """ Ensure target_storage and target_model_names are parsed and passed through. """ - from litellm.proxy._types import LitellmUserRoles import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) @@ -611,7 +625,9 @@ def test_create_file_for_each_model( assert openai_call_found, "OpenAI call not found with expected parameters" -def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_create_file_with_expires_after( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that expires_after is properly parsed and passed through when creating a file """ @@ -624,18 +640,25 @@ def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_ proxy_logging_obj._add_proxy_hooks(llm_router) class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Verify expires_after is in the request if isinstance(create_file_request, dict): expires_after = create_file_request.get("expires_after") else: expires_after = getattr(create_file_request, "expires_after", None) - + # Verify expires_after was passed correctly assert expires_after is not None, "expires_after should be in the request" assert expires_after["anchor"] == "created_at" assert expires_after["seconds"] == 2592000 - + # Return a dummy response return OpenAIFileObject( id="file-abc123", @@ -646,17 +669,21 @@ def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_ purpose="fine-tune", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles() @@ -688,7 +715,9 @@ def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_ assert result["purpose"] == "fine-tune" -def test_create_file_with_expires_after_missing_anchor(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_create_file_with_expires_after_missing_anchor( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that an error is returned when expires_after[anchor] is missing """ @@ -717,10 +746,15 @@ def test_create_file_with_expires_after_missing_anchor(mocker: MockerFixture, mo assert response.status_code == 400 error_detail = response.json() - assert "expires_after" in error_detail["error"]["message"].lower() or "both" in error_detail["error"]["message"].lower() + assert ( + "expires_after" in error_detail["error"]["message"].lower() + or "both" in error_detail["error"]["message"].lower() + ) -def test_create_file_with_expires_after_missing_seconds(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_create_file_with_expires_after_missing_seconds( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that an error is returned when expires_after[seconds] is missing """ @@ -749,10 +783,15 @@ def test_create_file_with_expires_after_missing_seconds(mocker: MockerFixture, m assert response.status_code == 400 error_detail = response.json() - assert "expires_after" in error_detail["error"]["message"].lower() or "both" in error_detail["error"]["message"].lower() + assert ( + "expires_after" in error_detail["error"]["message"].lower() + or "both" in error_detail["error"]["message"].lower() + ) -def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_create_file_with_expires_after_valid_values( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that expires_after works with valid anchor and seconds values """ @@ -765,18 +804,25 @@ def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monk proxy_logging_obj._add_proxy_hooks(llm_router) class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Verify expires_after is in the request if isinstance(create_file_request, dict): expires_after = create_file_request.get("expires_after") else: expires_after = getattr(create_file_request, "expires_after", None) - + # Verify expires_after was passed correctly assert expires_after is not None, "expires_after should be in the request" assert expires_after["anchor"] == "created_at" assert expires_after["seconds"] == 3600 - + return OpenAIFileObject( id="file-abc123", object="file", @@ -786,17 +832,21 @@ def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monk purpose="fine-tune", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles() @@ -827,7 +877,9 @@ def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monk assert result["purpose"] == "fine-tune" -def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_create_file_without_expires_after( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that file creation works normally without expires_after """ @@ -840,16 +892,25 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l proxy_logging_obj._add_proxy_hooks(llm_router) class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Verify expires_after is None when not provided if isinstance(create_file_request, dict): expires_after = create_file_request.get("expires_after") else: expires_after = getattr(create_file_request, "expires_after", None) - + # expires_after should be None when not provided - assert expires_after is None, "expires_after should be None when not provided" - + assert ( + expires_after is None + ), "expires_after should be None when not provided" + return OpenAIFileObject( id="file-abc123", object="file", @@ -859,17 +920,21 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l purpose="fine-tune", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles() @@ -898,11 +963,13 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l assert result["purpose"] == "fine-tune" -def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_managed_files_with_loadbalancing( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that managed files work with loadbalancing when both target_model_names and enable_loadbalancing_on_batch_endpoints are enabled. - + This ensures that the priority order is correct: - managed files should take precedence over deprecated loadbalancing - managed files internally use llm_router.acreate_file() which provides loadbalancing @@ -912,28 +979,34 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll # Enable loadbalancing on batch endpoints monkeypatch.setattr("litellm.enable_loadbalancing_on_batch_endpoints", True) - + proxy_logging_obj = ProxyLogging( user_api_key_cache=DualCache(default_in_memory_ttl=1) ) proxy_logging_obj._add_proxy_hooks(llm_router) - + # Track calls to verify loadbalancing through router router_acreate_file_calls = [] - + class ManagedFilesWithLoadbalancing(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Verify we receive the target model names - assert len(target_model_names_list) > 0, "Should have target_model_names_list" - + assert ( + len(target_model_names_list) > 0 + ), "Should have target_model_names_list" + # Simulate what managed files does - call llm_router.acreate_file for each model # This is where loadbalancing happens internally for model in target_model_names_list: - router_acreate_file_calls.append({ - "model": model, - "via_router": True - }) - + router_acreate_file_calls.append({"model": model, "via_router": True}) + # Return a managed file ID (base64 encoded) return OpenAIFileObject( id="litellm_managed_file_abc123", @@ -944,23 +1017,29 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll purpose="batch", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - + import litellm.proxy.proxy_server as ps from litellm.proxy._types import LitellmUserRoles - proxy_logging_obj.proxy_hook_mapping["managed_files"] = ManagedFilesWithLoadbalancing() + proxy_logging_obj.proxy_hook_mapping[ + "managed_files" + ] = ManagedFilesWithLoadbalancing() monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) monkeypatch.setattr( "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj @@ -971,12 +1050,12 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( api_key="test-key", user_role=LitellmUserRoles.PROXY_ADMIN ) - + try: # Create batch file content test_file_content = b'{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]}}' test_file = ("batch_data.jsonl", test_file_content, "application/jsonl") - + # Make request with both target_model_names AND enable_loadbalancing_on_batch_endpoints response = client.post( "/v1/files", @@ -987,7 +1066,7 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll }, headers={"Authorization": "Bearer test-key"}, ) - + # Verify success assert response.status_code == 200, response.text finally: @@ -995,13 +1074,17 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll result = response.json() assert result["id"] == "litellm_managed_file_abc123" assert result["purpose"] == "batch" - + # Verify that managed files was called (via router for loadbalancing) # This proves that managed files took precedence over deprecated loadbalancing - assert len(router_acreate_file_calls) == 2, "Should have called router for both models" + assert ( + len(router_acreate_file_calls) == 2 + ), "Should have called router for both models" assert router_acreate_file_calls[0]["model"] == "azure-gpt-3-5-turbo" assert router_acreate_file_calls[1]["model"] == "gpt-3.5-turbo" - assert all(call["via_router"] for call in router_acreate_file_calls), "All calls should go through router" + assert all( + call["via_router"] for call in router_acreate_file_calls + ), "All calls should go through router" def test_create_file_with_nested_litellm_metadata( @@ -1009,22 +1092,29 @@ def test_create_file_with_nested_litellm_metadata( ): """ Test that nested litellm_metadata is correctly parsed from form data in bracket notation. - + Regression test for: litellm_metadata[spend_logs_metadata][owner] format should be correctly parsed into nested dictionary structure. """ from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.types.llms.openai import OpenAIFileObject - + proxy_logging_obj = ProxyLogging( user_api_key_cache=DualCache(default_in_memory_ttl=1) ) proxy_logging_obj._add_proxy_hooks(llm_router) - + captured_litellm_metadata = {} - + class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Capture litellm_metadata for verification if isinstance(create_file_request, dict): captured_litellm_metadata.update( @@ -1034,7 +1124,7 @@ def test_create_file_with_nested_litellm_metadata( captured_litellm_metadata.update( getattr(create_file_request, "litellm_metadata", {}) ) - + return OpenAIFileObject( id="file-test-123", object="file", @@ -1044,28 +1134,32 @@ def test_create_file_with_nested_litellm_metadata( purpose="fine-tune", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - + proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles() monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) monkeypatch.setattr( "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj ) - + test_file_content = b'{"prompt": "Hello", "completion": "Hi"}' test_file = ("test.jsonl", test_file_content, "application/jsonl") - + # Test with nested litellm_metadata in bracket notation response = client.post( "/v1/files", @@ -1080,12 +1174,12 @@ def test_create_file_with_nested_litellm_metadata( }, headers={"Authorization": "Bearer test-key"}, ) - + # Verify success assert response.status_code == 200 result = response.json() assert result["id"] == "file-test-123" - + # Verify nested metadata was correctly parsed assert "spend_logs_metadata" in captured_litellm_metadata assert captured_litellm_metadata["spend_logs_metadata"]["owner"] == "john_doe" @@ -1099,26 +1193,33 @@ def test_create_file_with_deep_nested_litellm_metadata( ): """ Test that deeply nested litellm_metadata is correctly parsed from form data. - + Regression test for: litellm_metadata[a][b][c] format should be correctly parsed. """ + import litellm.proxy.proxy_server as ps from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.proxy._types import LitellmUserRoles - import litellm.proxy.proxy_server as ps from litellm.types.llms.openai import OpenAIFileObject - + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) - + proxy_logging_obj = ProxyLogging( user_api_key_cache=DualCache(default_in_memory_ttl=1) ) proxy_logging_obj._add_proxy_hooks(llm_router) - + captured_litellm_metadata = {} - + class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): if isinstance(create_file_request, dict): captured_litellm_metadata.update( create_file_request.get("litellm_metadata", {}) @@ -1127,7 +1228,7 @@ def test_create_file_with_deep_nested_litellm_metadata( captured_litellm_metadata.update( getattr(create_file_request, "litellm_metadata", {}) ) - + return OpenAIFileObject( id="file-test-456", object="file", @@ -1137,33 +1238,37 @@ def test_create_file_with_deep_nested_litellm_metadata( purpose="batch", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - + proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles() monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) monkeypatch.setattr( "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj ) - + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" ) - + try: test_file_content = b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo"}}' test_file = ("nested.jsonl", test_file_content, "application/jsonl") - + # Test with deeply nested metadata response = client.post( "/v1/files", @@ -1177,12 +1282,12 @@ def test_create_file_with_deep_nested_litellm_metadata( }, headers={"Authorization": "Bearer test-key"}, ) - + # Verify success assert response.status_code == 200, response.text result = response.json() assert result["id"] == "file-test-456" - + # Verify deeply nested metadata was correctly parsed assert "config" in captured_litellm_metadata assert "database" in captured_litellm_metadata["config"] @@ -1356,7 +1461,9 @@ def test_file_team_injects_when_caller_sends_nothing( # --------------------------------------------------------------------------- -def _post_file_raw(monkeypatch, llm_router: Router, team_metadata: dict, form_data: dict): +def _post_file_raw( + monkeypatch, llm_router: Router, team_metadata: dict, form_data: dict +): """POST /v1/files and return the raw response (no status assertion).""" from litellm.proxy.auth.user_api_key_auth import user_api_key_auth diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 223f0b335f..45ec4b726a 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1431,6 +1431,83 @@ class TestOverrideOpenAIResponseModel: assert response_obj.model == actual_model_used assert response_obj.model != requested_model + def test_override_model_uses_winning_model_for_fastest_response(self): + """ + Test that when fastest_response batch completion is used with a + comma-separated model list, the response model is set to the winning + model's group name (not the comma-separated list). + """ + requested_model = "openai/gpt-4o,gemini/gemini-2.5-flash" + winning_model_group = "gemini/gemini-2.5-flash" + downstream_model = "gemini-2.5-flash" + + response_obj = MagicMock() + response_obj.model = downstream_model + response_obj._hidden_params = { + "fastest_response_batch_completion": True, + "additional_headers": { + "x-litellm-model-group": winning_model_group, + }, + } + + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + + assert response_obj.model == winning_model_group + assert response_obj.model != requested_model + + def test_override_model_preserves_response_when_fastest_response_no_model_group( + self, + ): + """ + Test that when fastest_response is set but no model group header is + available, the actual downstream model is preserved. + """ + requested_model = "openai/gpt-4o,gemini/gemini-2.5-flash" + downstream_model = "gpt-4o-2024-08-06" + + response_obj = MagicMock() + response_obj.model = downstream_model + response_obj._hidden_params = { + "fastest_response_batch_completion": True, + "additional_headers": {}, + } + + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + + assert response_obj.model == downstream_model + + def test_override_model_normal_when_fastest_response_not_set(self): + """ + Test that when fastest_response_batch_completion is not set, the + normal override behavior applies (model is set to requested_model). + """ + requested_model = "openai/gpt-4o" + downstream_model = "gpt-4o-2024-08-06" + + response_obj = MagicMock() + response_obj.model = downstream_model + response_obj._hidden_params = { + "additional_headers": { + "x-litellm-model-group": "openai/gpt-4o", + }, + } + + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + + assert response_obj.model == requested_model + class TestIsAzureModelRouterRequest: """Tests for _is_azure_model_router_request helper""" diff --git a/tests/test_litellm/proxy/test_response_model_sanitization.py b/tests/test_litellm/proxy/test_response_model_sanitization.py index 22785bbcb9..b7253d9833 100644 --- a/tests/test_litellm/proxy/test_response_model_sanitization.py +++ b/tests/test_litellm/proxy/test_response_model_sanitization.py @@ -273,3 +273,61 @@ async def test_proxy_streaming_azure_model_router_preserves_actual_model(monkeyp # Azure Model Router: preserve actual model used, not the router model assert payload["model"] == actual_model_used assert payload["model"] != router_model + + +@pytest.mark.asyncio +async def test_proxy_streaming_fastest_response_preserves_winning_model(monkeypatch): + """ + Regression test for fastest_response streaming: + + When the client sends a comma-separated model list with fastest_response=True, + the streaming chunks should preserve the winning model's name from the + downstream response, NOT override to the comma-separated list. + """ + comma_separated_models = "openai/gpt-4o,gemini/gemini-2.5-flash" + winning_model = "gemini-2.5-flash" + + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + async def _iterator_hook( + user_api_key_dict: UserAPIKeyAuth, + response: AsyncGenerator, + request_data: dict, + ): + yield _make_model_response_stream_chunk(model=winning_model) + + monkeypatch.setattr( + proxy_server.proxy_logging_obj, + "async_post_call_streaming_iterator_hook", + _iterator_hook, + ) + monkeypatch.setattr( + proxy_server.proxy_logging_obj, + "async_post_call_streaming_hook", + AsyncMock(side_effect=lambda **kwargs: kwargs["response"]), + ) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-1234") + + gen = proxy_server.async_data_generator( + response=MagicMock(), + user_api_key_dict=user_api_key_dict, + request_data={ + "model": comma_separated_models, + "_litellm_client_requested_model": comma_separated_models, + "fastest_response": True, + }, + ) + + chunks = [] + async for item in gen: + chunks.append(item) + + assert len(chunks) >= 2 + first = chunks[0] + assert first.startswith("data: ") + + payload = json.loads(first[len("data: ") :].strip()) + assert payload["model"] == winning_model + assert payload["model"] != comma_separated_models diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index a931a9bc93..aaa131c23d 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1899,3 +1899,97 @@ class TestStreamingIDConsistency: else getattr(assistant_messages[0], "tool_calls", None) ) assert tool_calls is not None and len(tool_calls) == 1 + + +class TestEnsureOutputItemContentPartAdded: + """Test that _ensure_output_item_for_chunk emits content_part.added after + output_item.added for message items.""" + + def _make_iterator(self): + """Create a minimal LiteLLMCompletionStreamingIterator for testing.""" + from unittest.mock import MagicMock + + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + iterator = LiteLLMCompletionStreamingIterator.__new__( + LiteLLMCompletionStreamingIterator + ) + iterator.sent_output_item_added_event = False + iterator.sent_content_part_added_event = False + iterator._sequence_number = 0 + iterator._cached_item_id = None + iterator._cached_reasoning_item_id = None + iterator._reasoning_active = False + iterator._pending_response_events = [] + return iterator + + def _make_text_chunk(self): + """Create a mock ModelResponseStream with a text delta.""" + from unittest.mock import MagicMock + + chunk = MagicMock() + delta = MagicMock() + delta.reasoning_content = None + delta.tool_calls = None + chunk.choices = [MagicMock(delta=delta)] + return chunk + + def _make_reasoning_chunk(self): + """Create a mock ModelResponseStream with a reasoning delta.""" + from unittest.mock import MagicMock + + chunk = MagicMock() + delta = MagicMock() + delta.reasoning_content = "thinking..." + delta.tool_calls = None + chunk.choices = [MagicMock(delta=delta)] + return chunk + + def test_message_item_emits_content_part_added(self): + """content_part.added must follow output_item.added for message items.""" + from litellm.types.llms.openai import ( + ContentPartAddedEvent, + OutputItemAddedEvent, + ResponsesAPIStreamEvents, + ) + + iterator = self._make_iterator() + chunk = self._make_text_chunk() + + iterator._ensure_output_item_for_chunk(chunk) + + events = iterator._pending_response_events + assert len(events) == 2 + assert isinstance(events[0], OutputItemAddedEvent) + assert events[0].type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED + assert isinstance(events[1], ContentPartAddedEvent) + assert events[1].type == ResponsesAPIStreamEvents.CONTENT_PART_ADDED + assert events[1].part.type == "output_text" + assert iterator.sent_content_part_added_event is True + + def test_reasoning_item_does_not_emit_content_part_added(self): + """Reasoning items should not get a content_part.added event.""" + from litellm.types.llms.openai import OutputItemAddedEvent + + iterator = self._make_iterator() + chunk = self._make_reasoning_chunk() + + iterator._ensure_output_item_for_chunk(chunk) + + events = iterator._pending_response_events + assert len(events) == 1 + assert isinstance(events[0], OutputItemAddedEvent) + assert iterator.sent_content_part_added_event is False + + def test_only_emits_once(self): + """Calling _ensure_output_item_for_chunk twice should not duplicate events.""" + iterator = self._make_iterator() + chunk = self._make_text_chunk() + + iterator._ensure_output_item_for_chunk(chunk) + iterator._ensure_output_item_for_chunk(chunk) + + events = iterator._pending_response_events + assert len(events) == 2 diff --git a/tests/test_litellm/router_utils/test_health_check_routing.py b/tests/test_litellm/router_utils/test_router_health_check_routing.py similarity index 100% rename from tests/test_litellm/router_utils/test_health_check_routing.py rename to tests/test_litellm/router_utils/test_router_health_check_routing.py diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index f4a7c4f499..a02464e946 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -24,9 +24,9 @@ def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): "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"), }, } ], @@ -646,8 +646,15 @@ def test_arouter_responses_api_bridge(): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = {"id": "resp_test", "object": "response", "status": "completed", "output": []} - mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' + mock_response.json.return_value = { + "id": "resp_test", + "object": "response", + "status": "completed", + "output": [], + } + mock_response.text = ( + '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' + ) with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -2147,7 +2154,10 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ) assert credentials is not None - assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" + assert ( + credentials["aws_bedrock_runtime_endpoint"] + == "https://bedrock-runtime.us-east-1.amazonaws.com" + ) assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -2169,11 +2179,11 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): credential_values={ "api_key": "resolved-api-key", "api_base": "https://resolved.openai.azure.com", - "api_version": "2024-02-01" - } + "api_version": "2024-02-01", + }, ) ] - + router = litellm.Router( model_list=[ { @@ -2197,7 +2207,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): assert credentials["custom_llm_provider"] == "azure" # Ensure credential name is removed after resolution assert "litellm_credential_name" not in credentials - + # Cleanup litellm.credential_list = [] @@ -2302,7 +2312,10 @@ async def test_aguardrail_helper(): # Mock the original function async def mock_original_function(**kwargs): - return {"result": "success", "selected_guardrail": kwargs.get("selected_guardrail")} + return { + "result": "success", + "selected_guardrail": kwargs.get("selected_guardrail"), + } result = await router._aguardrail_helper( model="content-filter", @@ -2336,7 +2349,10 @@ async def test_aguardrail(): # Mock the original function async def mock_original_function(**kwargs): - return {"result": "success", "selected_guardrail": kwargs.get("selected_guardrail")} + return { + "result": "success", + "selected_guardrail": kwargs.get("selected_guardrail"), + } result = await router.aguardrail( guardrail_name="content-filter", @@ -2346,6 +2362,7 @@ async def test_aguardrail(): assert result["result"] == "success" assert result["selected_guardrail"]["id"] == "guardrail-1" + @pytest.mark.asyncio async def test_anthropic_messages_call_type_is_cached(): """ @@ -2417,36 +2434,33 @@ async def test_anthropic_messages_call_type_is_cached(): additional_headers=None, ), ) - + cache = DualCache() deployment_check = PromptCachingDeploymentCheck(cache=cache) prompt_cache = PromptCachingCache(cache=cache) - + # Create messages with enough tokens to pass the caching threshold test_messages = [ { - "role": "user", + "role": "user", "content": [ { - "type": "text", + "type": "text", "text": "test long message here" * 1024, - "cache_control": { - "type": "ephemeral", - "ttl": "5m" - } + "cache_control": {"type": "ephemeral", "ttl": "5m"}, } - ] + ], } ] test_model_id = "test-model-id-123" - + # Create a payload with anthropic_messages call type payload = create_standard_logging_payload() payload["call_type"] = CallTypes.anthropic_messages.value payload["messages"] = test_messages payload["model"] = "anthropic/claude-3-5-sonnet-20240620" payload["model_id"] = test_model_id - + # Log the success event (should cache the model_id) await deployment_check.async_log_success_event( kwargs={"standard_logging_object": payload}, @@ -2454,19 +2468,23 @@ async def test_anthropic_messages_call_type_is_cached(): start_time=1234567890.0, end_time=1234567891.0, ) - + # Small delay to ensure cache write completes await asyncio.sleep(0.1) - + # Verify that the model_id was actually cached cached_result = await prompt_cache.async_get_model_id( messages=test_messages, tools=None, ) - + # This assertion will FAIL if anthropic_messages is filtered out - assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" - assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" + assert ( + cached_result is not None + ), "Model ID should be cached for anthropic_messages call type" + assert ( + cached_result["model_id"] == test_model_id + ), f"Expected {test_model_id}, got {cached_result['model_id']}" def test_update_kwargs_with_deployment_propagates_model_tags(): @@ -2682,9 +2700,7 @@ def test_credential_name_injected_as_tag(): ) kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="xai-model" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert "Credential: xAI" in kwargs["metadata"]["tags"] @@ -2709,9 +2725,7 @@ def test_credential_name_not_duplicated_in_tags(): ) kwargs: dict = {"metadata": {"tags": ["Credential: xAI", "A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="xai-model" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["metadata"]["tags"].count("Credential: xAI") == 1 @@ -2733,9 +2747,7 @@ def test_credential_name_not_injected_when_absent(): ) kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-model" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-model") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["metadata"]["tags"] == ["A.101"] diff --git a/tests/test_litellm/test_setup_wizard.py b/tests/test_litellm/test_setup_wizard.py index e10bd893e3..c96d6d7ed6 100644 --- a/tests/test_litellm/test_setup_wizard.py +++ b/tests/test_litellm/test_setup_wizard.py @@ -58,7 +58,7 @@ _ANTHROPIC = { _AZURE = { "id": "azure", "name": "Azure OpenAI", - "env_key": "AZURE_API_KEY", + "env_key": "AZURE_AI_API_KEY", "models": [], "test_model": None, "needs_api_base": True, @@ -123,8 +123,8 @@ def test_build_config_master_key_quoted(): def test_build_config_does_not_mutate_env_vars(): """_build_config must not modify the caller's env_vars dict.""" env_vars = { - "AZURE_API_KEY": "az-key", - "_LITELLM_AZURE_API_BASE_AZURE": "https://my.azure.com", + "AZURE_AI_API_KEY": "az-key", + "_LITELLM_AZURE_AI_API_BASE_AZURE": "https://my.azure.com", "_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-deployment", } original_keys = set(env_vars.keys()) @@ -134,8 +134,8 @@ def test_build_config_does_not_mutate_env_vars(): def test_build_config_azure_uses_deployment_name(): env_vars = { - "AZURE_API_KEY": "az-key", - "_LITELLM_AZURE_API_BASE_AZURE": "https://my.azure.com", + "AZURE_AI_API_KEY": "az-key", + "_LITELLM_AZURE_AI_API_BASE_AZURE": "https://my.azure.com", "_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-gpt4o", } config = SetupWizard._build_config([_AZURE], env_vars, "sk-master") @@ -147,7 +147,7 @@ def test_build_config_azure_uses_deployment_name(): def test_build_config_azure_no_deployment_skipped(): """Azure without a deployment name should emit nothing (not fallback to gpt-4o).""" - env_vars = {"AZURE_API_KEY": "az-key"} # no deployment sentinel + env_vars = {"AZURE_AI_API_KEY": "az-key"} # no deployment sentinel config = SetupWizard._build_config([_AZURE], env_vars, "sk-master") # No azure model entry should be emitted when deployment name is absent assert "model: azure/" not in config @@ -157,7 +157,7 @@ def test_build_config_no_display_name_collision_openai_and_azure(): """OpenAI gpt-4o and azure gpt-4o should get distinct model_name values.""" env_vars = { "OPENAI_API_KEY": "sk-openai", - "AZURE_API_KEY": "az-key", + "AZURE_AI_API_KEY": "az-key", "_LITELLM_AZURE_DEPLOYMENT_AZURE": "gpt-4o", } config = SetupWizard._build_config([_OPENAI, _AZURE], env_vars, "sk-master") @@ -182,7 +182,7 @@ def test_build_config_internal_sentinel_keys_excluded(): """_LITELLM_ prefixed sentinel keys must not appear in environment_variables.""" env_vars = { "OPENAI_API_KEY": "sk-real", - "_LITELLM_AZURE_API_BASE_AZURE": "https://x.azure.com", + "_LITELLM_AZURE_AI_API_BASE_AZURE": "https://x.azure.com", } config = SetupWizard._build_config([_OPENAI], env_vars, "sk-master") assert "_LITELLM_" not in config diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index eb2f510e23..e984403a05 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -831,6 +831,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): }, }, "supports_native_streaming": {"type": "boolean"}, + "supports_native_structured_output": {"type": "boolean"}, "tiered_pricing": { "type": "array", "items": { diff --git a/tests/test_team_logging.py b/tests/test_team_logging.py index 913b1e1949..9e89d945ed 100644 --- a/tests/test_team_logging.py +++ b/tests/test_team_logging.py @@ -59,137 +59,3 @@ async def chat_completion(session, key, model="azure-gpt-3.5", request_metadata= if status != 200: raise Exception(f"Request did not return a 200 status code: {status}") - - -@pytest.mark.skip(reason="flaky test - covered by simpler unit testing.") -@pytest.mark.asyncio -@pytest.mark.flaky(retries=12, delay=2) -async def test_aaateam_logging(): - """ - -> Team 1 logs to project 1 - -> Create Key - -> Make chat/completions call - -> Fetch logs from langfuse - """ - try: - async with aiohttp.ClientSession() as session: - - key = await generate_key( - session, models=["fake-openai-endpoint"], team_id="team-1" - ) # team-1 logs to project 1 - - from litellm._uuid import uuid - - _trace_id = f"trace-{uuid.uuid4()}" - _request_metadata = { - "trace_id": _trace_id, - } - - await chat_completion( - session, - key["key"], - model="fake-openai-endpoint", - request_metadata=_request_metadata, - ) - - # Test - if the logs were sent to the correct team on langfuse - import langfuse - - print(f"langfuse_public_key: {os.getenv('LANGFUSE_PROJECT1_PUBLIC')}") - print(f"langfuse_secret_key: {os.getenv('LANGFUSE_HOST')}") - langfuse_client = langfuse.Langfuse( - public_key=os.getenv("LANGFUSE_PROJECT1_PUBLIC"), - secret_key=os.getenv("LANGFUSE_PROJECT1_SECRET"), - host="https://us.cloud.langfuse.com", - ) - - await asyncio.sleep(30) - - print(f"searching for trace_id={_trace_id} on langfuse") - - generations = langfuse_client.get_generations(trace_id=_trace_id).data - print(generations) - assert len(generations) == 1 - except Exception as e: - pytest.fail(f"Unexpected error: {str(e)}") - - -@pytest.mark.skip(reason="todo fix langfuse credential error") -@pytest.mark.asyncio -async def test_team_2logging(): - """ - -> Team 1 logs to project 2 - -> Create Key - -> Make chat/completions call - -> Fetch logs from langfuse - """ - langfuse_public_key = os.getenv("LANGFUSE_PROJECT2_PUBLIC") - - print(f"langfuse_public_key: {langfuse_public_key}") - langfuse_secret_key = os.getenv("LANGFUSE_PROJECT2_SECRET") - print(f"langfuse_secret_key: {langfuse_secret_key}") - langfuse_host = "https://us.cloud.langfuse.com" - - try: - assert langfuse_public_key is not None - assert langfuse_secret_key is not None - except Exception as e: - # skip test if langfuse credentials are not set - return - - try: - async with aiohttp.ClientSession() as session: - - key = await generate_key( - session, models=["fake-openai-endpoint"], team_id="team-2" - ) # team-1 logs to project 1 - - from litellm._uuid import uuid - - _trace_id = f"trace-{uuid.uuid4()}" - _request_metadata = { - "trace_id": _trace_id, - } - - await chat_completion( - session, - key["key"], - model="fake-openai-endpoint", - request_metadata=_request_metadata, - ) - - # Test - if the logs were sent to the correct team on langfuse - import langfuse - - langfuse_client = langfuse.Langfuse( - public_key=langfuse_public_key, - secret_key=langfuse_secret_key, - host=langfuse_host, - ) - - await asyncio.sleep(30) - - print(f"searching for trace_id={_trace_id} on langfuse") - - generations = langfuse_client.get_generations(trace_id=_trace_id).data - print("Team 2 generations", generations) - - # team-2 should have 1 generation with this trace id - assert len(generations) == 1 - - # team-1 should have 0 generations with this trace id - langfuse_client_1 = langfuse.Langfuse( - public_key=os.getenv("LANGFUSE_PROJECT1_PUBLIC"), - secret_key=os.getenv("LANGFUSE_PROJECT1_SECRET"), - host="https://us.cloud.langfuse.com", - ) - - generations_team_1 = langfuse_client_1.get_generations( - trace_id=_trace_id - ).data - print("Team 1 generations", generations_team_1) - - assert len(generations_team_1) == 0 - - except Exception as e: - pytest.fail("Team 2 logging failed: " + str(e)) diff --git a/tests/unified_google_tests/vertex_key.json b/tests/unified_google_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/unified_google_tests/vertex_key.json +++ b/tests/unified_google_tests/vertex_key.json @@ -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" } diff --git a/tests/vector_store_tests/test_azure_ai_vector_store.py b/tests/vector_store_tests/test_azure_ai_vector_store.py index 52eb6635a9..58e45f259a 100644 --- a/tests/vector_store_tests/test_azure_ai_vector_store.py +++ b/tests/vector_store_tests/test_azure_ai_vector_store.py @@ -17,10 +17,10 @@ async def test_basic_search_vector_store(sync_mode): "vector_store_id": "my-vector-index", "custom_llm_provider": "azure_ai", "azure_search_service_name": "azure-kb-search", - "litellm_embedding_model": "azure/text-embedding-3-large", + "litellm_embedding_model": "azure_ai/text-embedding-3-large", "litellm_embedding_config": { - "api_base": os.getenv("AZURE_AI_SEARCH_EMBEDDING_API_BASE"), - "api_key": os.getenv("AZURE_AI_SEARCH_EMBEDDING_API_KEY"), + "api_base": os.getenv("AZURE_AI_API_BASE"), + "api_key": os.getenv("AZURE_AI_API_KEY"), }, "api_key": os.getenv("AZURE_SEARCH_API_KEY"), }