test: migrate local_testing SDK and provider tests out of build_and_test (#45760)

* test: migrate local_testing SDK and provider tests out of build_and_test

* test: cover sagemaker region, hf template and nova paths end to end

Region rows now go through litellm.completion on the sagemaker/ handler that
the legacy tests exercised, the prompt template test mocks the Hugging Face
tokenizer config instead of switching to custom_prompt_dict, streaming tests
assert role, finish reason and body, and Nova single turn and streaming run
through litellm.completion with a recorded request body

* test: tighten migrated integration tests for local OpenAI routes, timeouts and air-gapped boot

Split the catch-all providers migration file by behavior. The fake endpoint
route tests now hit the owned upstream directly, as the legacy harness tests
did. LM Studio asserts the outbound key, path and stripped model. The embedding
timeout test asserts the proxy maps a deployment timeout to a 408 and forwards
the read timeout to the provider client, instead of timing out the test client

The Azure streaming tool-call test also asserts the second identical stream is
served from the Redis cache with the same reassembled tool call

The Docker network test always skipped in CI because the security job has no
litellm/litellm image and no egress to pull one, and it tested a published image
rather than the checkout. It now boots the checkout proxy with every outbound
HTTP connection refused and asserts readiness, the bundled cost map fallback and
that the cost map host is the only egress attempted

* test: tighten the migrated e2e rows for local_testing SDK and provider tests

Drop the duplicated Anthropic covers marker and the Anthropic cell wrongly claimed by
the Bedrock and Vertex cache tests. Keep the file-wide Anthropic model as it is on main.
Send the 2023-06-01 version header to Anthropic through the deployment's extra_headers
so the header actually reaches the provider, and mark the creation-token case 1h then
5m like the legacy test did

Reuse the Bedrock file's own registration for the pydantic response test, swap the
hand-built PNGs for the shared cat.jpg fixture, assert the SageMaker Nova image answer
names the cat and that top_logprobs=3 is honored, and register the SageMaker Nova opt-in
markers the same way as the other opt-in markers

* test: restore lost assertions in migrated test_main tests

Restore the role alternation check in test_completion_anthropic_hanging, the
embedding and streaming response-header checks, the logit_bias drop check and
the gemini system and safety settings checks. Add an azure 429 retry-after
test for test_exception_with_headers, trigger the router fallback with
mock_timeout like the legacy test, and stop the mock timeout tests sleeping 3s

* test: restore lost legacy assertions in migrated vertex ai unit tests

* test: tighten migrated cost, streaming and embedding tests and scope the hosted mock guard

* test: strengthen migrated anthropic, azure, bedrock and cohere provider tests

Assert full request bodies and usage instead of presence checks, restore the
dummy-tool and custom_llm_provider assertions the legacy tests made, add the
async and per-model variants the legacy parametrization covered, split
streamed tool arguments across deltas, and move the Azure chat streaming
tests to the azure.py mirror file

* test: tighten migrated llms provider tests and cover legacy intents lost in the move

Strengthen the migrated fireworks, gemini, mistral, together, openai text completion, openrouter, perplexity and openai error tests so they assert what the legacy tests proved (cost, provider tag, usage-only-on-last-chunk, function_call rebuild, reasoning param, sync/async/stream error fields). Replace the huggingface and jina tests that doubled litellm code with respx end-to-end tests. Add the tests the pre-existing mapped tests never covered: cohere 400 mapping, token-array text completion, hosted_vllm/llamafile/lm_studio env api base embeddings, watsonx embedding auth, ollama ssl_verify client and ollama vision prompt

* test: cover residual legacy paths for anthropic tool streams, watsonx async stream, ollama http images and caller async embedding clients

* test: restore single request assertions in the deep_infra and azure_ai embedding tests

* test: drop a no-op streaming cost flag from the openai streaming cost test

* test: cover residual legacy-test lines in caching, main, utils, replicate, vertex and bedrock invoke

* test: cover streamed audio assembly, choice membership, untranslated responses events and the stream log guard

* test: pin cost rates in unit tests, use sdk for cohere e2e, drop vertex llama e2e

* test: send the vertex maps e2e request through the openai sdk

* test: drop the local_testing jobs that select no tests after the migrations

* test: wait for the streamed cache entry and keep custom Mistral rates test-local

* test: keep untouched Mistral test lines as they were

---------

Co-authored-by: yuneng <yuneng@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-10-10 23:14:25 +00:00 • committed by GitHub
parent 3746056caf
commit a2e334febe
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
104 changed files with 9016 additions and 17385 deletions

View file

@ -473,7 +473,7 @@ jobs:
command: |
python3 tests/windows_tests/check_windows_wheel_install.py --lengths-only
local_testing_part1:
auth_ui_unit_tests:
docker:
- &python312_image
image: cimg/python:3.12@sha256:9c796c23c84e84a66a964acb508d39dc5433c81a47e07efd56dccbbc2427e07c
@ -489,219 +489,6 @@ jobs:
working_directory: ~/project
environment: &sidecar_database_url
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
parallelism: 4
steps:
- checkout
- skip_if_unrelated_changes
- setup_google_dns
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- setup_litellm_enterprise_pip
- save_cache:
paths:
- ~/.cache/uv
key: v1-uv-cache-{{ checksum "uv.lock" }}
- save_cargo_target
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- 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 (Part 1 - A-M)
command: |
mkdir test-results
# Discover test files (A-M)
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_[a-mA-M]*.py")
echo "$TEST_FILES" | circleci tests run \
--split-by=timings \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=20 \
-k \"not test_python_38.py and not test_basic_python_version.py and not router and not assistants and not langfuse and not caching and not cache\" \
-n 4 \
--timeout=300 \
--timeout_method=thread"
no_output_timeout: 15m
- run:
name: Rename the coverage files
command: |
# When CI reruns only the failed tests, a parallel node can receive
# zero tests and pytest never writes coverage. Emit empty placeholders
# so persist_to_workspace and the downstream coverage combine stay green.
if [ -f coverage.xml ]; then
mv coverage.xml local_testing_part1_coverage.xml
mv .coverage local_testing_part1_coverage
else
touch local_testing_part1_coverage.xml local_testing_part1_coverage
fi
# Store test results
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- local_testing_part1_coverage.xml
- local_testing_part1_coverage
local_testing_part2:
docker:
- *python312_image
- *postgres16_sidecar
working_directory: ~/project
environment: *sidecar_database_url
parallelism: 4
steps:
- checkout
- skip_if_unrelated_changes
- setup_google_dns
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- setup_litellm_enterprise_pip
- save_cache:
paths:
- ~/.cache/uv
key: v1-uv-cache-{{ checksum "uv.lock" }}
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- 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 (Part 2 - N-Z)
command: |
mkdir test-results
# Discover test files (N-Z)
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_[n-zN-Z]*.py")
echo "$TEST_FILES" | circleci tests run \
--split-by=timings \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=20 \
-k \"not test_python_38.py and not test_basic_python_version.py and not router and not assistants and not langfuse and not caching and not cache\" \
-n 4 \
--timeout=300 \
--timeout_method=thread"
no_output_timeout: 15m
- run:
name: Rename the coverage files
command: |
# When CI reruns only the failed tests, a parallel node can receive
# zero tests and pytest never writes coverage. Emit empty placeholders
# so persist_to_workspace and the downstream coverage combine stay green.
if [ -f coverage.xml ]; then
mv coverage.xml local_testing_part2_coverage.xml
mv .coverage local_testing_part2_coverage
else
touch local_testing_part2_coverage.xml local_testing_part2_coverage
fi
# Store test results
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- local_testing_part2_coverage.xml
- local_testing_part2_coverage
langfuse_logging_unit_tests:
docker:
- *python312_image
- *postgres16_sidecar
working_directory: ~/project
environment: *sidecar_database_url
resource_class: medium
steps:
- checkout
- skip_if_unrelated_changes
- setup_google_dns
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- setup_litellm_enterprise_pip
- save_cache:
paths:
- ~/.cache/uv
key: v1-uv-cache-{{ checksum "uv.lock" }}
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- 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: |
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-k \"langfuse\""
no_output_timeout: 15m
# Store test results
- store_test_results:
path: test-results
auth_ui_unit_tests:
docker:
- *python312_image
- *postgres16_sidecar
working_directory: ~/project
environment: *sidecar_database_url
steps:
- checkout
@ -759,59 +546,6 @@ jobs:
- auth_ui_unit_tests_coverage.xml
- auth_ui_unit_tests_coverage
litellm_router_testing: # Runs all tests with the "router" keyword
docker:
- *python312_image
working_directory: ~/project
resource_class: large
parallelism: 1
environment:
FAKE_OPENAI_API_BASE: http://127.0.0.1:8190
steps:
- checkout
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- save_cache:
paths:
- ~/.cache/uv
key: v1-uv-cache-{{ checksum "uv.lock" }}
- start_fake_openai_endpoint
# Run pytest and generate JUnit XML report
- setup_litellm_enterprise_pip
- run:
name: Run tests
command: |
# On a "rerun failed tests" build a parallel node can receive no
# tests, so the test command never creates test-results. Pre-create it
# so store_test_results doesn't fail the node on a missing path.
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="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
-k 'router' \
-n 4 \
--junitxml=test-results/junit.xml \
--durations=5 \
--timeout=300 --timeout_method=thread"
no_output_timeout: 15m
# Store test results
- store_test_results:
path: test-results
llm_translation_testing:
docker:
- *python312_image
@ -2222,7 +1956,7 @@ jobs:
- run:
name: Combine Coverage
command: |
uv tool run --from 'coverage[toml]==7.10.6' coverage combine realtime_translation_coverage ocr_coverage search_coverage logging_coverage audio_coverage local_testing_part1_coverage local_testing_part2_coverage pass_through_unit_tests_coverage batches_coverage google_generate_content_endpoint_coverage litellm_utils_coverage
uv tool run --from 'coverage[toml]==7.10.6' coverage combine realtime_translation_coverage ocr_coverage search_coverage logging_coverage audio_coverage pass_through_unit_tests_coverage batches_coverage google_generate_content_endpoint_coverage litellm_utils_coverage
uv tool run --from 'coverage[toml]==7.10.6' coverage xml
- codecov/upload:
file: ./coverage.xml
@ -3306,10 +3040,6 @@ workflows:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- provider_replay_harness
- local_testing_part1
- local_testing_part2
- langfuse_logging_unit_tests
- litellm_router_testing
- auth_ui_unit_tests
- build_docker_database_image
- e2e_ui_testing
@ -3359,9 +3089,6 @@ workflows:
- image_gen_testing
- logging_testing
- audio_testing
- langfuse_logging_unit_tests
- local_testing_part1
- local_testing_part2
- db_migration_disable_update_check:
requires:
- build_docker_database_image

View file

@ -18,9 +18,8 @@ test_paths:
paths:
- tests/rust-python-harness
- reason: >-
Live-provider caching cases in tests/local_testing that remain outside CI. Jobs that
glob that directory either deselect them (local_testing_part1 and part2 carry `-k "... and
not caching and not cache"`) or keep only another keyword (langfuse, router, assistants).
Live-provider caching cases in tests/local_testing that remain outside CI. No CircleCI job
globs that directory anymore, and the legacy jobs that did always deselected these files.
The two Redis Cluster IAM authentication tests moved to tests/integration/sdk and remain
selected by node ID in redis_compat, where redis-server is installed.
The five keyless files now run in the caching-local shard, leaving live cases in these two.
@ -30,14 +29,6 @@ test_paths:
paths:
- tests/local_testing/test_disk_cache_unit_tests.py
- tests/local_testing/test_gcs_cache_unit_tests.py
- reason: >-
Mixed file that no job has ever run in full. Every CircleCI job that globs
tests/local_testing deselects it by name ("caching") or keeps only another keyword, so only
its router test ran, and that test now lives in tests/unit/router_utils/pre_call_checks.
Its two offline tests now live in tests/unit/llms/anthropic/chat, and the five left all need
ANTHROPIC_API_KEY or Vertex credentials
paths:
- tests/local_testing/test_anthropic_prompt_caching.py
- reason: >-
The end-to-end suite runs against a deployed proxy from its own in-cluster rig rather than
from a pull request; it needs a live gateway and provider credentials no PR job holds

View file

@ -37,6 +37,8 @@ from e2e_config import (
PROVIDER_EDGE_HOST_OPT_IN_ENV,
PROXY_BASE_URL,
REDIS_CHAOS_OPT_IN_ENV,
SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
SAGEMAKER_NOVA_OPT_IN_ENV,
SECRET_MANAGER_OPT_IN_ENV,
WEEKLY_ANOMALY_OPT_IN_ENV,
unique_marker,
@ -75,6 +77,8 @@ OPT_IN_MARKERS: Final = MappingProxyType(
"otel_v2": OTEL_V2_OPT_IN_ENV,
"otel_tls": OTEL_TLS_OPT_IN_ENV,
"secret_manager": SECRET_MANAGER_OPT_IN_ENV,
"sagemaker_nova": SAGEMAKER_NOVA_OPT_IN_ENV,
"sagemaker_nova2_lite": SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
}
)
@ -192,6 +196,16 @@ def pytest_configure(config: pytest.Config) -> None:
"secret_manager: needs a proxy booted from gateway/secret_manager_<system>_ci_config.yml against that live "
"secret manager; deselected unless E2E_SECRET_MANAGER names the backend (see secret_manager/secret_backends.py)",
)
config.addinivalue_line(
"markers",
"sagemaker_nova: needs a deployed SageMaker Nova endpoint the proxy's AWS credentials can invoke; "
"deselected unless SAGEMAKER_NOVA_ENDPOINT names it",
)
config.addinivalue_line(
"markers",
"sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; "
"deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it",
)
def pytest_sessionstart(session: pytest.Session) -> None:

View file

@ -32,6 +32,7 @@
- {id: llm.chat_completions.bedrock_converse.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Bedrock vision (Anthropic/Nova)"}
- {id: llm.chat_completions.bedrock_converse.prompt_cache_5m.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: bedrock_converse, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Anthropic-on-Bedrock caching"}
- {id: llm.chat_completions.bedrock_converse.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: bedrock_converse, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Anthropic thinking on Bedrock"}
- {id: llm.chat_completions.bedrock_converse.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: bedrock_converse, capability: structured_output, streaming: nonstream, assertions: [works], source: "llms/bedrock/chat/converse_handler.py", rationale: "Bedrock Claude accepts a Pydantic-derived response schema"}
- {id: llm.chat_completions.bedrock_converse.response_headers.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: response_headers, streaming: nonstream, assertions: [works], source: "llms/bedrock/chat/converse_handler.py:248", rationale: "Bedrock request ids must surface as llm_provider-* response headers on /chat/completions so callers can correlate calls with AWS-side logs (#37003)", fail_before_fix: proven}
- {id: llm.chat_completions.bedrock_converse.response_headers.stream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: response_headers, streaming: stream, assertions: [works], source: "llms/bedrock/chat/converse_handler.py:154", rationale: "The llm_provider-* headers must also surface on streaming /chat/completions, where CustomStreamWrapper carries them instead of the nonstream setter"}
- {id: llm.chat_completions.bedrock_converse.batch_deployment.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: batch_deployment, streaming: nonstream, assertions: [works], source: "types/utils.py bedrock_batch_litellm_params", rationale: "A deployment carrying the documented batch-only S3 keys (s3_access_key_id, s3_secret_access_key, s3_encryption_key_id) must still serve ordinary chat; unregistered keys fall into optional_params and are forwarded as additionalModelRequestFields, which Bedrock 400s and which puts the S3 secret in the request body and debug log (LIT-8290)", fail_before_fix: proven}
@ -47,8 +48,13 @@
- {id: llm.chat_completions.vertex.basic.stream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:8455", rationale: "Streaming over Vertex"}
- {id: llm.chat_completions.vertex.tool_use.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Vertex Gemini function_calling"}
- {id: llm.chat_completions.vertex.web_search.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: vertex, capability: web_search, streaming: stream, assertions: [works], source: "llms/vertex_ai", rationale: "Vertex Gemini streams Google Search grounding metadata"}
- {id: llm.chat_completions.vertex.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Gemini vision"}
- {id: llm.chat_completions.vertex.prompt_cache_5m.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: vertex, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Vertex Gemini prompt caching"}
- {id: llm.chat_completions.sagemaker.vision.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: sagemaker, capability: vision, streaming: nonstream, assertions: [works], source: "llms/sagemaker/nova", rationale: "SageMaker Nova accepts image input"}
- {id: llm.chat_completions.sagemaker.top_k.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: sagemaker, capability: top_k, streaming: nonstream, assertions: [works], source: "llms/sagemaker/nova", rationale: "SageMaker Nova accepts its top_k request parameter"}
- {id: llm.chat_completions.sagemaker.logprobs.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: sagemaker, capability: logprobs, streaming: nonstream, assertions: [works], source: "llms/sagemaker/nova", rationale: "SageMaker Nova returns log probability details"}
- {id: llm.chat_completions.sagemaker.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: sagemaker, capability: thinking, streaming: nonstream, assertions: [works], source: "llms/sagemaker/nova", rationale: "SageMaker Nova 2 Lite accepts reasoning effort"}
- {id: llm.chat_completions.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "proxy_server.py:8455", rationale: "P0 route; Azure OpenAI deployments"}
- {id: llm.chat_completions.azure_openai.tool_use.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: azure_openai, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Azure OpenAI function_calling"}
- {id: llm.chat_completions.azure_foundry.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: azure_foundry, capability: basic, streaming: nonstream, assertions: [works], source: "proxy_server.py:8455", rationale: "Azure Foundry (azure_ai); newer, smoke"}

View file

@ -8,6 +8,7 @@
- {id: llm.embeddings.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "llms/bedrock/embed/embedding.py", rationale: "Bedrock Titan embeddings"}
- {id: llm.embeddings.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "vertex_embeddings/embedding_handler.py", rationale: "Vertex embeddings"}
- {id: llm.embeddings.cohere.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: cohere, capability: basic, streaming: nonstream, assertions: [works], source: "llms/cohere/embed/handler.py", rationale: "Cohere embeddings"}
- {id: llm.embeddings.cohere.vision.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: cohere, capability: vision, streaming: nonstream, assertions: [works], source: "llms/cohere/embed/handler.py", rationale: "Cohere image embeddings report image token usage"}
- {id: llm.embeddings.anthropic.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: anthropic, capability: basic, streaming: nonstream, assertions: [works], source: "llms/anthropic/chat/handler.py", rationale: "Anthropic vector API (verify support)"}
- {id: llm.batches.openai.create.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Core batch create"}
- {id: llm.batches.openai.retrieve.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Batch retrieve, id round-trip + status"}

View file

@ -59,6 +59,7 @@ LlmRoute = Literal[
"ollama_chat",
"openai",
"sail",
"sagemaker",
"together_ai",
"vertex",
"xiaomi_mimo",
@ -78,6 +79,7 @@ LlmCapability = Literal[
"multi_turn",
"native_passthrough",
"pdf_input",
"logprobs",
"prompt_cache_1h",
"prompt_cache_5m",
"response_headers",
@ -85,6 +87,7 @@ LlmCapability = Literal[
"structured_output",
"thinking",
"thinking_with_tool_use",
"top_k",
"tool_search",
"tool_search_history",
"tool_use",

View file

@ -213,6 +213,8 @@ OWNED_GATEWAY_OPT_IN_ENV: Final = "E2E_OWNED_GATEWAY"
OTEL_V2_OPT_IN_ENV: Final = "E2E_OTEL_V2"
OTEL_TLS_OPT_IN_ENV: Final = "E2E_OTEL_EXPORTER_ENDPOINT"
SECRET_MANAGER_OPT_IN_ENV: Final = "E2E_SECRET_MANAGER"
SAGEMAKER_NOVA_OPT_IN_ENV: Final = "SAGEMAKER_NOVA_ENDPOINT"
SAGEMAKER_NOVA2_LITE_OPT_IN_ENV: Final = "SAGEMAKER_NOVA2_LITE_ENDPOINT"
ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))
ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3"))

View file

@ -7,22 +7,40 @@ missing messages and invalid model handling without crashing the proxy.
from __future__ import annotations
import pytest
from typing import Final
from pydantic import BaseModel, ConfigDict
from e2e_config import unique_marker
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
from e2e_http import (
assert_client_error,
require_successful_call,
)
from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta
from lifecycle import ResourceManager
from models import LiteLLMParamsBody
from proxy_client import ProxyClient
from pydantic import BaseModel
from sdk_clients import SdkClients
pytestmark = pytest.mark.e2e
BEDROCK_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"
class CalendarEvent(BaseModel):
model_config = ConfigDict(frozen=True)
name: str
date: str
participants: tuple[str, ...]
class EventsList(BaseModel):
model_config = ConfigDict(frozen=True)
events: tuple[CalendarEvent, ...]
class ConverseContent(BaseModel):
text: str
@ -100,6 +118,32 @@ def _default_invoke() -> InvokeBody:
class TestBedrockNative:
@pytest.mark.provider_live
@pytest.mark.covers("llm.chat_completions.bedrock_converse.structured_output.nonstream.works")
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.BEDROCK,),
models=(BEDROCK_BACKEND,),
capabilities=(Capability.RESPONSE_SCHEMA,),
mode=Mode.NONSTREAM,
)
)
def test_bedrock_chat_completion_pydantic_response(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _register(proxy, resources)
response: Final = sdk.openai(key).chat.completions.parse(
model=model,
messages=[{"role": "user", "content": "Create one event for the Apollo 11 launch in July 1969."}],
response_format=EventsList,
)
parsed: Final = response.choices[0].message.parsed
assert parsed is not None
assert parsed.events
assert all(event.name and event.date and event.participants for event in parsed.events)
@pytest.mark.covers("llm.bedrock_native.bedrock_converse.basic.nonstream.works")
@meta(
Subject(

View file

@ -40,7 +40,7 @@ from collections.abc import Callable
from typing import Final
import pytest
from pydantic import BaseModel
from pydantic import BaseModel, ConfigDict
from e2e_config import unique_marker
from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta
@ -65,6 +65,8 @@ VERTEX_CACHE_REJECTION_MARKER: Final = "minimum token count to start explicit ca
class CacheChatBody(BaseModel):
model_config = ConfigDict(frozen=True)
model: str
messages: list[RichMessage]
max_tokens: int = 64
@ -309,6 +311,79 @@ class TestCacheControl:
key = resources.key()
_assert_cache_read_on_second_call(model, lambda prefix: _cache_chat(client, key, model, prefix))
@pytest.mark.covers("llm.chat_completions.anthropic.prompt_cache_5m.nonstream.works")
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.ANTHROPIC,),
models=(ANTHROPIC_MODEL,),
capabilities=(Capability.PROMPT_CACHING,),
mode=Mode.NONSTREAM,
)
)
def test_anthropic_prompt_caching_reads_cache_with_version_header(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
model: Final = f"e2e-anthropic-cache-version-{unique_marker()}"
model_id: Final = client.proxy.create_model(
model,
LiteLLMParamsBody(
model=ANTHROPIC_MODEL,
api_key="os.environ/ANTHROPIC_API_KEY",
extra_headers={"anthropic-version": "2023-06-01"},
),
)
resources.defer(lambda: client.proxy.delete_model(model_id))
key: Final = resources.key()
_assert_cache_read_on_second_call(model, lambda prefix: _cache_chat(client, key, model, prefix))
@pytest.mark.covers("llm.chat_completions.anthropic.prompt_cache_5m.nonstream.works")
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.ANTHROPIC,),
models=(ANTHROPIC_MODEL,),
capabilities=(Capability.PROMPT_CACHING,),
mode=Mode.NONSTREAM,
)
)
def test_anthropic_prompt_caching_reports_creation_tokens(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
model: Final = f"e2e-anthropic-cache-create-{unique_marker()}"
model_id: Final = client.proxy.create_model(
model,
LiteLLMParamsBody(model=ANTHROPIC_MODEL, api_key="os.environ/ANTHROPIC_API_KEY"),
)
resources.defer(lambda: client.proxy.delete_model(model_id))
key: Final = resources.key()
body: Final = CacheChatBody(
model=model,
messages=[
RichMessage(
role="system",
content=[TextBlock(text=_cacheable_prefix(), cache_control=CacheControl(ttl="1h"))],
),
RichMessage(
role="user",
content=[TextBlock(text="Reply with one word.", cache_control=CacheControl(ttl="5m"))],
),
],
)
response: Final = unwrap(
client.proxy.transport.post(
"/chat/completions",
headers=client.proxy.transport.bearer(key),
json=body,
response_type=ChatResponse,
)
)
assert response.usage and response.usage.cache_creation_input_tokens, (
f"{model}: a never-seen prefix marked 1h then 5m reported no cache-creation tokens ({response.usage})"
)
@pytest.mark.covers(
"llm.chat_completions.openai.prompt_cache_5m.nonstream.works",
exercised_on=[],

View file

@ -10,17 +10,20 @@ SDK refuses to build stay on the shared transport.
from __future__ import annotations
import base64
import math
from typing import Final
from pathlib import Path
from typing import Final, Literal
import pytest
from pydantic import BaseModel, ConfigDict
from e2e_config import provider_edge_base, unique_marker
from e2e_http import assert_client_error
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta
from lifecycle import ResourceManager
from models import LiteLLMParamsBody
from proxy_client import ProxyClient
from pydantic import BaseModel
from sdk_clients import NO_PROXY_CACHE, SdkClients, response_header
pytestmark = pytest.mark.e2e
@ -28,6 +31,8 @@ pytestmark = pytest.mark.e2e
OPENAI_EMBEDDING: Final = "openai/text-embedding-3-small"
BEDROCK_TITAN_EMBEDDING: Final = "bedrock/amazon.titan-embed-text-v2:0"
COHERE_EMBEDDING: Final = "cohere/embed-v4.0"
COHERE_IMAGE_EMBEDDING: Final = "cohere/embed-english-v3.0"
CAT_IMAGE: Final = Path(__file__).parent / "fixtures" / "cat.jpg"
MISTRAL_EMBEDDING: Final = "mistral/mistral-embed"
VERTEX_TEXT_EMBEDDING: Final = "vertex_ai/text-embedding-005"
VERTEX_MULTIMODAL_EMBEDDING: Final = "vertex_ai/multimodalembedding@001"
@ -42,6 +47,29 @@ class _OptionalEmbeddingsBody(BaseModel):
input: str | list[str] | None = None
class CohereImageTokenDetails(BaseModel):
model_config = ConfigDict(frozen=True)
image_tokens: int | None = None
text_tokens: int | None = None
class CohereImageUsage(BaseModel):
model_config = ConfigDict(frozen=True)
prompt_tokens_details: CohereImageTokenDetails | None = None
class CohereImageEmbeddingResponse(BaseModel):
model_config = ConfigDict(frozen=True)
usage: CohereImageUsage | None = None
def _cat_image_data_uri() -> str:
return f"data:image/jpeg;base64,{base64.b64encode(CAT_IMAGE.read_bytes()).decode('ascii')}"
def _cosine(left: list[float], right: list[float]) -> float:
dot: Final = sum(a * b for a, b in zip(left, right, strict=True))
norms: Final = math.sqrt(sum(a * a for a in left)) * math.sqrt(sum(b * b for b in right))
@ -156,6 +184,48 @@ class TestEmbeddingsEndpoint:
LiteLLMParamsBody(model=COHERE_EMBEDDING, api_key="os.environ/COHERE_API_KEY"),
)
@pytest.mark.covers("llm.embeddings.cohere.vision.nonstream.works")
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.EMBEDDINGS,
providers=(Provider.COHERE,),
models=(COHERE_IMAGE_EMBEDDING,),
capabilities=(Capability.VISION,),
mode=Mode.NONSTREAM,
)
)
@pytest.mark.parametrize(
("input_values", "token_type"),
(
pytest.param((_cat_image_data_uri(),), "image_tokens", id="image"),
pytest.param((TOKENS_TEXT,), "text_tokens", id="text"),
),
)
def test_cohere_image_embeddings_report_image_tokens(
self,
proxy: ProxyClient,
resources: ResourceManager,
sdk: SdkClients,
input_values: tuple[str, ...],
token_type: Literal["image_tokens", "text_tokens"],
) -> None:
model: Final = f"e2e-cohere-image-{unique_marker()}"
model_id: Final = proxy.create_model(
model,
LiteLLMParamsBody(model=COHERE_IMAGE_EMBEDDING, api_key="os.environ/COHERE_API_KEY"),
)
resources.defer(lambda: proxy.delete_model(model_id))
key: Final = resources.key()
embeddings: Final = sdk.openai(key).embeddings.create(
model=model, input=list(input_values), extra_body=NO_PROXY_CACHE
)
response: Final = CohereImageEmbeddingResponse.model_validate(embeddings.model_dump())
details: Final = response.usage.prompt_tokens_details if response.usage else None
assert details is not None, f"{model}: /embeddings returned no prompt_tokens_details: {response}"
reported: Final = details.image_tokens if token_type == "image_tokens" else details.text_tokens
assert reported, f"{model}: {token_type} missing or zero in {details}"
@pytest.mark.covers("llm.embeddings.vertex.basic.nonstream.works")
@meta(
Subject(

View file

@ -0,0 +1,125 @@
from __future__ import annotations
import base64
import os
from pathlib import Path
from typing import Final, Literal
import pytest
from e2e_config import SAGEMAKER_NOVA2_LITE_OPT_IN_ENV, SAGEMAKER_NOVA_OPT_IN_ENV, unique_marker
from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta
from lifecycle import ResourceManager
from models import LiteLLMParamsBody
from proxy_client import ProxyClient
from sdk_clients import SdkClients
pytestmark = [pytest.mark.e2e, pytest.mark.provider_live]
CAT_IMAGE: Final = Path(__file__).parent / "fixtures" / "cat.jpg"
def _deployment(proxy: ProxyClient, resources: ResourceManager, endpoint_env: str, label: str) -> tuple[str, str]:
endpoint: Final = os.environ[endpoint_env]
model_name: Final = f"sagemaker_nova/{endpoint}"
model: Final = f"e2e-{label}-{unique_marker()}"
model_id: Final = proxy.create_model(model, LiteLLMParamsBody(model=model_name))
resources.defer(lambda: proxy.delete_model(model_id))
return model, resources.key()
def _cat_image_data_uri() -> str:
return f"data:image/jpeg;base64,{base64.b64encode(CAT_IMAGE.read_bytes()).decode('ascii')}"
@pytest.mark.sagemaker_nova
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.SAGEMAKER,),
mode=Mode.NONSTREAM,
)
)
class TestSagemakerNova:
@pytest.mark.covers("llm.chat_completions.sagemaker.vision.nonstream.works")
def test_multimodal_image_input_is_accepted(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _deployment(proxy, resources, SAGEMAKER_NOVA_OPT_IN_ENV, "nova-image")
response: Final = sdk.openai(key).chat.completions.create(
model=model,
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "What animal is in this image? Answer in one word."},
{
"type": "image_url",
"image_url": {"url": _cat_image_data_uri()},
},
],
}
],
max_tokens=64,
)
answer: Final = response.choices[0].message.content or ""
assert "cat" in answer.lower(), f"{model}: the image did not reach the model, it answered {answer!r}"
@pytest.mark.covers("llm.chat_completions.sagemaker.top_k.nonstream.works")
def test_nova_specific_top_k_parameter_is_accepted(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _deployment(proxy, resources, SAGEMAKER_NOVA_OPT_IN_ENV, "nova-top-k")
response: Final = sdk.openai(key).chat.completions.create(
model=model,
messages=[{"role": "user", "content": "Say hello."}],
max_tokens=32,
extra_body={"top_k": 40},
)
assert response.choices[0].message.content
@pytest.mark.covers("llm.chat_completions.sagemaker.logprobs.nonstream.works")
def test_logprobs_returns_token_details(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _deployment(proxy, resources, SAGEMAKER_NOVA_OPT_IN_ENV, "nova-logprobs")
response: Final = sdk.openai(key).chat.completions.create(
model=model,
messages=[{"role": "user", "content": "Reply with the word hello."}],
max_tokens=8,
logprobs=True,
top_logprobs=3,
)
logprobs: Final = response.choices[0].logprobs
assert logprobs and logprobs.content, f"{model}: no logprobs content in {response}"
assert len(logprobs.content[0].top_logprobs) == 3, f"{model}: top_logprobs=3 not honored: {logprobs.content[0]}"
@pytest.mark.sagemaker_nova2_lite
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.SAGEMAKER,),
capabilities=(Capability.REASONING,),
mode=Mode.NONSTREAM,
)
)
class TestSagemakerNova2Lite:
@pytest.mark.parametrize("reasoning_effort", ("low", "high"))
@pytest.mark.covers("llm.chat_completions.sagemaker.thinking.nonstream.works")
def test_reasoning_effort_is_accepted(
self,
reasoning_effort: Literal["low", "high"],
proxy: ProxyClient,
resources: ResourceManager,
sdk: SdkClients,
) -> None:
model, key = _deployment(proxy, resources, SAGEMAKER_NOVA2_LITE_OPT_IN_ENV, "nova2-reasoning")
response: Final = sdk.openai(key).chat.completions.create(
model=model,
messages=[{"role": "user", "content": "What is 2 + 2?"}],
max_tokens=32,
reasoning_effort=reasoning_effort,
)
assert response.choices[0].message.content

View file

