mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
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:
parent
3746056caf
commit
a2e334febe
104 changed files with 9016 additions and 17385 deletions
|
|
@ -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
|
||||
|
|
|
|||
13
.github/ci-coverage-allowlist.yml
vendored
13
.github/ci-coverage-allowlist.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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=[],
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
125
tests/e2e/llm_translation/test_sagemaker_nova_e2e.py
Normal file
125
tests/e2e/llm_translation/test_sagemaker_nova_e2e.py
Normal 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
|
||||
116
tests/e2e/llm_translation/test_vertex_capabilities_e2e.py
Normal file
116
tests/e2e/llm_translation/test_vertex_capabilities_e2e.py
Normal 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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"]),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
58
tests/integration/providers/test_lm_studio_chat_wire.py
Normal file
58
tests/integration/providers/test_lm_studio_chat_wire.py
Normal 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,
|
||||
}
|
||||
|
|
@ -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"}
|
||||
39
tests/integration/security/test_air_gapped_boot.py
Normal file
39
tests/integration/security/test_air_gapped_boot.py
Normal 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"}
|
||||
|
|
@ -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
|
||||
|
|
@ -1,3 +0,0 @@
|
|||
{
|
||||
"gpt-3.5-turbo": 7.7e-05
|
||||
}
|
||||
|
|
@ -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,
|
||||
),
|
||||
]
|
||||
|
|
@ -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
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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
|
|
@ -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"},
|
||||
# ]
|
||||
# )
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
|
@ -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"
|
||||
|
|
@ -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"
|
||||
)
|
||||
|
|
@ -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])
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
@ -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}")
|
||||
|
|
@ -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)}")
|
||||
|
|
@ -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)
|
||||
|
|
@ -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"
|
||||
|
|
@ -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
|
||||
|
|
@ -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())
|
||||
|
|
@ -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()
|
||||
|
|
@ -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())
|
||||
|
|
@ -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"
|
||||
)
|
||||
|
|
@ -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()
|
||||
|
|
@ -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"
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
|
@ -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
|
|
@ -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))
|
||||
|
|
@ -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"))
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
@ -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,)
|
||||
|
|
|
|||
0
tests/unit/llms/fireworks_ai/embed/__init__.py
Normal file
0
tests/unit/llms/fireworks_ai/embed/__init__.py
Normal 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)
|
||||
|
|
@ -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]
|
||||
219
tests/unit/llms/gemini/test_gemini_streaming.py
Normal file
219
tests/unit/llms/gemini/test_gemini_streaming.py
Normal 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"}
|
||||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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}]
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
34
tests/unit/llms/mistral/test_mistral_embedding.py
Normal file
34
tests/unit/llms/mistral/test_mistral_embedding.py
Normal 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)
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
39
tests/unit/llms/ollama/test_ollama_text_completion.py
Normal file
39
tests/unit/llms/ollama/test_ollama_text_completion.py
Normal 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!"
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
79
tests/unit/llms/openai/test_openai_authentication_error.py
Normal file
79
tests/unit/llms/openai/test_openai_authentication_error.py
Normal 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
|
||||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
36
tests/unit/llms/together_ai/test_together_ai_embedding.py
Normal file
36
tests/unit/llms/together_ai/test_together_ai_embedding.py
Normal 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)
|
||||
|
|
@ -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
|
|
@ -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"
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
17
tests/unit/test_docker_no_network_on_deploy.py
Normal file
17
tests/unit/test_docker_no_network_on_deploy.py
Normal 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
Loading…
Add table
Reference in a new issue