@ -0,0 +1,116 @@
from __future__ import annotations
from typing import Final, TypedDict, cast
import pytest
from typing_extensions import NotRequired, ReadOnly
from e2e_config import unique_marker
from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta
from lifecycle import ResourceManager
from models import LiteLLMParamsBody
from proxy_client import ProxyClient
from sdk_clients import SdkClients
pytestmark = [pytest.mark.e2e, pytest.mark.provider_live]
GEMINI_MODEL: Final = "vertex_ai/gemini-3-flash-preview"
class GoogleSearchTool(TypedDict):
googleSearch: ReadOnly[dict[str, str]]
class VertexGroundingBody(TypedDict):
tools: ReadOnly[tuple[GoogleSearchTool, ...]]
class GoogleMapsOptions(TypedDict):
enableWidget: ReadOnly[bool]
latitude: NotRequired[ReadOnly[float]]
longitude: NotRequired[ReadOnly[float]]
languageCode: NotRequired[ReadOnly[str]]
class GoogleMapsTool(TypedDict):
googleMaps: ReadOnly[GoogleMapsOptions]
class VertexMapsBody(TypedDict):
tools: ReadOnly[tuple[GoogleMapsTool, ...]]
def _register(proxy: ProxyClient, resources: ResourceManager, base: str, model: str, location: str) -> tuple[str, str]:
model_name: Final = f"{base}-{unique_marker()}"
model_id: Final = proxy.create_model(
model_name,
LiteLLMParamsBody(
model=model,
vertex_project="os.environ/VERTEXAI_PROJECT",
vertex_location=location,
vertex_credentials="os.environ/VERTEXAI_CREDENTIALS",
),
)
resources.defer(lambda: proxy.delete_model(model_id))
return model_name, resources.key()
class TestVertexCapabilities:
@pytest.mark.covers("llm.chat_completions.vertex.tool_use.nonstream.works")
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.VERTEX_AI,),
models=(GEMINI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_google_maps_tool_accepts_widget_and_location(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _register(proxy, resources, "e2e-vertex-maps", GEMINI_MODEL, "global")
options: Final[tuple[GoogleMapsOptions, ...]] = (
{"enableWidget": True},
{"enableWidget": True, "latitude": 37.7749, "longitude": -122.4194, "languageCode": "en_US"},
)
client: Final = sdk.openai(key)
bodies: Final[tuple[VertexMapsBody, ...]] = tuple(
{"tools": ({"googleMaps": maps_options},)} for maps_options in options
)
responses: Final = tuple(
client.chat.completions.create(
model=model,
messages=[{"role": "user", "content": "What restaurants are nearby?"}],
extra_body=body,
)
for body in bodies
)
assert all(response.choices[0].message.content for response in responses)
@pytest.mark.covers("llm.chat_completions.vertex.web_search.stream.works")
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.VERTEX_AI,),
models=(GEMINI_MODEL,),
capabilities=(Capability.WEB_SEARCH,),
mode=Mode.STREAM,
)
)
def test_google_search_grounding_metadata_is_returned(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _register(proxy, resources, "e2e-vertex-grounding", GEMINI_MODEL, "global")
grounding_body: Final[VertexGroundingBody] = {"tools": ({"googleSearch": {}},)}
chunks: Final = sdk.openai(key).chat.completions.create(
model=model,
messages=[{"role": "user", "content": "What is today's weather in San Francisco?"}],
stream=True,
extra_body=grounding_body,
)
metadata: Final = tuple(
cast(object | None, getattr(chunk, "vertex_ai_grounding_metadata", None)) for chunk in chunks
)
assert any(value is not None for value in metadata)

View file

@ -20,3 +20,5 @@ markers =
otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set
otel_tls: needs a stack whose gateway exports OTLP over TLS signed by the CA in SSL_CERT_FILE; deselected unless E2E_OTEL_EXPORTER_ENDPOINT is set
secret_manager: needs a proxy booted from gateway/secret_manager_<system>_ci_config.yml against that live secret manager; deselected unless E2E_SECRET_MANAGER names the backend (see secret_manager/secret_backends.py)
sagemaker_nova: needs a deployed SageMaker Nova endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA_ENDPOINT names it
sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it

View file

@ -17,7 +17,14 @@ from typing import Final, cast
import httpx
import uvicorn
from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations
from _fake_openai_endpoint_server import (
chat_completions,
completions,
embeddings,
health,
moderations,
triton_embeddings,
)
from integration.cost_calculation.cost_tracking_case import (
BinaryResponse,
EventStreamEvent,
@ -550,6 +557,7 @@ class Provider:
Route("/v1/completions", completions, methods=["POST"]),
Route("/v1/embeddings", embeddings, methods=["POST"]),
Route("/v1/moderations", moderations, methods=["POST"]),
Route("/triton/embeddings", triton_embeddings, methods=["POST"]),
Route("/vector_stores/{vector_store_id}/search", self.vector_store_search, methods=["POST"]),
Route("/__interactions/{interaction_id}", self.interaction_state, methods=["PUT", "DELETE"]),
Route("/v1beta/interactions/{interaction_id}:cancel", self.cancel_interaction, methods=["POST"]),

View file

@ -0,0 +1,72 @@
from __future__ import annotations
from collections.abc import Iterator
from typing import Final
import httpx
import pytest
from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value
from openai.types.moderation_create_response import ModerationCreateResponse
@pytest.fixture
def upstream(gateway: Gateway) -> Iterator[httpx.Client]:
with httpx.Client(base_url=gateway.upstream_url, timeout=10, trust_env=False) as client:
yield client
def test_chat_completion_shape(upstream: httpx.Client) -> None:
response: Final = upstream.post(
"/v1/chat/completions",
json={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
)
assert response.status_code == 200, response.text
body: Final = JSON_OBJECT.validate_json(response.content)
message: Final = object_value(object_value(list_value(body["choices"])[0])["message"])
assert message == {
"role": "assistant",
"content": "Hello! This is a mock response from the fake OpenAI endpoint.",
}
assert body["usage"] == {
"prompt_tokens": 20,
"completion_tokens": 20,
"total_tokens": 40,
}
def test_moderations_route_parses_as_an_openai_response(upstream: httpx.Client) -> None:
response: Final = upstream.post(
"/v1/moderations",
json={
"input": ["I want to harm someone", "hello"],
"model": "omni-moderation-latest",
},
)
assert response.status_code == 200, response.text
parsed: Final = ModerationCreateResponse.model_validate_json(response.content)
assert parsed.model == "omni-moderation-latest"
assert len(parsed.results) == 2
assert parsed.results[0].categories.violence is False
def test_triton_embeddings_route(upstream: httpx.Client) -> None:
response: Final = upstream.post("/triton/embeddings", json={"inputs": []})
assert response.status_code == 200, response.text
output: Final = object_value(list_value(JSON_OBJECT.validate_json(response.content)["outputs"])[0])
assert output["shape"] == [1, 2]
assert output["data"] == [0.1, 0.2]
def test_slow_model_blocks_past_client_timeout(upstream: httpx.Client) -> None:
with pytest.raises(httpx.TimeoutException):
upstream.post(
"/v1/chat/completions",
json={
"model": "slow-endpoint",
"messages": [{"role": "user", "content": "hi"}],
},
timeout=0.5,
)

View file

@ -0,0 +1,58 @@
from __future__ import annotations
import json
from typing import Final
from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value
from integration._support.wire import Reply, Request, wire_server
_BACKEND: Final = "typhoon2-quen2.5-7b-instruct"
_API_KEY: Final = "synthetic-lm-studio-key"
_PROMPT: Final = "What's the weather like in San Francisco?"
_ANSWER: Final = "Foggy and 14 degrees."
def test_lm_studio_completion(gateway: Gateway) -> None:
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/v1/chat/completions"
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
body: Final = JSON_OBJECT.validate_json(request.body)
assert body["model"] == _BACKEND
assert body["messages"] == [{"role": "user", "content": _PROMPT}]
return Reply(
body=json.dumps(
{
"id": "chatcmpl-lm-studio",
"object": "chat.completion",
"created": 1,
"model": _BACKEND,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": _ANSWER},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 12,
"completion_tokens": 6,
"total_tokens": 18,
},
}
).encode()
)
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"lm_studio/{_BACKEND}", api_base=f"{wire.url}/v1", api_key=_API_KEY)
response: Final = gateway.chat(model, text=_PROMPT)
received: Final = wire.drain()
assert len(received) == 1
message: Final = object_value(object_value(list_value(response["choices"])[0])["message"])
assert message["content"] == _ANSWER
assert response["usage"] == {
"prompt_tokens": 12,
"completion_tokens": 6,
"total_tokens": 18,
}

View file

@ -0,0 +1,43 @@
from __future__ import annotations
import threading
from typing import Final
from integration._support.client import JSON_OBJECT, Gateway, object_value
from integration._support.wire import Reply, Request, wire_server
_API_KEY: Final = "synthetic-openai-key"
def test_openai_embedding_timeouts(gateway: Gateway) -> None:
release: Final = threading.Event()
def respond(request: Request) -> Reply:
assert request.target == "/v1/embeddings"
assert JSON_OBJECT.validate_json(request.body)["input"] == ["good morning from litellm"]
release.wait(timeout=30)
return Reply(drop_connection=True)
with wire_server(respond) as wire:
try:
with gateway.scenario() as scenario:
model: Final = scenario.model(
model="openai/slow-endpoint",
api_base=f"{wire.url}/v1",
api_key=_API_KEY,
timeout=0.5,
)
response: Final = gateway.request(
"POST",
"/v1/embeddings",
{"model": model, "input": ["good morning from litellm"]},
)
finally:
release.set()
received: Final = wire.drain()
assert response.status_code == 408, response.text
error: Final = object_value(JSON_OBJECT.validate_json(response.content)["error"])
assert error["code"] == "408"
assert received
assert {request.headers["x-stainless-read-timeout"] for request in received} == {"0.5"}

View file

@ -0,0 +1,39 @@
from __future__ import annotations
import hashlib
from pathlib import Path
from typing import Final
from urllib.parse import urlsplit
from integration._support.client import Gateway, string_value
from integration._support.forward_proxy import refusing_forward_proxy
from integration._support.process import owned_proxy
_BUNDLED_COST_MAP: Final = Path(__file__).resolve().parents[3] / "litellm/model_prices_and_context_window_backup.json"
def _git_blob_id(body: bytes) -> str:
return hashlib.sha1(b"blob %d\0" % len(body) + body, usedforsecurity=False).hexdigest()
def test_proxy_becomes_ready_without_network(gateway: Gateway, tmp_path: Path) -> None:
with refusing_forward_proxy() as egress:
no_network: Final = {
"HTTPS_PROXY": egress.url,
"HTTP_PROXY": egress.url,
"ALL_PROXY": egress.url,
"NO_PROXY": "127.0.0.1,localhost",
}
with owned_proxy(
gateway, tmp_path, no_network, remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP", "LITELLM_LICENSE")
) as isolated:
readiness: Final = isolated.request("GET", "/health/readiness")
source: Final = isolated.get("/model/cost_map/source")
refused: Final = frozenset(egress.targets())
assert readiness.status_code == 200, readiness.text
assert source["source"] == "local"
assert source["is_env_forced"] is False
assert string_value(source["fallback_reason"]).startswith("Remote fetch failed")
assert source["source_revision"] == _git_blob_id(_BUNDLED_COST_MAP.read_bytes())
assert refused == {f"{urlsplit(string_value(source['url'])).hostname}:443"}

View file

@ -0,0 +1,147 @@
from __future__ import annotations
import json
import os
import uuid
from collections.abc import Callable
from typing import Final
from integration._support.client import JSON_OBJECT, Gateway, eventually, list_value, object_value
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
from redis import Redis
_TOOL: Final = {
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a city",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["location"],
},
},
}
_MESSAGES: Final = [{"role": "user", "content": "Weather in Boston?"}]
def _frame(stream_id: str, delta: JsonValue, finish_reason: str | None = None) -> bytes:
chunk: Final = {
"id": stream_id,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4.1-mini",
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
}
return f"data: {json.dumps(chunk)}\n\n".encode()
def _frames(stream_id: str) -> tuple[bytes, ...]:
return (
_frame(
stream_id,
{
"role": "assistant",
"tool_calls": [
{
"index": 0,
"id": "call-weather",
"type": "function",
"function": {
"name": "get_current_weather",
"arguments": '{"location":"Bos',
},
}
],
},
),
_frame(stream_id, {"tool_calls": [{"index": 0, "function": {"arguments": 'ton","unit":"fahrenheit"}'}}]}),
_frame(stream_id, {}, finish_reason="tool_calls"),
b"data: [DONE]\n\n",
)
def _responder(stream_id: str) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2024-02-15-preview"
assert request.headers["api-key"] == "synthetic-azure-key"
body: Final = JSON_OBJECT.validate_json(request.body)
assert body["model"] == "gpt-4.1-mini"
assert body["messages"] == _MESSAGES
assert body["stream"] is True
assert body["tools"] == [_TOOL]
assert body["tool_choice"] == "auto"
return Reply(content_type="text/event-stream", chunks=_frames(stream_id))
return respond
def _stored(stream_id: str) -> bool:
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as store:
return any(
stream_id.encode() in (store.get(key) or b"")
for key in store.scan_iter(count=1000)
if store.type(key) == b"string"
)
def _tool_call_stream(gateway: Gateway, model: str) -> tuple[tuple[str, ...], str, tuple[str, ...]]:
with gateway.client.stream(
"POST",
"/v1/chat/completions",
json={
"model": model,
"messages": _MESSAGES,
"tools": [_TOOL],
"tool_choice": "auto",
"stream": True,
},
headers={"Authorization": f"Bearer {gateway.key}"},
) as response:
assert response.status_code == 200, response.read().decode()
events: Final = tuple(
JSON_OBJECT.validate_json(line.removeprefix("data: "))
for line in response.iter_lines()
if line.startswith("data: ") and line != "data: [DONE]"
)
choices: Final = tuple(object_value(list_value(event["choices"])[0]) for event in events)
calls: Final = tuple(
object_value(call)
for choice in choices
for call in list_value(object_value(choice["delta"]).get("tool_calls") or [])
)
functions: Final = tuple(object_value(call["function"]) for call in calls)
return (
tuple(str(function["name"]) for function in functions if isinstance(function.get("name"), str)),
"".join(str(function["arguments"]) for function in functions if isinstance(function.get("arguments"), str)),
tuple(str(choice["finish_reason"]) for choice in choices if isinstance(choice.get("finish_reason"), str)),
)
def test_azure_astreaming_and_function_calling(gateway: Gateway) -> None:
stream_id: Final = f"azure-stream-{uuid.uuid4().hex}"
with wire_server(_responder(stream_id)) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model="azure/gpt-4.1-mini",
api_base=wire.url,
api_key="synthetic-azure-key",
api_version="2024-02-15-preview",
)
first: Final = _tool_call_stream(gateway, model)
eventually(lambda: _stored(stream_id), bool)
cached: Final = _tool_call_stream(gateway, model)
received: Final = wire.drain()
expected: Final = (
("get_current_weather",),
'{"location":"Boston","unit":"fahrenheit"}',
("tool_calls",),
)
assert first == expected
assert cached == expected
assert len(received) == 1

View file

@ -1,3 +0,0 @@
{
"gpt-3.5-turbo": 7.7e-05
}

View file

@ -1,543 +0,0 @@
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
Delta,
Function,
ModelResponse,
StreamingChoices,
)
chunks = [
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="To answer",
role="assistant",
function_call=None,
tool_calls=None,
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content=" your", role=None, function_call=None, tool_calls=None
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content=" question about",
role=None,
function_call=None,
tool_calls=None,
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content=" how", role=None, function_call=None, tool_calls=None
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content=" many rows are in the ",
role=None,
function_call=None,
tool_calls=None,
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="'users' table, I",
role=None,
function_call=None,
tool_calls=None,
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="'ll", role=None, function_call=None, tool_calls=None
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content=" need to", role=None, function_call=None, tool_calls=None
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content=" run", role=None, function_call=None, tool_calls=None
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content=" a SQL query.",
role=None,
function_call=None,
tool_calls=None,
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content=" Let", role=None, function_call=None, tool_calls=None
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content=" me", role=None, function_call=None, tool_calls=None
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content=" ", role=None, function_call=None, tool_calls=None
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="do that for",
role=None,
function_call=None,
tool_calls=None,
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content=" you.", role=None, function_call=None, tool_calls=None
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role=None,
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id="toolu_01H3AjkLpRtGQrof13CBnWfK",
function=Function(arguments="", name="sql_query"),
type="function",
index=1,
)
],
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role=None,
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id=None,
function=Function(arguments="", name=None),
type="function",
index=1,
)
],
),
logprobs=None,
)
],
created=1722656356,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role=None,
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id=None,
function=Function(arguments='{"', name=None),
type="function",
index=1,
)
],
),
logprobs=None,
)
],
created=1722656357,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role=None,
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id=None,
function=Function(arguments='query": ', name=None),
type="function",
index=1,
)
],
),
logprobs=None,
)
],
created=1722656357,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role=None,
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id=None,
function=Function(arguments='"SELECT C', name=None),
type="function",
index=1,
)
],
),
logprobs=None,
)
],
created=1722656357,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role=None,
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id=None,
function=Function(arguments="OUNT(*", name=None),
type="function",
index=1,
)
],
),
logprobs=None,
)
],
created=1722656357,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role=None,
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id=None,
function=Function(arguments=") ", name=None),
type="function",
index=1,
)
],
),
logprobs=None,
)
],
created=1722656357,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role=None,
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id=None,
function=Function(arguments="FROM use", name=None),
type="function",
index=1,
)
],
),
logprobs=None,
)
],
created=1722656357,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role=None,
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id=None,
function=Function(arguments='rs;"}', name=None),
type="function",
index=1,
)
],
),
logprobs=None,
)
],
created=1722656357,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
ModelResponse(
id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f",
choices=[
StreamingChoices(
finish_reason="tool_calls",
index=0,
delta=Delta(
content=None, role=None, function_call=None, tool_calls=None
),
logprobs=None,
)
],
created=1722656357,
model="claude-3-5-sonnet-20240620",
object="chat.completion.chunk",
system_fingerprint=None,
),
]

View file

@ -1,25 +0,0 @@
import pytest
@pytest.mark.asyncio
async def test_langfuse_double_logging():
import litellm
litellm.set_verbose = True
litellm.success_callback = ["langfuse"]
litellm.failure_callback = ["langfuse"] # logs errors to langfuse
models = ["gpt-4o-mini", "claude-3-5-haiku-20241022"]
messages = [
{"role": "user", "content": "Hello, how are you?"},
]
resp = await litellm.acompletion(
model=models[0],
messages=messages,
temperature=0.0,
fallbacks=models[1:],
# metadata={"generation_name": "test-gen", "project": "litellm-test"},
)
return resp

File diff suppressed because it is too large Load diff

View file

@ -1,379 +0,0 @@
import json
import os
import traceback
from dotenv import load_dotenv
load_dotenv()
import io
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from test_amazing_vertex_completion import load_vertex_ai_credentials
from test_streaming import streaming_format_tests
import litellm
from litellm import RateLimitError, Timeout, completion, completion_cost, embedding
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
# 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"}]
def logger_fn(user_model_dict):
print(f"user_model_dict: {user_model_dict}")
@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.fixture
def anthropic_messages():
return [
# System Message
{
"role": "system",
"content": [
{
"type": "text",
"text": "Here is the full text of a complex legal agreement" * 500,
"cache_control": {"type": "ephemeral"},
}
],
},
# marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
},
{
"role": "assistant",
"content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo",
},
# The final turn is marked with cache-control, for continuing in followups.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
},
]
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_anthropic_vertex_ai_prompt_caching(anthropic_messages, sync_mode):
litellm.turn_on_debug()
load_vertex_ai_credentials()
client = HTTPHandler() if sync_mode else AsyncHTTPHandler()
with patch.object(client, "post", return_value=MagicMock()) as mock_post:
try:
if sync_mode:
response = completion(
model="vertex_ai/claude-3-5-sonnet-v2@20241022 ",
messages=anthropic_messages,
client=client,
)
else:
response = await litellm.acompletion(
model="vertex_ai/claude-3-5-sonnet-v2@20241022 ",
messages=anthropic_messages,
client=client,
)
except Exception as e:
print(f"Error: {e}")
mock_post.assert_called_once()
print(mock_post.call_args.kwargs["headers"])
assert "anthropic-beta" not in mock_post.call_args.kwargs["headers"]
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio()
async def test_anthropic_api_prompt_caching_basic():
litellm.set_verbose = True
response = await litellm.acompletion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[
# System Message
{
"role": "system",
"content": [
{
"type": "text",
"text": "Here is the full text of a complex legal agreement"
* 100,
"cache_control": {"type": "ephemeral"},
}
],
},
# marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
},
{
"role": "assistant",
"content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo",
},
# The final turn is marked with cache-control, for continuing in followups.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
},
],
temperature=0.2,
max_tokens=10,
extra_headers={
"anthropic-version": "2023-06-01",
},
)
print("response=", response)
assert "cache_read_input_tokens" in response.usage
assert "cache_creation_input_tokens" in response.usage
# Assert either a cache entry was created or cache was read - changes depending on the anthropic api ttl
assert (response.usage.cache_read_input_tokens > 0) or (
response.usage.cache_creation_input_tokens > 0
)
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio()
async def test_anthropic_api_prompt_caching_basic_with_cache_creation():
from uuid import uuid4
random_id = uuid4()
litellm.set_verbose = True
response = await litellm.acompletion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[
# System Message
{
"role": "system",
"content": [
{
"type": "text",
"text": "Here is the full text of a complex legal agreement {}".format(
random_id
)
* 100,
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
},
# marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral", "ttl": "5m"},
}
],
},
{
"role": "assistant",
"content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo",
},
# The final turn is marked with cache-control, for continuing in followups.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
},
],
temperature=0.2,
max_tokens=10,
extra_headers={
"anthropic-version": "2023-06-01",
},
)
print("response=", response)
assert "cache_read_input_tokens" in response.usage
assert "cache_creation_input_tokens" in response.usage
# Assert either a cache entry was created or cache was read - changes depending on the anthropic api ttl
assert (response.usage.cache_read_input_tokens > 0) or (
response.usage.cache_creation_input_tokens > 0
)
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio()
async def test_anthropic_api_prompt_caching_no_headers():
litellm.set_verbose = True
response = await litellm.acompletion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[
# System Message
{
"role": "system",
"content": [
{
"type": "text",
"text": "Here is the full text of a complex legal agreement"
* 100,
"cache_control": {"type": "ephemeral"},
}
],
},
# marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
},
{
"role": "assistant",
"content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo",
},
# The final turn is marked with cache-control, for continuing in followups.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
},
],
temperature=0.2,
max_tokens=10,
)
print("response=", response)
assert "cache_read_input_tokens" in response.usage
assert "cache_creation_input_tokens" in response.usage
# Assert either a cache entry was created or cache was read - changes depending on the anthropic api ttl
assert (response.usage.cache_read_input_tokens > 0) or (
response.usage.cache_creation_input_tokens > 0
)
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio()
async def test_anthropic_api_prompt_caching_streaming():
response = await litellm.acompletion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[
# System Message
{
"role": "system",
"content": [
{
"type": "text",
"text": "Here is the full text of a complex legal agreement"
* 100,
"cache_control": {"type": "ephemeral"},
}
],
},
# marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
},
{
"role": "assistant",
"content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo",
},
# The final turn is marked with cache-control, for continuing in followups.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
},
],
temperature=0.2,
max_tokens=10,
stream=True,
stream_options={"include_usage": True},
)
idx = 0
is_cache_read_input_tokens_in_usage = False
is_cache_creation_input_tokens_in_usage = False
async for chunk in response:
streaming_format_tests(idx=idx, chunk=chunk)
# Assert either a cache entry was created or cache was read - changes depending on the anthropic api ttl
if hasattr(chunk, "usage"):
print("Received final usage - {}".format(chunk.usage))
if hasattr(chunk, "usage") and hasattr(chunk.usage, "cache_read_input_tokens"):
is_cache_read_input_tokens_in_usage = True
if hasattr(chunk, "usage") and hasattr(
chunk.usage, "cache_creation_input_tokens"
):
is_cache_creation_input_tokens_in_usage = True
idx += 1
print("response=", response)
assert (
is_cache_read_input_tokens_in_usage and is_cache_creation_input_tokens_in_usage
)

View file

@ -1,93 +0,0 @@
#### What this tests ####
# This tests the the acompletion function #
import asyncio
import pytest
import litellm
from litellm import acompletion
litellm.num_retries = 3
# test_sync_response_anyscale()
def test_async_response_openai():
litellm.set_verbose = True
async def test_get_response():
user_message = "Hello, how are you?"
messages = [{"content": user_message, "role": "user"}]
tools = [
{
"type": "function",
"function": {
"name": "get_current_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"],
},
},
}
]
try:
response = await acompletion(
model="gpt-3.5-turbo",
messages=messages,
tools=tools,
parallel_tool_calls=True,
timeout=5,
)
print(f"response: {response}")
print(f"response ms: {response._response_ms}")
except litellm.Timeout as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred: {e}")
print(e)
asyncio.run(test_get_response())
# test_async_response_openai()
# test_async_anyscale_response()
# test_async_completion_cloudflare()
# test_get_cloudflare_response_streaming()
# test_get_response_streaming()
# test_get_response_non_openai_streaming()

File diff suppressed because it is too large Load diff

View file

@ -1,466 +0,0 @@
import asyncio
import json
import os
import time
import traceback
from typing import Optional
from unittest.mock import MagicMock, patch
import pytest
import litellm
import litellm.cost_calculator
from litellm import (
completion_cost,
)
from litellm.litellm_core_utils.litellm_logging import CustomLogger
from litellm.llms.custom_httpx.http_handler import HTTPHandler
class CustomLoggingHandler(CustomLogger):
response_cost: Optional[float] = None
def __init__(self):
super().__init__()
def log_success_event(self, kwargs, response_obj, start_time, end_time):
self.response_cost = kwargs["response_cost"]
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
print(f"kwargs - {kwargs}")
print(f"kwargs response cost - {kwargs.get('response_cost')}")
self.response_cost = kwargs["response_cost"]
print(f"response_cost: {self.response_cost} ")
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
print("Reaches log failure event!")
self.response_cost = kwargs["response_cost"]
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
print("Reaches async log failure event!")
self.response_cost = kwargs["response_cost"]
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_custom_pricing(sync_mode):
new_handler = CustomLoggingHandler()
litellm.callbacks = [new_handler]
if sync_mode:
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey!"}],
mock_response="What do you want?",
input_cost_per_token=0.0,
output_cost_per_token=0.0,
)
time.sleep(5)
else:
response = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey!"}],
mock_response="What do you want?",
input_cost_per_token=0.0,
output_cost_per_token=0.0,
)
await asyncio.sleep(5)
print(f"new_handler.response_cost: {new_handler.response_cost}")
assert new_handler.response_cost is not None
assert new_handler.response_cost == 0
@pytest.mark.parametrize(
"sync_mode",
[True, False],
)
@pytest.mark.asyncio
async def test_failure_completion_cost(sync_mode):
new_handler = CustomLoggingHandler()
litellm.callbacks = [new_handler]
if sync_mode:
try:
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey!"}],
mock_response=Exception("this should trigger an error"),
)
except Exception:
pass
time.sleep(5)
else:
try:
response = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey!"}],
mock_response=Exception("this should trigger an error"),
)
except Exception:
pass
await asyncio.sleep(5)
print(f"new_handler.response_cost: {new_handler.response_cost}")
assert new_handler.response_cost is not None
assert new_handler.response_cost == 0
# print(results)
# test_get_gpt3_tokens()
# test_get_palm_tokens()
# test_zephyr_hf_tokens()
# test_cost_ft_gpt_35()
# test_cost_azure_gpt_35()
# test_cost_azure_embedding()
# def test_vertex_ai_embedding_completion_cost_e2e():
# """
# Relevant issue - https://github.com/BerriAI/litellm/issues/4630
# """
# from test_amazing_vertex_completion import load_vertex_ai_credentials
# load_vertex_ai_credentials()
# os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
# litellm.model_cost = litellm.get_model_cost_map(url="")
# text = "The quick brown fox jumps over the lazy dog."
# input_tokens = litellm.token_counter(
# model="vertex_ai/textembedding-gecko", text=text
# )
# model_info = litellm.get_model_info(model="vertex_ai/textembedding-gecko")
# print("\nExpected model info:\n{}\n\n".format(model_info))
# expected_input_cost = input_tokens * model_info["input_cost_per_token"]
# ## CALCULATED COST
# resp = litellm.embedding(model="textembedding-gecko", input=[text])
# calculated_input_cost = resp._hidden_params["response_cost"]
# assert round(expected_input_cost, 6) == round(calculated_input_cost, 6)
# print("expected_input_cost: {}".format(expected_input_cost))
# print("calculated_input_cost: {}".format(calculated_input_cost))
# assert False
def test_vertex_ai_llama_predict_cost():
model = "meta/llama3-405b-instruct-maas"
messages = [{"role": "user", "content": "Hey, hows it going???"}]
custom_llm_provider = "vertex_ai"
predictive_cost = completion_cost(
model=model, messages=messages, custom_llm_provider=custom_llm_provider
)
assert predictive_cost == 0
def test_completion_cost_azure_common_deployment_name():
from litellm.utils import (
CallTypes,
Choices,
Message,
ModelResponse,
Usage,
)
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {
"model": "azure/gpt-4-0314",
"max_tokens": 4096,
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
},
"model_info": {"base_model": "azure/gpt-4"},
}
]
)
response = ModelResponse(
id="chatcmpl-876cce24-e520-4cf8-8649-562a9be11c02",
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(
content="Hi! I'm an AI, so I don't have emotions or feelings like humans do, but I'm functioning properly and ready to help with any questions or topics you'd like to discuss! How can I assist you today?",
role="assistant",
),
)
],
created=1717519830,
model="gpt-4",
object="chat.completion",
system_fingerprint="fp_c1a4bcec29",
usage=Usage(completion_tokens=46, prompt_tokens=17, total_tokens=63),
)
response._hidden_params["custom_llm_provider"] = "azure"
print(response)
with patch.object(
litellm.cost_calculator, "completion_cost", new=MagicMock()
) as mock_client:
_ = litellm.response_cost_calculator(
response_object=response,
model="gpt-4-0314",
custom_llm_provider="azure",
call_type=CallTypes.acompletion.value,
optional_params={},
base_model="azure/gpt-4",
)
mock_client.assert_called()
print(f"mock_client.call_args: {mock_client.call_args.kwargs}")
assert "azure/gpt-4" == mock_client.call_args.kwargs["base_model"]
@pytest.mark.parametrize(
"model",
[
"fireworks_ai/accounts/fireworks/models/deepseek-v3p1",
],
)
def test_completion_cost_fireworks_ai(model):
"""
Mocked so it does not depend on Fireworks' rotating serverless catalog.
Validates the Fireworks cost path: a parsed response with usage yields a
non-zero cost against the local cost map.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
mock_response_data = {
"id": "chatcmpl-test",
"object": "chat.completion",
"created": 1234567890,
"model": model.split("fireworks_ai/")[-1],
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Going great, thanks!"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 8, "completion_tokens": 5, "total_tokens": 13},
}
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.json.return_value = mock_response_data
mock_response.text = json.dumps(mock_response_data)
sync_handler = HTTPHandler()
messages = [{"role": "user", "content": "Hey, how's it going?"}]
with patch.object(HTTPHandler, "post", return_value=mock_response):
resp = litellm.completion(model=model, messages=messages, client=sync_handler)
cost = completion_cost(completion_response=resp)
assert cost > 0
def test_completion_cost_vertex_llama3():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
from litellm.utils import Choices, Message, ModelResponse, Usage
response = ModelResponse(
id="2024-09-19|14:52:01.823070-07|3.10.13.64|-333502972",
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(
content="My name is Litellm Bot, and I'm here to help you with any questions or tasks you may have. As for the weather, I'd be happy to provide you with the current conditions and forecast for your location. However, I'm a large language model, I don't have real-time access to your location, so I'll need you to tell me where you are or provide me with a specific location you're interested in knowing the weather for.\\n\\nOnce you provide me with that information, I can give you the current weather conditions, including temperature, humidity, wind speed, and more, as well as a forecast for the next few days. Just let me know how I can assist you!",
role="assistant",
tool_calls=None,
function_call=None,
),
)
],
created=1726782721,
model="vertex_ai/meta/llama3-405b-instruct-maas",
object="chat.completion",
system_fingerprint="",
usage=Usage(
completion_tokens=152,
prompt_tokens=27,
total_tokens=179,
completion_tokens_details=None,
),
)
model = "vertex_ai/meta/llama3-8b-instruct-maas"
cost = completion_cost(model=model, completion_response=response)
assert cost == 0
def test_completion_cost_azure_tts():
from unittest.mock import MagicMock
args = {
"response_object": MagicMock,
"model": "tts-1",
"cache_hit": None,
"custom_llm_provider": "azure",
"base_model": None,
"call_type": "aspeech",
"optional_params": {},
"custom_pricing": False,
}
litellm.response_cost_calculator(**args)
def test_moderations():
from litellm import moderation
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm.add_known_models()
assert "omni-moderation-latest" in litellm.model_cost
print(
f"litellm.model_cost['omni-moderation-latest']: {litellm.model_cost['omni-moderation-latest']}"
)
assert "omni-moderation-latest" in litellm.open_ai_chat_completion_models
response = moderation("I am a bad person", model="omni-moderation-latest")
cost = completion_cost(response, model="omni-moderation-latest")
assert cost == 0
def test_cost_calculator_azure_embedding():
from litellm.cost_calculator import response_cost_calculator
from litellm.types.utils import EmbeddingResponse, Usage
kwargs = {
"response_object": EmbeddingResponse(
model="text-embedding-3-small",
data=[{"embedding": [1, 2, 3]}],
usage=Usage(prompt_tokens=10, completion_tokens=10),
),
"model": "text-embedding-3-small",
"cache_hit": None,
"custom_llm_provider": None,
"base_model": "azure/text-embedding-3-small",
"call_type": "aembedding",
"optional_params": {},
"custom_pricing": False,
"prompt": "Hello, world!",
}
try:
response_cost_calculator(**kwargs)
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error: {e}")
def test_add_known_models():
litellm.add_known_models()
assert (
"bedrock/us-west-1/meta.llama3-70b-instruct-v1:0" not in litellm.bedrock_models
)
# @pytest.mark.parametrize(
# "base_model_arg", [
# {"base_model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"},
# {"model_info": "anthropic.claude-3-sonnet-20240229-v1:0"},
# ]
# )

View file

@ -1,115 +0,0 @@
import traceback
from dotenv import load_dotenv
load_dotenv()
import io
from typing import Literal
import pytest
from pydantic import BaseModel, ConfigDict
import litellm
from litellm import Router, completion_cost, stream_chunk_builder
models = [
dict(
model_name="openai/gpt-3.5-turbo",
),
dict(
model_name="anthropic/claude-3-haiku-20240307",
),
dict(
model_name="together_ai/meta-llama/Llama-2-7b-chat-hf",
),
]
router = Router(
model_list=[
{
"model_name": m["model_name"],
"litellm_params": {
"model": m.get("model", m["model_name"]),
},
}
for m in models
],
routing_strategy="simple-shuffle",
num_retries=3,
retry_after=1,
timeout=60.0,
allowed_fails=2,
cooldown_time=0,
debug_level="INFO",
)
@pytest.mark.parametrize(
"model",
[
"openai/gpt-3.5-turbo",
# "anthropic/claude-3-haiku-20240307",
# "together_ai/meta-llama/Llama-2-7b-chat-hf",
],
)
def test_run(model: str):
"""
Relevant issue - https://github.com/BerriAI/litellm/issues/4965
"""
litellm.set_verbose = True
prompt = "Hi"
kwargs = dict(
model=model,
messages=[{"role": "user", "content": prompt}],
temperature=0.001,
top_p=0.001,
max_tokens=20,
input_cost_per_token=2,
output_cost_per_token=2,
)
print(f"--------- {model} ---------")
print(f"Prompt: {prompt}")
response = router.completion(**kwargs) # type: ignore
non_stream_output = response.choices[0].message.content.replace("\n", "") # type: ignore
non_stream_cost_calc = response._hidden_params["response_cost"] * 100
print(f"Non-stream output: {non_stream_output}")
print(f"Non-stream usage : {response.usage}") # type: ignore
non_stream_usage = response.usage
try:
print(
f"Non-stream cost : {response._hidden_params['response_cost'] * 100:.4f}"
)
except TypeError:
print("Non-stream cost : NONE")
print(f"Non-stream cost : {completion_cost(response) * 100:.4f} (response)")
response = router.completion(**kwargs, stream=True, stream_options={"include_usage": True}) # type: ignore
response = stream_chunk_builder(list(response), messages=kwargs["messages"]) # type: ignore
output = response.choices[0].message.content.replace("\n", "") # type: ignore
if response.usage.completion_tokens != non_stream_usage.completion_tokens:
pytest.skip(
"LLM API returning inconsistent usage"
) # handles transient openai errors
streaming_cost_calc = (
completion_cost(
response,
custom_cost_per_token={
"input_cost_per_token": kwargs["input_cost_per_token"],
"output_cost_per_token": kwargs["output_cost_per_token"],
},
)
* 100
)
print(f"Stream output : {output}")
print(f"Stream usage : {response.usage}") # type: ignore
print(f"Stream cost : {streaming_cost_calc} (response)")
print("")
if output == non_stream_output:
# assert cost is the same
assert streaming_cost_calc == non_stream_cost_calc

View file

@ -1,385 +0,0 @@
"""
Test to ensure Docker container does not go out to network on deploy.
This test verifies that the LiteLLM proxy container does not make outbound
network requests during startup. This is important for:
1. Air-gapped environments where outbound network is restricted
2. Security compliance requiring no unexpected network calls
3. Fast container startup without network dependencies
The test works by:
1. Building/running the container with network disabled
2. Verifying the container starts successfully without network
3. Checking for any errors related to network failures during startup
"""
import os
import re
import subprocess
import time
import pytest
from tests._master_key import MASTER_KEY
def is_docker_available() -> bool:
"""Check if Docker is available and running."""
try:
result = subprocess.run(
["docker", "info"],
capture_output=True,
text=True,
timeout=10,
)
return result.returncode == 0
except (subprocess.TimeoutExpired, FileNotFoundError):
return False
@pytest.mark.skipif(
not is_docker_available(),
reason="Docker not available",
)
class TestDockerNoNetworkOnDeploy:
"""
Test suite for verifying Docker container starts without network access.
"""
# Container and image names for testing
TEST_IMAGE_NAME = "litellm-no-network-test"
TEST_CONTAINER_NAME = "litellm-no-network-test-container"
# Timeout for container operations
CONTAINER_START_TIMEOUT = 60 # seconds
# Patterns that indicate network-related failures during startup
NETWORK_ERROR_PATTERNS = [
r"connection refused",
r"network is unreachable",
r"could not resolve host",
r"name resolution failed",
r"dns lookup failed",
r"failed to establish.*connection",
r"socket.gaierror",
r"urllib.error.URLError.*Errno",
r"requests.exceptions.ConnectionError",
r"httpx.*ConnectError",
r"aiohttp.*ClientConnectorError",
]
# Patterns that indicate EXPECTED local-only startup
LOCAL_STARTUP_PATTERNS = [
r"starting.*(server|proxy)",
r"listening on",
r"uvicorn running",
r"application startup complete",
]
@pytest.fixture(autouse=True)
def cleanup(self):
"""Cleanup any existing test containers before and after each test."""
self._cleanup_container()
yield
self._cleanup_container()
def _cleanup_container(self):
"""Remove test container if it exists."""
subprocess.run(
["docker", "rm", "-f", self.TEST_CONTAINER_NAME],
capture_output=True,
timeout=30,
)
def _build_test_image(self) -> bool:
"""
Build the Docker image if needed.
Returns True if image is available (built or already exists).
"""
# Check if main litellm image exists
result = subprocess.run(
["docker", "images", "-q", "litellm/litellm"],
capture_output=True,
text=True,
timeout=10,
)
if result.stdout.strip():
# Image exists, use it
return True
# Try to build from Dockerfile
dockerfile_path = os.path.join(
os.path.dirname(__file__),
"..",
"..",
"..",
"Dockerfile",
)
if os.path.exists(dockerfile_path):
result = subprocess.run(
[
"docker",
"build",
"-t",
self.TEST_IMAGE_NAME,
"-f",
dockerfile_path,
os.path.dirname(dockerfile_path),
],
capture_output=True,
text=True,
timeout=600, # 10 minutes for build
)
return result.returncode == 0
return False
def test_container_starts_without_network(self):
"""
Test that the container can start with network completely disabled.
This test runs the container with --network=none to ensure no outbound
network requests are required during startup.
"""
# Use a minimal config that doesn't require external services
minimal_config = f"""
model_list:
- model_name: fake-model
litellm_params:
model: fake/fake-model
general_settings:
master_key: {MASTER_KEY}
database_url: null
environment_variables: {{}}
"""
# Create a temporary config file
config_path = "/tmp/litellm_test_config.yaml"
with open(config_path, "w") as f:
f.write(minimal_config)
image_to_use = "litellm/litellm"
# Check if image exists
result = subprocess.run(
["docker", "images", "-q", image_to_use],
capture_output=True,
text=True,
timeout=10,
)
if not result.stdout.strip():
pytest.skip(f"Docker image {image_to_use} not available")
# Run container with network disabled
run_cmd = [
"docker",
"run",
"--name",
self.TEST_CONTAINER_NAME,
"--network=none", # Disable all network access
"-v",
f"{config_path}:/app/config.yaml:ro",
"-e",
f"LITELLM_MASTER_KEY={MASTER_KEY}",
"-e",
"DATABASE_URL=", # Empty to disable DB
"-e",
"STORE_MODEL_IN_DB=false",
"-e",
"LITELLM_LOG=DEBUG",
"-d", # Detached mode
image_to_use,
"--config",
"/app/config.yaml",
]
result = subprocess.run(
run_cmd,
capture_output=True,
text=True,
timeout=30,
)
if result.returncode != 0:
pytest.fail(f"Failed to start container: {result.stderr}")
# Wait for container to start up
time.sleep(5)
# Check container logs for any network-related errors
logs_result = subprocess.run(
["docker", "logs", self.TEST_CONTAINER_NAME],
capture_output=True,
text=True,
timeout=30,
)
logs = logs_result.stdout + logs_result.stderr
# Check for network error patterns (case-insensitive)
network_errors = []
for pattern in self.NETWORK_ERROR_PATTERNS:
matches = re.findall(pattern, logs, re.IGNORECASE)
if matches:
network_errors.extend(matches)
# Check if container is still running
inspect_result = subprocess.run(
[
"docker",
"inspect",
"-f",
"{{.State.Running}}",
self.TEST_CONTAINER_NAME,
],
capture_output=True,
text=True,
timeout=10,
)
is_running = inspect_result.stdout.strip() == "true"
# If container crashed, get exit code and reason
# If container crashed, get exit code and reason
if not is_running:
exit_result = subprocess.run(
[
"docker",
"inspect",
"-f",
"{{.State.ExitCode}}",
self.TEST_CONTAINER_NAME,
],
capture_output=True,
text=True,
timeout=10,
)
exit_code = exit_result.stdout.strip()
# Container not running is OK if it didn't crash due to network issues
# Check if exit was due to network errors
if network_errors:
pytest.fail(
f"Container failed with network errors (exit code {exit_code}): "
f"{network_errors}\n\nFull logs:\n{logs}"
)
# Assert no network errors were found
assert len(network_errors) == 0, (
f"Container made network requests during startup that failed: "
f"{network_errors}"
)
def test_no_external_urls_in_startup_code(self):
"""
Static analysis test: check that startup code doesn't contain
hardcoded external URLs that would be called during import/startup.
This is a complementary test to catch issues without needing Docker.
"""
# Directories to check for startup code
startup_dirs = [
"litellm/proxy",
"litellm/__init__.py",
"litellm/main.py",
]
# Patterns that indicate external URL calls during startup (not in functions)
problematic_patterns = [
# Immediate HTTP calls (not inside functions)
r'^requests\.get\(["\'](https?://)',
r'^httpx\.get\(["\'](https?://)',
r'^urllib\.request\.urlopen\(["\'](https?://)',
]
# Files that are OK to have URLs (they're called on-demand, not startup)
allowed_files = [
"model_prices_and_context_window.json", # Static data file
"test_", # Test files
"_test.py",
"conftest.py",
]
workspace_root = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
issues_found = []
for startup_dir in startup_dirs:
full_path = os.path.join(workspace_root, startup_dir)
if not os.path.exists(full_path):
continue
if os.path.isfile(full_path):
files_to_check = [full_path]
else:
files_to_check = []
for root, _dirs, files in os.walk(full_path):
for f in files:
if f.endswith(".py"):
files_to_check.append(os.path.join(root, f))
for filepath in files_to_check:
# Skip allowed files
if any(allowed in filepath for allowed in allowed_files):
continue
try:
with open(filepath, "r") as f:
content = f.read()
for pattern in problematic_patterns:
matches = re.findall(pattern, content, re.MULTILINE)
if matches:
issues_found.append(f"{filepath}: {pattern} matched")
except Exception:
pass # Skip unreadable files
# This test is informational - we document but don't fail
if issues_found:
pass
@pytest.mark.skipif(
not is_docker_available(),
reason="Docker not available",
)
def test_container_build_no_network_fetch():
"""
Test that the Docker build process doesn't require network for runtime.
This verifies that all dependencies are properly bundled and no
runtime network calls are made during container initialization.
Note: Build itself may need network to resolve dependencies, but runtime should not.
"""
# This is a simplified version - full test would need to:
# 1. Build image with --network=none (requires pre-cached deps)
# 2. Or run built image in isolated network
# For now, just verify the Dockerfile doesn't have wget/curl in CMD/ENTRYPOINT
workspace_root = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
dockerfile_path = os.path.join(workspace_root, "Dockerfile")
if not os.path.exists(dockerfile_path):
pytest.skip("Dockerfile not found")
with open(dockerfile_path, "r") as f:
content = f.read()
# Check for network calls in CMD/ENTRYPOINT
problematic = []
lines = content.split("\n")
for i, line in enumerate(lines, 1):
line_upper = line.strip().upper()
if line_upper.startswith(("CMD", "ENTRYPOINT")):
if any(
cmd in line.lower()
for cmd in ["curl", "wget", "fetch", "http://", "https://"]
):
problematic.append(f"Line {i}: {line.strip()}")
assert (
len(problematic) == 0
), f"Dockerfile CMD/ENTRYPOINT contains network calls: {problematic}"

File diff suppressed because one or more lines are too long

View file

@ -1,770 +0,0 @@
import asyncio
import os
import traceback
from typing import Any
from unittest.mock import MagicMock, patch
import httpx
import pytest
from openai import (
AsyncAzureOpenAI,
AsyncOpenAI,
AuthenticationError,
AzureOpenAI,
BadRequestError,
OpenAIError,
)
import litellm
from litellm import ( # AuthenticationError,; RateLimitError,; ServiceUnavailableError,; OpenAIError,
completion,
)
litellm.vertex_project = "litellm-ci-cd"
litellm.vertex_location = "us-central1"
litellm.num_retries = 0
# litellm.failure_callback = ["sentry"]
#### What this tests ####
# This tests exception mapping -> trigger an exception from an llm provider -> assert if output is of the expected type
# 5 providers -> OpenAI, Azure, Anthropic, Cohere, Replicate
# 3 main types of exceptions -> - Rate Limit Errors, Context Window Errors, Auth errors (incorrect/rotated key, etc.)
# Approach: Run each model through the test -> assert if the correct error (always the same one) is triggered
exception_models = [
"sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
"bedrock/anthropic.claude-instant-v1",
]
# Test 1: Context Window Errors
models = ["command-nightly"]
# for model in litellm.models_by_provider["bedrock"]:
# test_context_window(model=model)
# test_context_window(model="chat-bison")
# test_context_window_with_fallbacks(model="command-nightly")
# Test 2: InvalidAuth Errors
@pytest.mark.parametrize("model", models)
def invalid_auth(model): # set the model key to an invalid key, depending on the model
messages = [{"content": "Hello, how are you?", "role": "user"}]
temporary_key = None
try:
if model == "gpt-3.5-turbo" or model == "gpt-3.5-turbo-instruct":
temporary_key = os.environ["OPENAI_API_KEY"]
os.environ["OPENAI_API_KEY"] = "bad-key"
elif "bedrock" in model:
temporary_aws_access_key = os.environ["AWS_ACCESS_KEY_ID"]
os.environ["AWS_ACCESS_KEY_ID"] = "bad-key"
temporary_aws_region_name = os.environ["AWS_REGION_NAME"]
os.environ["AWS_REGION_NAME"] = "bad-key"
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_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"
elif model == "command-nightly":
temporary_key = os.environ["COHERE_API_KEY"]
os.environ["COHERE_API_KEY"] = "bad-key"
elif "j2" in model:
temporary_key = os.environ["AI21_API_KEY"]
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"
elif model in litellm.openrouter_models:
temporary_key = os.environ["OPENROUTER_API_KEY"]
os.environ["OPENROUTER_API_KEY"] = "bad-key"
elif model in litellm.aleph_alpha_models:
temporary_key = os.environ["ALEPH_ALPHA_API_KEY"]
os.environ["ALEPH_ALPHA_API_KEY"] = "bad-key"
elif model in litellm.nlp_cloud_models:
os.environ["NLP_CLOUD_API_KEY"] = "bad-key"
elif (
model
== "replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1"
):
temporary_key = os.environ["REPLICATE_API_KEY"]
os.environ["REPLICATE_API_KEY"] = "bad-key"
print(f"model: {model}")
response = completion(model=model, messages=messages)
print(f"response: {response}")
except AuthenticationError as e:
print(f"AuthenticationError Caught Exception - {str(e)}")
except (
OpenAIError
) as e: # is at least an openai error -> in case of random model errors - e.g. overloaded server
print(f"OpenAIError Caught Exception - {e}")
except Exception as e:
print(type(e))
print(type(AuthenticationError))
print(e.__class__.__name__)
print(f"Uncaught Exception - {e}")
pytest.fail(f"Error occurred: {e}")
if temporary_key != None: # reset the key
if model == "gpt-3.5-turbo":
os.environ["OPENAI_API_KEY"] = temporary_key
elif model == "chatgpt-test":
os.environ["AZURE_AI_API_KEY"] = temporary_key
azure = True
elif model == "claude-3-5-haiku-20241022":
os.environ["ANTHROPIC_API_KEY"] = temporary_key
elif model == "command-nightly":
os.environ["COHERE_API_KEY"] = temporary_key
elif (
model
== "replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1"
):
os.environ["REPLICATE_API_KEY"] = temporary_key
elif "j2" in model:
os.environ["AI21_API_KEY"] = temporary_key
elif "togethercomputer" in model:
os.environ["TOGETHERAI_API_KEY"] = temporary_key
elif model in litellm.aleph_alpha_models:
os.environ["ALEPH_ALPHA_API_KEY"] = temporary_key
elif model in litellm.nlp_cloud_models:
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
os.environ["AWS_SECRET_ACCESS_KEY"] = temporary_secret_key
return
# for model in litellm.models_by_provider["bedrock"]:
# invalid_auth(model=model)
# invalid_auth(model="command-nightly")
# Test 3: Invalid Request Error
@pytest.mark.parametrize("model", models)
def test_invalid_request_error(model):
messages = [{"content": "hey, how's it going?", "role": "user"}]
with pytest.raises(BadRequestError):
completion(model=model, messages=messages, max_tokens="hello world")
def test_completion_azure_exception():
try:
import openai
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = True
## Test azure call
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_AI_API_KEY"] = old_azure_key
print(f"response: {response}")
print(response)
except openai.AuthenticationError as e:
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}")
# test_completion_azure_exception()
async def asynctest_completion_azure_exception():
try:
import openai
import litellm
print("azure gpt-3.5 test\n\n")
litellm.set_verbose = True
## Test azure call
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"}],
)
print(f"response: {response}")
print(response)
except openai.AuthenticationError as e:
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:
print("Got wrong exception")
print("exception", e)
pytest.fail(f"Error occurred: {e}")
# import asyncio
# asyncio.run(
# asynctest_completion_azure_exception()
# )
def asynctest_completion_openai_exception_bad_model():
try:
import asyncio
import openai
import litellm
print("azure exception bad model\n\n")
litellm.set_verbose = True
## Test azure call
async def test():
response = await litellm.acompletion(
model="openai/gpt-6",
messages=[{"role": "user", "content": "hello"}],
)
asyncio.run(test())
except openai.NotFoundError:
print("Good job this is a NotFoundError for a model that does not exist!")
print("Passed")
except Exception as e:
print("Raised wrong type of exception", type(e))
pytest.fail(f"Error occurred: {e}")
# asynctest_completion_openai_exception_bad_model()
def asynctest_completion_azure_exception_bad_model():
try:
import asyncio
import openai
import litellm
print("azure exception bad model\n\n")
litellm.set_verbose = True
## Test azure call
async def test():
response = await litellm.acompletion(
model="azure/gpt-12",
messages=[{"role": "user", "content": "hello"}],
)
asyncio.run(test())
except openai.NotFoundError:
print("Good job this is a NotFoundError for a model that does not exist!")
print("Passed")
except Exception as e:
print("Raised wrong type of exception", type(e))
pytest.fail(f"Error occurred: {e}")
# asynctest_completion_azure_exception_bad_model()
def test_completion_openai_exception():
# test if openai:gpt raises openai.AuthenticationError
try:
import openai
print("openai gpt-3.5 test\n\n")
litellm.set_verbose = True
## Test azure call
old_azure_key = os.environ["OPENAI_API_KEY"]
os.environ["OPENAI_API_KEY"] = "good morning"
response = completion(
model="gpt-4",
messages=[{"role": "user", "content": "hello"}],
)
print(f"response: {response}")
print(response)
except openai.AuthenticationError as e:
os.environ["OPENAI_API_KEY"] = old_azure_key
print("OpenAI: good job got the correct error for openai when key not set")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_openai_exception()
def test_completion_mistral_exception():
# test if mistral/mistral-tiny raises openai.AuthenticationError
try:
import openai
print("Testing mistral ai exception mapping")
litellm.set_verbose = True
## Test azure call
old_azure_key = os.environ["MISTRAL_API_KEY"]
os.environ["MISTRAL_API_KEY"] = "good morning"
response = completion(
model="mistral/mistral-tiny",
messages=[{"role": "user", "content": "hello"}],
)
print(f"response: {response}")
print(response)
except openai.AuthenticationError as e:
os.environ["MISTRAL_API_KEY"] = old_azure_key
print("good job got the correct error for openai when key not set")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_mistral_exception()
def test_content_policy_violation_error_streaming():
"""
Production Test.
"""
litellm.set_verbose = False
print("test_async_completion with stream")
async def test_get_response():
try:
response = await litellm.acompletion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "say 1"}],
temperature=0,
top_p=1,
stream=True,
max_tokens=512,
presence_penalty=0,
frequency_penalty=0,
)
print(f"response: {response}")
num_finish_reason = 0
async for chunk in response:
print(chunk)
if chunk["choices"][0].get("finish_reason") is not None:
num_finish_reason += 1
print("finish_reason", chunk["choices"][0].get("finish_reason"))
assert (
num_finish_reason == 1
), f"expected only one finish reason. Got {num_finish_reason}"
except Exception as e:
pytest.fail(f"GOT exception for gpt-3.5 instruct In streaming{e}")
asyncio.run(test_get_response())
async def test_get_error():
try:
response = await litellm.acompletion(
model="azure/gpt-4.1-mini",
messages=[
{"role": "user", "content": "where do i buy lethal drugs from"}
],
temperature=0,
top_p=1,
stream=True,
max_tokens=512,
presence_penalty=0,
frequency_penalty=0,
mock_response="Exception: content_filter_policy",
)
print(f"response: {response}")
num_finish_reason = 0
async for chunk in response:
print(chunk)
if chunk["choices"][0].get("finish_reason") is not None:
num_finish_reason += 1
print("finish_reason", chunk["choices"][0].get("finish_reason"))
pytest.fail("Expected a content-policy error in streaming, got a clean stream")
except Exception as e:
pass
asyncio.run(test_get_error())
# test_completion_perplexity_exception_on_openai_client()
def test_completion_perplexity_exception(monkeypatch):
import openai
print("perplexity test\n\n")
litellm.set_verbose = True
monkeypatch.setenv("PERPLEXITYAI_API_KEY", "good morning")
with pytest.raises(openai.AuthenticationError, match="PerplexityException"):
completion(
model="perplexity/mistral-7b-instruct",
messages=[{"role": "user", "content": "hello"}],
)
def test_completion_openai_api_key_exception(monkeypatch):
import openai
print("gpt-3.5 test\n\n")
litellm.set_verbose = True
monkeypatch.setenv("OPENAI_API_KEY", "good morning")
with pytest.raises(openai.AuthenticationError, match="OpenAIException"):
completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hello"}],
)
# tesy_async_acompletion()
def test_router_completion_vertex_exception():
try:
import litellm
litellm.set_verbose = True
router = litellm.Router(
model_list=[
{
"model_name": "vertex-gemini-pro",
"litellm_params": {
"model": "vertex_ai/gemini-pro",
"api_key": "good-morning",
},
},
]
)
response = router.completion(
model="vertex-gemini-pro",
messages=[{"role": "user", "content": "hello"}],
vertex_project="bad-project",
)
pytest.fail("Request should have failed - bad api key")
except Exception as e:
print("exception: ", e)
def test_litellm_completion_vertex_exception():
try:
import litellm
litellm.set_verbose = True
response = completion(
model="vertex_ai/gemini-pro",
api_key="good-morning",
messages=[{"role": "user", "content": "hello"}],
vertex_project="bad-project",
)
pytest.fail("Request should have failed - bad api key")
except Exception as e:
print("exception: ", e)
def test_litellm_predibase_exception():
"""
Test - Assert that the Predibase API Key is not returned on Authentication Errors
"""
try:
import litellm
litellm.set_verbose = True
response = completion(
model="predibase/llama-3-8b-instruct",
messages=[{"role": "user", "content": "What is the meaning of life?"}],
tenant_id="c4768f95",
api_key="hf-rawapikey",
)
pytest.fail("Request should have failed - bad api key")
except Exception as e:
if "hf-rawapikey" in str(e):
pytest.fail("predibase error leaked the raw api key")
print("exception: ", e)
# # test_invalid_request_error(model="command-nightly")
# # Test 3: Rate Limit Errors
# def test_model_call(model):
# try:
# sample_text = "how does a court case get to the Supreme Court?"
# messages = [{ "content": sample_text,"role": "user"}]
# print(f"model: {model}")
# response = completion(model=model, messages=messages)
# except RateLimitError as e:
# print(f"headers: {e.response.headers}")
# return True
# # except OpenAIError: # is at least an openai error -> in case of random model errors - e.g. overloaded server
# # return True
# except Exception as e:
# print(f"Uncaught Exception {model}: {type(e).__name__} - {e}")
# traceback.print_exc()
# pass
# return False
# # Repeat each model 500 times
# # extended_models = [model for model in models for _ in range(250)]
# extended_models = ["azure/gpt-4.1-mini" for _ in range(250)]
# def worker(model):
# return test_model_call(model)
# # Create a dictionary to store the results
# counts = {True: 0, False: 0}
# # Use Thread Pool Executor
# with ThreadPoolExecutor(max_workers=500) as executor:
# # Use map to start the operation in thread pool
# results = executor.map(worker, extended_models)
# # Iterate over results and count True/False
# for result in results:
# counts[result] += 1
# accuracy_score = counts[True]/(counts[True] + counts[False])
# print(f"accuracy_score: {accuracy_score}")
def test_anthropic_tool_calling_exception():
"""
Related - https://github.com/BerriAI/litellm/issues/4348
"""
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {},
},
}
]
try:
litellm.completion(
model="claude-haiku-4-5-20251001",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
tools=tools,
)
except litellm.BadRequestError:
pass
from typing import Optional, Union
from openai import OpenAI
def _pre_call_utils(
call_type: str,
data: dict,
client: Union[OpenAI, AsyncOpenAI],
sync_mode: bool,
streaming: Optional[bool],
):
if call_type == "embedding":
data["input"] = "Hello world!"
if isinstance(client, (AzureOpenAI, AsyncAzureOpenAI)):
mapped_target: Any = client.embeddings.with_raw_response
patched_attr = "create"
else:
mapped_target = client
patched_attr = "post"
if sync_mode:
original_function = litellm.embedding
else:
original_function = litellm.aembedding
elif call_type == "chat_completion":
data["messages"] = [{"role": "user", "content": "Hello world"}]
if streaming is True:
data["stream"] = True
mapped_target = client.chat.completions.with_raw_response # type: ignore
patched_attr = "create"
if sync_mode:
original_function = litellm.completion
else:
original_function = litellm.acompletion
elif call_type == "completion":
data["prompt"] = "Hello world"
if streaming is True:
data["stream"] = True
mapped_target = client.completions.with_raw_response # type: ignore
patched_attr = "create"
if sync_mode:
original_function = litellm.text_completion
else:
original_function = litellm.atext_completion
return data, original_function, mapped_target, patched_attr
@pytest.mark.parametrize(
"sync_mode",
[True, False],
)
@pytest.mark.parametrize(
"provider, model, call_type, streaming",
[
("azure", "azure/text-embedding-ada-002", "embedding", True),
("azure", "azure_text/gpt-3.5-turbo-instruct", "completion", True),
],
)
@pytest.mark.asyncio
async def test_exception_with_headers(sync_mode, provider, model, call_type, streaming):
"""
User feedback: litellm says "No deployments available for selected model, Try again in 60 seconds"
but Azure says to retry in at most 9s
```
{"message": "litellm.proxy.proxy_server.embeddings(): Exception occured - No deployments available for selected model, Try again in 60 seconds. Passed model=text-embedding-ada-002. pre-call-checks=False, allowed_model_region=n/a, cooldown_list=[('b49cbc9314273db7181fe69b1b19993f04efb88f2c1819947c538bac08097e4c', {'Exception Received': 'litellm.RateLimitError: AzureException RateLimitError - Requests to the Embeddings_Create Operation under Azure OpenAI API version 2023-09-01-preview have exceeded call rate limit of your current OpenAI S0 pricing tier. Please retry after 9 seconds. Please go here: https://aka.ms/oai/quotaincrease if you would like to further increase the default rate limit.', 'Status Code': '429'})]", "level": "ERROR", "timestamp": "2024-08-22T03:25:36.900476"}
```
"""
print(f"Received args: {locals()}")
import openai
if sync_mode:
if provider == "openai":
openai_client = openai.OpenAI(api_key="")
elif provider == "azure":
openai_client = openai.AzureOpenAI(
api_key="", base_url="", api_version=litellm.AZURE_DEFAULT_API_VERSION
)
else:
if provider == "openai":
openai_client = openai.AsyncOpenAI(api_key="")
elif provider == "azure":
openai_client = openai.AsyncAzureOpenAI(
api_key="", base_url="", api_version=litellm.AZURE_DEFAULT_API_VERSION
)
data = {"model": model}
data, original_function, mapped_target, patched_attr = _pre_call_utils(
call_type=call_type,
data=data,
client=openai_client,
sync_mode=sync_mode,
streaming=streaming,
)
cooldown_time = 30.0
def _return_exception(*args, **kwargs):
from httpx import Headers, Request, Response
kwargs = {
"request": Request("POST", "https://www.google.com"),
"message": "Error code: 429 - Rate Limit Error!",
"body": {"detail": "Rate Limit Error!"},
"code": None,
"param": None,
"type": None,
"response": Response(
status_code=429,
headers=Headers(
{
"date": "Sat, 21 Sep 2024 22:56:53 GMT",
"server": "uvicorn",
"retry-after": "30",
"content-length": "30",
"content-type": "application/json",
}
),
request=Request("POST", "http://0.0.0.0:9000/chat/completions"),
),
"status_code": 429,
"request_id": None,
}
exception = Exception()
for k, v in kwargs.items():
setattr(exception, k, v)
raise exception
with patch.object(
mapped_target,
patched_attr,
side_effect=_return_exception,
):
new_retry_after_mock_client = MagicMock(return_value=-1)
litellm.utils._get_retry_after_from_exception_header = (
new_retry_after_mock_client
)
async def call_and_drain():
if sync_mode:
resp = original_function(**data, client=openai_client)
if streaming:
for chunk in resp:
continue
else:
resp = await original_function(**data, client=openai_client)
if streaming:
async for chunk in resp:
continue
with pytest.raises(litellm.RateLimitError) as exc_info:
await call_and_drain()
assert exc_info.value.litellm_response_headers is not None
assert int(exc_info.value.litellm_response_headers["retry-after"]) == cooldown_time
@pytest.mark.asyncio
@pytest.mark.parametrize("model", ["azure/gpt-4.1-mini", "openai/gpt-3.5-turbo"])
async def test_bad_request_error_contains_httpx_response(model):
"""
Test that the BadRequestError contains the httpx response
Relevant issue: https://github.com/BerriAI/litellm/issues/6732
"""
with pytest.raises(litellm.BadRequestError) as exc_info:
await litellm.acompletion(
model=model,
messages=[{"role": "user", "content": "Hello world"}],
bad_arg="bad_arg",
)
e = exc_info.value
print("e.response", e.response)
print("vars(e.response)", vars(e.response))
assert e.response is not None
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.parametrize("stream_mode", [True, False])
@pytest.mark.parametrize("model", ["gpt-4.1-nano"]) # "gpt-4o-mini",
@pytest.mark.asyncio
async def test_exception_bubbling_up(sync_mode, stream_mode, model):
"""
make sure code, param, and type are bubbled up
"""
import litellm
litellm.set_verbose = True
async def _call_with_bad_role():
if sync_mode:
litellm.completion(
model=model,
messages=[{"role": "usera", "content": "hi"}],
stream=stream_mode,
sync_stream=sync_mode,
)
else:
await litellm.acompletion(
model=model,
messages=[{"role": "usera", "content": "hi"}],
stream=stream_mode,
sync_stream=sync_mode,
)
with pytest.raises(Exception, match='litellm\\.BadRequestError: OpenAIException - Invalid value') as exc_info:
await _call_with_bad_role()
assert exc_info.value.code == "invalid_value"
assert exc_info.value.param is not None
assert exc_info.value.type == "invalid_request_error"

View file

@ -1,112 +0,0 @@
"""Regression tests for the canned OpenAI mock that decouples the suite from a
hosted endpoint (see ``tests/fake_openai_endpoint.py`` and
``tests/_fake_openai_endpoint_server.py``).
Before this, router/completion/triton tests hardcoded a shared Railway mock as
their ``api_base``; when that single deployment went down, unrelated CI jobs
failed with ``404 Application not found``. These tests pin the two behaviors
those tests now rely on from the local stand-in (the Triton embeddings shape and
the ``slow-endpoint`` delay) and guard against the dead host creeping back in.
"""
from __future__ import annotations
import re
from pathlib import Path
from typing import Final
import httpx
import pytest
from openai.types import ModerationCreateResponse
from tests.fake_openai_endpoint import (
_LOCAL_DEFAULT,
_resolve_base,
ensure_fake_openai_endpoint,
)
_REPO_ROOT = Path(__file__).resolve().parents[2]
_MIGRATED_FILES = (
"tests/llm_translation/test_triton.py",
"tests/local_testing/test_completion.py",
)
# A live ``railway.app/`` or ``railway.app"`` api_base reappearing in a migrated
# file would re-couple CI to an external service's uptime. The deliberately
# broken ``...railway.appzzzzz`` fallback URL is not a live host, so it does not
# match.
_LIVE_HOSTED_MOCK = re.compile(r"railway\.app(?:/|\")")
def test_chat_completion_shape():
base = ensure_fake_openai_endpoint()
response = httpx.post(
f"{base}/v1/chat/completions",
json={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
timeout=10,
)
assert response.status_code == 200
body = response.json()
assert body["choices"][0]["message"]["content"]
assert body["usage"]["total_tokens"] == 40
def test_moderations_route_parses_as_an_openai_response():
base: Final = ensure_fake_openai_endpoint()
response: Final = httpx.post(
f"{base}/v1/moderations",
json={"input": ["I want to harm someone", "hello"], "model": "omni-moderation-latest"},
timeout=10,
)
assert response.status_code == 200
parsed: Final = ModerationCreateResponse.model_validate(response.json())
assert parsed.model == "omni-moderation-latest"
assert len(parsed.results) == 2
assert parsed.results[0].categories.violence is False
def test_triton_embeddings_route():
base = ensure_fake_openai_endpoint()
response = httpx.post(f"{base}/triton/embeddings", json={"inputs": []}, timeout=10)
assert response.status_code == 200
output = response.json()["outputs"][0]
assert output["shape"] == [1, 2]
assert output["data"] == [0.1, 0.2]
def test_slow_model_blocks_past_client_timeout():
base = ensure_fake_openai_endpoint()
with pytest.raises(httpx.TimeoutException):
httpx.post(
f"{base}/v1/chat/completions",
json={
"model": "slow-endpoint",
"messages": [{"role": "user", "content": "hi"}],
},
timeout=0.5,
)
def test_remote_env_base_resolves_to_local(monkeypatch):
monkeypatch.setenv(
"FAKE_OPENAI_API_BASE",
"https://exampleopenaiendpoint-production.up.railway.app",
)
assert _resolve_base() == _LOCAL_DEFAULT
@pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "[::1]"])
def test_loopback_env_base_is_honored(monkeypatch, host):
base = f"http://{host}:9191"
monkeypatch.setenv("FAKE_OPENAI_API_BASE", base)
assert _resolve_base() == base
@pytest.mark.parametrize("relative_path", _MIGRATED_FILES)
def test_migrated_files_have_no_live_hosted_mock(relative_path):
source = (_REPO_ROOT / relative_path).read_text()
assert not _LIVE_HOSTED_MOCK.search(source), (
f"{relative_path} references the hosted Railway mock; point api_base at "
"FAKE_OPENAI_API_BASE so CI does not depend on an external service"
)

View file

@ -1,144 +0,0 @@
# What is this?
## Test to make sure function call response always works with json.loads() -> no extra parsing required. Relevant issue - https://github.com/BerriAI/litellm/issues/2654
import traceback
from dotenv import load_dotenv
load_dotenv()
import io
import json
import warnings
from typing import List
import pytest
import litellm
from litellm import completion
# Just a stub to keep the sample code simple
class Trade:
def __init__(self, order: dict):
self.order = order
@staticmethod
def buy(order: dict):
return Trade(order)
@staticmethod
def sell(order: dict):
return Trade(order)
def trade(model_name: str) -> List[Trade]: # type: ignore
def parse_order(order: dict) -> Trade:
action = order["action"]
if action == "buy":
return Trade.buy(order)
elif action == "sell":
return Trade.sell(order)
else:
raise ValueError(f"Invalid action {action}")
def parse_call(call) -> List[Trade]:
arguments = json.loads(call.function.arguments)
trades = [parse_order(order) for order in arguments["orders"]]
return trades
tool_spec = {
"type": "function",
"function": {
"name": "trade",
"description": "Execute orders to manage the portfolio. Orders will be executed immediately at the stated prices.",
"parameters": {
"type": "object",
"properties": {
"orders": {
"type": "array",
"items": {
"type": "object",
"properties": {
"action": {"type": "string", "enum": ["buy", "sell"]},
"asset": {"type": "string"},
"amount": {
"type": "number",
"description": "Amount of asset to buy or sell.",
},
},
"required": ["action", "asset", "amount"],
},
},
},
},
},
}
try:
response = completion(
model_name,
[
{
"role": "system",
"content": """You are an expert asset manager, managing a portfolio.
Always use the `trade` function. Make sure that you call it correctly. For example, the following is a valid call:
```
trade({
"orders": [
{"action": "buy", "asset": "BTC", "amount": 0.1},
{"action": "sell", "asset": "ETH", "amount": 0.2}
]
})
```
If there are no trades to make, call `trade` with an empty array:
```
trade({ "orders": [] })
```
""",
},
{
"role": "user",
"content": """Manage the portfolio.
Don't jabber.
This is the current market data:
```
{market_data}
```
Your portfolio is as follows:
```
{portfolio}
```
""".replace(
"{market_data}", "BTC: 64,000 USD\nETH: 3,500 USD"
).replace(
"{portfolio}", "USD: 1000, BTC: 0.1, ETH: 0.2"
),
},
],
tools=[tool_spec],
tool_choice={
"type": "function",
"function": {"name": tool_spec["function"]["name"]}, # type: ignore
},
)
calls = response.choices[0].message.tool_calls
trades = [trade for call in calls for trade in parse_call(call)]
return trades
except litellm.InternalServerError:
pass
@pytest.mark.parametrize(
"model", ["us.anthropic.claude-haiku-4-5-20251001-v1:0"]
)
@pytest.mark.flaky(retries=6, delay=10)
def test_function_call_parsing(model):
trades = trade(model)
print([trade.order for trade in trades if trade is not None])

View file

@ -1,302 +0,0 @@
import traceback
from dotenv import load_dotenv
load_dotenv()
import io
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
from litellm import RateLimitError, Timeout, completion_cost, embedding
litellm.num_retries = 0
litellm.cache = None
# litellm.set_verbose=True
import json
# litellm.success_callback = ["langfuse"]
def get_current_weather(location, unit="fahrenheit"):
"""Get the current weather in a given location"""
if "tokyo" in location.lower():
return json.dumps({"location": "Tokyo", "temperature": "10", "unit": "celsius"})
elif "san francisco" in location.lower():
return json.dumps(
{"location": "San Francisco", "temperature": "72", "unit": "fahrenheit"}
)
elif "paris" in location.lower():
return json.dumps({"location": "Paris", "temperature": "22", "unit": "celsius"})
else:
return json.dumps({"location": location, "temperature": "unknown"})
# Example dummy function hard coded to return the same weather
# In production, this could be your backend API or an external API
# test_parallel_function_call()
from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message
_PARALLEL_TOOL_HISTORY_MESSAGES = [
{
"role": "user",
"content": "What's the weather like in San Francisco, Tokyo, and Paris? - give me 3 responses",
},
Message(
content="Here are the current weather conditions for San Francisco, Tokyo, and Paris:",
role="assistant",
tool_calls=[
ChatCompletionMessageToolCall(
index=1,
function=Function(
arguments='{"location": "San Francisco, CA", "unit": "fahrenheit"}',
name="get_current_weather",
),
id="tooluse_Jj98qn6xQlOP_PiQr-w9iA",
type="function",
)
],
function_call=None,
),
{
"tool_call_id": "tooluse_Jj98qn6xQlOP_PiQr-w9iA",
"role": "tool",
"name": "get_current_weather",
"content": '{"location": "San Francisco", "temperature": "72", "unit": "fahrenheit"}',
},
]
@pytest.mark.parametrize(
"model, messages",
[
# Anthropic Messages API: a dummy tool is injected without modify_params,
# so tool history with no tools= completes instead of raising.
("claude-haiku-4-5-20251001", _PARALLEL_TOOL_HISTORY_MESSAGES),
(
"us.anthropic.claude-sonnet-4-5-20250929-v1:0",
[
{
"role": "user",
"content": "What's the weather like in San Francisco, Tokyo, and Paris? - give me 3 responses",
}
],
),
(
"claude-haiku-4-5-20251001",
[
{
"role": "user",
"content": "What's the weather like in San Francisco, Tokyo, and Paris? - give me 3 responses",
}
],
),
],
)
def test_parallel_function_call_anthropic_error_msg(model, messages):
"""
Tool history without an explicit ``tools`` param must complete, not raise.
Anthropic (and Bedrock Invoke via ``AnthropicConfig.transform_request``)
inject a dummy tool so CLIs work with ``modify_params`` left off. Bedrock
Converse's no-raise behavior is covered offline in
``tests/unit/llms/bedrock/chat/test_converse_transformation.py``
(see #24158, #27138), which needs no live credentials.
"""
# Force modify_params off as a clean baseline: it exercises the Anthropic
# dummy-tool path, which injects regardless of modify_params
# (other tests in this file set it to True and don't reset it)
original_modify_params = litellm.modify_params
litellm.modify_params = False
try:
litellm.set_verbose = True
second_response = litellm.completion(
model=model,
messages=messages,
temperature=0.2,
seed=22,
drop_params=True,
) # get a new response from the model where it can see the function response
print("second response\n", second_response)
except litellm.InternalServerError as e:
print(e)
except litellm.RateLimitError as e:
print(e)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
finally:
litellm.modify_params = original_modify_params
def test_parallel_function_call_stream():
try:
litellm.set_verbose = True
# Step 1: send the conversation and available functions to the model
messages = [
{
"role": "user",
"content": "What's the weather like in San Francisco, Tokyo, and Paris?",
}
]
tools = [
{
"type": "function",
"function": {
"name": "get_current_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"],
},
},
}
]
response = litellm.completion(
model="gpt-6-luna",
messages=messages,
tools=tools,
stream=True,
tool_choice="auto", # auto is default, but we'll be explicit
complete_response=True,
)
print("Response\n", response)
# for chunk in response:
# print(chunk)
response_message = response.choices[0].message
tool_calls = response_message.tool_calls
print("length of tool calls", len(tool_calls))
print("Expecting there to be 3 tool calls")
assert (
len(tool_calls) > 1
) # this has to call the function for SF, Tokyo and parise
# Step 2: check if the model wanted to call a function
if tool_calls:
# Step 3: call the function
# Note: the JSON response may not always be valid; be sure to handle errors
available_functions = {
"get_current_weather": get_current_weather,
} # only one function in this example, but you can have multiple
messages.append(
response_message
) # extend conversation with assistant's reply
print("Response message\n", response_message)
# Step 4: send the info for each function call and function response to the model
for tool_call in tool_calls:
function_name = tool_call.function.name
function_to_call = available_functions[function_name]
function_args = json.loads(tool_call.function.arguments)
function_response = function_to_call(
location=function_args.get("location"),
unit=function_args.get("unit"),
)
messages.append(
{
"tool_call_id": tool_call.id,
"role": "tool",
"name": function_name,
"content": function_response,
}
) # extend conversation with function response
print(f"messages: {messages}")
second_response = litellm.completion(
model="gpt-6-luna", messages=messages, temperature=0.2, seed=22, reasoning_effort="none"
) # get a new response from the model where it can see the function response
print("second response\n", second_response)
return second_response
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_parallel_function_call_stream()
@pytest.mark.parametrize("sync_mode", [False])
@pytest.mark.asyncio
@pytest.mark.flaky(retries=6, delay=1)
async def test_watsonx_tool_choice(sync_mode, monkeypatch):
import json
from litellm import acompletion, completion
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
# Mock the IAM token generation to avoid actual API calls
monkeypatch.setenv("WATSONX_API_KEY", "mock-api-key")
monkeypatch.setenv("WATSONX_TOKEN", "mock-watsonx-token")
monkeypatch.setenv("WATSONX_API_BASE", "https://us-south.ml.cloud.ibm.com")
monkeypatch.setenv("WATSONX_PROJECT_ID", "mock-project-id")
monkeypatch.setattr(litellm, "set_verbose", True)
tools = [
{
"type": "function",
"function": {
"name": "get_current_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"],
},
},
}
]
messages = [{"role": "user", "content": "What is the weather in San Francisco?"}]
client = HTTPHandler() if sync_mode else AsyncHTTPHandler()
with patch.object(client, "post", return_value=MagicMock()) as mock_completion:
try:
if sync_mode:
resp = completion(
model="watsonx/meta-llama/llama-3-1-8b-instruct",
messages=messages,
tools=tools,
tool_choice="auto",
client=client,
)
else:
resp = await acompletion(
model="watsonx/meta-llama/llama-3-1-8b-instruct",
messages=messages,
tools=tools,
tool_choice="auto",
client=client,
stream=True,
)
print(resp)
mock_completion.assert_called_once()
print(mock_completion.call_args.kwargs)
json_data = json.loads(mock_completion.call_args.kwargs["data"])
assert json_data["tool_choice_option"] == "auto"
except Exception as e:
print(e)
if "The read operation timed out" in str(e):
pytest.skip("Skipping test due to timeout")
else:
raise e

View file

@ -1,29 +0,0 @@
import litellm
from litellm import get_optional_params
litellm.add_function_to_prompt = True
optional_params = get_optional_params(
model="",
tools=[
{
"type": "function",
"function": {
"description": "Get the current weather in a given location",
"name": "get_current_weather",
"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"],
},
},
}
],
tool_choice="auto",
)
assert optional_params is not None

View file

@ -1,36 +0,0 @@
import os, sys, traceback
import litellm
from dotenv import load_dotenv
def generate_text():
try:
litellm.set_verbose = True
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this image?"},
{
"type": "image_url",
"image_url": {
"url": "https://avatars.githubusercontent.com/u/17561003?v=4"
},
},
],
}
]
response = litellm.completion(
model="gemini/gemini-pro-vision",
messages=messages,
stop="Hello world",
num_retries=3,
)
print(response)
assert isinstance(response.choices[0].message.content, str) == True
except Exception as exception:
raise Exception("An error occurred during text generation:", exception)
# generate_text()

View file

@ -1,78 +0,0 @@
from typing import Literal
def calculate_img_tokens(
width,
height,
mode: Literal["low", "high", "auto"] = "auto",
base_tokens: int = 85, # openai default - https://openai.com/pricing
):
if mode == "low":
return base_tokens
elif mode == "high" or mode == "auto":
resized_width, resized_height = resize_image_high_res(
width=width, height=height
)
tiles_needed_high_res = calculate_tiles_needed(resized_width, resized_height)
tile_tokens = (base_tokens * 2) * tiles_needed_high_res
total_tokens = base_tokens + tile_tokens
return total_tokens
def resize_image_high_res(width, height):
# Maximum dimensions for high res mode
max_short_side = 768
max_long_side = 2000
# Determine the longer and shorter sides
longer_side = max(width, height)
shorter_side = min(width, height)
# Calculate the aspect ratio
aspect_ratio = longer_side / shorter_side
# Resize based on the short side being 768px
if width <= height: # Portrait or square
resized_width = max_short_side
resized_height = int(resized_width * aspect_ratio)
# if the long side exceeds the limit after resizing, adjust both sides accordingly
if resized_height > max_long_side:
resized_height = max_long_side
resized_width = int(resized_height / aspect_ratio)
else: # Landscape
resized_height = max_short_side
resized_width = int(resized_height * aspect_ratio)
# if the long side exceeds the limit after resizing, adjust both sides accordingly
if resized_width > max_long_side:
resized_width = max_long_side
resized_height = int(resized_width / aspect_ratio)
return resized_width, resized_height
# Test the function with the given example
def calculate_tiles_needed(
resized_width, resized_height, tile_width=512, tile_height=512
):
tiles_across = (resized_width + tile_width - 1) // tile_width
tiles_down = (resized_height + tile_height - 1) // tile_height
total_tiles = tiles_across * tiles_down
return total_tiles
# Test high res mode with 1875 x 768 image
resized_width_high_res = 1875
resized_height_high_res = 768
tiles_needed_high_res = calculate_tiles_needed(
resized_width_high_res, resized_height_high_res
)
print(
f"Tiles needed for high res image ({resized_width_high_res}x{resized_height_high_res}): {tiles_needed_high_res}"
)
# If you had the original size and needed to resize and then calculate tiles:
original_size = (10000, 4096)
resized_size_high_res = resize_image_high_res(*original_size)
print(f"Resized dimensions in high res mode: {resized_size_high_res}")
tiles_needed = calculate_tiles_needed(*resized_size_high_res)
print(f"Tiles needed for high res image {resized_size_high_res}: {tiles_needed}")

View file

@ -1,29 +0,0 @@
# #### What this tests ####
# # This tests calling litellm.max_budget by making back-to-back gpt-4 calls
# # commenting out this test for circle ci, as it causes other tests to fail, since litellm.max_budget would impact other litellm imports
# import sys, os, json
# import traceback
# import pytest
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import litellm
# # litellm.set_verbose = True
# from litellm import completion, BudgetExceededError
# def test_max_budget():
# try:
# litellm.max_budget = 0.001 # sets a max budget of $0.001
# messages = [{"role": "user", "content": "Hey, how's it going"}]
# response = completion(model="gpt-4", messages=messages, stream=True)
# for chunk in response:
# continue
# print(litellm._current_cost)
# completion(model="gpt-4", messages=messages, stream=True)
# litellm.max_budget = float('inf')
# except BudgetExceededError as e:
# pass
# except Exception as e:
# pytest.fail(f"An error occured: {str(e)}")

View file

@ -1,11 +0,0 @@
#### What this tests ####
# This tests context fallback dict
import sys, os
import traceback
import pytest
import litellm
from litellm import longer_context_model_fallback_dict
print(longer_context_model_fallback_dict)

View file

@ -1,90 +0,0 @@
#### What this tests ####
# This tests mock request calls to litellm
import os
import pytest
import litellm
import time
def test_mock_request_with_mock_timeout():
"""
Allow user to set 'mock_timeout = True', this allows for testing if fallbacks/retries are working on timeouts.
"""
start_time = time.time()
with pytest.raises(litellm.Timeout):
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, I'm a mock request"}],
timeout=3,
mock_timeout=True,
)
end_time = time.time()
assert end_time - start_time >= 3, f"Time taken: {end_time - start_time}"
def test_router_mock_request_with_mock_timeout():
"""
Allow user to set 'mock_timeout = True', this allows for testing if fallbacks/retries are working on timeouts.
"""
start_time = time.time()
router = litellm.Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
],
)
with pytest.raises(litellm.Timeout):
router.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, I'm a mock request"}],
timeout=3,
mock_timeout=True,
)
end_time = time.time()
assert end_time - start_time >= 3, f"Time taken: {end_time - start_time}"
def test_router_mock_request_with_mock_timeout_with_fallbacks():
"""
Allow user to set 'mock_timeout = True', this allows for testing if fallbacks/retries are working on timeouts.
"""
litellm.set_verbose = True
start_time = time.time()
router = litellm.Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
{
"model_name": "gpt-4.1-nano",
"litellm_params": {
"model": "gpt-4.1-nano",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
],
fallbacks=[{"gpt-3.5-turbo": ["gpt-4.1-nano"]}],
)
response = router.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, I'm a mock request"}],
timeout=3,
num_retries=1,
mock_timeout=True,
)
print(response)
end_time = time.time()
assert end_time - start_time >= 3, f"Time taken: {end_time - start_time}"
assert "gpt-4.1-nano" in response.model, "Model should be gpt-4.1-nano"

View file

@ -1,125 +0,0 @@
import json
from dotenv import load_dotenv
load_dotenv()
import pytest
import litellm
## for ollama we can't test making the completion call
def test_ollama_vision_model():
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
from unittest.mock import patch
with patch.object(client, "post") as mock_post:
try:
litellm.completion(
model="ollama/llama3.2-vision:11b",
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "Whats in this image?"},
{
"type": "image_url",
"image_url": {
"url": "https://dummyimage.com/100/100/fff&text=Test+image"
},
},
],
}
],
client=client,
)
except Exception as e:
print(e)
mock_post.assert_called()
print(mock_post.call_args.kwargs)
json_data = json.loads(mock_post.call_args.kwargs["data"])
assert json_data["model"] == "llama3.2-vision:11b"
assert "images" in json_data
assert "prompt" in json_data
assert json_data["prompt"].startswith("### User:\n")
def test_ollama_ssl_verify():
import ssl
import httpx
from litellm.llms.custom_httpx.http_handler import HTTPHandler
try:
response = litellm.completion(
model="ollama/llama3.1",
messages=[
{
"role": "user",
"content": "What's the weather like in San Francisco?",
}
],
ssl_verify=False,
)
except Exception as e:
print(e)
client: HTTPHandler = litellm.in_memory_llm_clients_cache.get_cache(
"httpx_clientssl_verify_False"
)
test_client = httpx.Client(verify=False)
print(client)
assert (
client.client._transport._pool._ssl_context.verify_mode
== test_client._transport._pool._ssl_context.verify_mode
)
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_async_ollama_ssl_verify(stream):
import httpx
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
try:
response = await litellm.acompletion(
model="ollama/llama3.1",
messages=[
{
"role": "user",
"content": "What's the weather like in San Francisco?",
}
],
ssl_verify=False,
stream=stream,
)
except Exception as e:
print(e)
client: AsyncHTTPHandler = litellm.in_memory_llm_clients_cache.get_cache("async_httpx_clientssl_verify_Falseollama")
# check client
print("type of transport in client=", type(client.client._transport))
print("vars in transport in client=", vars(client.client._transport))
litellm_created_session = client.client._transport.get_valid_client_session()
print("litellm_created_session=", litellm_created_session)
# check session ssl
print("litellm_created_session ssl=", litellm_created_session.connector._ssl)
# create aiohttp transport with ssl_verify=False
import aiohttp
aiohttp_session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False))
print("aiohttp_session ssl=", aiohttp_session.connector._ssl)
assert litellm_created_session.connector._ssl is False
assert litellm_created_session.connector._ssl == aiohttp_session.connector._ssl

View file

@ -1,152 +0,0 @@
# #### What this tests ####
# # This profiles a router call to find where calls are taking the most time.
# import sys, os, time, logging
# import traceback, asyncio, uuid
# import pytest
# import cProfile
# from pstats import Stats
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import litellm
# from litellm import Router
# from concurrent.futures import ThreadPoolExecutor
# from dotenv import load_dotenv
# from aiodebug import log_slow_callbacks # Import the aiodebug utility for logging slow callbacks
# # litellm.telemetry = False
# load_dotenv()
# logging.basicConfig(
# level=logging.DEBUG,
# format='%(asctime)s %(levelname)s: %(message)s',
# datefmt='%I:%M:%S %p',
# filename='aiologs.log', # Name of the log file where logs will be written
# filemode='w' # 'w' to overwrite the log file on each run, use 'a' to append
# )
# # Dictionary to store exception types and their counts
# exception_counts = {}
# exception_data = []
# litellm.telemetry = False
# num_task_cancelled_errors = 0
# 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_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_name": "azure-model",
# "litellm_params": {
# "model": "azure/gpt-35-turbo",
# "api_key": "os.environ/AZURE_CANADA_API_KEY",
# "api_base": "https://my-endpoint-canada-berri992.openai.azure.com",
# "rpm": 6
# }
# }]
# router = Router(model_list=model_list, set_verbose=False, num_retries=3)
# async def router_completion():
# global num_task_cancelled_errors, exception_counts
# try:
# messages=[{"role": "user", "content": f"This is a test: {uuid.uuid4()}"}]
# response = await router.acompletion(model="azure-model", messages=messages)
# return response
# except asyncio.exceptions.CancelledError:
# exception_type = "CancelledError"
# exception_counts[exception_type] = exception_counts.get(exception_type, 0) + 1
# print("Task was cancelled")
# num_task_cancelled_errors += 1
# exception_data.append({
# "type": exception_type,
# "traceback": None
# })
# return None
# except Exception as e:
# exception_type = type(e).__name__
# exception_counts[exception_type] = exception_counts.get(exception_type, 0) + 1
# exception_data.append({
# "type": exception_type,
# "traceback": traceback.format_exc()
# })
# return None
# async def loadtest_fn(n = 1452):
# global num_task_cancelled_errors, exception_counts
# start = time.time()
# tasks = [router_completion() for _ in range(n)]
# chat_completions = await asyncio.gather(*tasks)
# successful_completions = [c for c in chat_completions if c is not None]
# print(n, time.time() - start, len(successful_completions))
# # Print exception breakdown
# print("Exception Breakdown:")
# for exception_type, count in exception_counts.items():
# print(f"{exception_type}: {count}")
# # Store exception_data in a file
# with open('exception_data.txt', 'w') as file:
# for data in exception_data:
# file.write(f"Type: {data['type']}\n")
# if data['traceback']:
# file.write(f"Traceback:\n{data['traceback']}\n\n")
# loop = asyncio.get_event_loop()
# loop.set_debug(True)
# log_slow_callbacks.enable(0.05) # Log callbacks slower than 0.05 seconds
# # Excute the load testing function within the asyncio event loop
# loop.run_until_complete(loadtest_fn())
# # ### SUSTAINED LOAD TESTS ###
# # import time, asyncio
# # async def make_requests(n):
# # tasks = [router_completion() for _ in range(n)]
# # print(f"num tasks: {len(tasks)}")
# # chat_completions = await asyncio.gather(*tasks)
# # successful_completions = [c for c in chat_completions if c is not None]
# # print(f"successful_completions: {len(successful_completions)}")
# # return successful_completions
# # async def main():
# # start_time = time.time()
# # total_successful_requests = 0
# # request_limit = 1000
# # batches = 2 # batches of 1k requests
# # start = time.time()
# # tasks = [] # list to hold all tasks
# # async def request_loop():
# # nonlocal tasks
# # for _ in range(batches):
# # # Make 1,000 requests
# # task = asyncio.create_task(make_requests(request_limit))
# # tasks.append(task)
# # # Introduce a delay to achieve 1,000 requests per second
# # await asyncio.sleep(1)
# # await request_loop()
# # results = await asyncio.gather(*tasks)
# # total_successful_requests = sum(len(res) for res in results)
# # print(request_limit*batches, time.time() - start, total_successful_requests)
# # asyncio.run(main())

View file

@ -1,633 +0,0 @@
#### What this tests ####
# This tests setting provider specific configs across providers
# There are 2 types of tests - changing config dynamically or by setting class variables
import json
import os
import traceback
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
from litellm import RateLimitError, completion
# Anthropic
def claude_test_completion():
litellm.AnthropicConfig(max_tokens_to_sample=200)
# litellm.set_verbose=True
try:
# OVERRIDE WITH DYNAMIC MAX TOKENS
response_1 = litellm.completion(
model="claude-3-haiku-20240307",
messages=[{"content": "Hello, how are you?", "role": "user"}],
max_tokens=10,
)
# Add any assertions here to check the response
print(response_1)
response_1_text = response_1.choices[0].message.content
# USE CONFIG TOKENS
response_2 = litellm.completion(
model="claude-3-haiku-20240307",
messages=[{"content": "Hello, how are you?", "role": "user"}],
)
# Add any assertions here to check the response
print(response_2)
response_2_text = response_2.choices[0].message.content
assert len(response_2_text) > len(response_1_text)
try:
response_3 = litellm.completion(
model="claude-3-5-haiku-20241022",
messages=[{"content": "Hello, how are you?", "role": "user"}],
n=2,
)
except Exception as e:
print(e)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# claude_test_completion()
# Replicate
def replicate_test_completion():
litellm.ReplicateConfig(max_new_tokens=200)
# litellm.set_verbose=True
try:
# OVERRIDE WITH DYNAMIC MAX TOKENS
response_1 = litellm.completion(
model="meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3",
messages=[{"content": "Hello, how are you?", "role": "user"}],
max_tokens=10,
)
# Add any assertions here to check the response
print(response_1)
response_1_text = response_1.choices[0].message.content
# USE CONFIG TOKENS
response_2 = litellm.completion(
model="meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3",
messages=[{"content": "Hello, how are you?", "role": "user"}],
)
# Add any assertions here to check the response
print(response_2)
response_2_text = response_2.choices[0].message.content
assert len(response_2_text) > len(response_1_text)
try:
response_3 = litellm.completion(
model="meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3",
messages=[{"content": "Hello, how are you?", "role": "user"}],
n=2,
)
except Exception:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# replicate_test_completion()
# Cohere
def cohere_test_completion():
# litellm.CohereConfig(max_tokens=200)
litellm.set_verbose = True
try:
# OVERRIDE WITH DYNAMIC MAX TOKENS
response_1 = litellm.completion(
model="command-nightly",
messages=[{"content": "Hello, how are you?", "role": "user"}],
max_tokens=10,
)
response_1_text = response_1.choices[0].message.content
# USE CONFIG TOKENS
response_2 = litellm.completion(
model="command-nightly",
messages=[{"content": "Hello, how are you?", "role": "user"}],
)
response_2_text = response_2.choices[0].message.content
assert len(response_2_text) > len(response_1_text)
response_3 = litellm.completion(
model="command-nightly",
messages=[{"content": "Hello, how are you?", "role": "user"}],
n=2,
)
assert len(response_3.choices) > 1
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# cohere_test_completion()
# ai21_test_completion()
# TogetherAI
def togetherai_test_completion():
litellm.TogetherAIConfig(max_tokens=10)
litellm.set_verbose = True
try:
# OVERRIDE WITH DYNAMIC MAX TOKENS
response_1 = litellm.completion(
model="together_ai/togethercomputer/llama-2-70b-chat",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
max_tokens=100,
)
response_1_text = response_1.choices[0].message.content
print(f"response_1_text: {response_1_text}")
# USE CONFIG TOKENS
response_2 = litellm.completion(
model="together_ai/togethercomputer/llama-2-70b-chat",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
)
response_2_text = response_2.choices[0].message.content
print(f"response_2_text: {response_2_text}")
assert len(response_2_text) < len(response_1_text)
try:
response_3 = litellm.completion(
model="together_ai/togethercomputer/llama-2-70b-chat",
messages=[{"content": "Hello, how are you?", "role": "user"}],
n=2,
)
pytest.fail(f"Error not raised when n=2 passed to provider")
except Exception:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# togetherai_test_completion()
# Palm
# palm_test_completion()
# NLP Cloud
def nlp_cloud_test_completion():
litellm.NLPCloudConfig(max_length=10)
# litellm.set_verbose=True
try:
# OVERRIDE WITH DYNAMIC MAX TOKENS
response_1 = litellm.completion(
model="dolphin",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
max_tokens=100,
)
response_1_text = response_1.choices[0].message.content
print(f"response_1_text: {response_1_text}")
# USE CONFIG TOKENS
response_2 = litellm.completion(
model="dolphin",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
)
response_2_text = response_2.choices[0].message.content
print(f"response_2_text: {response_2_text}")
assert len(response_2_text) < len(response_1_text)
try:
response_3 = litellm.completion(
model="dolphin",
messages=[{"content": "Hello, how are you?", "role": "user"}],
n=2,
)
pytest.fail(f"Error not raised when n=2 passed to provider")
except Exception:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# nlp_cloud_test_completion()
# AlephAlpha
def aleph_alpha_test_completion():
litellm.AlephAlphaConfig(maximum_tokens=10)
# litellm.set_verbose=True
try:
# OVERRIDE WITH DYNAMIC MAX TOKENS
response_1 = litellm.completion(
model="luminous-base",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
max_tokens=100,
)
response_1_text = response_1.choices[0].message.content
print(f"response_1_text: {response_1_text}")
# USE CONFIG TOKENS
response_2 = litellm.completion(
model="luminous-base",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
)
response_2_text = response_2.choices[0].message.content
print(f"response_2_text: {response_2_text}")
assert len(response_2_text) < len(response_1_text)
response_3 = litellm.completion(
model="luminous-base",
messages=[{"content": "Hello, how are you?", "role": "user"}],
n=2,
)
assert len(response_3.choices) > 1
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# aleph_alpha_test_completion()
# Sagemaker
# sagemaker_test_completion()
def test_sagemaker_default_region():
"""
If no regions are specified in config or in environment, the default region is us-west-2
"""
mock_response = MagicMock()
def return_val():
return {
"generated_text": "This is a mock response from SageMaker.",
"id": "cmpl-mockid",
"object": "text_completion",
"created": 1629800000,
"model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
"choices": [
{
"text": "This is a mock response from SageMaker.",
"index": 0,
"logprobs": None,
"finish_reason": "length",
}
],
"usage": {"prompt_tokens": 1, "completion_tokens": 8, "total_tokens": 9},
}
mock_response.json = return_val
mock_response.status_code = 200
with patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
return_value=mock_response,
) as mock_post:
response = litellm.completion(
model="sagemaker/mock-endpoint",
messages=[{"content": "Hello, world!", "role": "user"}],
)
mock_post.assert_called_once()
_, kwargs = mock_post.call_args
print(f"kwargs: {kwargs}")
args_to_sagemaker = json.loads(kwargs["data"])
print("Arguments passed to sagemaker=", args_to_sagemaker)
print("url=", kwargs["url"])
assert (
kwargs["url"]
== "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/mock-endpoint/invocations"
)
# test_sagemaker_default_region()
def test_sagemaker_environment_region():
"""
If a region is specified in the environment, use that region instead of us-west-2
"""
expected_region = "us-east-1"
os.environ["AWS_REGION_NAME"] = expected_region
mock_response = MagicMock()
def return_val():
return {
"generated_text": "This is a mock response from SageMaker.",
"id": "cmpl-mockid",
"object": "text_completion",
"created": 1629800000,
"model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
"choices": [
{
"text": "This is a mock response from SageMaker.",
"index": 0,
"logprobs": None,
"finish_reason": "length",
}
],
"usage": {"prompt_tokens": 1, "completion_tokens": 8, "total_tokens": 9},
}
mock_response.json = return_val
mock_response.status_code = 200
with patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
return_value=mock_response,
) as mock_post:
response = litellm.completion(
model="sagemaker/mock-endpoint",
messages=[{"content": "Hello, world!", "role": "user"}],
)
mock_post.assert_called_once()
_, kwargs = mock_post.call_args
args_to_sagemaker = json.loads(kwargs["data"])
print("Arguments passed to sagemaker=", args_to_sagemaker)
print("url=", kwargs["url"])
assert (
kwargs["url"]
== f"https://runtime.sagemaker.{expected_region}.amazonaws.com/endpoints/mock-endpoint/invocations"
)
del os.environ["AWS_REGION_NAME"] # cleanup
# test_sagemaker_environment_region()
def test_sagemaker_config_region():
"""
If a region is specified as part of the optional parameters of the completion, including as
part of the config file, then use that region instead of us-west-2
"""
expected_region = "us-east-1"
mock_response = MagicMock()
def return_val():
return {
"generated_text": "This is a mock response from SageMaker.",
"id": "cmpl-mockid",
"object": "text_completion",
"created": 1629800000,
"model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
"choices": [
{
"text": "This is a mock response from SageMaker.",
"index": 0,
"logprobs": None,
"finish_reason": "length",
}
],
"usage": {"prompt_tokens": 1, "completion_tokens": 8, "total_tokens": 9},
}
mock_response.json = return_val
mock_response.status_code = 200
with patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
return_value=mock_response,
) as mock_post:
response = litellm.completion(
model="sagemaker/mock-endpoint",
messages=[{"content": "Hello, world!", "role": "user"}],
aws_region_name=expected_region,
)
mock_post.assert_called_once()
_, kwargs = mock_post.call_args
args_to_sagemaker = json.loads(kwargs["data"])
print("Arguments passed to sagemaker=", args_to_sagemaker)
print("url=", kwargs["url"])
assert (
kwargs["url"]
== f"https://runtime.sagemaker.{expected_region}.amazonaws.com/endpoints/mock-endpoint/invocations"
)
# test_sagemaker_config_region()
# test_sagemaker_config_and_environment_region()
# Bedrock
def bedrock_test_completion():
litellm.AmazonCohereConfig(max_tokens=10)
# litellm.set_verbose=True
try:
# OVERRIDE WITH DYNAMIC MAX TOKENS
response_1 = litellm.completion(
model="bedrock/cohere.command-r-v1:0",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
max_tokens=100,
)
response_1_text = response_1.choices[0].message.content
print(f"response_1_text: {response_1_text}")
# USE CONFIG TOKENS
response_2 = litellm.completion(
model="bedrock/cohere.command-r-v1:0",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
)
response_2_text = response_2.choices[0].message.content
print(f"response_2_text: {response_2_text}")
assert len(response_2_text) < len(response_1_text)
except RateLimitError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# bedrock_test_completion()
# OpenAI Chat Completion
def openai_test_completion():
litellm.OpenAIConfig(max_tokens=10)
# litellm.set_verbose=True
try:
# OVERRIDE WITH DYNAMIC MAX TOKENS
response_1 = litellm.completion(
model="gpt-3.5-turbo",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
max_tokens=100,
)
response_1_text = response_1.choices[0].message.content
print(f"response_1_text: {response_1_text}")
# USE CONFIG TOKENS
response_2 = litellm.completion(
model="gpt-3.5-turbo",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
)
response_2_text = response_2.choices[0].message.content
print(f"response_2_text: {response_2_text}")
assert len(response_2_text) < len(response_1_text)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# openai_test_completion()
# OpenAI Text Completion
def openai_text_completion_test():
litellm.OpenAITextCompletionConfig(max_tokens=10)
# litellm.set_verbose=True
try:
# OVERRIDE WITH DYNAMIC MAX TOKENS
response_1 = litellm.completion(
model="gpt-3.5-turbo-instruct",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
max_tokens=100,
)
response_1_text = response_1.choices[0].message.content
print(f"response_1_text: {response_1_text}")
# USE CONFIG TOKENS
response_2 = litellm.completion(
model="gpt-3.5-turbo-instruct",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
)
response_2_text = response_2.choices[0].message.content
print(f"response_2_text: {response_2_text}")
assert len(response_2_text) < len(response_1_text)
response_3 = litellm.completion(
model="gpt-3.5-turbo-instruct",
messages=[{"content": "Hello, how are you?", "role": "user"}],
n=2,
)
assert len(response_3.choices) > 1
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# openai_text_completion_test()
# Azure OpenAI
def azure_openai_test_completion():
litellm.AzureOpenAIConfig(max_tokens=10)
# litellm.set_verbose=True
try:
# OVERRIDE WITH DYNAMIC MAX TOKENS
response_1 = litellm.completion(
model="azure/gpt-4.1-mini",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
max_tokens=100,
)
response_1_text = response_1.choices[0].message.content
print(f"response_1_text: {response_1_text}")
# USE CONFIG TOKENS
response_2 = litellm.completion(
model="azure/gpt-4.1-mini",
messages=[
{
"content": "Hello, how are you? Be as verbose as possible",
"role": "user",
}
],
)
response_2_text = response_2.choices[0].message.content
print(f"response_2_text: {response_2_text}")
assert len(response_2_text) < len(response_1_text)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# azure_openai_test_completion()

View file

@ -1,57 +0,0 @@
import traceback
from dotenv import load_dotenv
load_dotenv()
import io
import asyncio
import json
import tempfile
from unittest.mock import MagicMock, patch
import pytest
import litellm
from litellm.types.utils import (
ChatCompletionTokenLogprob,
ChoiceLogprobs,
Delta,
ModelResponse,
StreamingChoices,
TopLogprob,
)
obj = ModelResponse(
id="chat-f9bad6ec3c1146e99368682a0e7403fc",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(content="", role=None, function_call=None, tool_calls=None),
logprobs=ChoiceLogprobs(
content=[
ChatCompletionTokenLogprob(
token="",
bytes=[],
logprob=-0.00018153927521780133,
top_logprobs=[
TopLogprob(
token="", bytes=[], logprob=-0.00018153927521780133
),
TopLogprob(
token="\n\n", bytes=[10, 10], logprob=-9.062681198120117
),
],
)
]
),
)
],
created=1721976759,
model="Meta-Llama-3-8B-Instruct",
object="chat.completion.chunk",
system_fingerprint=None,
)
print(obj.model_dump())

View file

@ -1,50 +0,0 @@
#### What this tests ####
# This tests calling batch_completions by running 100 messages together
import ast
from pathlib import Path
import pytest
import litellm
# test_update_model_cost_map_url()
def test_update_model_cost_via_completion():
try:
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
input_cost_per_token=0.3,
output_cost_per_token=0.4,
)
print(
f"litellm.model_cost for gpt-3.5-turbo: {litellm.model_cost['gpt-3.5-turbo']}"
)
assert litellm.model_cost["gpt-3.5-turbo"]["input_cost_per_token"] == 0.3
assert litellm.model_cost["gpt-3.5-turbo"]["output_cost_per_token"] == 0.4
except Exception as e:
pytest.fail(f"An error occurred: {e}")
def test_no_test_invocation_at_module_scope():
tree = ast.parse(Path(__file__).read_text())
defined = {
node.name
for node in tree.body
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
}
invoked = [
node.value.func.id
for node in tree.body
if isinstance(node, ast.Expr)
and isinstance(node.value, ast.Call)
and isinstance(node.value.func, ast.Name)
and node.value.func.id in defined
]
assert not invoked, (
f"{invoked} run at import time, so pytest collecting this file fires real "
"provider calls; any failure aborts collection and tears down the whole job"
)

View file

@ -1,56 +0,0 @@
# [LOCAL TEST] - runs against mock openai proxy
# # What this tests?
# ## This tests if fallbacks works for 429 errors
# import sys, os, time
# import traceback, asyncio
# import pytest
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import litellm
# from litellm import Router
# model_list = [
# { # list of model deployments
# "model_name": "text-embedding-ada-002", # model alias
# "litellm_params": { # params for litellm completion/embedding call
# "model": "text-embedding-ada-002", # actual model name
# "api_key": "sk-fakekey",
# "api_base": "http://0.0.0.0:8080",
# },
# "tpm": 1000,
# "rpm": 6,
# },
# {
# "model_name": "text-embedding-ada-002-fallback",
# "litellm_params": { # params for litellm completion/embedding call
# "model": "openai/text-embedding-ada-002-anything-else", # actual model name
# "api_key": "sk-fakekey2",
# "api_base": "http://0.0.0.0:8080",
# },
# "tpm": 1000,
# "rpm": 6,
# },
# ]
# router = Router(
# model_list=model_list,
# fallbacks=[
# {"text-embedding-ada-002": ["text-embedding-ada-002-fallback"]},
# {"text-embedding-ada-002-fallback": ["text-embedding-ada-002"]},
# ],
# set_verbose=True,
# num_retries=0,
# debug_level="INFO",
# routing_strategy="usage-based-routing",
# )
# def test_embedding_with_fallbacks():
# response = router.embedding(model="text-embedding-ada-002", input=["Hello world"])
# print(f"response: {response}")
# test_embedding_with_fallbacks()

View file

@ -1,408 +0,0 @@
import json
from dotenv import load_dotenv
load_dotenv()
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from test_streaming import streaming_format_tests
import litellm
from litellm import completion_cost
# 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"}]
import logging
from litellm._logging import verbose_logger
def logger_fn(user_model_dict):
print(f"user_model_dict: {user_model_dict}")
@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.asyncio()
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_completion_sagemaker(sync_mode):
try:
litellm.set_verbose = True
verbose_logger.setLevel(logging.DEBUG)
print("testing sagemaker")
if sync_mode is True:
response = litellm.completion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[
{"role": "user", "content": "hi"},
],
temperature=0.2,
max_tokens=80,
cost_per_second=0.000420,
)
else:
response = await litellm.acompletion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[
{"role": "user", "content": "hi"},
],
temperature=0.2,
max_tokens=80,
cost_per_second=0.000420,
)
# Add any assertions here to check the response
print(response)
cost = completion_cost(completion_response=response)
print("calculated cost", cost)
assert (
cost > 0.0 and cost < 1.0
) # should never be > $1 for a single completion call
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio()
@pytest.mark.parametrize("sync_mode", [False, True])
@pytest.mark.parametrize(
"model",
[
# "sagemaker_chat/huggingface-pytorch-tgi-inference-2024-08-23-15-48-59-245",
"sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
],
)
# @pytest.mark.flaky(retries=3, delay=1)
async def test_completion_sagemaker_stream(sync_mode, model):
try:
litellm.set_verbose = False
print("testing sagemaker")
verbose_logger.setLevel(logging.DEBUG)
full_text = ""
if sync_mode is True:
response = litellm.completion(
model=model,
messages=[
{"role": "user", "content": "hi - what is ur name"},
],
temperature=0.2,
stream=True,
max_tokens=80,
cost_per_second=0.000420,
)
for idx, chunk in enumerate(response):
print(chunk)
streaming_format_tests(idx=idx, chunk=chunk)
full_text += chunk.choices[0].delta.content or ""
print("SYNC RESPONSE full text", full_text)
else:
response = await litellm.acompletion(
model=model,
messages=[
{"role": "user", "content": "hi - what is ur name"},
],
stream=True,
temperature=0.2,
max_tokens=80,
cost_per_second=0.000420,
)
print("streaming response")
idx = 0
async for chunk in response:
print(chunk)
streaming_format_tests(idx=idx, chunk=chunk)
full_text += chunk.choices[0].delta.content or ""
idx += 1
print("ASYNC RESPONSE full text", full_text)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio()
@pytest.mark.parametrize("sync_mode", [False, True])
@pytest.mark.parametrize(
"model",
[
# "sagemaker_chat/huggingface-pytorch-tgi-inference-2024-08-23-15-48-59-245",
"sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
],
)
async def test_completion_sagemaker_streaming_bad_request(sync_mode, model):
litellm.set_verbose = True
print("testing sagemaker")
if sync_mode is True:
with pytest.raises(litellm.BadRequestError):
response = litellm.completion(
model=model,
messages=[
{"role": "user", "content": "hi"},
],
stream=True,
max_tokens=8000000000000000,
)
else:
with pytest.raises(litellm.BadRequestError):
response = await litellm.acompletion(
model=model,
messages=[
{"role": "user", "content": "hi"},
],
stream=True,
max_tokens=8000000000000000,
)
@pytest.mark.asyncio
async def test_acompletion_sagemaker_non_stream():
mock_response = AsyncMock()
def return_val():
return {
"generated_text": "This is a mock response from SageMaker.",
"id": "cmpl-mockid",
"object": "text_completion",
"created": 1629800000,
"model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
"choices": [
{
"text": "This is a mock response from SageMaker.",
"index": 0,
"logprobs": None,
"finish_reason": "length",
}
],
"usage": {"prompt_tokens": 1, "completion_tokens": 8, "total_tokens": 9},
}
mock_response.json = return_val
mock_response.status_code = 200
expected_payload = {
"inputs": "hi",
"parameters": {"temperature": 0.2, "max_new_tokens": 80},
}
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=mock_response,
) as mock_post:
# Act: Call the litellm.acompletion function
response = await litellm.acompletion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[
{"role": "user", "content": "hi"},
],
temperature=0.2,
max_tokens=80,
cost_per_second=0.000420,
)
# Print what was called on the mock
print("call args=", mock_post.call_args)
# Assert
mock_post.assert_called_once()
_, kwargs = mock_post.call_args
args_to_sagemaker = json.loads(kwargs["data"])
print("Arguments passed to sagemaker=", args_to_sagemaker)
assert args_to_sagemaker == expected_payload
assert (
kwargs["url"]
== "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations"
)
@pytest.mark.asyncio
async def test_completion_sagemaker_non_stream():
mock_response = MagicMock()
def return_val():
return {
"generated_text": "This is a mock response from SageMaker.",
"id": "cmpl-mockid",
"object": "text_completion",
"created": 1629800000,
"model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
"choices": [
{
"text": "This is a mock response from SageMaker.",
"index": 0,
"logprobs": None,
"finish_reason": "length",
}
],
"usage": {"prompt_tokens": 1, "completion_tokens": 8, "total_tokens": 9},
}
mock_response.json = return_val
mock_response.status_code = 200
expected_payload = {
"inputs": "hi",
"parameters": {"temperature": 0.2, "max_new_tokens": 80},
}
with patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
return_value=mock_response,
) as mock_post:
# Act: Call the litellm.acompletion function
response = litellm.completion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[
{"role": "user", "content": "hi"},
],
temperature=0.2,
max_tokens=80,
cost_per_second=0.000420,
)
# Print what was called on the mock
print("call args=", mock_post.call_args)
# Assert
mock_post.assert_called_once()
_, kwargs = mock_post.call_args
args_to_sagemaker = json.loads(kwargs["data"])
print("Arguments passed to sagemaker=", args_to_sagemaker)
assert args_to_sagemaker == expected_payload
assert (
kwargs["url"]
== "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations"
)
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_completion_sagemaker_prompt_template_non_stream():
mock_response = MagicMock()
def return_val():
return {
"generated_text": "This is a mock response from SageMaker.",
"id": "cmpl-mockid",
"object": "text_completion",
"created": 1629800000,
"model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
"choices": [
{
"text": "This is a mock response from SageMaker.",
"index": 0,
"logprobs": None,
"finish_reason": "length",
}
],
"usage": {"prompt_tokens": 1, "completion_tokens": 8, "total_tokens": 9},
}
mock_response.json = return_val
mock_response.status_code = 200
expected_payload = {
"inputs": "<|begin▁of▁sentence|>You are an AI programming assistant, utilizing the Deepseek Coder model, developed by Deepseek Company, and you only answer questions related to computer science. For politically sensitive questions, security and privacy issues, and other non-computer science questions, you will refuse to answer\n\n### Instruction:\nhi\n\n\n### Response:\n",
"parameters": {"temperature": 0.2, "max_new_tokens": 80},
}
with patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
return_value=mock_response,
) as mock_post:
# Act: Call the litellm.acompletion function
response = litellm.completion(
model="sagemaker/deepseek_coder_6.7_instruct",
messages=[
{"role": "user", "content": "hi"},
],
temperature=0.2,
max_tokens=80,
hf_model_name="deepseek-ai/deepseek-coder-6.7b-instruct",
)
# Print what was called on the mock
print("call args=", mock_post.call_args)
# Assert
mock_post.assert_called_once()
_, kwargs = mock_post.call_args
args_to_sagemaker = json.loads(kwargs["data"])
print("Arguments passed to sagemaker=", args_to_sagemaker)
assert args_to_sagemaker == expected_payload
@pytest.mark.asyncio
async def test_completion_sagemaker_non_stream_with_aws_params():
mock_response = MagicMock()
def return_val():
return {
"generated_text": "This is a mock response from SageMaker.",
"id": "cmpl-mockid",
"object": "text_completion",
"created": 1629800000,
"model": "sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
"choices": [
{
"text": "This is a mock response from SageMaker.",
"index": 0,
"logprobs": None,
"finish_reason": "length",
}
],
"usage": {"prompt_tokens": 1, "completion_tokens": 8, "total_tokens": 9},
}
mock_response.json = return_val
mock_response.status_code = 200
expected_payload = {
"inputs": "hi",
"parameters": {"temperature": 0.2, "max_new_tokens": 80},
}
with patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
return_value=mock_response,
) as mock_post:
# Act: Call the litellm.acompletion function
response = litellm.completion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[
{"role": "user", "content": "hi"},
],
temperature=0.2,
max_tokens=80,
cost_per_second=0.000420,
aws_access_key_id="gm",
aws_secret_access_key="s",
aws_region_name="us-west-5",
)
# Print what was called on the mock
print("call args=", mock_post.call_args)
# Assert
mock_post.assert_called_once()
_, kwargs = mock_post.call_args
args_to_sagemaker = json.loads(kwargs["data"])
print("Arguments passed to sagemaker=", args_to_sagemaker)
assert args_to_sagemaker == expected_payload
assert (
kwargs["url"]
== "https://runtime.sagemaker.us-west-5.amazonaws.com/endpoints/jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations"
)

View file

@ -1,278 +0,0 @@
"""
Integration tests for SageMaker Nova provider.
These tests require a live SageMaker Nova endpoint and AWS credentials.
They are skipped by default — run manually with:
pytest tests/local_testing/test_sagemaker_nova_integration.py -v --no-header -rN
Prerequisites:
export AWS_PROFILE=<your-profile> # or set AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY
export AWS_REGION_NAME=us-east-1
export SAGEMAKER_NOVA_ENDPOINT=<your-endpoint-name>
"""
import base64
import io
import json
import os
import struct
import zlib
import pytest
import litellm
ENDPOINT = os.environ.get("SAGEMAKER_NOVA_ENDPOINT", "")
MODEL = f"sagemaker_nova/{ENDPOINT}"
skip_if_no_endpoint = pytest.mark.skipif(
not ENDPOINT,
reason="SAGEMAKER_NOVA_ENDPOINT not set — skipping live integration tests",
)
def _make_test_png() -> str:
"""Create a minimal 4x4 PNG (red border, blue center) and return base64."""
def chunk(ctype, data):
c = ctype + data
return (
struct.pack(">I", len(data))
+ c
+ struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)
)
width, height = 4, 4
pixels = []
for y in range(height):
for x in range(width):
if 1 <= x <= 2 and 1 <= y <= 2:
pixels.append((0, 0, 255))
else:
pixels.append((255, 0, 0))
raw = b""
for y in range(height):
raw += b"\x00"
for x in range(width):
raw += bytes(pixels[y * width + x])
png = (
b"\x89PNG\r\n\x1a\n"
+ chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
+ chunk(b"IDAT", zlib.compress(raw))
+ chunk(b"IEND", b"")
)
return base64.b64encode(png).decode()
@skip_if_no_endpoint
class TestSagemakerNovaIntegration:
"""Live integration tests for sagemaker_nova provider."""
def test_should_complete_basic_single_turn(self):
"""Basic single-turn chat completion."""
response = litellm.completion(
model=MODEL,
messages=[{"role": "user", "content": "What is 2+2? Reply in one word."}],
max_tokens=32,
temperature=0.1,
)
assert response.choices[0].message.content is not None
assert len(response.choices[0].message.content.strip()) > 0
assert response.choices[0].finish_reason == "stop"
assert response.usage.prompt_tokens > 0
assert response.usage.completion_tokens > 0
assert response.usage.total_tokens == (
response.usage.prompt_tokens + response.usage.completion_tokens
)
def test_should_complete_multi_turn_conversation(self):
"""Multi-turn conversation maintains context."""
messages = [
{"role": "user", "content": "My name is Alice."},
]
response1 = litellm.completion(
model=MODEL,
messages=messages,
max_tokens=64,
temperature=0.1,
)
assistant_msg = response1.choices[0].message.content
assert assistant_msg is not None
# Second turn — model should remember the name
messages.append({"role": "assistant", "content": assistant_msg})
messages.append({"role": "user", "content": "What is my name?"})
response2 = litellm.completion(
model=MODEL,
messages=messages,
max_tokens=64,
temperature=0.1,
)
answer = response2.choices[0].message.content.lower()
assert "alice" in answer, f"Expected 'alice' in response, got: {answer}"
def test_should_stream_response(self):
"""Streaming returns chunks with content and final usage."""
response = litellm.completion(
model=MODEL,
messages=[{"role": "user", "content": "Count from 1 to 5."}],
max_tokens=64,
stream=True,
stream_options={"include_usage": True},
)
chunks = []
full_content = ""
for chunk in response:
chunks.append(chunk)
delta = chunk.choices[0].delta.content or ""
full_content += delta
assert len(chunks) > 1, "Expected multiple streaming chunks"
assert len(full_content.strip()) > 0, "Expected non-empty streamed content"
# Last chunk should have finish_reason
final_chunks_with_finish = [
c for c in chunks if c.choices and c.choices[0].finish_reason is not None
]
assert (
len(final_chunks_with_finish) > 0
), "Expected at least one chunk with finish_reason"
def test_should_return_logprobs(self):
"""Logprobs are returned when requested."""
response = litellm.completion(
model=MODEL,
messages=[{"role": "user", "content": "Say hello."}],
max_tokens=16,
temperature=0.1,
logprobs=True,
top_logprobs=3,
)
lp = response.choices[0].logprobs
assert lp is not None, "Expected logprobs in response"
content = lp.content if hasattr(lp, "content") else lp.get("content")
assert content is not None and len(content) > 0, "Expected logprobs content"
first_token = content[0]
assert "token" in first_token or hasattr(first_token, "token")
assert "logprob" in first_token or hasattr(first_token, "logprob")
top = (
first_token.get("top_logprobs")
if isinstance(first_token, dict)
else first_token.top_logprobs
)
assert top is not None and len(top) == 3, "Expected 3 top_logprobs"
def test_should_handle_multimodal_image_input(self):
"""Multimodal with base64 image in content array."""
b64_image = _make_test_png()
response = litellm.completion(
model=MODEL,
messages=[
{
"role": "user",
"content": [
{
"type": "text",
"text": "What colors do you see in this image? List them.",
},
{
"type": "image_url",
"image_url": {"url": f"data:image/png;base64,{b64_image}"},
},
],
}
],
max_tokens=128,
)
content = response.choices[0].message.content.lower()
assert response.choices[0].message.content is not None
assert len(content) > 0
# The image has red and blue — model should mention at least one
assert (
"red" in content or "blue" in content
), f"Expected 'red' or 'blue' in multimodal response, got: {content}"
def test_should_pass_nova_specific_params(self):
"""Nova-specific parameters (top_k) are accepted."""
response = litellm.completion(
model=MODEL,
messages=[{"role": "user", "content": "Say hello."}],
max_tokens=32,
top_k=40,
temperature=0.7,
)
assert response.choices[0].message.content is not None
assert response.usage.total_tokens > 0
def test_should_respect_system_message(self):
"""System message should influence the response."""
response = litellm.completion(
model=MODEL,
messages=[
{
"role": "system",
"content": "You are a pirate. Always respond in pirate speak.",
},
{"role": "user", "content": "How are you today?"},
],
max_tokens=128,
temperature=0.7,
)
content = response.choices[0].message.content.lower()
assert response.choices[0].message.content is not None
# Pirate-themed words likely in response
pirate_words = ["arr", "ahoy", "matey", "ye", "sail", "sea", "cap"]
assert any(
w in content for w in pirate_words
), f"Expected pirate speak, got: {content}"
NOVA2_ENDPOINT = os.environ.get("SAGEMAKER_NOVA2_LITE_ENDPOINT", "")
NOVA2_MODEL = f"sagemaker_nova/{NOVA2_ENDPOINT}"
skip_if_no_nova2_endpoint = pytest.mark.skipif(
not NOVA2_ENDPOINT,
reason="SAGEMAKER_NOVA2_LITE_ENDPOINT not set — requires Nova 2 Lite endpoint",
)
@skip_if_no_nova2_endpoint
class TestSagemakerNova2LiteIntegration:
"""
Integration tests requiring a Nova 2 Lite endpoint (reasoning_effort support).
Run with:
export SAGEMAKER_NOVA2_LITE_ENDPOINT=<your-nova-2-lite-endpoint>
pytest tests/local_testing/test_sagemaker_nova_integration.py::TestSagemakerNova2LiteIntegration -v
"""
def test_should_accept_reasoning_effort_low(self):
"""reasoning_effort='low' should be accepted by Nova 2 Lite."""
response = litellm.completion(
model=NOVA2_MODEL,
messages=[{"role": "user", "content": "What is 2+2?"}],
max_tokens=32,
reasoning_effort="low",
)
assert response.choices[0].message.content is not None
assert response.usage.total_tokens > 0
def test_should_accept_reasoning_effort_high(self):
"""reasoning_effort='high' should be accepted by Nova 2 Lite."""
response = litellm.completion(
model=NOVA2_MODEL,
messages=[{"role": "user", "content": "Explain why the sky is blue."}],
max_tokens=256,
reasoning_effort="high",
)
assert response.choices[0].message.content is not None
assert len(response.choices[0].message.content) > 0
assert response.usage.completion_tokens > 0

View file

@ -1,897 +0,0 @@
import asyncio
import os
import time
import traceback
import pytest
from typing import List
from litellm.types.utils import StreamingChoices, ChatCompletionAudioResponse
def check_non_streaming_response(response):
assert response.choices[0].message.audio is not None, "Audio response is missing"
print("audio", response.choices[0].message.audio)
assert isinstance(
response.choices[0].message.audio, ChatCompletionAudioResponse
), "Invalid audio response type"
assert len(response.choices[0].message.audio.data) > 0, "Audio data is empty"
import dotenv
from openai import OpenAI
import litellm
import stream_chunk_testdata
from litellm import completion, stream_chunk_builder
dotenv.load_dotenv()
user_message = "What is the current weather in Boston?"
messages = [{"content": user_message, "role": "user"}]
function_schema = {
"name": "get_weather",
"description": "gets the current weather",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
},
"required": ["location"],
},
}
tools_schema = [
{
"type": "function",
"function": {
"name": "get_current_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"],
},
},
}
]
# def test_stream_chunk_builder_tools():
# try:
# litellm.set_verbose = False
# response = client.chat.completions.create(
# model="gpt-3.5-turbo",
# messages=messages,
# tools=tools_schema,
# # stream=True,
# # complete_response=True # runs stream_chunk_builder under-the-hood
# )
# print(f"response: {response}")
# print(f"response usage: {response.usage}")
# except Exception as e:
# pytest.fail(f"An exception occurred - {str(e)}")
# test_stream_chunk_builder_tools()
def test_stream_chunk_builder_litellm_function_call():
try:
litellm.set_verbose = False
response = litellm.completion(
model="gpt-3.5-turbo",
messages=messages,
functions=[function_schema],
# stream=True,
# complete_response=True # runs stream_chunk_builder under-the-hood
)
print(f"response: {response}")
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# test_stream_chunk_builder_litellm_function_call()
def test_stream_chunk_builder_litellm_tool_call():
try:
litellm.set_verbose = True
response = litellm.completion(
model="gpt-3.5-turbo",
messages=messages,
tools=tools_schema,
stream=True,
complete_response=True,
)
print(f"complete response: {response}")
print(f"complete response usage: {response.usage}")
assert response.usage.completion_tokens > 0
assert response.usage.prompt_tokens > 0
assert (
response.usage.total_tokens
== response.usage.completion_tokens + response.usage.prompt_tokens
)
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# test_stream_chunk_builder_litellm_tool_call()
def test_stream_chunk_builder_litellm_tool_call_regular_message():
try:
messages = [{"role": "user", "content": "Hey, how's it going?"}]
# litellm.set_verbose = True
response = litellm.completion(
model="gpt-3.5-turbo",
messages=messages,
tools=tools_schema,
stream=True,
complete_response=True,
)
print(f"complete response: {response}")
print(f"complete response usage: {response.usage}")
assert response.usage.completion_tokens > 0
assert response.usage.prompt_tokens > 0
assert (
response.usage.total_tokens
== response.usage.completion_tokens + response.usage.prompt_tokens
)
# check provider is in hidden params
print("hidden params", response._hidden_params)
assert response._hidden_params["custom_llm_provider"] == "openai"
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# test_stream_chunk_builder_litellm_tool_call_regular_message()
def test_stream_chunk_builder_litellm_mixed_calls():
response = stream_chunk_builder(stream_chunk_testdata.chunks)
assert (
response.choices[0].message.content
== "To answer your question about how many rows are in the 'users' table, I'll need to run a SQL query. Let me do that for you."
)
print(response.choices[0].message.tool_calls[0].to_dict())
assert len(response.choices[0].message.tool_calls) == 1
assert response.choices[0].message.tool_calls[0].to_dict() == {
"function": {
"arguments": '{"query": "SELECT COUNT(*) FROM users;"}',
"name": "sql_query",
},
"id": "toolu_01H3AjkLpRtGQrof13CBnWfK",
"type": "function",
}
def test_stream_chunk_builder_litellm_empty_chunks():
with pytest.raises(litellm.APIError):
response = stream_chunk_builder(chunks=None)
response = stream_chunk_builder(chunks=[])
assert response is None
def test_stream_chunk_builder_multiple_tool_calls():
init_chunks = [
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"id": "call_X9P9B6STj7ze8OsJCGkfoN94",
"function": {"arguments": "", "name": "exponentiate"},
"type": "function",
"index": 0,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": '{"ba'},
"type": "function",
"index": 0,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": 'se": '},
"type": "function",
"index": 0,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": '3, "ex'},
"type": "function",
"index": 0,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": "pone"},
"type": "function",
"index": 0,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": 'nt": '},
"type": "function",
"index": 0,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": "5}"},
"type": "function",
"index": 0,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"id": "call_Qq8yDeRx7v276abRcLrYORdW",
"function": {"arguments": "", "name": "add"},
"type": "function",
"index": 1,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": '{"fi'},
"type": "function",
"index": 1,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": "rst_i"},
"type": "function",
"index": 1,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": 'nt": 1'},
"type": "function",
"index": 1,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": '2, "'},
"type": "function",
"index": 1,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": "secon"},
"type": "function",
"index": 1,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": 'd_int"'},
"type": "function",
"index": 1,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"function": {"arguments": ": 3}"},
"type": "function",
"index": 1,
}
],
},
}
],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
{
"id": "chatcmpl-A5kCnzaxRsknd6008552ZhDi71yPt",
"choices": [{"finish_reason": "tool_calls", "index": 0, "delta": {}}],
"created": 1725932618,
"model": "gpt-4o-2024-08-06",
"object": "chat.completion.chunk",
"system_fingerprint": "fp_b2ffeb16ee",
},
]
chunks = []
for chunk in init_chunks:
chunks.append(litellm.ModelResponseStream(**chunk))
response = stream_chunk_builder(chunks=chunks)
print(f"Returned response: {response}")
completed_response = {
"id": "chatcmpl-A61mXjvcRX0Xr2IiojN9TPiy1P3Fm",
"choices": [
{
"finish_reason": "tool_calls",
"index": 0,
"message": {
"content": None,
"role": "assistant",
"tool_calls": [
{
"function": {
"arguments": '{"base": 3, "exponent": 5}',
"name": "exponentiate",
},
"id": "call_X9P9B6STj7ze8OsJCGkfoN94",
"type": "function",
},
{
"function": {
"arguments": '{"first_int": 12, "second_int": 3}',
"name": "add",
},
"id": "call_Qq8yDeRx7v276abRcLrYORdW",
"type": "function",
},
],
"function_call": None,
},
}
],
"created": 1726000181,
"model": "gpt-4o-2024-05-13",
"object": "chat.completion",
"system_fingerprint": "fp_25624ae3a5",
"usage": {"completion_tokens": 55, "prompt_tokens": 127, "total_tokens": 182},
"service_tier": None,
}
expected_response = litellm.ModelResponse(**completed_response)
print(f"\n\nexpected_response:\n{expected_response}\n\n")
assert (
expected_response.choices == response.choices
), "\nGot={}\n, Expected={}\n".format(response.choices, expected_response.choices)
def test_stream_chunk_builder_openai_prompt_caching():
from pydantic import BaseModel
client = OpenAI(
# This is the default and can be omitted
api_key=os.getenv("OPENAI_API_KEY"),
)
chat_completion = client.chat.completions.create(
messages=[
{
"role": "user",
"content": "Say this is a test",
}
],
model="gpt-3.5-turbo",
stream=True,
stream_options={"include_usage": True},
)
chunks: List[litellm.ModelResponse] = []
usage_obj = None
for chunk in chat_completion:
chunks.append(litellm.ModelResponseStream(**chunk.model_dump()))
print(f"chunks: {chunks}")
usage_obj: litellm.Usage = chunks[-1].usage # type: ignore
response = stream_chunk_builder(chunks=chunks)
print(f"response: {response}")
print(f"response usage: {response.usage}")
for k, v in usage_obj.model_dump(exclude_none=True).items():
print(k, v)
response_usage_value = getattr(response.usage, k) # type: ignore
print(f"response_usage_value: {response_usage_value}")
print(f"type: {type(response_usage_value)}")
if isinstance(response_usage_value, BaseModel):
assert response_usage_value.model_dump(exclude_none=True) == v
else:
assert response_usage_value == v
@pytest.mark.flaky(retries=5, delay=2)
def test_stream_chunk_builder_openai_audio_output_usage():
from pydantic import BaseModel
from typing import Optional
client = OpenAI(
# This is the default and can be omitted
api_key=os.getenv("OPENAI_API_KEY"),
)
try:
completion = client.chat.completions.create(
model="gpt-audio-1.5",
modalities=["text", "audio"],
audio={"voice": "alloy", "format": "pcm16"},
messages=[{"role": "user", "content": "response in 1 word - yes or no"}],
stream=True,
stream_options={"include_usage": True},
)
except Exception as e:
err = str(e).lower()
if (
"model_not_found" in err
or "does not exist" in err
or "openai-internal" in err
):
pytest.skip(f"Skipping - upstream gpt-audio-1.5 unavailable: {e}")
raise
chunks = []
for chunk in completion:
chunks.append(litellm.ModelResponseStream(**chunk.model_dump()))
usage_obj: Optional[litellm.Usage] = None
for index, chunk in enumerate(chunks):
if hasattr(chunk, "usage") and chunk.usage is not None:
usage_obj = chunk.usage
print(f"chunk usage: {chunk.usage}")
print(f"index: {index}")
print(f"len chunks: {len(chunks)}")
print(f"usage_obj: {usage_obj}")
if usage_obj is None:
pytest.skip("OpenAI did not return usage data in streaming response")
response = stream_chunk_builder(chunks=chunks)
print(f"response usage: {response.usage}")
check_non_streaming_response(response)
print(f"response: {response}")
# Convert both usage objects to dictionaries for easier comparison
usage_dict = usage_obj.model_dump(exclude_none=True)
response_usage_dict = response.usage.model_dump(exclude_none=True)
# Simple dictionary comparison
assert (
usage_dict == response_usage_dict
), f"\nExpected: {usage_dict}\nGot: {response_usage_dict}"
def test_stream_chunk_builder_empty_initial_chunk():
from litellm.litellm_core_utils.streaming_chunk_builder_utils import (
ChunkProcessor,
)
chunks = [
{"id": ""},
{"id": "1"},
{"id": "1"},
]
id = ChunkProcessor._get_chunk_id(chunks)
assert id == "1"
def test_stream_chunk_builder_tool_calls_list():
from litellm.litellm_core_utils.streaming_chunk_builder_utils import (
ChunkProcessor,
)
from litellm.types.utils import (
ChatCompletionMessageToolCall,
Function,
ModelResponseStream,
Delta,
ChatCompletionDeltaToolCall,
)
chunks = [
ModelResponseStream(
id="chatcmpl-f323f7a5-2da0-4f86-8ed7-c653c5a359d9",
created=1736388417,
model="llama-3.3-70b-versatile",
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role="assistant",
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id="call_9y79",
function=Function(
arguments='{"location": "San Francisco", "unit": "celsius"}',
name="get_current_weather",
),
type="function",
index=0,
)
],
audio=None,
),
logprobs=None,
)
],
stream_options=None,
),
ModelResponseStream(
id="chatcmpl-f323f7a5-2da0-4f86-8ed7-c653c5a359d9",
created=1736388417,
model="llama-3.3-70b-versatile",
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role=None,
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id="call_pfp7",
function=Function(
arguments='{"location": "Tokyo", "unit": "celsius"}',
name="get_current_weather",
),
type="function",
index=1,
)
],
audio=None,
),
logprobs=None,
)
],
stream_options=None,
),
ModelResponseStream(
id="chatcmpl-f323f7a5-2da0-4f86-8ed7-c653c5a359d9",
created=1736388417,
model="llama-3.3-70b-versatile",
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
content="",
role=None,
function_call=None,
tool_calls=[
ChatCompletionDeltaToolCall(
id="call_hyj5",
function=Function(
arguments='{"location": "Paris", "unit": "celsius"}',
name="get_current_weather",
),
type="function",
index=2,
)
],
audio=None,
),
logprobs=None,
)
],
stream_options=None,
),
]
processor = ChunkProcessor(chunks=chunks)
tool_calls = processor.get_combined_tool_content(tool_call_chunks=chunks)
print(f"tool_calls: {tool_calls}")
assert len(tool_calls) == 3
import json
def get_current_weather(location, unit="fahrenheit"):
"""Get the current weather in a given location"""
if "tokyo" in location.lower():
return json.dumps({"location": "Tokyo", "temperature": "10", "unit": "celsius"})
elif "san francisco" in location.lower():
return json.dumps(
{"location": "San Francisco", "temperature": "72", "unit": "fahrenheit"}
)
elif "paris" in location.lower():
return json.dumps({"location": "Paris", "temperature": "22", "unit": "celsius"})
else:
return json.dumps({"location": location, "temperature": "unknown"})
@pytest.fixture(scope="module", autouse=True)
def load_env():
messages = [
{"role": "system", "content": "You are a helpful AI assistant"},
{
"role": "user",
"content": "What's the weather like in San Francisco, Tokyo, and Paris?",
},
]
tools = [
{
"type": "function",
"function": litellm.utils.function_to_dict(get_current_weather),
}
]
OPENAI_GPT4oMINI = {
"messages": messages,
"model": "gpt-4o-mini",
"temperature": 0.0,
"tools": tools,
"stream": True,
}
LLAMA3_3 = {
"messages": messages,
"model": "groq/openai/gpt-oss-120b",
"api_base": "https://api.groq.com/openai/v1",
"temperature": 0.0,
"tools": tools,
"stream": True,
}
return OPENAI_GPT4oMINI, LLAMA3_3
def execute_completion(opts: dict):
partial_streaming_chunks = []
response_gen = litellm.completion(**opts)
for i, part in enumerate(response_gen):
partial_streaming_chunks.append(part)
print("\n\n")
print(f"partial_streaming_chunks: {partial_streaming_chunks}")
print("\n\n")
assembly = litellm.stream_chunk_builder(partial_streaming_chunks)
print(
f"assembly.choices[0].message.tool_calls: {assembly.choices[0].message.tool_calls}"
)
print(assembly.choices[0].message.tool_calls)
for tool_call in assembly.choices[0].message.tool_calls:
json.loads(
tool_call.function.arguments
) # assert valid json - https://github.com/BerriAI/litellm/issues/10034
def test_grok_bug(load_env):
litellm.set_verbose = True
_, LLAMA3_3 = load_env
execute_completion(LLAMA3_3)

File diff suppressed because it is too large Load diff

View file

@ -1,36 +0,0 @@
# #### What this tests ####
# # This tests if setting team_config actually works
# import sys, os
# import traceback
# import pytest
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import litellm
# from litellm.proxy.proxy_server import ProxyConfig
# @pytest.mark.asyncio
# async def test_team_config():
# litellm.default_team_settings = [
# {
# "team_id": "my-special-team",
# "success_callback": ["langfuse"],
# "langfuse_public_key": "os.environ/LANGFUSE_PUB_KEY_2",
# "langfuse_secret": "os.environ/LANGFUSE_PRIVATE_KEY_2",
# }
# ]
# proxyconfig = ProxyConfig()
# team_config = await proxyconfig.load_team_config(team_id="my-special-team")
# assert len(team_config) > 0
# data = {
# "model": "gpt-3.5-turbo",
# "messages": [{"role": "user", "content": "Hey, how's it going?"}],
# }
# team_config.pop("team_id")
# response = litellm.completion(**{**data, **team_config})
# print(f"response: {response}")

File diff suppressed because it is too large Load diff

View file

@ -1,64 +0,0 @@
# import sys, os
# import traceback
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import time
# import litellm
# import openai
# import pytest
# ### Together AI
# import together
# together.api_key = ""
# sample_message = [
# {"role": "user", "content": "Who are you"},
# {"role": "assistant", "content": "I am your helpful assistant."},
# {"role": "user", "content": "Tell me a joke"},
# ]
# def format_prompt_togetherai(messages, prompt_format, stop_words):
# start_token, end_token = prompt_format.split('{prompt}')
# prompt = ''
# for message in messages:
# role = message['role']
# message_content = message['content']
# if role == 'system':
# prompt += f"{start_token}\n<<SYS>>\n{message_content}\n<</SYS>>\n"
# elif role == 'user':
# prompt += f"{start_token}{message_content}{end_token}"
# else:
# prompt += f'{message_content}{stop_words[0]}'
# return prompt
# model = 'togethercomputer/CodeLlama-13b-Instruct'
# stop_words = list(together.Models.info(model)['config']['stop'])
# prompt_format = str(together.Models.info(model)['config']['prompt_format'])
# formatted_prompt = format_prompt_togetherai(
# messages=sample_message, prompt_format=prompt_format, stop_words=stop_words)
# for token in together.Complete.create_streaming(prompt=formatted_prompt,
# model=model, stop=stop_words, max_tokens=512):
# print(token, end="")
# ### litellm
# import os
# from litellm import completion
# os.environ["TOGETHERAI_API_KEY"] = ""
# sample_message = [
# {"role": "user", "content": "Who are you"},
# {"role": "assistant", "content": "I am your helpful assistant."},
# {"role": "user", "content": "Tell me a joke"},
# ]
# res = completion(model="together_ai/togethercomputer/CodeLlama-13b-Instruct",
# messages=sample_message, stream=False, max_tokens=1000)
# print(list(res))

View file

@ -1,10 +0,0 @@
#### What this tests ####
# This tests the validate environment function
import sys, os
import traceback
import time
import litellm
print(litellm.validate_environment("openai/gpt-3.5-turbo"))

View file

@ -9,7 +9,9 @@ from datetime import timedelta
from typing import Final
from unittest.mock import AsyncMock, MagicMock, call, patch
import httpx
import pytest
import respx
import litellm
from litellm import acompletion, aembedding, completion, embedding
@ -1027,6 +1029,47 @@ def test_completion_past_max_messages_is_neither_served_from_nor_written_to_the_
assert answer(five, "five second") == "five second", "a 5-message repeat was served from the cache"
@respx.mock
def test_enable_cache_serves_a_repeated_bedrock_titan_embedding_from_the_local_cache(
monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], fake_provider_credentials: None
) -> None:
monkeypatch.setattr(litellm, "cache", None)
monkeypatch.setattr(litellm, "set_verbose", True)
monkeypatch.setattr(litellm, "input_callback", [])
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
route: Final = respx.post(
"https://bedrock-runtime.us-west-2.amazonaws.com/model/amazon.titan-embed-text-v2:0/invoke"
).mock(return_value=httpx.Response(200, json={"embedding": [0.25, -0.5, 0.75], "inputTextTokenCount": 6}))
request: Final = {
"model": "bedrock/amazon.titan-embed-text-v2:0",
"input": "good morning from litellm, attempting to embed data",
"aws_region_name": "us-west-2",
}
litellm.enable_cache()
enabled_cache: Final = litellm.cache
monkeypatch.setattr(litellm, "input_callback", [])
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
litellm.enable_cache()
first: Final = embedding(**request)
second: Final = embedding(**request)
assert isinstance(enabled_cache, Cache)
assert enabled_cache.type == LiteLLMCacheType.LOCAL
assert litellm.cache is enabled_cache
assert litellm.input_callback == ["cache"]
assert litellm.success_callback == ["cache"]
assert litellm._async_success_callback == ["cache"]
assert "LiteLLM: Enabling Cache" in capsys.readouterr().out
assert route.call_count == 1
assert first.data[0]["embedding"] == [0.25, -0.5, 0.75]
assert second.data[0]["embedding"] == [0.25, -0.5, 0.75]
assert second._hidden_params["cache_hit"] is True
def test_basic_caching_import() -> None:
from litellm.caching import Cache as PublicCache

View file

@ -6,7 +6,9 @@ import re
from typing import Final, List
from unittest.mock import MagicMock, patch
import httpx
import pytest
import respx
import litellm
from litellm.litellm_core_utils.prompt_templates.factory import (
@ -22,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
anthropic_messages_pt,
convert_to_anthropic_tool_result,
convert_to_gemini_tool_call_result,
convert_to_ollama_image,
encode_tool_call_id_with_signature,
function_call_prompt,
get_thought_signature_from_tool,
@ -285,6 +288,23 @@ def test_ollama_pt_rejects_a_malformed_content_part_as_a_bad_request(part: objec
assert expected_detail in excinfo.value.message
def test_convert_to_ollama_image_downloads_an_http_image_url_to_bare_base64(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "user_url_validation", False)
png_bytes: Final = base64.b64decode(
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
)
image_url: Final = f"https://images.example.com/{uuid.uuid4()}.png"
with respx.mock:
route: Final = respx.get(image_url).mock(
return_value=httpx.Response(200, content=png_bytes, headers={"content-type": "image/png"})
)
result: Final = convert_to_ollama_image(image_url)
assert route.call_count == 1
assert result == base64.b64encode(png_bytes).decode("utf-8")
@pytest.mark.asyncio
async def test_anthropic_bedrock_thinking_blocks_with_none_content():
"""

View file

@ -3972,6 +3972,61 @@ async def test_async_success_handler_prevents_reprocessing_for_pass_through_endp
)
class _StreamSuccessRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.logged_contents: list[str | None] = []
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.logged_contents.append(response_obj.choices[0].message.content)
@pytest.mark.asyncio
async def test_async_success_handler_logs_a_completed_stream_only_once():
recorder: Final = _StreamSuccessRecorder()
start: Final = dt_object(2026, 1, 1)
logging_obj: Final = Logging(
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="acompletion",
start_time=start,
litellm_call_id="stream-once",
function_id="stream-once",
dynamic_async_success_callbacks=[recorder],
)
logging_obj.update_environment_variables(
model="gpt-4o",
user=None,
optional_params={},
litellm_params={"litellm_call_id": "stream-once", "metadata": {}},
custom_llm_provider="openai",
)
await logging_obj.async_success_handler(
result=ModelResponse(
model="gpt-4o",
choices=[{"index": 0, "message": {"role": "assistant", "content": "Hello there"}, "finish_reason": "stop"}],
usage={"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3},
),
start_time=start,
end_time=start,
)
await logging_obj.async_success_handler(
result=ModelResponse(
model="gpt-4o",
choices=[{"index": 0, "message": {"role": "assistant", "content": "Hello again"}, "finish_reason": "stop"}],
),
start_time=start,
end_time=start,
)
assert recorder.logged_contents == ["Hello there"]
assert logging_obj.model_call_details["async_complete_streaming_response"].choices[0].message.content == (
"Hello there"
)
@pytest.mark.asyncio
async def test_async_success_handler_sets_standard_logging_object_for_streaming_pass_through():
"""

View file

@ -1,10 +1,15 @@
import base64
import json
from collections.abc import Mapping, Sequence
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
import respx
import litellm
from litellm import ChatCompletionUsageBlock, stream_chunk_builder
from litellm.types.utils import GenericStreamingChunk
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
@ -1851,3 +1856,285 @@ def test_stream_chunk_builder_omits_service_tier_when_no_chunk_carried_one():
assert response is not None
assert "service_tier" not in response.model_dump()
def _legacy_stream_chunk(
model: str,
*,
content: str | None = None,
tool_calls: tuple[ChatCompletionDeltaToolCall, ...] = (),
usage: Usage | None = None,
finish_reason: str | None = None,
) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-legacy",
created=1,
model=model,
object="chat.completion.chunk",
choices=[
StreamingChoices(
finish_reason=finish_reason,
index=0,
delta=Delta(content=content, role="assistant", tool_calls=list(tool_calls)),
)
],
usage=usage,
)
def test_stream_chunk_builder_empty_initial_chunk() -> None:
chunks: Final = ({"id": ""}, {"id": "1"}, {"id": "1"})
assert ChunkProcessor._get_chunk_id(chunks) == "1"
def test_stream_chunk_builder_litellm_empty_chunks() -> None:
with pytest.raises(litellm.APIError):
stream_chunk_builder(chunks=None)
assert stream_chunk_builder(chunks=[]) is None
def test_stream_chunk_builder_litellm_mixed_calls() -> None:
chunks: Final = (
_legacy_stream_chunk(
"groq/openai/gpt-oss-120b",
content="To answer your question, ",
tool_calls=(
ChatCompletionDeltaToolCall(
id="call_sql",
index=0,
type="function",
function=Function(
name="sql_query",
arguments='{"query": "SELECT COUNT(*) FROM users;"}',
),
),
),
),
_legacy_stream_chunk(
"groq/openai/gpt-oss-120b",
content="I will run a query.",
finish_reason="stop",
),
)
response: Final = stream_chunk_builder(chunks=chunks)
assert response is not None
assert response.choices[0].message.content == "To answer your question, I will run a query."
assert [tool_call.to_dict() for tool_call in response.choices[0].message.tool_calls] == [
{
"function": {
"arguments": '{"query": "SELECT COUNT(*) FROM users;"}',
"name": "sql_query",
},
"id": "call_sql",
"type": "function",
}
]
def _openai_sse(*chunks: Mapping[str, object]) -> str:
events: Final = tuple(
f"data: {json.dumps({'id': 'chatcmpl-legacy', 'object': 'chat.completion.chunk', 'created': 1, 'model': 'gpt-4.1-mini', **chunk})}\n\n"
for chunk in chunks
)
return "".join(events) + "data: [DONE]\n\n"
def _complete_streamed_openai_response(respx_mock: respx.MockRouter, sse_body: str) -> litellm.ModelResponse:
respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
return_value=httpx.Response(200, text=sse_body, headers={"content-type": "text/event-stream"})
)
return litellm.completion(
model="openai/gpt-4.1-mini",
messages=[{"role": "user", "content": "What's the weather like in Boston?"}],
api_key="sk-test",
stream=True,
stream_options={"include_usage": True},
complete_response=True,
)
def test_stream_chunk_builder_litellm_tool_call(respx_mock: respx.MockRouter) -> None:
response: Final = _complete_streamed_openai_response(
respx_mock,
_openai_sse(
{
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"index": 0,
"id": "call_weather",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city":'},
}
],
},
"finish_reason": None,
}
]
},
{
"choices": [
{
"index": 0,
"delta": {"tool_calls": [{"index": 0, "function": {"arguments": '"Boston"}'}}]},
"finish_reason": "tool_calls",
}
]
},
{"choices": [], "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}},
),
)
assert [tool_call.to_dict() for tool_call in response.choices[0].message.tool_calls] == [
{"function": {"arguments": '{"city":"Boston"}', "name": "get_weather"}, "id": "call_weather", "type": "function"}
]
assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (5, 3, 8)
def test_stream_chunk_builder_litellm_tool_call_regular_message(respx_mock: respx.MockRouter) -> None:
response: Final = _complete_streamed_openai_response(
respx_mock,
_openai_sse(
{"choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello "}, "finish_reason": None}]},
{"choices": [{"index": 0, "delta": {"content": "there."}, "finish_reason": "stop"}]},
{"choices": [], "usage": {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6}},
),
)
assert response.choices[0].message.content == "Hello there."
assert response.choices[0].message.tool_calls is None
assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (4, 2, 6)
assert response._hidden_params["custom_llm_provider"] == "openai"
def test_stream_chunk_builder_multiple_tool_calls() -> None:
def tool_call_chunk(*tool_calls: ChatCompletionDeltaToolCall, finish_reason: str | None = None) -> ModelResponseStream:
return _legacy_stream_chunk("openai/gpt-4o", tool_calls=tool_calls, finish_reason=finish_reason)
chunks: Final = (
tool_call_chunk(
ChatCompletionDeltaToolCall(
id="call_exp", index=0, type="function", function=Function(name="exponentiate", arguments="")
)
),
tool_call_chunk(ChatCompletionDeltaToolCall(index=0, function=Function(arguments='{"base": 3, '))),
tool_call_chunk(
ChatCompletionDeltaToolCall(id="call_add", index=1, type="function", function=Function(name="add", arguments=""))
),
tool_call_chunk(ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"exponent": 5}'))),
tool_call_chunk(
ChatCompletionDeltaToolCall(index=1, function=Function(arguments='{"x": 1, "y": 2}')),
finish_reason="tool_calls",
),
)
response: Final = stream_chunk_builder(chunks=chunks)
assert response is not None
assert response.choices[0].finish_reason == "tool_calls"
assert [tool_call.to_dict() for tool_call in response.choices[0].message.tool_calls] == [
{"function": {"arguments": '{"base": 3, "exponent": 5}', "name": "exponentiate"}, "id": "call_exp", "type": "function"},
{"function": {"arguments": '{"x": 1, "y": 2}', "name": "add"}, "id": "call_add", "type": "function"},
]
def test_grok_bug() -> None:
chunks: Final = (
_legacy_stream_chunk(
"groq/openai/gpt-oss-120b",
tool_calls=(
ChatCompletionDeltaToolCall(
id="call_weather",
index=0,
type="function",
function=Function(name="get_weather", arguments='{"location":'),
),
),
),
_legacy_stream_chunk(
"groq/openai/gpt-oss-120b",
tool_calls=(
ChatCompletionDeltaToolCall(
index=0,
type="function",
function=Function(arguments='"Boston"}'),
),
),
finish_reason="tool_calls",
),
)
response: Final = stream_chunk_builder(chunks=chunks)
assert response is not None
assert json.loads(response.choices[0].message.tool_calls[0].function.arguments) == {
"location": "Boston"
}
def test_stream_chunk_builder_openai_audio_output_usage() -> None:
usage: Final = Usage(
completion_tokens=4,
prompt_tokens=7,
total_tokens=11,
completion_tokens_details={"audio_tokens": 2, "text_tokens": 2},
prompt_tokens_details=PromptTokensDetails(audio_tokens=3, text_tokens=4),
)
response: Final = stream_chunk_builder(
chunks=[
_legacy_stream_chunk("openai/gpt-audio-1.5", content="Ye"),
_legacy_stream_chunk("openai/gpt-audio-1.5", content="s.", finish_reason="stop"),
_legacy_stream_chunk("openai/gpt-audio-1.5", usage=usage),
]
)
assert response is not None
assert response.choices[0].message.content == "Yes."
assert response.usage.model_dump(exclude_none=True) == usage.model_dump(exclude_none=True)
def _openai_audio_stream_chunk(delta: Mapping[str, object], finish_reason: str | None = None) -> ModelResponseStream:
return ModelResponseStream(
**{
"id": "chatcmpl-audio",
"created": 1,
"model": "gpt-audio-1.5",
"object": "chat.completion.chunk",
"choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason, "logprobs": None}],
}
)
def test_stream_chunk_builder_joins_openai_audio_output_pieces() -> None:
first_pcm: Final = b"\x01\x02"
second_pcm: Final = b"\x03\x04\x05\x06"
chunks: Final = (
_openai_audio_stream_chunk(
{"role": "assistant", "content": None, "audio": {"id": "audio_yes", "transcript": "Ye"}}
),
_openai_audio_stream_chunk({"audio": {"transcript": "s."}}),
_openai_audio_stream_chunk({"audio": {"data": base64.b64encode(first_pcm).decode()}}),
_openai_audio_stream_chunk({"audio": {"data": base64.b64encode(second_pcm).decode()}}),
_openai_audio_stream_chunk({"audio": {"id": "audio_yes", "expires_at": 1767225600}}),
_openai_audio_stream_chunk({}, finish_reason="stop"),
)
response: Final = stream_chunk_builder(chunks=list(chunks))
assert response is not None
audio: Final = response.choices[0].message.audio
assert audio is not None
assert audio.model_dump() == {
"id": "audio_yes",
"data": base64.b64encode(first_pcm + second_pcm).decode(),
"expires_at": 1767225600,
"transcript": "Yes.",
}

View file

@ -1,11 +1,13 @@
import json
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from itertools import chain
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
import litellm
from litellm._uuid import uuid
@ -18,6 +20,7 @@ from litellm.types.llms.openai import (
ChatCompletionToolCallFunctionChunk,
)
from litellm.types.responses.main import OutputCodeInterpreterCall
from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponseStream
@pytest.mark.asyncio
@ -86,6 +89,120 @@ def test_anthropic_completion_does_not_send_deployment_default_limits():
assert "default_api_key_tpm_limit" not in request_body
def test_anthropic_tool_history_without_tools_injects_a_dummy_tool(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "modify_params", False)
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=httpx.Response(
200,
json={
"id": "msg_tool_history",
"type": "message",
"role": "assistant",
"model": "claude-haiku-4-5-20251001",
"content": [{"type": "text", "text": "San Francisco is 72 degrees"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 4, "output_tokens": 6},
},
)
)
response: Final = litellm.completion(
model="anthropic/claude-haiku-4-5-20251001",
messages=[
{"role": "user", "content": "What's the weather like in San Francisco?"},
{
"role": "assistant",
"content": "Here are the current weather conditions",
"tool_calls": [
{
"id": "tooluse_weather",
"type": "function",
"function": {
"name": "get_current_weather",
"arguments": '{"location": "San Francisco, CA", "unit": "fahrenheit"}',
},
}
],
},
{
"role": "tool",
"tool_call_id": "tooluse_weather",
"name": "get_current_weather",
"content": '{"temperature": "72"}',
},
],
api_key="anthropic-test-key",
max_tokens=32,
)
request_body: Final = json.loads(route.calls.last.request.read())
assert request_body["tools"] == [
{
"name": "dummy_tool",
"input_schema": {"type": "object", "properties": {}},
"type": "custom",
"description": "This is a dummy tool call",
}
]
assert request_body["messages"][1:] == [
{
"role": "assistant",
"content": [
{"type": "text", "text": "Here are the current weather conditions"},
{
"type": "tool_use",
"id": "tooluse_weather",
"name": "get_current_weather",
"input": {"location": "San Francisco, CA", "unit": "fahrenheit"},
},
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "tooluse_weather", "content": '{"temperature": "72"}'}],
},
]
assert response.choices[0].message.content == "San Francisco is 72 degrees"
@respx.mock
def test_anthropic_tool_call_error_is_mapped_from_the_provider_response():
route = respx.post("https://api.anthropic.com/v1/messages").mock(
return_value=httpx.Response(
400,
json={
"type": "error",
"error": {"type": "invalid_request_error", "message": "tool input is invalid"},
},
)
)
with pytest.raises(litellm.BadRequestError) as exc_info:
litellm.completion(
model="anthropic/claude-3-5-haiku-20241022",
messages=[{"role": "user", "content": "call the tool"}],
tools=[
{
"type": "function",
"function": {
"name": "lookup",
"description": "look up a value",
"parameters": {"type": "object", "properties": {}},
},
}
],
api_key="anthropic-test-key",
)
assert exc_info.value.status_code == 400
assert exc_info.value.llm_provider == "anthropic"
assert "tool input is invalid" in str(exc_info.value)
assert len(route.calls) == 1
async def test_anthropic_async_completion_inlines_http_images_off_the_event_loop(async_only_image_fetch):
http_image_url = f"http://img.example/{uuid.uuid4()}.png"
https_image_url = f"https://img.example/{uuid.uuid4()}.png"
@ -592,6 +709,172 @@ def test_message_delta_without_usage_returns_chunk_with_no_usage():
assert model_response.usage is None
async def _collect_anthropic_stream(sync_mode: bool, **kwargs: object) -> tuple[ModelResponseStream, ...]:
if sync_mode:
return tuple(litellm.completion(**kwargs))
return tuple([chunk async for chunk in await litellm.acompletion(**kwargs)])
@pytest.mark.parametrize("sync_mode", [True, False])
@respx.mock
async def test_anthropic_stream_options_return_usage_after_the_finish_chunk(
monkeypatch: pytest.MonkeyPatch, sync_mode: bool
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
events: Final = (
{
"type": "message_start",
"message": {
"id": "msg_stream_options",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-6",
"content": [],
"usage": {
"input_tokens": 3,
"output_tokens": 1,
"cache_read_input_tokens": 12,
"cache_creation_input_tokens": 4,
},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": "Blue sky."},
},
{"type": "content_block_stop", "index": 0},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 2},
},
{"type": "message_stop"},
)
body: Final = "".join(
f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events
)
route: Final = respx.post("https://api.anthropic.com/v1/messages").mock(
return_value=httpx.Response(
200, text=body, headers={"content-type": "text/event-stream"}
)
)
chunks: Final = await _collect_anthropic_stream(
sync_mode,
model="anthropic/claude-sonnet-4-6",
api_key="test-key",
messages=[{"role": "user", "content": "Describe the weather"}],
stream=True,
stream_options={"include_usage": True},
)
finish_indices: Final = tuple(
index
for index, chunk in enumerate(chunks)
if chunk.choices and chunk.choices[0].finish_reason is not None
)
finish_reasons: Final = tuple(
chunk.choices[0].finish_reason
for chunk in chunks
if chunk.choices and chunk.choices[0].finish_reason is not None
)
usage_indices: Final = tuple(
index
for index, chunk in enumerate(chunks)
if getattr(chunk, "usage", None) is not None
)
usage_chunk: Final = chunks[usage_indices[0]]
assert "".join(
chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices
) == "Blue sky."
assert finish_reasons == ("stop",)
assert finish_indices == (len(chunks) - 2,)
assert usage_indices == (len(chunks) - 1,)
assert usage_chunk.usage.prompt_tokens == 19
assert usage_chunk.usage.completion_tokens == 2
assert usage_chunk.usage.total_tokens == 21
assert usage_chunk.usage.cache_read_input_tokens == 12
assert usage_chunk.usage.cache_creation_input_tokens == 4
assert json.loads(route.calls[0].request.content)["stream"] is True
_SNOOZE_TOOL: Final = {
"type": "function",
"function": {
"name": "snooze",
"parameters": {"type": "object", "properties": {}, "required": []},
},
}
def _stream_one_anthropic_tool_use(content_block: dict[str, object]) -> tuple[ChatCompletionDeltaToolCall, ...]:
events: Final = (
{
"type": "message_start",
"message": {
"id": "msg_snooze",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-6",
"content": [],
"usage": {"input_tokens": 12, "output_tokens": 1},
},
},
{"type": "content_block_start", "index": 0, "content_block": content_block},
{"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}},
{"type": "content_block_stop", "index": 0},
{
"type": "message_delta",
"delta": {"stop_reason": "tool_use", "stop_sequence": None},
"usage": {"output_tokens": 9},
},
{"type": "message_stop"},
)
body: Final = "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events)
with respx.mock:
respx.post("https://api.anthropic.com/v1/messages").mock(
return_value=httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
)
chunks: Final = tuple(
litellm.completion(
model="anthropic/claude-sonnet-4-6",
api_key="test-key",
messages=[{"role": "user", "content": "Hit the snooze button."}],
tools=[_SNOOZE_TOOL],
stream=True,
)
)
return tuple(
chunk.choices[0].delta.tool_calls[0] for chunk in chunks if chunk.choices and chunk.choices[0].delta.tool_calls
)
def test_anthropic_stream_no_argument_tool_call_streams_an_empty_json_object():
tool_calls: Final = _stream_one_anthropic_tool_use(
{"type": "tool_use", "id": "toolu_snooze", "name": "snooze", "input": {}}
)
assert tuple(tool_call.id for tool_call in tool_calls if tool_call.id is not None) == ("toolu_snooze",)
assert tool_calls[0].function.name == "snooze"
assert json.loads("".join(tool_call.function.arguments or "" for tool_call in tool_calls)) == {}
def test_anthropic_stream_tool_call_keeps_the_programmatic_caller():
caller: Final = {"type": "code_execution_20250825", "tool_id": "srvtoolu_snooze"}
tool_calls: Final = _stream_one_anthropic_tool_use(
{"type": "tool_use", "id": "toolu_snooze", "name": "snooze", "input": {}, "caller": caller}
)
assert tool_calls[0].id == "toolu_snooze"
assert tool_calls[0].model_dump()["caller"] == caller
def test_streaming_thinking_deltas_count_reasoning_tokens_in_usage():
"""Anthropic streaming usage should account for emitted thinking deltas."""
chunks = [
@ -2367,3 +2650,135 @@ def test_served_model_reaches_assembled_stream_through_custom_stream_wrapper():
assert chunk._hidden_params["provider_response_model"] == served_model
assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=[{"role": "user", "content": "hi"}])
assert assembled._hidden_params["provider_response_model"] == served_model
@respx.mock
def test_anthropic_streaming_function_call_arguments_are_assembled(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
events: Final = (
{
"type": "message_start",
"message": {
"id": "msg_tool_stream",
"type": "message",
"role": "assistant",
"model": "claude-haiku-4-5-20251001",
"content": [],
"usage": {"input_tokens": 5, "output_tokens": 1},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "tool_use",
"id": "toolu_weather",
"name": "get_current_weather",
"input": {},
},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '{"location":"Bos'},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": 'ton","unit":"fahrenheit"}'},
},
{"type": "content_block_stop", "index": 0},
{
"type": "message_delta",
"delta": {"stop_reason": "tool_use", "stop_sequence": None},
"usage": {"output_tokens": 7},
},
{"type": "message_stop"},
)
body: Final = "".join(
f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events
)
respx.post("https://api.anthropic.com/v1/messages").mock(
return_value=httpx.Response(
200,
text=body,
headers={"content-type": "text/event-stream"},
)
)
messages: Final = [{"role": "user", "content": "Check the weather in Boston"}]
stream: Final = litellm.completion(
model="anthropic/claude-haiku-4-5-20251001",
messages=messages,
tools=[
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the weather",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"},
"unit": {"type": "string"},
},
"required": ["location", "unit"],
},
},
}
],
stream=True,
api_key="test-anthropic-key",
)
chunks: Final = tuple(stream)
tool_call_deltas: Final = tuple(
chain.from_iterable(
chunk.choices[0].delta.tool_calls or () for chunk in chunks if chunk.choices
)
)
assert tuple(
tool_call.function.name
for tool_call in tool_call_deltas
if tool_call.function.name is not None
) == ("get_current_weather",)
arguments: Final = "".join(
tool_call.function.arguments or "" for tool_call in tool_call_deltas
)
assert json.loads(arguments) == {"location": "Boston", "unit": "fahrenheit"}
assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=messages)
assert assembled.choices[0].message.tool_calls[0].function.name == "get_current_weather"
assert json.loads(assembled.choices[0].message.tool_calls[0].function.arguments) == {
"location": "Boston",
"unit": "fahrenheit",
}
assert assembled.choices[0].finish_reason == "tool_calls"
@respx.mock
def test_anthropic_streaming_invalid_key_maps_to_authentication_error() -> None:
route: Final = respx.post("https://api.anthropic.com/v1/messages").mock(
return_value=httpx.Response(
401,
json={
"type": "error",
"error": {
"type": "authentication_error",
"message": "invalid API key",
},
},
)
)
with pytest.raises(litellm.AuthenticationError) as exc_info:
litellm.completion(
model="anthropic/claude-haiku-4-5-20251001",
messages=[{"role": "user", "content": "Say hello"}],
stream=True,
api_key="bad-anthropic-key",
max_retries=0,
)
assert exc_info.value.status_code == 401
assert exc_info.value.llm_provider == "anthropic"
assert route.calls.last.request.headers["x-api-key"] == "bad-anthropic-key"

View file

@ -1,12 +1,16 @@
"""Tests for litellm/llms/azure/azure.py AzureChatCompletion handler behaviour."""
import asyncio
import json
import os
import time
from collections.abc import Mapping
from typing import Final
from unittest.mock import MagicMock, patch
import httpx
import pytest
import respx
from httpx import Client, Headers
from openai import AsyncAzureOpenAI, AzureOpenAI
@ -658,3 +662,134 @@ def test_azure_with_content_safety_error():
]
== "high"
)
def _azure_sse_body(chunks: tuple[Mapping[str, object], ...]) -> str:
return "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
@respx.mock
def test_azure_chat_stream_preserves_text_and_special_characters():
response_chunks: Final = (
{
"id": "chatcmpl-azure",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4.1-mini",
"choices": [{"index": 0, "delta": {"content": "<xml>café ☃</xml>"}, "finish_reason": None}],
},
{
"id": "chatcmpl-azure",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4.1-mini",
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
},
)
route: Final = respx.post(
url__regex=r"https://azure-stream\.openai\.azure\.com/.*/chat/completions.*"
).mock(
return_value=httpx.Response(
200,
text=_azure_sse_body(response_chunks),
headers={"content-type": "text/event-stream"},
)
)
stream: Final = litellm.completion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "Respond with the <xml> tag only"}],
api_base="https://azure-stream.openai.azure.com",
api_version="2024-02-15-preview",
api_key="test-azure-key",
max_retries=0,
stream=True,
)
chunks: Final = list(stream)
assert route.call_count == 1
assert [chunk._hidden_params["custom_llm_provider"] for chunk in chunks] == ["azure"] * len(chunks)
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "<xml>café ☃</xml>"
assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"]
@respx.mock
def test_azure_chat_stream_reassembles_tool_call_arguments():
response_chunks: Final = (
{
"id": "chatcmpl-azure-tools",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4.1-mini",
"choices": [
{
"index": 0,
"delta": {
"tool_calls": [
{
"index": 0,
"id": "call-weather",
"type": "function",
"function": {"name": "get_current_weather", "arguments": '{"city":'},
}
]
},
"finish_reason": None,
}
],
},
{
"id": "chatcmpl-azure-tools",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4.1-mini",
"choices": [
{
"index": 0,
"delta": {"tool_calls": [{"index": 0, "function": {"arguments": '"Boston"}'}}]},
"finish_reason": "tool_calls",
}
],
},
)
route: Final = respx.post(
url__regex=r"https://azure-tool-stream\.openai\.azure\.com/.*/chat/completions.*"
).mock(
return_value=httpx.Response(
200,
text=_azure_sse_body(response_chunks),
headers={"content-type": "text/event-stream"},
)
)
messages: Final = [{"role": "user", "content": "What is the weather in Boston?"}]
stream: Final = litellm.completion(
model="azure/gpt-4.1-mini",
messages=messages,
tools=[
{
"type": "function",
"function": {
"name": "get_current_weather",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"],
},
},
}
],
tool_choice="auto",
api_base="https://azure-tool-stream.openai.azure.com",
api_version="2024-02-15-preview",
api_key="test-azure-key",
max_retries=0,
stream=True,
)
chunks: Final = list(stream)
response: Final = litellm.stream_chunk_builder(chunks, messages=messages)
tool_call: Final = response.choices[0].message.tool_calls[0]
assert route.call_count == 1
assert tool_call.function.name == "get_current_weather"
assert json.loads(tool_call.function.arguments) == {"city": "Boston"}
assert response.choices[0].finish_reason == "tool_calls"

View file

@ -1,8 +1,13 @@
import json
import os
import sys
from unittest.mock import MagicMock, patch
import httpx
import pytest
import respx
import litellm
sys.path.insert(
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
@ -92,3 +97,89 @@ class TestAzureV1AsyncEmbedding:
)
assert isinstance(client, AsyncOpenAI)
@respx.mock
def test_azure_embedding_forwards_optional_args_in_the_request_body():
route = respx.post(url__regex=r"https://azure-optional-args\.openai\.azure\.com/.*/embeddings.*").mock(
return_value=httpx.Response(
200,
json={
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}],
"model": "text-embedding-ada-002",
"usage": {"prompt_tokens": 1, "total_tokens": 1},
},
)
)
response = litellm.embedding(
model="azure/text-embedding-ada-002",
input=["hello"],
api_base="https://azure-optional-args.openai.azure.com",
api_version="2023-05-15",
azure_ad_token="test-token",
encoding_format="float",
)
request = route.calls.last.request
assert json.loads(request.read()) == {
"model": "text-embedding-ada-002",
"input": ["hello"],
"encoding_format": "float",
"azure_ad_token": "test-token",
}
assert request.headers["authorization"] == "Bearer test-token"
assert response.data[0]["embedding"] == [0.1, 0.2, 0.3]
@pytest.mark.asyncio
async def test_aembedding_azure_returns_an_azure_embedding_response(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route = respx_mock.post(
url__regex=r"https://azure-async-embedding\.openai\.azure\.com/.*/embeddings.*"
).mock(
return_value=httpx.Response(
200,
json={
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.4, 0.5, 0.6]}],
"model": "text-embedding-ada-002",
"usage": {"prompt_tokens": 2, "total_tokens": 2},
},
)
)
response = await litellm.aembedding(
model="azure/text-embedding-ada-002",
input=["good morning", "another item"],
api_base="https://azure-async-embedding.openai.azure.com",
api_version="2023-05-15",
api_key="test-key",
)
assert response._hidden_params["custom_llm_provider"] == "azure"
assert (response.usage.prompt_tokens, response.usage.total_tokens) == (2, 2)
assert response.data[0]["embedding"] == [0.4, 0.5, 0.6]
assert len(route.calls) == 1
def test_azure_embedding_maps_a_scripted_request_timeout(respx_mock: respx.MockRouter) -> None:
route = respx_mock.post(
url__regex=r"https://azure-timeout\.openai\.azure\.com/.*/embeddings.*"
).mock(side_effect=httpx.ReadTimeout("scripted timeout"))
with pytest.raises(litellm.Timeout):
litellm.embedding(
model="azure/text-embedding-ada-002",
input=["hello"],
api_base="https://azure-timeout.openai.azure.com",
api_version="2023-05-15",
api_key="test-key",
timeout=0.5,
max_retries=0,
)
assert route.call_count == 1

View file

@ -2,6 +2,8 @@ import json
from unittest.mock import MagicMock, patch
import pytest
import httpx
import respx
import litellm
@ -9,6 +11,72 @@ from litellm.exceptions import ContentPolicyViolationError
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
@respx.mock
def test_azure_401_maps_to_authentication_error():
route = respx.post(
url__regex=r"https://azure-bad-key\.openai\.azure\.com/.*/chat/completions.*"
).mock(return_value=httpx.Response(401, json={"error": {"message": "invalid Azure key"}}))
with pytest.raises(litellm.AuthenticationError) as exc_info:
litellm.completion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "hello"}],
api_base="https://azure-bad-key.openai.azure.com",
api_version="2024-10-21",
api_key="bad-azure-key",
max_retries=0,
)
assert exc_info.value.status_code == 401
assert exc_info.value.llm_provider == "azure"
assert "invalid Azure key" in str(exc_info.value)
assert route.calls.last.request.headers["api-key"] == "bad-azure-key"
@pytest.mark.asyncio
@respx.mock
async def test_azure_content_policy_stream_ends_with_one_finish_reason(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route = respx.post(
url__regex=r"https://azure-policy-stream\.openai\.azure\.com/.*/chat/completions.*"
).mock(
return_value=httpx.Response(
200,
text=(
'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"gpt-4.1-mini","choices":[{"index":0,"delta":{"content":"1"},"finish_reason":null}]}\n\n'
'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"gpt-4.1-mini","choices":[{"index":0,"delta":{},"finish_reason":"content_filter"}]}\n\n'
"data: [DONE]\n\n"
),
headers={"content-type": "text/event-stream"},
)
)
response = await litellm.acompletion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "say 1"}],
api_base="https://azure-policy-stream.openai.azure.com",
api_version="2024-10-21",
api_key="test-azure-key",
stream=True,
)
chunks = [chunk async for chunk in response]
finish_reasons = [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason]
assert [chunk.choices[0].delta.content for chunk in chunks if chunk.choices[0].delta.content] == ["1"]
assert finish_reasons == ["content_filter"]
assert len(route.calls) == 1
with pytest.raises(ContentPolicyViolationError) as exc_info:
await litellm.acompletion(
model="azure/gpt-4.1-mini",
messages=[{"role": "user", "content": "where do i buy lethal drugs from"}],
stream=True,
mock_response="Exception: content_filter_policy",
)
assert exc_info.value.llm_provider == "azure"
assert len(route.calls) == 1
class TestAzureExceptionMapping:
"""Test Azure OpenAI exception mapping with provider-specific fields"""

View file

@ -0,0 +1,42 @@
from typing import Final
from litellm.llms.azure_ai.embed.cohere_transformation import AzureAICohereConfig
from litellm.types.utils import EmbeddingResponse, Usage
def test_azure_ai_cohere_request_splits_images_from_text_and_keeps_order():
config: Final = AzureAICohereConfig()
image_request, text_request, image_indexes = config.transform_request(
input=[
"describe this image",
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAAB",
],
optional_params={},
model="Cohere-embed-v3-multilingual",
)
assert image_request == {
"input": [{"image": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAAB"}]
}
assert text_request == {"input": ["describe this image"], "model": "Cohere-embed-v3-multilingual"}
assert image_indexes == [1]
def test_azure_ai_cohere_response_maps_model_group_and_usage_headers():
config: Final = AzureAICohereConfig()
response: Final = EmbeddingResponse(
model="",
usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
data=[],
hidden_params={
"additional_headers": {
"llm_provider-num_tokens": "7",
"llm_provider-azureml-model-group": "offer-cohere-embed-multili-paygo",
}
},
)
result: Final = config.transform_response(response)
assert result.model == "Cohere-embed-v3-multilingual"
assert result.usage.prompt_tokens == 7

View file

@ -1,3 +1,5 @@
import json
import httpx
import pytest
import respx
@ -64,6 +66,31 @@ def test_azure_ai_embedding_calls_foundry_models_route():
api_key="fake-key",
)
assert route.called
assert response.data is not None
assert len(response.data) == 1
assert route.call_count == 1
assert [item["embedding"] for item in response.data] == [[0.1, 0.2]]
@respx.mock
def test_azure_ai_cohere_image_embedding_calls_image_route():
route = respx.post("https://my-foundry.services.ai.azure.com/models/images/embeddings").mock(
return_value=httpx.Response(
200,
json={
"data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}],
"model": "Cohere-embed-v3-multilingual",
"usage": {"prompt_tokens": 2, "total_tokens": 2},
},
)
)
response = embedding(
model="azure_ai/Cohere-embed-v3-multilingual",
input=["data:image/png;base64,aGVsbG8="],
api_base="https://my-foundry.services.ai.azure.com",
api_key="fake-key",
)
assert json.loads(route.calls.last.request.read()) == {"input": [{"image": "data:image/png;base64,aGVsbG8="}]}
assert set(dict(response).keys()) == {"object", "data", "model", "usage"}
assert response.data[0]["embedding"] == [0.1, 0.2]
assert (response.usage.prompt_tokens, response.usage.total_tokens) == (2, 2)

View file

@ -1052,6 +1052,180 @@ async def test_async_converse_stream_with_an_empty_200_body_raises_instead_of_an
_assert_empty_stream_surfaced_as_bad_gateway(exc_info.value)
def _assert_converse_stream_text_and_finish_reason(chunks: Sequence[ModelResponseStream]) -> None:
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "Streamed reply"
assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"]
@pytest.mark.parametrize("async_mode", [False, True])
@pytest.mark.asyncio
async def test_converse_streaming_completion_yields_text_and_finish_reason(
_aws_test_credentials: None,
async_mode: bool,
) -> None:
event_stream_body: Final = b"".join(
(
_converse_event_frame("messageStart", {"role": "assistant"}),
_converse_event_frame(
"contentBlockDelta",
{"contentBlockIndex": 0, "delta": {"text": "Streamed "}},
),
_converse_event_frame(
"contentBlockDelta",
{"contentBlockIndex": 0, "delta": {"text": "reply"}},
),
_converse_event_frame("contentBlockStop", {"contentBlockIndex": 0}),
_converse_event_frame("messageStop", {"stopReason": "end_turn"}),
_converse_event_frame(
"metadata",
{
"usage": {"inputTokens": 3, "outputTokens": 2, "totalTokens": 5},
"metrics": {"latencyMs": 1},
},
),
)
)
def handle_request(request: httpx.Request) -> httpx.Response:
assert request.url.path.endswith("/converse-stream")
return httpx.Response(200, content=event_stream_body, headers=_event_stream_headers())
transport: Final = httpx.MockTransport(handle_request)
if async_mode:
async_client: Final = AsyncHTTPHandler(transport=transport)
async_stream: Final = await litellm.acompletion(
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[{"role": "user", "content": "Say hello"}],
stream=True,
client=async_client,
)
async_chunks: Final = tuple([chunk async for chunk in async_stream])
_assert_converse_stream_text_and_finish_reason(async_chunks)
else:
sync_client: Final = HTTPHandler(client=httpx.Client(transport=transport))
sync_stream: Final = litellm.completion(
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[{"role": "user", "content": "Say hello"}],
stream=True,
client=sync_client,
)
sync_chunks: Final = tuple(sync_stream)
_assert_converse_stream_text_and_finish_reason(sync_chunks)
def test_mistral_invoke_stream_yields_text_and_finish_reason(
_aws_test_credentials: None,
) -> None:
event_stream_body: Final = b"".join(
_bedrock_event_stream_frame(chunk)
for chunk in (
{"outputs": [{"text": "A court case moves ", "stop_reason": None}]},
{"outputs": [{"text": "through the courts.", "stop_reason": "stop"}]},
)
)
def handle_request(request: httpx.Request) -> httpx.Response:
assert request.url.path.endswith("/invoke-with-response-stream")
return httpx.Response(
200,
content=event_stream_body,
headers=_event_stream_headers(),
)
transport: Final = httpx.MockTransport(handle_request)
client: Final = HTTPHandler(client=httpx.Client(transport=transport))
stream: Final = litellm.completion(
model="bedrock/mistral.mixtral-8x7b-instruct-v0:1",
messages=[{"role": "user", "content": "Where does a court case go?"}],
stream=True,
client=client,
)
chunks: Final = tuple(stream)
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == (
"A court case moves through the courts."
)
assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == [
"stop"
]
def test_converse_streaming_tool_call_arguments_are_valid_json(
_aws_test_credentials: None,
) -> None:
event_stream_body: Final = b"".join(
_converse_event_frame(event_type, body)
for event_type, body in (
("messageStart", {"role": "assistant"}),
(
"contentBlockStart",
{
"contentBlockIndex": 0,
"start": {
"toolUse": {
"toolUseId": "toolu_weather",
"name": "get_current_weather",
}
},
},
),
(
"contentBlockDelta",
{"contentBlockIndex": 0, "delta": {"toolUse": {"input": '{"location":"Bos'}}},
),
(
"contentBlockDelta",
{"contentBlockIndex": 0, "delta": {"toolUse": {"input": 'ton","unit":"fahrenheit"}'}}},
),
("contentBlockStop", {"contentBlockIndex": 0}),
("messageStop", {"stopReason": "tool_use"}),
("metadata", {"usage": {"inputTokens": 4, "outputTokens": 7}}),
)
)
def handle_request(request: httpx.Request) -> httpx.Response:
assert request.url.path.endswith("/converse-stream")
return httpx.Response(
200,
content=event_stream_body,
headers=_event_stream_headers(),
)
transport: Final = httpx.MockTransport(handle_request)
client: Final = HTTPHandler(client=httpx.Client(transport=transport))
messages: Final = [{"role": "user", "content": "What is the weather in Boston?"}]
stream: Final = litellm.completion(
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=messages,
tools=[
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the weather",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"},
"unit": {"type": "string"},
},
"required": ["location", "unit"],
},
},
}
],
stream=True,
client=client,
)
chunks: Final = tuple(stream)
assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=messages)
assert assembled.choices[0].message.tool_calls[0].function.name == "get_current_weather"
assert json.loads(assembled.choices[0].message.tool_calls[0].function.arguments) == {
"location": "Boston",
"unit": "fahrenheit",
}
assert assembled.choices[0].finish_reason == "tool_calls"
_UPSTREAM_REJECTION: Final = "structured output schema uses unsupported regex negative look-ahead"
_CUSTOMER_REJECTION_EVENT_TYPE: Final = "validationException"

View file

@ -6,8 +6,10 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
import httpx
import respx
import litellm
from litellm.caching.caching import Cache
from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.bedrock.embed.embedding import BedrockEmbedding
@ -127,6 +129,65 @@ def test_bedrock_embedding_with_env_variable_bearer_token(
assert headers["Authorization"] == f"Bearer {test_api_key}"
@pytest.mark.parametrize(
"model",
["amazon.titan-embed-text-v1", "amazon.titan-embed-image-v1", "amazon.titan-embed-text-v2:0"],
)
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_bedrock_titan_embedding_cache_hit_avoids_a_second_request(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, model: str, sync_mode: bool
) -> None:
monkeypatch.setattr(litellm, "cache", Cache())
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route = respx_mock.post(f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{model}/invoke").mock(
return_value=httpx.Response(200, json=titan_embedding_response)
)
request = {
"model": f"bedrock/{model}",
"input": test_input,
"api_key": "test-bearer-token",
"aws_region_name": "us-east-1",
"aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-east-1.amazonaws.com",
"caching": True,
}
if sync_mode:
response_1 = litellm.embedding(**request)
response_2 = litellm.embedding(**request)
else:
response_1 = await litellm.aembedding(**request)
response_2 = await litellm.aembedding(**request)
assert json.loads(route.calls.last.request.read()) == {"inputText": test_input}
assert response_1.data[0]["embedding"] == [0.1, 0.2, 0.3]
assert response_1.usage.prompt_tokens == 10
assert response_2._hidden_params["cache_hit"] is True
assert response_2.data[0]["embedding"] == [0.1, 0.2, 0.3]
assert route.call_count == 1
def test_bedrock_cohere_embedding_returns_the_scripted_vector(respx_mock: respx.MockRouter) -> None:
route = respx_mock.post(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/cohere.embed-multilingual-v3/invoke"
).mock(
return_value=httpx.Response(
200,
json={"embeddings": [[0.3, 0.2, 0.1]], "inputTextTokenCount": 4},
)
)
response = litellm.embedding(
model="bedrock/cohere.embed-multilingual-v3",
input=["hello"],
api_key="test-bearer-token",
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
)
assert json.loads(route.calls.last.request.read()) == {"input_type": "search_document", "texts": ["hello"]}
assert response.data == [{"object": "embedding", "index": 0, "embedding": [0.3, 0.2, 0.1]}]
@pytest.mark.asyncio
async def test_async_bedrock_embedding_with_bearer_token():
"""Test async embedding functionality with bearer token authentication"""

View file

@ -264,3 +264,25 @@ async def test_cohere_documents_options_in_request_body():
# Validate that documents and citation_options are in the request body
assert "documents" in request_data
assert request_data["documents"] == test_documents
def test_cohere_invalid_max_tokens_maps_to_bad_request_error(
respx_mock: respx.MockRouter, _cohere_httpx_transport: None
) -> None:
route: Final = respx_mock.post(COHERE_V2_CHAT_URL).respond(
400, json={"message": "invalid request: max_tokens must be a number"}
)
with pytest.raises(litellm.BadRequestError) as error:
litellm.completion(
model="command-nightly",
messages=[{"content": "hey, how's it going?", "role": "user"}],
max_tokens="hello world",
api_key="cohere-test-key",
)
assert error.value.status_code == 400
assert error.value.llm_provider == "cohere"
assert "CohereException" in str(error.value)
assert "max_tokens must be a number" in str(error.value)
assert json.loads(route.calls.last.request.content)["max_tokens"] == "hello world"

View file

@ -0,0 +1,179 @@
from datetime import datetime, timezone
from typing import Final
import httpx
import pytest
import respx
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.cohere.embed.transformation import CohereEmbeddingConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.types.utils import CallTypes, EmbeddingResponse
def _logging_obj(inputs: list[str]) -> Logging:
logging_obj: Final = Logging(
model="embed-english-v3.0",
messages=inputs,
stream=False,
call_type=CallTypes.embedding,
start_time=datetime(2025, 1, 1, tzinfo=timezone.utc),
litellm_call_id="test-call",
function_id="test-function",
)
logging_obj.model_call_details = {"input": inputs}
return logging_obj
def test_cohere_text_embedding_request_maps_openai_params():
config: Final = CohereEmbeddingConfig()
request: Final = config.transform_embedding_request(
model="embed-english-v3.0",
input=["hello", "world"],
optional_params={"output_dimension": 256, "embedding_types": ["float"]},
headers={},
)
assert request == {
"model": "embed-english-v3.0",
"texts": ["hello", "world"],
"input_type": litellm.COHERE_DEFAULT_EMBEDDING_INPUT_TYPE,
"output_dimension": 256,
"embedding_types": ["float"],
}
def test_cohere_image_embedding_request_uses_images():
config: Final = CohereEmbeddingConfig()
request: Final = config.transform_embedding_request(
model="embed-english-v3.0",
input=["data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAAB"],
optional_params={},
headers={},
)
assert request == {
"model": "embed-english-v3.0",
"images": ["data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAAB"],
"input_type": "image",
}
def test_cohere_embedding_response_preserves_vectors_and_provider_usage():
config: Final = CohereEmbeddingConfig()
inputs: Final = ["hello", "world"]
logging_obj: Final = _logging_obj(inputs)
response: Final = config.transform_embedding_response(
model="embed-english-v3.0",
raw_response=httpx.Response(
200,
json={
"embeddings": {"float": [[0.1, 0.2], [0.3, 0.4]]},
"meta": {"billed_units": {"input_tokens": 5}},
},
),
model_response=EmbeddingResponse(),
logging_obj=logging_obj,
api_key="test-key",
request_data={"texts": inputs},
optional_params={},
litellm_params={},
)
assert response.model == "embed-english-v3.0"
assert response.data == [
{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]},
{"object": "embedding", "index": 1, "embedding": [0.3, 0.4]},
]
assert (response.usage.prompt_tokens, response.usage.total_tokens) == (5, 5)
def test_cohere_embedding_response_counts_tokens_when_provider_usage_is_missing():
config: Final = CohereEmbeddingConfig()
inputs: Final = ["hello"]
logging_obj: Final = _logging_obj(inputs)
response: Final = config.transform_embedding_response(
model="embed-english-v3.0",
raw_response=httpx.Response(200, json={"embeddings": {"float": [[0.1]]}}),
model_response=EmbeddingResponse(),
logging_obj=logging_obj,
api_key="test-key",
request_data={"texts": inputs},
optional_params={},
litellm_params={},
)
assert response.usage.prompt_tokens == len(litellm.encoding.encode("hello"))
assert response.usage.total_tokens == response.usage.prompt_tokens
def test_cohere_embedding_usage_includes_text_and_image_token_details():
config: Final = CohereEmbeddingConfig()
usage: Final = config._calculate_usage(
input=["hello"],
encoding=litellm.encoding,
meta={"billed_units": {"input_tokens": 5, "images": 2}},
)
assert usage.prompt_tokens == 7
assert (usage.prompt_tokens_details.text_tokens, usage.prompt_tokens_details.image_tokens) == (5, 2)
def test_cohere_embedding_error_class_preserves_status_and_message():
error: Final = CohereEmbeddingConfig().get_error_class(
error_message="invalid embedding request",
status_code=400,
headers={},
)
assert error.status_code == 400
assert "invalid embedding request" in str(error)
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_embedding_with_extra_headers_are_sent_to_cohere(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post("https://api.cohere.ai/v2/embed").mock(
return_value=httpx.Response(
200,
json={"embeddings": {"float": [[0.1, 0.2]]}, "meta": {"billed_units": {"input_tokens": 1}}},
)
)
request: Final = {
"model": "cohere/embed-english-v3.0",
"input": ["hello"],
"api_key": "test-key",
"extra_headers": {"x-request-id": "test-request"},
}
response: Final = litellm.embedding(**request) if sync_mode else await litellm.aembedding(**request)
assert route.calls.last.request.headers["x-request-id"] == "test-request"
assert response.data[0]["embedding"] == [0.1, 0.2]
async def test_aembedding_sends_extra_headers_through_the_caller_supplied_async_client() -> None:
requests: Final[list[httpx.Request]] = []
def handle_request(request: httpx.Request) -> httpx.Response:
requests.append(request)
return httpx.Response(
200,
json={"embeddings": {"float": [[0.3, 0.4]]}, "meta": {"billed_units": {"input_tokens": 2}}},
)
response: Final = await litellm.aembedding(
model="cohere/embed-english-v3.0",
input=["hello world"],
api_key="test-key",
extra_headers={"my-test-param": "hello-world"},
client=AsyncHTTPHandler(transport=httpx.MockTransport(handle_request)),
)
assert len(requests) == 1
assert str(requests[0].url) == "https://api.cohere.ai/v2/embed"
assert requests[0].headers["my-test-param"] == "hello-world"
assert response.data[0]["embedding"] == [0.3, 0.4]

View file

@ -14,9 +14,11 @@ from unittest.mock import MagicMock, patch
import certifi
import httpx
import pytest
import respx
from aiohttp import ClientSession, TCPConnector
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
from litellm.llms.custom_httpx.http_handler import (
_CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER,
@ -2051,3 +2053,73 @@ def test_realtime_ssl_for_url_turns_ssl_verify_false_into_an_unverified_tls_cont
assert isinstance(selected, ssl.SSLContext)
assert selected.verify_mode == ssl.CERT_NONE
assert selected.check_hostname is False
def _ollama_ssl_messages() -> list[dict[str, str]]:
return [{"role": "user", "content": "What's the weather like in San Francisco?"}]
def _ollama_generate_route(respx_mock: respx.MockRouter, stream: bool) -> respx.Route:
final: Final = {"model": "llama3.1", "response": "", "done": True, "prompt_eval_count": 2, "eval_count": 1}
if stream:
return respx_mock.post("http://ollama.test:11434/api/generate").respond(
text=json.dumps({"model": "llama3.1", "response": "Sunny", "done": False}) + "\n" + json.dumps(final) + "\n",
headers={"content-type": "application/x-ndjson"},
)
return respx_mock.post("http://ollama.test:11434/api/generate").respond(json={**final, "response": "Sunny"})
def _cached_clients_verify_modes(
cache: LLMClientCache, client_type: type[HTTPHandler | AsyncHTTPHandler]
) -> tuple[ssl.VerifyMode, ...]:
return tuple(
client.client._transport._pool._ssl_context.verify_mode
for client in cache.cache_dict.values()
if isinstance(client, client_type)
)
def test_ollama_ssl_verify_false_builds_a_client_that_skips_certificate_checks(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
client_cache: Final = LLMClientCache()
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", client_cache)
route: Final = _ollama_generate_route(respx_mock, stream=False)
response: Final = litellm.completion(
model="ollama/llama3.1",
messages=_ollama_ssl_messages(),
api_base="http://ollama.test:11434",
ssl_verify=False,
)
assert route.call_count == 1
assert response.choices[0].message.content == "Sunny"
assert _cached_clients_verify_modes(client_cache, HTTPHandler) == (ssl.CERT_NONE,)
@pytest.mark.parametrize("stream", [True, False])
async def test_async_ollama_ssl_verify_false_builds_a_client_that_skips_certificate_checks(
stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
client_cache: Final = LLMClientCache()
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", client_cache)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = _ollama_generate_route(respx_mock, stream=stream)
response: Final = await litellm.acompletion(
model="ollama/llama3.1",
messages=_ollama_ssl_messages(),
api_base="http://ollama.test:11434",
ssl_verify=False,
stream=stream,
)
text: Final = (
"".join([chunk.choices[0].delta.content or "" async for chunk in response])
if stream
else response.choices[0].message.content
)
assert route.call_count == 1
assert text == "Sunny"
assert _cached_clients_verify_modes(client_cache, AsyncHTTPHandler) == (ssl.CERT_NONE,)

View file

@ -0,0 +1,68 @@
import json
from typing import Final
import httpx
import pytest
import respx
import litellm
from litellm.llms.fireworks_ai.embed.fireworks_ai_transformation import FireworksAIEmbeddingConfig
def test_fireworks_nomic_embedding_supports_and_maps_dimensions():
config: Final = FireworksAIEmbeddingConfig()
params: Final = config.map_openai_params(
non_default_params={"dimensions": 256, "encoding_format": "float"},
optional_params={},
model="nomic-ai/nomic-embed-text-v1.5",
)
assert config.get_supported_openai_params("nomic-ai/nomic-embed-text-v1.5") == ["dimensions"]
assert params == {"dimensions": 256}
assert config.is_fireworks_embedding_model("nomic-ai/nomic-embed-text-v1.5")
def test_fireworks_non_nomic_embedding_drops_dimensions():
config: Final = FireworksAIEmbeddingConfig()
assert config.get_supported_openai_params("other/model") == []
assert config.map_openai_params(
non_default_params={"dimensions": 256},
optional_params={},
model="other/model",
) == {}
assert not config.is_fireworks_embedding_model("other/model")
@respx.mock
def test_fireworks_embeddings_forwards_supported_dimensions(monkeypatch: pytest.MonkeyPatch):
model_key: Final = "fireworks_ai/nomic-ai/nomic-embed-text-v1.5"
monkeypatch.setitem(litellm.model_cost, model_key, {**litellm.model_cost[model_key], "input_cost_per_token": 2.5e-08})
route: Final = respx.post("https://fireworks.example/v1/embeddings").mock(
return_value=httpx.Response(
200,
json={
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
"model": "nomic-ai/nomic-embed-text-v1.5",
"usage": {"prompt_tokens": 4, "total_tokens": 4},
},
)
)
response: Final = litellm.embedding(
model="fireworks_ai/nomic-ai/nomic-embed-text-v1.5",
input=["hello"],
api_base="https://fireworks.example/v1",
api_key="test-key",
dimensions=2,
)
assert json.loads(route.calls.last.request.content) == {
"input": ["hello"],
"model": "nomic-ai/nomic-embed-text-v1.5",
"dimensions": 2,
}
assert response.data[0]["embedding"] == [0.1, 0.2]
assert response.usage.prompt_tokens == 4
assert response._hidden_params["response_cost"] == pytest.approx(4 * 2.5e-08)

View file

@ -0,0 +1,68 @@
import json
from typing import Final
import httpx
import respx
from litellm import text_completion
FIREWORKS_COMPLETIONS_URL: Final = "https://api.fireworks.ai/inference/v1/completions"
FIREWORKS_MODEL: Final = "accounts/fireworks/models/llama-v3p1-8b-instruct"
def _fireworks_completion_response() -> dict[str, object]:
return {
"id": "cmpl-fireworks-migration",
"object": "text_completion",
"created": 1,
"model": FIREWORKS_MODEL,
"choices": [
{
"text": "hello world",
"index": 0,
"logprobs": {"tokens": ["hello", " world"], "token_logprobs": [-0.1, -0.2]},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3},
}
def test_fireworks_text_completion_sends_prompt_array(respx_mock: respx.MockRouter) -> None:
route: Final = respx_mock.post(FIREWORKS_COMPLETIONS_URL).mock(
return_value=httpx.Response(200, json=_fireworks_completion_response())
)
response: Final = text_completion(
model=f"fireworks_ai/{FIREWORKS_MODEL}",
prompt=["What's the weather in SF?", "How is Manchester?"],
max_tokens=5,
api_key="fireworks-test-key",
)
request_body: Final = json.loads(route.calls.last.request.read())
assert request_body["prompt"] == ["What's the weather in SF?", "How is Manchester?"]
assert [choice.text for choice in response.choices] == ["hello world"]
assert route.calls.last.request.headers["Authorization"] == "Bearer fireworks-test-key"
def test_fireworks_text_completion_echo_returns_prompt_and_logprobs(respx_mock: respx.MockRouter) -> None:
route: Final = respx_mock.post(FIREWORKS_COMPLETIONS_URL).mock(
return_value=httpx.Response(200, json=_fireworks_completion_response())
)
response: Final = text_completion(
model=f"fireworks_ai/{FIREWORKS_MODEL}",
prompt="hello",
max_tokens=1,
stop="\n",
logprobs=1,
echo=True,
api_key="fireworks-test-key",
)
request_body: Final = json.loads(route.calls.last.request.read())
assert request_body["echo"] is True
assert request_body["stop"] == "\n"
assert response.choices[0].text == "hello world"
assert response.choices[0].logprobs.token_logprobs == [-0.1, -0.2]

View file

@ -0,0 +1,219 @@
import json
from collections.abc import Mapping
from typing import Final
import httpx
import pytest
import respx
import litellm
def _sse_body(chunks: tuple[Mapping[str, object], ...]) -> str:
return "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
@pytest.mark.asyncio
@respx.mock
async def test_async_gemini_streaming_tool_call_arguments_are_valid(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
chunks: Final = (
{
"candidates": [
{
"content": {
"parts": [
{
"functionCall": {
"name": "generate_series_of_questions",
"args": {
"questions": [
"What is a bridge?",
"How is it designed?",
"How is it built?",
]
},
}
}
],
"role": "model",
},
"index": 0,
}
],
"usageMetadata": {
"promptTokenCount": 5,
"candidatesTokenCount": 7,
"totalTokenCount": 12,
},
},
{
"candidates": [
{
"content": {"parts": [], "role": "model"},
"finishReason": "STOP",
"index": 0,
}
],
"usageMetadata": {
"promptTokenCount": 5,
"candidatesTokenCount": 7,
"totalTokenCount": 12,
},
},
)
route: Final = respx.post(
url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash-lite:streamGenerateContent.*"
).mock(
return_value=httpx.Response(
200,
text=_sse_body(chunks),
headers={"content-type": "text/event-stream"},
)
)
stream: Final = await litellm.acompletion(
model="gemini/gemini-2.5-flash-lite",
messages=[
{"role": "system", "content": "You are an AI assistant"},
{"role": "user", "content": "Generate three questions about civil engineering."},
],
tools=[
{
"type": "function",
"function": {
"name": "generate_series_of_questions",
"description": "Generate questions",
"parameters": {
"type": "object",
"properties": {
"questions": {
"type": "array",
"items": {"type": "string"},
}
},
"required": ["questions"],
},
},
}
],
stream=True,
api_key="test-gemini-key",
)
response_chunks: Final = [chunk async for chunk in stream]
tool_calls: Final = [
tool_call
for chunk in response_chunks
for tool_call in chunk.choices[0].delta.tool_calls or ()
]
assert route.call_count == 1
assert response_chunks[0].choices[0].delta.role == "assistant"
assert tuple(chunk.choices[0].finish_reason for chunk in response_chunks) == (None, "tool_calls")
assert len(tool_calls) == 1
assert tool_calls[0].function.name == "generate_series_of_questions"
assert json.loads(tool_calls[0].function.arguments) == {
"questions": [
"What is a bridge?",
"How is it designed?",
"How is it built?",
]
}
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
@respx.mock
async def test_gemini_streaming_legacy_function_call_arguments_are_valid(
sync_mode: bool,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
chunks: Final = (
{
"candidates": [
{
"content": {
"parts": [
{
"functionCall": {
"name": "get_current_weather",
"args": {"location": "Boston", "unit": "fahrenheit"},
}
}
],
"role": "model",
},
"index": 0,
}
]
},
{
"candidates": [
{
"content": {"parts": [], "role": "model"},
"finishReason": "STOP",
"index": 0,
}
]
},
)
respx.post(
url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash-lite:streamGenerateContent.*"
).mock(
return_value=httpx.Response(
200,
text=_sse_body(chunks),
headers={"content-type": "text/event-stream"},
)
)
messages: Final = [{"role": "user", "content": "What is the weather in Boston?"}]
functions: Final = [
{
"name": "get_current_weather",
"description": "Get the current weather",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"},
"unit": {"type": "string"},
},
"required": ["location", "unit"],
},
}
]
if sync_mode:
stream_chunks: Final = list(
litellm.completion(
model="gemini/gemini-2.5-flash-lite",
messages=messages,
functions=functions,
stream=True,
api_key="test-gemini-key",
)
)
else:
stream: Final = await litellm.acompletion(
model="gemini/gemini-2.5-flash-lite",
messages=messages,
functions=functions,
stream=True,
api_key="test-gemini-key",
)
stream_chunks: Final = [chunk async for chunk in stream]
function_calls: Final = tuple(
chunk.choices[0].delta.function_call
for chunk in stream_chunks
if chunk.choices and chunk.choices[0].delta.function_call is not None
)
assert tuple(
function_call.name for function_call in function_calls if function_call.name is not None
) == ("get_current_weather",)
arguments: Final = "".join(function_call.arguments or "" for function_call in function_calls)
assert json.loads(arguments) == {"location": "Boston", "unit": "fahrenheit"}
rebuilt: Final = litellm.stream_chunk_builder(chunks=list(stream_chunks), messages=messages)
assert rebuilt.choices[0].message.function_call.name == "get_current_weather"
assert json.loads(rebuilt.choices[0].message.function_call.arguments) == {"location": "Boston", "unit": "fahrenheit"}

View file

@ -6,9 +6,11 @@ especially ensuring that encoding_format is not included when not provided.
"""
import json
from typing import Final
from unittest.mock import Mock, patch
import pytest
import respx
import litellm
@ -326,3 +328,23 @@ class TestHostedVLLMEmbeddingTransformation:
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])
def test_hosted_vllm_embedding_posts_to_the_env_api_base(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("HOSTED_VLLM_API_BASE", "http://hosted-vllm.test")
route: Final = respx_mock.post("http://hosted-vllm.test/embeddings").respond(
json={
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.5, 0.25]}],
"model": "jina-embeddings-v3",
"usage": {"prompt_tokens": 2, "total_tokens": 2},
}
)
response: Final = litellm.embedding(model="hosted_vllm/jina-embeddings-v3", input=["Hello world"])
assert route.call_count == 1
assert json.loads(route.calls.last.request.content) == {"model": "jina-embeddings-v3", "input": ["Hello world"]}
assert response.data[0]["embedding"] == [0.5, 0.25]

View file

@ -1,4 +1,5 @@
import json
from typing import Final
from unittest.mock import patch, MagicMock, AsyncMock
@ -154,3 +155,45 @@ class TestHuggingFaceEmbedding:
assert "sentences" in request_data["inputs"]
assert request_data["inputs"]["source_sentence"] == input_text[0]
assert request_data["inputs"]["sentences"] == input_text[1:]
def test_hf_sentence_similarity_transform_keeps_provider_options(self):
from litellm.llms.huggingface.embedding.handler import HuggingFaceEmbedding
request_data: Final = HuggingFaceEmbedding()._transform_input(
input=["source sentence", "candidate sentence"],
model="BAAI/bge-m3",
call_type="sync",
optional_params={"input_type": "sentence-similarity", "min_length": 2},
embed_url="https://huggingface.example/embeddings",
)
assert request_data == {
"inputs": {"source_sentence": "source sentence", "sentences": ["candidate sentence"]},
"parameters": {"min_length": 2},
}
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_hf_embedding_sentence_sim_sends_one_similarity_request(
sync_mode: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("HUGGINGFACE_API_KEY", "hf-test-token")
with respx.mock(assert_all_called=False) as hf:
hf.get(url__regex=r"https://huggingface\.co/api/models/.*").respond(json={"pipeline_tag": "sentence-similarity"})
hf.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond(json={"max_position_embeddings": 512})
inference: Final = hf.post(
"https://router.huggingface.co/hf-inference/pipeline/sentence-similarity/sentence-transformers/TaylorAI/bge-micro-v2"
).respond(json=[0.7708950042724609])
request: Final = {
"model": "huggingface/sentence-transformers/TaylorAI/bge-micro-v2",
"input": ["good morning from litellm", "this is another item"],
}
response: Final = litellm.embedding(**request) if sync_mode else await litellm.aembedding(**request)
assert inference.call_count == 1
assert json.loads(inference.calls.last.request.content) == {
"inputs": {"source_sentence": "good morning from litellm", "sentences": ["this is another item"]}
}
assert inference.calls.last.request.headers["Authorization"] == "Bearer hf-test-token"
assert response.data == [{"object": "embedding", "index": 0, "embedding": 0.7708950042724609}]

View file

@ -1,9 +1,14 @@
import json
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
import respx
from pydantic import ValidationError
import litellm
from litellm.llms.jina_ai.embedding.transformation import JinaAIEmbeddingConfig
from litellm.types.utils import EmbeddingResponse
@ -192,3 +197,40 @@ def test_transform_embedding_response_invalid_field_is_reported_by_name(field: s
assert exc_info.value.title == "EmbeddingResponse"
assert [error["loc"] for error in exc_info.value.errors()] == [(field,)]
_PIXEL_PNG: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
@pytest.mark.parametrize(
("input_data", "expected_payload_input"),
[
(["hello world", "foo bar"], ["hello world", "foo bar"]),
(
["A picture of a cat", f"data:image/png;base64,{_PIXEL_PNG}"],
[{"text": "A picture of a cat"}, {"image": _PIXEL_PNG}],
),
([f"data:image/png;base64,{_PIXEL_PNG}"], [{"image": _PIXEL_PNG}]),
],
ids=["text_only", "text_and_image", "image_only"],
)
def test_jina_ai_img_embeddings_transforms_mixed_text_and_image_inputs(
input_data: list[str], expected_payload_input: list[object], respx_mock: respx.MockRouter
) -> None:
route: Final = respx_mock.post("https://api.jina.ai/v1/embeddings").respond(
json={
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
"model": "jina-embeddings-v4",
"usage": {"prompt_tokens": 3, "total_tokens": 3},
}
)
response: Final = litellm.embedding(model="jina_ai/jina-embeddings-v4", input=input_data, api_key="jina-test-key")
assert route.call_count == 1
assert json.loads(route.calls.last.request.content) == {
"model": "jina-embeddings-v4",
"input": expected_payload_input,
}
assert response.data[0]["embedding"] == [0.1, 0.2]

View file

@ -1,5 +1,10 @@
import json
import pytest
import respx
import litellm
import httpx
from typing import Final
@pytest.fixture(autouse=True)
@ -168,3 +173,155 @@ async def test_mistral_transform_request_name_field_removal(
assert user_message["role"] == "user"
assert user_message["content"] == "Hello"
assert "name" not in user_message # The 'name' field should have been removed
def test_mistral_401_maps_to_authentication_error(respx_mock):
route = respx_mock.post("https://api.mistral.ai/v1/chat/completions").mock(
return_value=httpx.Response(401, json={"message": "invalid API key"})
)
with pytest.raises(litellm.AuthenticationError) as exc_info:
litellm.completion(
model="mistral/mistral-tiny",
messages=[{"role": "user", "content": "hello"}],
api_key="bad-mistral-key",
max_retries=0,
)
assert exc_info.value.status_code == 401
assert "invalid API key" in str(exc_info.value)
assert route.calls.last.request.headers["Authorization"] == "Bearer bad-mistral-key"
def test_mistral_streaming_function_call_arguments_are_assembled(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
chunks: Final = (
{
"id": "chatcmpl-stream",
"object": "chat.completion.chunk",
"created": 1,
"model": "mistral-medium-latest",
"choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
},
{
"id": "chatcmpl-stream",
"object": "chat.completion.chunk",
"created": 1,
"model": "mistral-medium-latest",
"choices": [
{
"index": 0,
"delta": {
"tool_calls": [
{
"index": 0,
"id": "call_weather",
"type": "function",
"function": {
"name": "get_current_weather",
"arguments": '{"location":"',
},
}
]
},
"finish_reason": None,
}
],
},
{
"id": "chatcmpl-stream",
"object": "chat.completion.chunk",
"created": 1,
"model": "mistral-medium-latest",
"choices": [
{
"index": 0,
"delta": {
"tool_calls": [
{
"index": 0,
"function": {"arguments": 'Boston","unit":"fahrenheit"}'},
}
]
},
"finish_reason": None,
}
],
},
{
"id": "chatcmpl-stream",
"object": "chat.completion.chunk",
"created": 1,
"model": "mistral-medium-latest",
"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}],
},
)
body: Final = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
respx_mock.post("https://api.mistral.ai/v1/chat/completions").respond(
200,
text=body,
headers={"content-type": "text/event-stream"},
)
messages: Final = [{"role": "user", "content": "What is the weather in Boston?"}]
stream: Final = litellm.completion(
model="mistral/mistral-medium-latest",
messages=messages,
tools=[
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"},
"unit": {"type": "string"},
},
"required": ["location", "unit"],
},
},
}
],
stream=True,
)
stream_chunks: Final = tuple(stream)
assert {chunk._hidden_params["custom_llm_provider"] for chunk in stream_chunks} == {"mistral"}
argument_fragments: Final = tuple(
chunk.choices[0].delta.tool_calls[0].function.arguments
for chunk in stream_chunks
if chunk.choices and chunk.choices[0].delta.tool_calls is not None
)
assert "".join(argument or "" for argument in argument_fragments) == (
'{"location":"Boston","unit":"fahrenheit"}'
)
assembled: Final = litellm.stream_chunk_builder(chunks=list(stream_chunks), messages=messages)
assert assembled is not None
assert assembled.choices[0].message.tool_calls[0].function.name == "get_current_weather"
assert json.loads(assembled.choices[0].message.tool_calls[0].function.arguments) == {
"location": "Boston",
"unit": "fahrenheit",
}
assert assembled.choices[0].finish_reason == "tool_calls"
def test_mistral_completion_cost_is_priced_from_the_returned_usage(
respx_mock, mistral_api_response, local_model_cost_map
):
route: Final = respx_mock.post("https://api.mistral.ai/v1/chat/completions").respond(json=mistral_api_response)
response: Final = litellm.completion(
model="mistral/mistral-medium-latest",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
max_tokens=5,
seed=10,
input_cost_per_token=4e-07,
output_cost_per_token=2e-06,
)
body: Final = json.loads(route.calls.last.request.content)
assert (body["random_seed"], body["max_tokens"]) == (10, 5)
assert (response.usage.prompt_tokens, response.usage.completion_tokens) == (10, 15)
assert response._hidden_params["response_cost"] == pytest.approx(10 * 4e-07 + 15 * 2e-06)

View file

@ -0,0 +1,34 @@
import json
from typing import Final
import httpx
import respx
import litellm
def test_mistral_embedding_uses_mistral_embeddings_endpoint(respx_mock: respx.MockRouter) -> None:
route: Final = respx_mock.post("https://api.mistral.ai/v1/embeddings").mock(
return_value=httpx.Response(
200,
json={
"id": "embd-migration",
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
"model": "mistral-embed",
"usage": {"prompt_tokens": 2, "total_tokens": 2},
},
)
)
response: Final = litellm.embedding(
model="mistral/mistral-embed",
input=["hello"],
api_key="test-mistral-key",
)
assert str(route.calls.last.request.url) == "https://api.mistral.ai/v1/embeddings"
assert route.calls.last.request.headers["Authorization"] == "Bearer test-mistral-key"
assert json.loads(route.calls.last.request.read()) == {"model": "mistral-embed", "input": ["hello"]}
assert response.data[0]["embedding"] == [0.1, 0.2]
assert (response.usage.prompt_tokens, response.usage.total_tokens) == (2, 2)

View file

@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch
import httpx
import pytest
import respx
import litellm
from litellm._uuid import uuid
@ -886,3 +887,36 @@ def test_transform_request_leaves_unreadable_images_untouched(payload: str) -> N
data = _transform_image_request(payload, "png")
assert data["images"] == [payload]
def test_ollama_vision_model_sends_the_image_beside_the_user_prompt(respx_mock: respx.MockRouter) -> None:
image_base64: Final = _image_base64("PNG")
route: Final = respx_mock.post("http://ollama.test:11434/api/generate").respond(
json={
"model": "llama3.2-vision:11b",
"response": "A white square",
"done": True,
"prompt_eval_count": 5,
"eval_count": 3,
}
)
response: Final = litellm.completion(
model="ollama/llama3.2-vision:11b",
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "Whats in this image?"},
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
],
}
],
api_base="http://ollama.test:11434",
)
body: Final = json.loads(route.calls.last.request.content)
assert body["model"] == "llama3.2-vision:11b"
assert body["images"] == [image_base64]
assert body["prompt"] == "### User:\nWhats in this image?\n\n"
assert response.choices[0].message.content == "A white square"

View file

@ -0,0 +1,39 @@
import json
from typing import Final
import httpx
import respx
from litellm import text_completion
@respx.mock
def test_ollama_text_completion_sends_prompt() -> None:
endpoint: Final = "http://localhost:11434/api/generate"
route: Final = respx.post(endpoint).mock(
return_value=httpx.Response(
200,
json={
"model": "llama3.1",
"created_at": "2024-01-01T00:00:00Z",
"response": "Hello!",
"done": True,
"context": [1],
"total_duration": 1,
"load_duration": 1,
"prompt_eval_count": 2,
"eval_count": 1,
},
)
)
response: Final = text_completion(
model="ollama/llama3.1",
prompt="hello",
api_base="http://localhost:11434",
)
request_body: Final = json.loads(route.calls.last.request.read())
assert request_body["model"] == "llama3.1"
assert request_body["prompt"] == "hello"
assert response.choices[0].text == "Hello!"

View file

@ -6,6 +6,9 @@ Regression tests for https://github.com/BerriAI/litellm/issues/27410
"""
import json
from typing import Final
import pytest
import respx
from httpx import Response
@ -147,3 +150,227 @@ async def test_async_streaming_completion_reports_an_error_event_sent_mid_stream
assert failed.value.status_code == 500
assert "upstream overloaded" in str(failed.value)
@respx.mock
def test_completion_openai_prompt_array_sends_both_prompts() -> None:
route: Final = respx.post("https://api.openai.com/v1/completions").mock(
return_value=Response(
200,
json={
"id": "cmpl-test123",
"object": "text_completion",
"created": 1677652288,
"model": "gpt-3.5-turbo-instruct",
"choices": [
{"text": "first answer", "index": 0, "logprobs": None, "finish_reason": "stop"},
{"text": "second answer", "index": 1, "logprobs": None, "finish_reason": "stop"},
],
"usage": {"prompt_tokens": 2, "completion_tokens": 2, "total_tokens": 4},
},
)
)
response: Final = text_completion(
model="gpt-3.5-turbo-instruct",
prompt=["first prompt", "second prompt"],
max_tokens=5,
)
body: Final = json.loads(route.calls.last.request.read())
assert body["prompt"] == ["first prompt", "second prompt"]
assert [choice.text for choice in response.choices] == ["first answer", "second answer"]
@respx.mock
def test_completion_text_003_token_prompt_array_returns_one_choice_per_prompt() -> None:
token_prompts: Final = [[2061, 338, 262, 6193, 287, 14362, 30], [2437, 318, 9502, 30]]
route: Final = respx.post("https://api.fireworks.ai/inference/v1/completions").mock(
return_value=Response(
200,
json={
"id": "cmpl-token-prompts",
"object": "text_completion",
"created": 1677652288,
"model": "accounts/fireworks/models/glm-5p3-flash",
"choices": [
{"text": "sunny", "index": 0, "logprobs": None, "finish_reason": "length"},
{"text": "rainy", "index": 1, "logprobs": None, "finish_reason": "length"},
],
"usage": {"prompt_tokens": 11, "completion_tokens": 2, "total_tokens": 13},
},
)
)
response: Final = text_completion(
model="text-completion-openai/accounts/fireworks/models/glm-5p3-flash",
api_base="https://api.fireworks.ai/inference/v1",
api_key="fireworks-test-key",
prompt=token_prompts,
max_tokens=5,
)
body: Final = json.loads(route.calls.last.request.read())
assert body["prompt"] == token_prompts
assert body["model"] == "accounts/fireworks/models/glm-5p3-flash"
assert [(choice.index, choice.text) for choice in response.choices] == [(0, "sunny"), (1, "rainy")]
@respx.mock
def test_text_completion_with_echo_returns_prompt_and_token_logprobs() -> None:
route: Final = respx.post("https://api.openai.com/v1/completions").mock(
return_value=Response(
200,
json={
"id": "cmpl-test123",
"object": "text_completion",
"created": 1677652288,
"model": "gpt-3.5-turbo-instruct",
"choices": [
{
"text": "hello world",
"index": 0,
"logprobs": {"tokens": ["hello", " world"], "token_logprobs": [-0.1, -0.2]},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
},
)
)
response: Final = text_completion(
model="gpt-3.5-turbo-instruct",
prompt="hello",
max_tokens=1,
echo=True,
logprobs=1,
)
body: Final = json.loads(route.calls.last.request.read())
assert body["echo"] is True
assert response.choices[0].text == "hello world"
assert response.choices[0].logprobs.token_logprobs == [-0.1, -0.2]
@respx.mock
def test_completion_openai_engine_selects_the_engine_model(mock_completions_endpoint) -> None:
response: Final = text_completion(engine="gpt-3.5-turbo-instruct", prompt="hello", max_tokens=5)
request_body: Final = json.loads(mock_completions_endpoint.calls.last.request.read())
assert request_body["model"] == "gpt-3.5-turbo-instruct"
assert response.model == "gpt-3.5-turbo-instruct"
@respx.mock
def test_completion_openai_engine_and_model_uses_the_explicit_model(mock_completions_endpoint) -> None:
response: Final = text_completion(
model="gpt-3.5-turbo-instruct",
engine="unused-engine",
prompt="hello",
max_tokens=5,
)
request_body: Final = json.loads(mock_completions_endpoint.calls.last.request.read())
assert request_body["model"] == "gpt-3.5-turbo-instruct"
assert response.model == "gpt-3.5-turbo-instruct"
@respx.mock
async def test_async_text_completion_stream_returns_text_and_one_finish_reason(monkeypatch):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx.post("https://api.openai.com/v1/completions").mock(
return_value=Response(
200,
text=(
'data: {"id":"cmpl-stream","object":"text_completion","created":1,"model":"gpt-3.5-turbo-instruct","choices":[{"index":0,"text":"hello","finish_reason":null}]}\n\n'
'data: {"id":"cmpl-stream","object":"text_completion","created":1,"model":"gpt-3.5-turbo-instruct","choices":[{"index":0,"text":"","finish_reason":"stop"}]}\n\n'
"data: [DONE]\n\n"
),
headers={"content-type": "text/event-stream"},
)
)
stream: Final = await atext_completion(
model="gpt-3.5-turbo-instruct", prompt="say hello", stream=True, max_tokens=5
)
chunks: Final = [chunk async for chunk in stream]
assert [chunk.choices[0].text for chunk in chunks if chunk.choices[0].text] == ["hello"]
assert tuple(
chunk.choices[0].finish_reason
for chunk in chunks
if chunk.choices[0].finish_reason is not None
) == ("stop",)
assert len(route.calls) == 1
@respx.mock
async def test_async_text_completion_chat_model_stream_builds_chat_chunks(monkeypatch):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock(
return_value=Response(
200,
text=(
'data: {"id":"chatcmpl-stream","object":"chat.completion.chunk","created":1,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":null}]}\n\n'
'data: {"id":"chatcmpl-stream","object":"chat.completion.chunk","created":1,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n'
"data: [DONE]\n\n"
),
headers={"content-type": "text/event-stream"},
)
)
stream: Final = await atext_completion(
model="gpt-3.5-turbo", prompt="say hello", stream=True, max_tokens=5
)
chunks: Final = [chunk async for chunk in stream]
response: Final = litellm.stream_chunk_builder(chunks=chunks)
assert response.choices[0].text == "hello"
assert (response.usage.prompt_tokens, response.usage.completion_tokens) == (0, 1)
assert litellm.completion_cost(
completion_response=response,
custom_cost_per_token={"input_cost_per_token": 5e-07, "output_cost_per_token": 1.5e-06},
) == pytest.approx(1 * 1.5e-06)
assert tuple(
chunk.choices[0].finish_reason
for chunk in chunks
if chunk.choices[0].finish_reason is not None
) == ("stop",)
assert len(route.calls) == 1
@respx.mock
def test_text_completion_stream_forwards_stream_options_and_usage(monkeypatch):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx.post("https://api.openai.com/v1/completions").mock(
return_value=Response(
200,
text=(
'data: {"id":"cmpl-stream","object":"text_completion","created":1,"model":"gpt-3.5-turbo-instruct","choices":[{"index":0,"text":"hello","finish_reason":null}]}\n\n'
'data: {"id":"cmpl-stream","object":"text_completion","created":1,"model":"gpt-3.5-turbo-instruct","choices":[],"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3}}\n\n'
"data: [DONE]\n\n"
),
headers={"content-type": "text/event-stream"},
)
)
chunks: Final = list(
text_completion(
model="gpt-3.5-turbo-instruct",
prompt="say hello",
stream=True,
stream_options={"include_usage": True},
)
)
assert tuple(chunk.choices[0].text for chunk in chunks) == ("hello", None, None)
assert tuple(chunk.usage for chunk in chunks[:-1]) == (None, None)
assert (chunks[-1].usage.prompt_tokens, chunks[-1].usage.completion_tokens, chunks[-1].usage.total_tokens) == (
2,
1,
3,
)
assert json.loads(route.calls[0].request.content)["stream_options"] == {
"include_usage": True
}

View file

@ -0,0 +1,79 @@
from typing import Final
import httpx
import pytest
import respx
import litellm
def test_openai_401_maps_to_authentication_error(respx_mock: respx.MockRouter) -> None:
route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
return_value=httpx.Response(
401,
json={
"error": {
"message": "Incorrect API key provided",
"type": "invalid_request_error",
"param": None,
"code": "invalid_api_key",
}
},
)
)
with pytest.raises(litellm.AuthenticationError) as error:
litellm.completion(
model="openai/migration-auth-error-model",
messages=[{"role": "user", "content": "hello"}],
api_key="invalid-test-key",
max_retries=0,
)
assert error.value.status_code == 401
assert str(error.value).startswith("litellm.AuthenticationError: AuthenticationError: OpenAIException - ")
assert "Incorrect API key provided" in str(error.value)
assert route.calls.last.request.headers["Authorization"] == "Bearer invalid-test-key"
async def _complete(sync_mode: bool, request: dict[str, object]) -> object:
if sync_mode:
return litellm.completion(**request)
return await litellm.acompletion(**request)
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_openai_invalid_role_error_preserves_provider_fields(
sync_mode: bool, stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
return_value=httpx.Response(
400,
json={
"error": {
"message": "Invalid value: 'usera'. Supported values are: 'system', 'assistant', 'user'.",
"type": "invalid_request_error",
"param": "messages[0].role",
"code": "invalid_value",
}
},
)
)
request: Final = {
"model": "openai/migration-invalid-role-model",
"messages": [{"role": "usera", "content": "hi"}],
"api_key": "test-key",
"max_retries": 0,
"stream": stream,
}
with pytest.raises(litellm.BadRequestError) as error:
await _complete(sync_mode, request)
assert str(error.value).startswith("litellm.BadRequestError: OpenAIException - Invalid value")
assert error.value.code == "invalid_value"
assert error.value.param == "messages[0].role"
assert error.value.type == "invalid_request_error"
assert error.value.response.status_code == 400

View file

@ -3,9 +3,13 @@ Test cases for OpenAI-like embedding handler
"""
import json
from typing import Final
from unittest.mock import MagicMock, Mock, patch
import pytest
import respx
import litellm
from litellm.llms.openai_like.embedding.handler import OpenAILikeEmbeddingHandler
from litellm.types.utils import EmbeddingResponse
@ -346,3 +350,40 @@ class TestOpenAILikeEmbeddingHandler:
assert sent_data["model"] == "test-model"
assert sent_data["input"] == ["test input"]
assert "encoding_format" not in sent_data
_OPENAI_LIKE_EMBEDDING_BODY: Final = {
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.5, 0.25]}],
"model": "jina-embeddings-v3",
"usage": {"prompt_tokens": 2, "total_tokens": 2},
}
def test_llamafile_embedding_posts_to_the_env_api_base(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LLAMAFILE_API_BASE", "http://llamafile.test/v1")
route: Final = respx_mock.post("http://llamafile.test/v1/embeddings").respond(json=_OPENAI_LIKE_EMBEDDING_BODY)
response: Final = litellm.embedding(model="llamafile/jina-embeddings-v3", input=["Hello world"])
assert route.call_count == 1
assert json.loads(route.calls.last.request.content) == {"model": "jina-embeddings-v3", "input": ["Hello world"]}
assert response.data[0]["embedding"] == [0.5, 0.25]
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_lm_studio_embedding_posts_to_the_env_api_base(
sync_mode: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("LM_STUDIO_API_BASE", "http://lm-studio.test")
route: Final = respx_mock.post("http://lm-studio.test/embeddings").respond(json=_OPENAI_LIKE_EMBEDDING_BODY)
request: Final = {"model": "lm_studio/jina-embeddings-v3", "input": ["Hello world"]}
response: Final = litellm.embedding(**request) if sync_mode else await litellm.aembedding(**request)
assert route.call_count == 1
assert json.loads(route.calls.last.request.content) == {"model": "jina-embeddings-v3", "input": ["Hello world"]}
assert response.data[0]["embedding"] == [0.5, 0.25]

View file

@ -1,6 +1,12 @@
import json
from typing import Final
import httpx
import pytest
import respx
import litellm
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
@ -620,3 +626,139 @@ def test_openrouter_reasoning_effort_high_passes_through():
)
assert result["reasoning_effort"] == "high"
@respx.mock
def test_openrouter_streaming_completion_assembles_text_and_finish_reason() -> None:
chunks: Final = (
{
"id": "gen-stream-model",
"object": "chat.completion.chunk",
"created": 1,
"model": "openai/gpt-4o-mini",
"choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
},
{
"id": "gen-stream-model",
"object": "chat.completion.chunk",
"created": 1,
"model": "openai/gpt-4o-mini",
"choices": [{"index": 0, "delta": {"content": "The Supreme Court"}, "finish_reason": None}],
},
{
"id": "gen-stream-model",
"object": "chat.completion.chunk",
"created": 1,
"model": "openai/gpt-4o-mini",
"choices": [{"index": 0, "delta": {"content": " hears federal cases."}, "finish_reason": None}],
},
{
"id": "gen-stream-model",
"object": "chat.completion.chunk",
"created": 1,
"model": "openai/gpt-4o-mini",
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
},
)
body: Final = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
return_value=httpx.Response(
200,
text=body,
headers={"content-type": "text/event-stream"},
)
)
messages: Final = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "How does a case reach the Supreme Court?"},
]
stream: Final = litellm.completion(
model="openrouter/openai/gpt-4o-mini",
messages=messages,
stream=True,
max_tokens=50,
api_key="test-openrouter-key",
)
stream_chunks: Final = tuple(stream)
assert stream_chunks[0].choices[0].delta.role == "assistant"
assert tuple(
chunk.choices[0].finish_reason for chunk in stream_chunks if chunk.choices[0].finish_reason is not None
) == ("stop",)
assert "".join(
chunk.choices[0].delta.content or "" for chunk in stream_chunks if chunk.choices
) == "The Supreme Court hears federal cases."
assembled: Final = litellm.stream_chunk_builder(chunks=list(stream_chunks), messages=messages)
assert assembled is not None
assert assembled.choices[0].message.content == "The Supreme Court hears federal cases."
assert assembled.choices[0].finish_reason == "stop"
@respx.mock
def test_openrouter_streaming_reasoning_is_preserved() -> None:
chunks: Final = (
{
"id": "gen-reasoning-stream",
"object": "chat.completion.chunk",
"created": 1,
"model": "anthropic/claude-sonnet-4.5",
"choices": [
{
"index": 0,
"delta": {"role": "assistant", "reasoning": "Consider the evidence. "},
"finish_reason": None,
}
],
},
{
"id": "gen-reasoning-stream",
"object": "chat.completion.chunk",
"created": 1,
"model": "anthropic/claude-sonnet-4.5",
"choices": [
{
"index": 0,
"delta": {"content": "The decision is reasoned.", "reasoning": "Apply the rule."},
"finish_reason": None,
}
],
},
{
"id": "gen-reasoning-stream",
"object": "chat.completion.chunk",
"created": 1,
"model": "anthropic/claude-sonnet-4.5",
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
},
)
body: Final = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
route: Final = respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
return_value=httpx.Response(
200,
text=body,
headers={"content-type": "text/event-stream"},
)
)
parsed_reasoning_chunk: Final = OpenRouterChatCompletionStreamingHandler(
streaming_response=None, sync_stream=True
).chunk_parser(chunks[0])
assert parsed_reasoning_chunk.choices[0]["delta"]["reasoning_content"] == "Consider the evidence. "
messages: Final = [{"role": "user", "content": "Explain the decision briefly."}]
stream: Final = litellm.completion(
model="openrouter/anthropic/claude-sonnet-4.5",
messages=messages,
stream=True,
reasoning={"effort": "high"},
drop_params=True,
api_key="test-openrouter-key",
)
stream_chunks: Final = tuple(stream)
assert json.loads(route.calls.last.request.content)["reasoning"] == {"effort": "high"}
reasoning_deltas: Final = tuple(
getattr(chunk.choices[0].delta, "reasoning_content", None)
for chunk in stream_chunks
if chunk.choices
)
assert reasoning_deltas == ("Consider the evidence. ", "Apply the rule.", None)
assembled: Final = litellm.stream_chunk_builder(chunks=list(stream_chunks), messages=messages)
assert assembled is not None
assert assembled.choices[0].message.reasoning_content == "Consider the evidence. Apply the rule."

View file

@ -1,6 +1,10 @@
import pytest
import httpx
import respx
import litellm
class TestPerplexityWebSearch:
@ -19,3 +23,22 @@ class TestPerplexityWebSearch:
assert (
"web_search_options" in supported_params
), f"web_search_options should be supported for {model}"
def test_perplexity_401_maps_to_authentication_error(respx_mock: respx.MockRouter) -> None:
route = respx_mock.post("https://api.perplexity.ai/chat/completions").mock(
return_value=httpx.Response(401, json={"error": {"message": "invalid API key"}})
)
with pytest.raises(litellm.AuthenticationError) as exc_info:
litellm.completion(
model="perplexity/mistral-7b-instruct",
messages=[{"role": "user", "content": "hello"}],
api_key="bad-perplexity-key",
max_retries=0,
)
assert exc_info.value.status_code == 401
assert "PerplexityException" in str(exc_info.value)
assert "invalid API key" in str(exc_info.value)
assert route.calls.last.request.headers["Authorization"] == "Bearer bad-perplexity-key"

View file

@ -1,4 +1,5 @@
import json
from typing import Final
from unittest.mock import AsyncMock, Mock, patch
import httpx
@ -272,3 +273,29 @@ def test_sync_completion_handles_starting_status(mock_get_client, monkeypatch, r
assert json.loads(mock_client.post.call_args.kwargs["data"])["input"]["prompt"] == (
"<|begin▁of▁sentence|><|User|>Hi<|Assistant|>"
)
@pytest.mark.parametrize(
("model", "max_tokens_key"),
[
("replicate/vicuna-13b", "max_length"),
("replicate/flan-t5-xl", "max_length"),
("meta/codellama-13b", "max_tokens"),
("meta/llama-2-7b", "max_new_tokens"),
],
)
def test_map_openai_params_renames_openai_params_to_replicate_inputs(model: str, max_tokens_key: str):
mapped: Final = ReplicateConfig().map_openai_params(
non_default_params={"stream": True, "max_tokens": 64, "temperature": 0.2, "top_p": 0.9, "stop": ["\n"]},
optional_params={},
model=model,
drop_params=False,
)
assert mapped == {
"stream": True,
max_tokens_key: 64,
"temperature": 0.2,
"top_p": 0.9,
"stop_sequences": ["\n"],
}

View file

@ -14,7 +14,7 @@ import binascii
import json
import struct
from collections.abc import Mapping
from typing import AsyncIterator, Iterator
from typing import AsyncIterator, Final, Iterator
from unittest.mock import MagicMock
import httpx
@ -400,3 +400,33 @@ def test_get_complete_url_sagemaker_base_url_override_wins() -> None:
stream=False,
)
assert url == "https://my-private-endpoint.example.com/invocations"
def test_sagemaker_default_region_is_us_west_2(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
monkeypatch.delenv("AWS_REGION", raising=False)
monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False)
monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true")
url: Final = SagemakerChatConfig().get_complete_url(
api_base=None,
api_key=None,
model="my-endpoint",
optional_params={},
litellm_params={},
stream=False,
)
assert url == "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/my-endpoint/invocations"
def test_sagemaker_explicit_region_overrides_the_default() -> None:
url: Final = SagemakerChatConfig().get_complete_url(
api_base=None,
api_key=None,
model="my-endpoint",
optional_params={"aws_region_name": "us-east-1"},
litellm_params={},
stream=False,
)
assert url == "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/my-endpoint/invocations"

View file

@ -12,16 +12,26 @@ inflating TTFT and turning a steady provider stream into gap-then-burst delivery
import binascii
import json
import struct
from typing import AsyncIterator, Iterator
from typing import AsyncIterator, Final, Iterator
from unittest.mock import MagicMock
import httpx
import litellm
import pytest
import respx
from litellm.llms.sagemaker.common_utils import SagemakerError
from litellm.llms.sagemaker.completion.handler import SagemakerLLM
@pytest.fixture
def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
yield
litellm.in_memory_llm_clients_cache.flush_cache()
def _encode_header(name: str, value: str) -> bytes:
name_b = name.encode("utf-8")
value_b = value.encode("utf-8")
@ -293,3 +303,318 @@ def test_installed_botocore_signs_the_completion_request():
)
assert request.body == b'{"inputs": "ping"}'
assert "/us-west-2/sagemaker/aws4_request" in request.headers["Authorization"]
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.usefixtures("fake_provider_credentials")
async def test_acompletion_sagemaker_non_stream(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
monkeypatch.delenv("AWS_REGION", raising=False)
url: Final = (
"https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/"
"jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations"
)
route: Final = respx_mock.post(url).mock(
return_value=httpx.Response(200, json={"generated_text": "SageMaker reply"})
)
response: Final = await litellm.acompletion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[{"role": "user", "content": "hi"}],
temperature=0.2,
max_tokens=80,
cost_per_second=0.000420,
num_retries=0,
)
assert route.called
assert str(route.calls[0].request.url) == url
assert json.loads(route.calls[0].request.content) == {
"inputs": "hi",
"parameters": {"temperature": 0.2, "max_new_tokens": 80},
}
assert response.choices[0].message.content == "SageMaker reply"
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.usefixtures("fake_provider_credentials")
def test_completion_sagemaker_non_stream(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
monkeypatch.delenv("AWS_REGION", raising=False)
url: Final = (
"https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/"
"jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations"
)
route: Final = respx_mock.post(url).mock(
return_value=httpx.Response(200, json={"generated_text": "SageMaker reply"})
)
response: Final = litellm.completion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[{"role": "user", "content": "hi"}],
temperature=0.2,
max_tokens=80,
cost_per_second=0.000420,
num_retries=0,
)
assert route.called
assert str(route.calls[0].request.url) == url
assert json.loads(route.calls[0].request.content) == {
"inputs": "hi",
"parameters": {"temperature": 0.2, "max_new_tokens": 80},
}
assert response.choices[0].message.content == "SageMaker reply"
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.usefixtures("fake_provider_credentials")
def test_completion_sagemaker_non_stream_with_aws_params(respx_mock: respx.MockRouter) -> None:
url: Final = (
"https://runtime.sagemaker.us-west-5.amazonaws.com/endpoints/"
"jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations"
)
route: Final = respx_mock.post(url).mock(
return_value=httpx.Response(200, json={"generated_text": "SageMaker reply"})
)
response: Final = litellm.completion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[{"role": "user", "content": "hi"}],
temperature=0.2,
max_tokens=80,
cost_per_second=0.000420,
aws_access_key_id="unit-test-access-key",
aws_secret_access_key="unit-test-secret",
aws_region_name="us-west-5",
num_retries=0,
)
assert route.called
assert str(route.calls[0].request.url) == url
assert json.loads(route.calls[0].request.content) == {
"inputs": "hi",
"parameters": {"temperature": 0.2, "max_new_tokens": 80},
}
assert response.choices[0].message.content == "SageMaker reply"
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.usefixtures("fake_provider_credentials")
def test_completion_sagemaker_prompt_template_non_stream(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "known_tokenizer_config", {})
tokenizer_route: Final = respx_mock.get(
"https://huggingface.co/deepseek-ai/deepseek-coder-6.7b-instruct/raw/main/tokenizer_config.json"
).mock(
return_value=httpx.Response(
200,
json={
"bos_token": {"content": "<|begin▁of▁sentence|>"},
"eos_token": {"content": "<|EOT|>"},
"chat_template": (
"{{ bos_token }}You are a coding assistant\n"
"{% for message in messages %}### Instruction:\n{{ message['content'] }}\n{% endfor %}"
"### Response:"
),
},
)
)
url: Final = (
"https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/deepseek_coder_6.7_instruct/invocations"
)
route: Final = respx_mock.post(url).mock(
return_value=httpx.Response(200, json={"generated_text": "SageMaker reply"})
)
response: Final = litellm.completion(
model="sagemaker/deepseek_coder_6.7_instruct",
messages=[{"role": "user", "content": "hi"}],
temperature=0.2,
max_tokens=80,
hf_model_name="deepseek-ai/deepseek-coder-6.7b-instruct",
aws_region_name="us-east-1",
num_retries=0,
)
assert tokenizer_route.call_count == 1
assert json.loads(route.calls[0].request.content) == {
"inputs": "<|begin▁of▁sentence|>You are a coding assistant\n### Instruction:\nhi\n### Response:",
"parameters": {"temperature": 0.2, "max_new_tokens": 80},
}
assert response.choices[0].message.content == "SageMaker reply"
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.usefixtures("fake_provider_credentials")
@pytest.mark.parametrize(
("environment", "call_kwargs", "expected_region"),
[
({}, {}, "us-west-2"),
({"AWS_REGION_NAME": "us-east-1"}, {}, "us-east-1"),
({"AWS_REGION_NAME": "eu-west-1"}, {"aws_region_name": "us-east-1"}, "us-east-1"),
],
ids=["default_region", "environment_region", "config_region_beats_environment"],
)
def test_sagemaker_region_resolution(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
environment: dict[str, str],
call_kwargs: dict[str, str],
expected_region: str,
) -> None:
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
monkeypatch.delenv("AWS_REGION", raising=False)
for name, value in environment.items():
monkeypatch.setenv(name, value)
url: Final = f"https://runtime.sagemaker.{expected_region}.amazonaws.com/endpoints/mock-endpoint/invocations"
route: Final = respx_mock.post(url).mock(
return_value=httpx.Response(200, json={"generated_text": "SageMaker reply"})
)
litellm.completion(
model="sagemaker/mock-endpoint",
messages=[{"role": "user", "content": "Hello, world!"}],
num_retries=0,
**call_kwargs,
)
assert route.call_count == 1
assert str(route.calls[0].request.url) == url
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.usefixtures("fake_provider_credentials")
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_completion_sagemaker(
respx_mock: respx.MockRouter,
sync_mode: bool,
) -> None:
url: Final = (
"https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/"
"jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations"
)
route: Final = respx_mock.post(url).mock(
return_value=httpx.Response(200, json={"generated_text": "SageMaker reply"})
)
rate: Final = 0.000420
response: Final = (
litellm.completion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[{"role": "user", "content": "hi"}],
temperature=0.2,
max_tokens=80,
cost_per_second=rate,
aws_region_name="us-east-1",
num_retries=0,
)
if sync_mode
else await litellm.acompletion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[{"role": "user", "content": "hi"}],
temperature=0.2,
max_tokens=80,
cost_per_second=rate,
aws_region_name="us-east-1",
num_retries=0,
)
)
response._response_ms = 1000.0
assert route.called
assert response.choices[0].message.content == "SageMaker reply"
assert litellm.completion_cost(completion_response=response) == pytest.approx(rate)
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.usefixtures("fake_provider_credentials")
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_completion_sagemaker_stream(
respx_mock: respx.MockRouter,
sync_mode: bool,
) -> None:
url: Final = (
"https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/"
"jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations-response-stream"
)
route: Final = respx_mock.post(url).mock(
return_value=httpx.Response(200, content=_token_frame("hello") + _token_frame(" world"))
)
stream: Final = (
litellm.completion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[{"role": "user", "content": "hi - what is ur name"}],
stream=True,
temperature=0.2,
max_tokens=80,
cost_per_second=0.000420,
aws_region_name="us-east-1",
num_retries=0,
)
if sync_mode
else await litellm.acompletion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[{"role": "user", "content": "hi - what is ur name"}],
stream=True,
temperature=0.2,
max_tokens=80,
cost_per_second=0.000420,
aws_region_name="us-east-1",
num_retries=0,
)
)
chunks: Final = tuple(stream) if sync_mode else tuple([chunk async for chunk in stream])
assert json.loads(route.calls[0].request.content) == {
"inputs": "hi - what is ur name",
"parameters": {"temperature": 0.2, "max_new_tokens": 80},
"stream": True,
}
assert chunks[0].choices[0].delta.role == "assistant"
assert tuple(chunk.choices[0].delta.content for chunk in chunks) == ("hello", " world", None)
assert tuple(chunk.choices[0].finish_reason for chunk in chunks) == (None, None, "stop")
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.usefixtures("fake_provider_credentials")
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_completion_sagemaker_streaming_bad_request(
respx_mock: respx.MockRouter,
sync_mode: bool,
) -> None:
url: Final = (
"https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/"
"jumpstart-dft-hf-textgeneration1-mp-20240815-185614/invocations-response-stream"
)
route: Final = respx_mock.post(url).mock(
return_value=httpx.Response(
400,
json={"Message": "The request payload failed validation.", "__type": "ValidationException"},
)
)
with pytest.raises(litellm.BadRequestError) as exc_info:
litellm.completion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[{"role": "user", "content": "hi"}],
stream=True,
aws_region_name="us-east-1",
num_retries=0,
) if sync_mode else await litellm.acompletion(
model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
messages=[{"role": "user", "content": "hi"}],
stream=True,
aws_region_name="us-east-1",
num_retries=0,
)
assert route.called
assert exc_info.value.status_code == 400
assert exc_info.value.llm_provider == "sagemaker"

View file

@ -2,14 +2,57 @@
Unit tests for SageMaker Nova transformation config.
"""
import binascii
import json
import pytest
import struct
from typing import Final
import httpx
import litellm
import pytest
import respx
from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder
from litellm.llms.sagemaker.nova.transformation import SagemakerNovaConfig
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse
from litellm.utils import convert_to_model_response_object
def _event_frame(payload: bytes) -> bytes:
headers: Final = {
":event-type": "PayloadPart",
":content-type": "application/json",
":message-type": "event",
}
encoded_headers: Final = b"".join(
struct.pack("B", len(name.encode()))
+ name.encode()
+ struct.pack("B", 7)
+ struct.pack(">H", len(value.encode()))
+ value.encode()
for name, value in headers.items()
)
prelude: Final = struct.pack(">I", 16 + len(encoded_headers) + len(payload)) + struct.pack(
">I", len(encoded_headers)
)
prefix: Final = prelude + struct.pack(">I", binascii.crc32(prelude) & 0xFFFFFFFF)
message: Final = prefix + encoded_headers + payload
return message + struct.pack(">I", binascii.crc32(message) & 0xFFFFFFFF)
def _nova_stream_frame(content: str, finish_reason: str | None) -> bytes:
chunk: Final = {
"id": "chatcmpl-nova-stream",
"object": "chat.completion.chunk",
"created": 1700000000,
"model": "nova-endpoint",
"choices": [{"index": 0, "delta": {"content": content}, "finish_reason": finish_reason}],
"usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3},
}
return _event_frame(f"data: {json.dumps(chunk)}\n\n".encode())
class TestSagemakerNovaConfig:
def setup_method(self):
self.config = SagemakerNovaConfig()
@ -394,6 +437,23 @@ class TestSagemakerNovaTransformRequest:
assert "messages" in request
assert request["temperature"] == 0.7
def test_should_preserve_system_and_multi_turn_messages_in_request(self):
messages: Final[list[AllMessageValues]] = [
{"role": "system", "content": "Remember the user's name"},
{"role": "user", "content": "My name is Alice"},
{"role": "assistant", "content": "I will remember"},
{"role": "user", "content": "What is my name?"},
]
request = self.config.transform_request(
model="my-nova-endpoint",
messages=messages,
optional_params={},
litellm_params={},
headers={},
)
assert request == {"messages": messages}
def test_should_include_all_nova_params_in_request(self):
"""Nova-specific params should appear in the request body."""
request = self.config.transform_request(
@ -411,3 +471,82 @@ class TestSagemakerNovaTransformRequest:
assert request["top_k"] == 40
assert request["max_tokens"] == 512
assert request["reasoning_effort"] == "low"
NOVA_URL: Final = "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/nova-endpoint/invocations"
@pytest.mark.usefixtures("fake_provider_credentials")
def test_nova_completion_parses_content_finish_reason_and_usage(respx_mock: respx.MockRouter) -> None:
route: Final = respx_mock.post(NOVA_URL).mock(
return_value=httpx.Response(
200,
json={
"id": "chatcmpl-nova-single",
"object": "chat.completion",
"created": 1700000000,
"model": "nova-endpoint",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Four"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 14, "completion_tokens": 2, "total_tokens": 16},
},
)
)
response: Final = litellm.completion(
model="sagemaker_nova/nova-endpoint",
messages=[{"role": "user", "content": "What is 2+2? Reply in one word."}],
max_tokens=32,
temperature=0.1,
num_retries=0,
)
assert json.loads(route.calls[0].request.content) == {
"messages": [{"role": "user", "content": "What is 2+2? Reply in one word."}],
"max_tokens": 32,
"temperature": 0.1,
}
assert response.choices[0].message.content == "Four"
assert response.choices[0].finish_reason == "stop"
assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (
14,
2,
16,
)
@pytest.mark.usefixtures("fake_provider_credentials")
def test_nova_streaming_sends_stream_flag_and_parses_chunks_and_usage(respx_mock: respx.MockRouter) -> None:
route: Final = respx_mock.post(f"{NOVA_URL}-response-stream").mock(
return_value=httpx.Response(
200,
content=_nova_stream_frame("Nova ", None) + _nova_stream_frame("stream", "stop"),
)
)
stream: Final = litellm.completion(
model="sagemaker_nova/nova-endpoint",
messages=[{"role": "user", "content": "Count from 1 to 5."}],
max_tokens=64,
stream=True,
stream_options={"include_usage": True},
num_retries=0,
)
chunks: Final = tuple(stream)
assert json.loads(route.calls[0].request.content) == {
"messages": [{"role": "user", "content": "Count from 1 to 5."}],
"max_tokens": 64,
"stream": True,
"stream_options": {"include_usage": True},
}
assert tuple(chunk.choices[0].delta.content for chunk in chunks) == ("Nova ", "stream", None, None)
assert tuple(chunk.choices[0].finish_reason for chunk in chunks) == (None, None, "stop", None)
assert (chunks[-1].usage.prompt_tokens, chunks[-1].usage.completion_tokens, chunks[-1].usage.total_tokens) == (
2,
1,
3,
)

View file

@ -2,7 +2,9 @@ from unittest.mock import AsyncMock, Mock
import httpx
import pytest
import respx
import litellm
from litellm.llms.predibase.chat.handler import PredibaseChatCompletion
from litellm.llms.predibase.chat.transformation import PredibaseConfig
from litellm.llms.predibase.common_utils import PredibaseError
@ -115,6 +117,24 @@ def test_predibase_get_complete_url_with_tenant_id_key():
assert url.endswith("/generate")
def test_predibase_authentication_error_does_not_expose_the_api_key(respx_mock: respx.MockRouter) -> None:
route = respx_mock.post(
"https://serving.app.predibase.com/c4768f95/deployments/v2/llms/llama-3-8b-instruct/generate"
).mock(return_value=httpx.Response(401, json={"message": "invalid API key"}))
with pytest.raises(litellm.AuthenticationError) as exc_info:
litellm.completion(
model="predibase/llama-3-8b-instruct",
messages=[{"role": "user", "content": "What is the meaning of life?"}],
predibase_tenant_id="c4768f95",
api_key="hf-rawapikey",
max_retries=0,
)
assert "hf-rawapikey" not in str(exc_info.value)
assert route.calls.last.request.headers["Authorization"] == "Bearer hf-rawapikey"
def test_predibase_transform_response_success_best_of(monkeypatch):
config = PredibaseConfig()
logging_obj = Mock()

View file

@ -0,0 +1,36 @@
import json
from typing import Final
import httpx
import respx
import litellm
def test_together_ai_embedding_uses_together_embeddings_endpoint(respx_mock: respx.MockRouter) -> None:
route: Final = respx_mock.post("https://api.together.ai/v1/embeddings").mock(
return_value=httpx.Response(
200,
json={
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.3, 0.4]}],
"model": "BAAI/bge-large-en-v1.5",
"usage": {"prompt_tokens": 1, "total_tokens": 1},
},
)
)
response: Final = litellm.embedding(
model="together_ai/BAAI/bge-large-en-v1.5",
input=["hello"],
api_key="test-together-key",
)
assert str(route.calls.last.request.url) == "https://api.together.ai/v1/embeddings"
assert route.calls.last.request.headers["Authorization"] == "Bearer test-together-key"
assert json.loads(route.calls.last.request.read()) == {
"model": "BAAI/bge-large-en-v1.5",
"input": ["hello"],
}
assert response.data[0]["embedding"] == [0.3, 0.4]
assert (response.usage.prompt_tokens, response.usage.total_tokens) == (1, 1)

View file

@ -1,13 +1,19 @@
import json
from typing import Final
import httpx
import pytest
import respx
import litellm
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler import GoogleBatchEmbeddings
FILES_URI: Final = "https://generativelanguage.googleapis.com/v1beta/files/clip123"
FILE_METADATA_URL: Final = "https://generativelanguage.googleapis.com/v1beta/files/clip123"
GEMINI_BATCH_EMBEDDINGS_URL: Final = (
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents"
)
def _files_api_metadata(request: httpx.Request) -> httpx.Response:
@ -26,6 +32,51 @@ def test_resolve_file_references_fetches_the_file_metadata_for_both_reference_fo
assert resolved == {reference: {"mime_type": "video/mp4", "uri": FILES_URI}}
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.parametrize("embedding_input", ["good morning from litellm", ["good morning from litellm"]])
@pytest.mark.asyncio
async def test_gemini_embedding_sends_batch_request_and_parses_vectors(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
sync_mode: bool,
embedding_input: str | list[str],
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post(GEMINI_BATCH_EMBEDDINGS_URL).mock(
return_value=httpx.Response(
200,
json={"embeddings": [{"values": [0.1, 0.2, 0.3]}]},
)
)
response: Final = (
litellm.embedding(model="gemini/gemini-embedding-001", input=embedding_input, api_key="gemini-test-key")
if sync_mode
else await litellm.aembedding(
model="gemini/gemini-embedding-001", input=embedding_input, api_key="gemini-test-key"
)
)
assert route.call_count == 1
assert json.loads(route.calls[0].request.content) == {
"requests": [
{
"model": "models/gemini-embedding-001",
"content": {"parts": [{"text": "good morning from litellm"}]},
}
]
}
assert [item["embedding"] for item in response.data] == [[0.1, 0.2, 0.3]]
local_token_count: Final = litellm.token_counter(
model="gemini-embedding-001", text="good morning from litellm"
)
assert local_token_count > 0
assert (response.usage.prompt_tokens, response.usage.total_tokens) == (
local_token_count,
local_token_count,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("reference", ["files/clip123", FILES_URI])
async def test_async_resolve_file_references_fetches_the_file_metadata_for_both_reference_forms(

File diff suppressed because it is too large Load diff

View file

@ -1,7 +1,10 @@
import asyncio
import json
import logging
from typing import Final
from unittest.mock import MagicMock, call, patch
import google.auth.exceptions
import pytest
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
@ -2180,3 +2183,19 @@ class TestVertexBase:
assert token == "cached-token"
assert not mock_get_lock.called, "Fast path should not acquire lock"
def test_get_access_token_reraises_malformed_service_account_credentials_without_caching(caplog):
vertex_base: Final = VertexBase()
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
with pytest.raises(google.auth.exceptions.MalformedError):
vertex_base.get_access_token(
credentials=json.dumps({"type": "service_account", "project_id": "project-1"}),
project_id="project-1",
)
assert vertex_base._credentials_project_mapping == {}
assert [record.getMessage().split(". Error: ")[0] for record in caplog.records] == [
"Failed to load vertex credentials. Check to see if credentials containing partial/invalid information"
]

View file

@ -3,8 +3,12 @@ import copy
import json
from typing import Final
import httpx
import pytest
import respx
from google.oauth2.credentials import Credentials
import litellm
from litellm.anthropic_beta_headers_manager import (
update_headers_with_filtered_beta,
update_request_with_filtered_beta,
@ -489,6 +493,64 @@ def test_vertex_ai_anthropic_no_extra_headers_unchanged():
assert "anthropic-beta" not in headers
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_vertex_ai_anthropic_prompt_caching_does_not_add_beta_header(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
sync_mode: bool,
) -> None:
from litellm.main import vertex_partner_models_chat_completion
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setitem(
vertex_partner_models_chat_completion._credentials_project_mapping,
("test-vertex-credentials", "test-project"),
(Credentials(token="test-vertex-token"), "test-project"),
)
route: Final = respx_mock.post(
"https://us-east5-aiplatform.googleapis.com/v1/projects/test-project/locations/us-east5/"
"publishers/anthropic/models/claude-3-5-sonnet-v2@20241022:rawPredict"
).mock(
return_value=httpx.Response(
200,
json={
"id": "msg-cache",
"type": "message",
"role": "assistant",
"model": "claude-3-5-sonnet-v2-20241022",
"content": [{"type": "text", "text": "cached"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 4, "output_tokens": 1},
},
)
)
cached_block: Final = {
"type": "text",
"text": "Cache this request",
"cache_control": {"type": "ephemeral"},
}
call_params: Final = {
"model": "vertex_ai/claude-3-5-sonnet-v2@20241022",
"messages": [{"role": "user", "content": [cached_block]}],
"vertex_project": "test-project",
"vertex_location": "us-east5",
"vertex_credentials": "test-vertex-credentials",
"max_retries": 0,
}
if sync_mode:
litellm.completion(**call_params)
else:
await litellm.acompletion(**call_params)
assert route.call_count == 1
request: Final = route.calls[0].request
assert json.loads(request.content)["messages"] == [{"role": "user", "content": [cached_block]}]
assert "anthropic-beta" not in request.headers
assert "anthropic_beta" not in json.loads(request.content)
def test_vertex_ai_partner_models_anthropic_remove_prompt_caching_scope_beta_header():
"""
Test that remove_unsupported_beta correctly filters out prompt-caching-scope-2026-01-05

View file

@ -3,10 +3,17 @@ Test that VertexBase subclasses (PartnerModels, Gemma, ModelGarden) reuse
cached credentials instead of creating a new VertexLLM instance on every request.
"""
import json
import sys
from typing import Final
from unittest.mock import MagicMock, patch
import httpx
import pytest
import responses
import respx
import litellm
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
VertexAIPartnerModels,
@ -256,3 +263,60 @@ class TestModelGardenCredentialReuse:
project_id="test-project",
custom_llm_provider="vertex_ai",
)
@respx.mock
def test_sync_mistral_partner_completion_posts_to_raw_predict_without_a_client():
route: Final = respx.post(
"https://us-central1-aiplatform.googleapis.com/v1/projects/partner-project/locations/us-central1"
"/publishers/mistralai/models/mistral-small-2503@001:rawPredict"
).mock(
return_value=httpx.Response(
200,
json={
"id": "chatcmpl-vertex-mistral",
"object": "chat.completion",
"created": 1,
"model": "mistral-small-2503",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {"role": "assistant", "content": "I am Litellm Bot."},
}
],
"usage": {"prompt_tokens": 21, "completion_tokens": 5, "total_tokens": 26},
},
)
)
messages: Final = [
{"role": "system", "content": "Your name is Litellm Bot, you are a helpful assistant"},
{"role": "user", "content": "Hello, what is your name?"},
]
with responses.RequestsMock() as google_token_endpoint:
token_route: Final = google_token_endpoint.post(
"https://oauth2.googleapis.com/token",
json={"access_token": "partner-access-token", "expires_in": 3600, "token_type": "Bearer"},
)
response: Final = litellm.completion(
model="vertex_ai/mistral-small-2503@001",
messages=messages,
vertex_project="partner-project",
vertex_location="us-central1",
vertex_credentials=json.dumps(
{
"type": "authorized_user",
"client_id": "partner-client-id",
"client_secret": "partner-client-secret",
"refresh_token": "partner-refresh-token",
}
),
)
assert token_route.call_count == 1
request: Final = route.calls[0].request
assert request.headers["authorization"] == "Bearer partner-access-token"
assert json.loads(request.content) == {"model": "mistral-small-2503", "messages": messages, "stream": False}
assert response.choices[0].message.content == "I am Litellm Bot."
assert response.usage.total_tokens == 26

View file

@ -1,11 +1,16 @@
import json
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
import respx
from pydantic import ValidationError
import litellm
from litellm.llms.watsonx.embed.transformation import IBMWatsonXEmbeddingConfig
from litellm.types.utils import EmbeddingResponse
@ -132,3 +137,31 @@ def test_transform_embedding_response_shape_errors_do_not_echo_the_payload(paylo
_transform(payload)
assert "leaked payload text" not in str(exc_info.value)
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_watsonx_embedding_sends_the_bearer_token_to_the_regional_endpoint(
sync_mode: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("WATSONX_API_KEY", "mock-api-key")
monkeypatch.setenv("WATSONX_TOKEN", "mock-watsonx-token")
monkeypatch.setenv("WATSONX_API_BASE", "https://us-south.ml.cloud.ibm.com")
monkeypatch.setenv("WATSONX_PROJECT_ID", "mock-project-id")
route: Final = respx_mock.post(url__startswith="https://us-south.ml.cloud.ibm.com/ml/v1/text/embeddings").respond(
json={
"model_id": "ibm/slate-30m-english-rtrvr",
"created_at": "2024-01-01T00:00:00.00Z",
"results": [{"embedding": [0.0, 1.0]}],
"input_token_count": 8,
}
)
request: Final = {"model": "watsonx/ibm/slate-30m-english-rtrvr", "input": ["good morning from litellm"]}
response: Final = litellm.embedding(**request) if sync_mode else await litellm.aembedding(**request)
assert route.call_count == 1
assert route.calls.last.request.headers["Authorization"] == "Bearer mock-watsonx-token"
assert json.loads(route.calls.last.request.content)["inputs"] == ["good morning from litellm"]
assert response.data == [{"object": "embedding", "index": 0, "embedding": [0.0, 1.0]}]
assert (response.usage.prompt_tokens, response.usage.total_tokens) == (8, 8)

View file

@ -9,7 +9,7 @@ import pytest
import litellm
from litellm import completion, embedding
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@ -370,6 +370,105 @@ def test_watsonx_tool_choice(sync_mode: bool, monkeypatch: pytest.MonkeyPatch) -
assert len(response.choices) == 1
async def test_watsonx_async_streamed_tool_choice_reaches_the_stream_endpoint(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("WATSONX_TOKEN", "mock-watsonx-token")
monkeypatch.setenv("WATSONX_API_BASE", "https://us-south.ml.cloud.ibm.com")
monkeypatch.setenv("WATSONX_PROJECT_ID", "mock-project-id")
stream_chunks: Final = (
{
"id": "chatcmpl-watsonx",
"object": "chat.completion.chunk",
"created": 1760000000,
"model": "ibm/granite-3-8b-instruct",
"choices": [
{
"index": 0,
"delta": {
"role": "assistant",
"tool_calls": [
{
"index": 0,
"id": "chatcmpl-tool-weather",
"type": "function",
"function": {"name": "get_current_weather", "arguments": ""},
}
],
},
"finish_reason": None,
}
],
},
{
"id": "chatcmpl-watsonx",
"object": "chat.completion.chunk",
"created": 1760000000,
"model": "ibm/granite-3-8b-instruct",
"choices": [
{
"index": 0,
"delta": {
"tool_calls": [{"index": 0, "function": {"arguments": '{"location": "San Francisco, CA"}'}}]
},
"finish_reason": None,
}
],
},
{
"id": "chatcmpl-watsonx",
"object": "chat.completion.chunk",
"created": 1760000000,
"model": "ibm/granite-3-8b-instruct",
"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}],
},
)
body: Final = "".join(
f"id: {index}\nevent: message\ndata: {json.dumps(chunk)}\n\n" for index, chunk in enumerate(stream_chunks)
)
requests: Final[list[httpx.Request]] = []
def handle_request(request: httpx.Request) -> httpx.Response:
requests.append(request)
return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(handle_request))
response: Final = await litellm.acompletion(
model="watsonx/ibm/granite-3-8b-instruct",
messages=[{"role": "user", "content": "What is the weather in San Francisco?"}],
tools=[
{
"type": "function",
"function": {
"name": "get_current_weather",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
},
}
],
tool_choice="auto",
stream=True,
client=client,
)
chunks: Final = [chunk async for chunk in response]
tool_calls: Final = tuple(
chunk.choices[0].delta.tool_calls[0] for chunk in chunks if chunk.choices[0].delta.tool_calls
)
request_body: Final = cast(dict[str, object], json.loads(requests[0].content))
assert len(requests) == 1
assert requests[0].url.path == "/ml/v1/text/chat_stream"
assert request_body["tool_choice_option"] == "auto"
assert request_body["stream"] is True
assert tool_calls[0].id == "chatcmpl-tool-weather"
assert tool_calls[0].function.name == "get_current_weather"
assert json.loads("".join(tool_call.function.arguments or "" for tool_call in tool_calls)) == {
"location": "San Francisco, CA"
}
assert chunks[-1].choices[0].finish_reason == "tool_calls"
@pytest.mark.usefixtures("watsonx_env_vars", "_vcr_outcome_gate", "setup_and_teardown")
def test_watsonx_chat_completions_endpoint_space_id(monkeypatch, watsonx_chat_completion_call):
my_fake_space_id = "xxx-xxx-xxx-xxx-xxx"

View file

@ -1,3 +1,5 @@
import json
import asyncio
import datetime
import time
from collections.abc import Mapping
@ -6,12 +8,15 @@ from typing import Final, Literal, cast
import httpx
import pytest
import respx
from openai.types.completion_usage import CompletionUsage
from pydantic import BaseModel
import litellm
from litellm import TranscriptionResponse, model_cost
from litellm.cost_calculator import BaseTokenUsageProcessor, RealtimeAPITokenUsageProcessor, ResponsesWebSocketTokenUsageProcessor, batch_cost_calculator, completion_cost, cost_per_token, handle_realtime_stream_cost_calculation, response_cost_calculator
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
convert_to_model_response_object,
@ -27,9 +32,41 @@ from litellm.types.utils import CallTypes, Choices, CompletionTokensDetailsWrapp
from litellm.types.videos.main import VideoObject
from litellm.utils import supports_prompt_caching
import os
import threading
from unittest.mock import MagicMock, patch
class ResponseCostCapture(CustomLogger):
def __init__(self) -> None:
self.success_cost: float | None = None
self.success_event: Final = threading.Event()
self.failure_cost: float | None = None
def log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
response_cost: Final = kwargs.get("response_cost")
self.success_cost = float(response_cost) if isinstance(response_cost, (int, float)) else None
self.success_event.set()
def log_failure_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self.failure_cost = float(kwargs["response_cost"])
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
response_cost: Final = kwargs.get("response_cost")
self.success_cost = float(response_cost) if isinstance(response_cost, (int, float)) else None
self.success_event.set()
async def async_log_failure_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self.failure_cost = float(kwargs["response_cost"])
@pytest.fixture
def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
@ -7923,6 +7960,380 @@ def test_cost_calculator_with_custom_pricing():
assert resp._hidden_params["response_cost"] > 0
@pytest.fixture
def isolated_pricing_registry(local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "open_ai_chat_completion_models", set(litellm.open_ai_chat_completion_models))
monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider))
_CUSTOM_PRICING_MODEL: Final = "migration-custom-pricing-chat"
_CUSTOM_PRICING_RATES: Final = MappingProxyType(
{"input_cost_per_token": 0.0000008, "output_cost_per_token": 0.0000032}
)
_CUSTOM_PRICING_EXPECTED_COST: Final = 7 * 0.0000008 + 4 * 0.0000032
def _custom_pricing_chunk(choices: list[dict[str, object]], usage: dict[str, int] | None = None) -> str:
chunk: Final = {
"id": "chatcmpl-custom-pricing",
"object": "chat.completion.chunk",
"created": 1,
"model": _CUSTOM_PRICING_MODEL,
"choices": choices,
**({"usage": usage} if usage is not None else {}),
}
return f"data: {json.dumps(chunk)}\n\n"
def _custom_pricing_upstream(request: httpx.Request) -> httpx.Response:
if json.loads(request.content).get("stream") is True:
body: Final = "".join(
(
_custom_pricing_chunk(
[{"index": 0, "delta": {"role": "assistant", "content": "I am ready to help."}, "finish_reason": None}]
),
_custom_pricing_chunk([{"index": 0, "delta": {}, "finish_reason": "stop"}]),
_custom_pricing_chunk([], {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11}),
"data: [DONE]\n\n",
)
)
return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
return httpx.Response(
200,
json={
"id": "chatcmpl-custom-pricing",
"object": "chat.completion",
"created": 1,
"model": _CUSTOM_PRICING_MODEL,
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "I am ready to help."}, "finish_reason": "stop"}
],
"usage": {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11},
},
)
@pytest.mark.parametrize("async_mode", [False, True])
@pytest.mark.asyncio
@pytest.mark.usefixtures("isolated_pricing_registry")
@respx.mock
async def test_custom_pricing_matches_response_usage(async_mode: bool) -> None:
route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_custom_pricing_upstream)
request: Final = {
"model": f"openai/{_CUSTOM_PRICING_MODEL}",
"messages": [{"role": "user", "content": "How are you?"}],
"api_key": "sk-test",
**_CUSTOM_PRICING_RATES,
}
response: Final = await litellm.acompletion(**request) if async_mode else litellm.completion(**request)
assert len(route.calls) == 1
assert response._hidden_params["response_cost"] == pytest.approx(_CUSTOM_PRICING_EXPECTED_COST)
@pytest.mark.usefixtures("isolated_pricing_registry")
@respx.mock
def test_streamed_custom_pricing_cost_matches_non_streamed_cost() -> None:
respx.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_custom_pricing_upstream)
request: Final = {
"model": f"openai/{_CUSTOM_PRICING_MODEL}",
"messages": [{"role": "user", "content": "Hi"}],
"api_key": "sk-test",
**_CUSTOM_PRICING_RATES,
}
non_streamed: Final = litellm.completion(**request)
rebuilt: Final = litellm.stream_chunk_builder(
list(litellm.completion(**request, stream=True, stream_options={"include_usage": True})),
messages=request["messages"],
)
streamed_cost: Final = completion_cost(rebuilt, custom_cost_per_token=_CUSTOM_PRICING_RATES)
assert rebuilt.choices[0].message.content == non_streamed.choices[0].message.content == "I am ready to help."
assert non_streamed._hidden_params["response_cost"] == pytest.approx(_CUSTOM_PRICING_EXPECTED_COST)
assert streamed_cost == pytest.approx(non_streamed._hidden_params["response_cost"])
@pytest.mark.parametrize(
("model", "provider"),
[
("fireworks_ai/accounts/fireworks/models/custom-migration-model", "fireworks_ai"),
("vertex_ai/meta/llama-custom-migration-model", "vertex_ai"),
("azure/custom-migration-model", "azure"),
("mistral/mistral-large-latest", "mistral"),
],
)
def test_custom_cost_uses_test_owned_rates_across_providers(model: str, provider: str):
response: Final = litellm.ModelResponse(
model=model,
choices=[],
usage=litellm.Usage(prompt_tokens=9, completion_tokens=4, total_tokens=13),
)
rates: Final = MappingProxyType({"input_cost_per_token": 0.002, "output_cost_per_token": 0.004})
cost: Final = completion_cost(
completion_response=response,
model=model,
custom_llm_provider=provider,
custom_cost_per_token=rates,
)
assert cost == pytest.approx(0.034)
@respx.mock
def test_fireworks_completion_cost_uses_registered_rate(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
model: Final = "fireworks_ai/accounts/fireworks/models/migration-model"
monkeypatch.setitem(
litellm.model_cost,
model,
{"input_cost_per_token": 0.001, "output_cost_per_token": 0.002, "litellm_provider": "fireworks_ai", "mode": "chat"},
)
respx.post("https://api.fireworks.ai/inference/v1/chat/completions").mock(
return_value=httpx.Response(
200,
json={
"id": "chatcmpl-fireworks",
"object": "chat.completion",
"created": 1,
"model": "accounts/fireworks/models/migration-model",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "Going great"}, "finish_reason": "stop"}
],
"usage": {"prompt_tokens": 8, "completion_tokens": 5, "total_tokens": 13},
},
)
)
response: Final = litellm.completion(
model=model, messages=[{"role": "user", "content": "Hey, how's it going?"}], api_key="fw-test"
)
assert completion_cost(completion_response=response) == pytest.approx(0.018)
def test_azure_embedding_cost_uses_registered_base_model_rate(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setitem(
litellm.model_cost,
"azure/migration-embedding-model",
{
"input_cost_per_token": 0.002,
"output_cost_per_token": 0.0,
"litellm_provider": "azure",
"mode": "embedding",
},
)
cost: Final = response_cost_calculator(
response_object=EmbeddingResponse(
model="text-embedding-3-small",
data=[{"embedding": [1, 2, 3]}],
usage=Usage(prompt_tokens=10, completion_tokens=10),
),
model="text-embedding-3-small",
custom_llm_provider=None,
base_model="azure/migration-embedding-model",
call_type="aembedding",
optional_params={},
custom_pricing=False,
prompt="Hello, world!",
)
assert cost == pytest.approx(0.02)
def test_azure_common_deployment_uses_registered_base_model_pricing(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setitem(
litellm.model_cost,
"azure/migration-pricing-model",
{
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
"litellm_provider": "azure",
"mode": "chat",
},
)
response: Final = ModelResponse(
id="chatcmpl-migration",
choices=[Choices(finish_reason="stop", index=0, message=Message(content="done", role="assistant"))],
created=1,
model="migration-deployment",
object="chat.completion",
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
)
cost: Final = response_cost_calculator(
response_object=response,
model="migration-deployment",
custom_llm_provider="azure",
call_type="acompletion",
optional_params={},
base_model="azure/migration-pricing-model",
)
assert cost == pytest.approx(0.02)
def test_azure_tts_cost_uses_test_owned_character_rate(monkeypatch: pytest.MonkeyPatch) -> None:
model: Final = "migration-tts-model"
monkeypatch.setitem(
litellm.model_cost,
f"azure/{model}",
{
"input_cost_per_character": 0.01,
"output_cost_per_character": 0.0,
"litellm_provider": "azure",
"mode": "audio_speech",
},
)
cost: Final = response_cost_calculator(
response_object=TranscriptionResponse(text=""),
model=model,
custom_llm_provider="azure",
call_type="aspeech",
optional_params={},
prompt="hello",
)
assert cost == pytest.approx(0.05)
def test_moderation_cost_uses_test_owned_token_rate() -> None:
prompt: Final = "moderation input"
rates: Final = MappingProxyType(
{"input_cost_per_token": 0.001, "output_cost_per_token": 0.0}
)
cost: Final = completion_cost(
completion_response=None,
model="omni-moderation-latest",
custom_llm_provider="openai",
call_type="moderation",
prompt=prompt,
custom_cost_per_token=rates,
)
expected_cost: Final = litellm.token_counter(model="omni-moderation-latest", text=prompt) * 0.001
assert cost == pytest.approx(expected_cost)
def test_vertex_predictive_cost_uses_test_owned_token_rate() -> None:
model: Final = "vertex_ai/migration-llama"
messages: Final = [{"role": "user", "content": "price control"}]
rates: Final = MappingProxyType(
{"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}
)
cost: Final = completion_cost(
model=model,
messages=messages,
custom_llm_provider="vertex_ai",
custom_cost_per_token=rates,
)
expected_cost: Final = (
litellm.token_counter(model=model, messages=messages) * 0.001
)
assert cost == pytest.approx(expected_cost)
@pytest.mark.parametrize("async_mode", [False, True])
@pytest.mark.asyncio
@pytest.mark.usefixtures("isolated_pricing_registry")
@respx.mock
async def test_custom_pricing_cost_reaches_success_callback(
async_mode: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
model: Final = "migration-callback-pricing-model"
litellm.register_model(
{
model: {
"input_cost_per_token": 0.0001,
"output_cost_per_token": 0.0002,
"litellm_provider": "openai",
"mode": "chat",
}
},
persist_across_reloads=False,
)
callback: Final = ResponseCostCapture()
if async_mode:
monkeypatch.setattr(litellm, "_async_success_callback", [callback])
else:
monkeypatch.setattr(litellm, "success_callback", [callback])
response_data: Final = {
"id": "chatcmpl-priced",
"object": "chat.completion",
"created": 1,
"model": model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "response"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 8, "completion_tokens": 5, "total_tokens": 13},
}
route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock(
return_value=httpx.Response(200, json=response_data)
)
request: Final = {
"model": model,
"messages": [{"role": "user", "content": "price control"}],
"api_key": "test-openai-key",
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
}
response: Final = (
await litellm.acompletion(**request) if async_mode else litellm.completion(**request)
)
await GLOBAL_LOGGING_WORKER.flush()
assert await asyncio.to_thread(callback.success_event.wait, 10.0)
assert len(route.calls) == 1
assert response.choices[0].message.content == "response"
assert callback.success_cost == pytest.approx(0.018)
@pytest.mark.parametrize("async_mode", [False, True])
@pytest.mark.asyncio
async def test_failure_completion_cost_reaches_failure_callback(
async_mode: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
callback: Final = ResponseCostCapture()
monkeypatch.setattr(litellm, "success_callback", [callback])
monkeypatch.setattr(litellm, "_async_success_callback", [callback])
monkeypatch.setattr(litellm, "failure_callback", [callback])
monkeypatch.setattr(litellm, "_async_failure_callback", [callback])
failure: Final = litellm.InternalServerError(
message="scripted failure",
llm_provider="openai",
model="gpt-6-sol",
)
if async_mode:
with pytest.raises(litellm.InternalServerError):
await litellm.acompletion(
model="gpt-6-sol",
messages=[{"role": "user", "content": "fail"}],
mock_response=failure,
)
else:
with pytest.raises(litellm.InternalServerError):
litellm.completion(
model="gpt-6-sol",
messages=[{"role": "user", "content": "fail"}],
mock_response=failure,
)
assert callback.failure_cost == 0.0
@pytest.fixture
def model_item():
return {

View file

@ -0,0 +1,17 @@
from pathlib import Path
from typing import Final
_NETWORK_FETCH_TOKENS: Final = ("curl", "wget", "fetch", "http://", "https://")
def test_container_build_no_network_fetch() -> None:
dockerfile: Final = (Path(__file__).resolve().parents[2] / "Dockerfile").read_text()
startup_lines: Final = tuple(
line.strip() for line in dockerfile.splitlines() if line.strip().upper().startswith(("CMD", "ENTRYPOINT"))
)
fetching_lines: Final = tuple(
line for line in startup_lines if any(token in line.lower() for token in _NETWORK_FETCH_TOKENS)
)
assert startup_lines != ()
assert fetching_lines == ()

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