scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 98

Merge remote-tracking branch 'origin/main' into litellm_heuristic_first_context_escalation
Some checks are pending
LiteLLM Rust / rust-lint (push) Waiting to run
Details
LiteLLM Rust / rust-test (push) Waiting to run
Details
LiteLLM Rust / rust-wheel (push) Waiting to run
Details
Terraform Modules / fmt, validate, test (aws) (push) Waiting to run
Details
Terraform Modules / fmt, validate, test (gcp) (push) Waiting to run
Details
Terraform Provider / gofmt, vet, build, test (push) Waiting to run
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Waiting to run
Details

Browse source
This commit is contained in:
Devin AI 2026-09-26 00:12:08 +00:00
commit 443468d530
1378 changed files with 33878 additions and 23199 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.circleci/scripts/classify_changes.sh
View file

@ -14,12 +14,12 @@ while IFS= read -r file || [ -n "$file" ]; do
[ -n "$file" ] || continue
case "$file" in
*.md | *.mdx) : ;;
pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/test_litellm/test_circleci_path_filter.py | tests/test_litellm/test_detect_changes.py)
pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/unit/test_circleci_path_filter.py | tests/unit/test_detect_changes.py)
has_mcp_dependencies=true ;;
esac
case "$file" in
tests/e2e/*/*.py) : ;;
tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/test_litellm/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/unit/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
has_provider_harness=true ;;
esac
case "$file" in

37
.circleci/scripts/unit_selection.sh
View file

@ -5,9 +5,14 @@ flag="${1:?usage: unit_selection.sh <codecov flag>}"
legacy_flags=(
caching-local
core-utils
enterprise-package
enterprise-routing
integrations
llm-other-providers
llm-vertex-ai
mcp-integration
misc
proxy-db-auth-checks
proxy-db-budgets
proxy-db-custom-logging
@ -22,11 +27,13 @@ legacy_flags=(
proxy-db-proxy-utils
proxy-extras
proxy-infra
responses-caching-types
)
legacy_paths() {
case "$1" in
caching-local) echo tests/unit/caching ;;
core-utils) echo tests/unit/litellm_core_utils ;;
enterprise-package)
echo tests/unit/enterprise/integrations
echo tests/unit/enterprise/proxy/auth
@ -36,6 +43,9 @@ legacy_paths() {
echo tests/unit/enterprise/proxy/test_audit_logging_endpoints.py
echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;;
enterprise-routing)
echo tests/unit/google_genai
echo tests/unit/router_strategy
echo tests/unit/router_utils
echo tests/unit/enterprise/enterprise_callbacks/send_emails
echo tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py
echo tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py
@ -47,10 +57,34 @@ legacy_paths() {
echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py
echo tests/unit/enterprise/proxy/test_managed_files_access_check.py
echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;;
integrations) echo tests/unit/integrations ;;
llm-other-providers) find tests/unit/llms -name 'test_*.py' -not -path 'tests/unit/llms/vertex_ai/*' ;;
llm-vertex-ai) echo tests/unit/llms/vertex_ai ;;
mcp-integration)
echo tests/unit/experimental_mcp_client
echo tests/unit/proxy/_experimental/mcp_server
echo tests/unit/responses/mcp
echo tests/mcp_tests/test_proxy_mcp_e2e.py ;;
misc)
find tests/unit -maxdepth 1 -name 'test_*.py'
echo tests/unit/test_router
echo tests/unit/a2a_protocol
echo tests/unit/batches
echo tests/unit/chat_completions
echo tests/unit/completion_extras
echo tests/unit/containers
echo tests/unit/embeddings
echo tests/unit/endpoints
echo tests/unit/files
echo tests/unit/images
echo tests/unit/interactions
echo tests/unit/messages
echo tests/unit/rag
echo tests/unit/rerank_api
echo tests/unit/rust_bridge
echo tests/unit/secret_managers
echo tests/unit/vector_stores
echo tests/unit/videos ;;
proxy-db-auth-checks)
echo tests/unit/proxy/auth/test_auth_checks.py
echo tests/unit/proxy/auth/test_user_api_key_auth.py
@ -113,6 +147,9 @@ legacy_paths() {
proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;;
proxy-extras) echo tests/unit/litellm_proxy_extras ;;
proxy-infra) echo tests/unit/gateway ;;
responses-caching-types)
find tests/unit/responses -name 'test_*.py' -not -path 'tests/unit/responses/mcp/*'
echo tests/unit/types ;;
*) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;;
esac
}

37
.circleci/tests.yml
View file

@ -341,6 +341,7 @@ workflows:
flag:
- enterprise-package
- proxy-infra
- responses-caching-types
- proxy-db-auth-checks
- proxy-db-jwt-and-keys
- proxy-db-proxy-server-core
@ -353,6 +354,42 @@ workflows:
- proxy-db-endpoints-and-responses
base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
- unit:
name: unit-llm-vertex-ai
flag: llm-vertex-ai
shards: 2
workers: 1
reruns: 2
base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
- unit:
name: unit-llm-other-providers
flag: llm-other-providers
shards: 3
reruns: 2
base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
- unit:
name: unit-core-utils
flag: core-utils
shards: 2
reruns: 1
base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
- unit:
name: unit-integrations
flag: integrations
shards: 2
reruns: 3
base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
- unit:
name: unit-misc
flag: misc
shards: 2
reruns: 2
base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
- unit:
name: unit-proxy-db-proxy-utils
flag: proxy-db-proxy-utils

18
.github/merge-smoke-tests.json vendored
View file

@ -1,15 +1,15 @@
{
"cases": {
"CHAT-JSON": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_returns_json_reply_over_injected_transport",
"CHAT-TEXT-STREAM": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_streams_text_deltas_over_injected_transport",
"CHAT-TOOL-STREAM": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_streams_tool_call_arguments_over_injected_transport",
"CHAT-JSON": "tests/unit/llms/openai/test_openai.py::test_acompletion_returns_json_reply_over_injected_transport",
"CHAT-TEXT-STREAM": "tests/unit/llms/openai/test_openai.py::test_acompletion_streams_text_deltas_over_injected_transport",
"CHAT-TOOL-STREAM": "tests/unit/llms/openai/test_openai.py::test_acompletion_streams_tool_call_arguments_over_injected_transport",
"MODEL-ALLOW": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_allows_listed_model_for_key",
"MODEL-DENY": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_denials_return_forbidden[key-key_model_access_denied]",
"COST-EXPLICIT": "tests/test_litellm/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones",
"COST-ZERO": "tests/test_litellm/test_cost_calculator.py::test_completion_cost_is_zero_when_explicit_rates_are_zero",
"LOG-CONTENT-ON": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_keeps_message_content_when_message_logging_is_on",
"LOG-CONTENT-OFF": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_redacts_message_content_when_message_logging_is_off",
"CALLBACK-SUCCESS": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_async_success_handler_delivers_standard_logging_payload_to_custom_logger",
"CALLBACK-FAILURE": "tests/test_litellm/litellm_core_utils/test_litellm_logging.py::test_async_failure_handler_delivers_failure_payload_to_custom_logger"
"COST-EXPLICIT": "tests/unit/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones",
"COST-ZERO": "tests/unit/test_cost_calculator.py::test_completion_cost_is_zero_when_explicit_rates_are_zero",
"LOG-CONTENT-ON": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_keeps_message_content_when_message_logging_is_on",
"LOG-CONTENT-OFF": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_standard_logging_payload_redacts_message_content_when_message_logging_is_off",
"CALLBACK-SUCCESS": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_async_success_handler_delivers_standard_logging_payload_to_custom_logger",
"CALLBACK-FAILURE": "tests/unit/litellm_core_utils/test_litellm_logging.py::test_async_failure_handler_delivers_failure_payload_to_custom_logger"
}
}

18
.github/workflows/_test-unit-base.yml vendored
View file

@ -13,12 +13,13 @@ on:
have its path existence-checked like any other token.
required: true
type: string
fork-flag:
unit-flag:
description: >-
Codecov flag of the `.circleci/tests.yml` job that now owns part of
this shard. CircleCI does not run on pull requests from forks, so on
those events this shard also runs the files
`.circleci/scripts/unit_selection.sh` lists for the flag.
this shard. The shard also runs the files
`.circleci/scripts/unit_selection.sh` lists for the flag, on every
event, because the CircleCI pipeline is manual-only while the tests
migrate.
required: false
type: string
default: ""
@ -175,8 +176,7 @@ jobs:
timeout-minutes: ${{ inputs.timeout-minutes }}
env:
TEST_PATH: ${{ inputs.test-path }}
FORK_FLAG: ${{ inputs.fork-flag }}
IS_FORK: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository }}
UNIT_FLAG: ${{ inputs.unit-flag }}
MAX_FAILURES: ${{ inputs.max-failures }}
WORKERS: ${{ inputs.workers }}
RERUNS: ${{ inputs.reruns }}
@ -186,11 +186,11 @@ jobs:
run: |
echo "has-coverage=false" >> "$GITHUB_OUTPUT"
selection="${TEST_PATH}"
if [ "${IS_FORK}" = "true" ] && [ -n "${FORK_FLAG}" ]; then
selection="${TEST_PATH} $(bash .circleci/scripts/unit_selection.sh "${FORK_FLAG}" | tr '\n' ' ')"
if [ -n "${UNIT_FLAG}" ]; then
selection="${TEST_PATH} $(bash .circleci/scripts/unit_selection.sh "${UNIT_FLAG}" | tr '\n' ' ')"
fi
if [ -z "${selection// /}" ]; then
echo "shard selection is empty on this event (CircleCI flag ${FORK_FLAG:-none} owns it); nothing to run"
echo "shard selection is empty; nothing to run"
exit 0
fi
pytest_args=()

16
.github/workflows/test-redis-compat.yml vendored
View file

@ -10,11 +10,11 @@ on:
- "litellm/_redis_credential_provider.py"
- "litellm/caching/redis_cache.py"
- "litellm/caching/evicted_client_closer.py"
- "tests/test_litellm/test_redis.py"
- "tests/unit/test_redis.py"
- "tests/local_testing/test_caching.py"
- "tests/test_litellm/caching/test_redis_connection_pool.py"
- "tests/test_litellm/caching/test_redis_cluster_cache.py"
- "tests/test_litellm/caching/test_evicted_client_closer.py"
- "tests/unit/caching/test_redis_connection_pool.py"
- "tests/unit/caching/test_redis_cluster_cache.py"
- "tests/unit/caching/test_evicted_client_closer.py"
- ".github/workflows/test-redis-compat.yml"
- "pyproject.toml"
- "uv.lock"
@ -84,10 +84,10 @@ jobs:
run: |
redis-server --version
uv run --no-sync pytest \
tests/test_litellm/test_redis.py \
tests/test_litellm/caching/test_redis_connection_pool.py \
tests/test_litellm/caching/test_redis_cluster_cache.py \
tests/test_litellm/caching/test_evicted_client_closer.py \
tests/unit/test_redis.py \
tests/unit/caching/test_redis_connection_pool.py \
tests/unit/caching/test_redis_cluster_cache.py \
tests/unit/caching/test_evicted_client_closer.py \
tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_azure_credentials \
tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_gcp_credentials \
--tb=short -vv \

6
.github/workflows/test-rust.yml vendored
View file

@ -24,7 +24,7 @@ on:
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
- "tests/unit/rust_bridge/native_route_wheel_test.py"
- ".github/workflows/test-rust.yml"
pull_request:
branches:
@ -52,7 +52,7 @@ on:
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
- "tests/unit/rust_bridge/native_route_wheel_test.py"
- ".github/workflows/test-rust.yml"
permissions:
@ -171,7 +171,7 @@ jobs:
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
- run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- run: python tests/unit/rust_bridge/native_route_wheel_test.py dist/*.whl
- name: Run pytest tests/test_litellm_rust with the compiled extension
run: make test-rust-extension

33
.github/workflows/test-unit-proxy-db.yml vendored
View file

@ -22,9 +22,10 @@ concurrency:
#
# `.circleci/tests.yml` runs each group's files on same-repo events under the
# `proxy-db-<group>` Codecov flag; `.circleci/scripts/unit_selection.sh` holds
# the file lists. CircleCI does not build pull requests from forks, so `fork-flag`
# makes the shard run that list there. `test-path` keeps the files that still
# reach real providers and never left tests/proxy_unit_tests.
# the file lists. That pipeline is manual-only while the tests migrate, so
# `unit-flag` makes the shard run that list on every event. `test-path` keeps
# the files that still reach real providers and never left
# tests/proxy_unit_tests.
#
# Design targets:
# * Every shard runs in <= 7 minutes of wall-clock on the default runner.
@ -78,7 +79,7 @@ jobs:
# Must run serially — event-loop conflict with the logging worker.
- test-group: key-generation
test-path: ""
fork-flag: proxy-db-key-generation
unit-flag: proxy-db-key-generation
workers: 0
dist: loadscope
timeout: 20
@ -86,13 +87,13 @@ jobs:
# ---- auth: split into 2 shards ----
- test-group: auth-checks
test-path: ""
fork-flag: proxy-db-auth-checks
unit-flag: proxy-db-auth-checks
workers: 4
dist: loadscope
timeout: 15
- test-group: jwt-and-keys
test-path: ""
fork-flag: proxy-db-jwt-and-keys
unit-flag: proxy-db-jwt-and-keys
workers: 4
dist: loadscope
timeout: 15
@ -100,7 +101,7 @@ jobs:
# ---- test_proxy_utils.py, single shard, worksteal distribution ----
- test-group: proxy-utils
test-path: ""
fork-flag: proxy-db-proxy-utils
unit-flag: proxy-db-proxy-utils
workers: 4
dist: worksteal
timeout: 15
@ -108,13 +109,13 @@ jobs:
# ---- proxy server: split into 2 shards ----
- test-group: proxy-server-core
test-path: "tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py"
fork-flag: proxy-db-proxy-server-core
unit-flag: proxy-db-proxy-server-core
workers: 4
dist: loadscope
timeout: 15
- test-group: proxy-runtime
test-path: ""
fork-flag: proxy-db-proxy-runtime
unit-flag: proxy-db-proxy-runtime
workers: 4
dist: loadscope
timeout: 15
@ -122,20 +123,20 @@ jobs:
# ---- logging: split into 2 shards ----
- test-group: custom-logging
test-path: "tests/proxy_unit_tests/test_proxy_custom_logger.py"
fork-flag: proxy-db-custom-logging
unit-flag: proxy-db-custom-logging
workers: 4
dist: loadscope
timeout: 15
- test-group: logging-misc
test-path: ""
fork-flag: proxy-db-logging-misc
unit-flag: proxy-db-logging-misc
workers: 4
dist: loadscope
timeout: 15
- test-group: db-and-spend
test-path: ""
fork-flag: proxy-db-db-and-spend
unit-flag: proxy-db-db-and-spend
workers: 4
dist: loadscope
timeout: 15
@ -143,27 +144,27 @@ jobs:
# ---- guardrails + budget + hooks: split into 2 ----
- test-group: guardrails-hooks
test-path: ""
fork-flag: proxy-db-guardrails-hooks
unit-flag: proxy-db-guardrails-hooks
workers: 4
dist: loadscope
timeout: 15
- test-group: budgets
test-path: ""
fork-flag: proxy-db-budgets
unit-flag: proxy-db-budgets
workers: 4
dist: loadscope
timeout: 15
- test-group: endpoints-and-responses
test-path: "tests/proxy_unit_tests/test_proxy_exception_mapping.py"
fork-flag: proxy-db-endpoints-and-responses
unit-flag: proxy-db-endpoints-and-responses
workers: 4
dist: loadscope
timeout: 15
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: ${{ matrix.test-path }}
fork-flag: ${{ matrix.fork-flag }}
unit-flag: ${{ matrix.unit-flag }}
workers: ${{ matrix.workers }}
reruns: 2
timeout-minutes: ${{ matrix.timeout }}

56
.github/workflows/test-unit.yml vendored
View file

@ -36,9 +36,9 @@ concurrency:
# Folding it in here is a follow-up, together with generalising that guard into
# assert_ci_coverage.py.
#
# `fork-flag` names the `.circleci/tests.yml` job that now runs part of the
# shard under the same Codecov flag. CircleCI does not build pull requests from
# forks, so the shard still runs those files there and skips them elsewhere.
# `unit-flag` names the `.circleci/tests.yml` job that now runs part of the
# shard under the same Codecov flag. That pipeline is manual-only while the
# tests migrate, so the shard also runs those files on every event.
jobs:
unit:
name: ${{ matrix.shard }}
@ -52,8 +52,8 @@ jobs:
include:
- shard: mcp-integration
artifact-name: mcp-integration
test-path: "tests/mcp_tests tests/test_litellm/experimental_mcp_client"
fork-flag: mcp-integration
test-path: "tests/mcp_tests"
unit-flag: mcp-integration
workers: 2
reruns: 0
timeout-minutes: 20
@ -62,6 +62,7 @@ jobs:
- shard: core-utils
artifact-name: core-utils
test-path: "tests/test_litellm/litellm_core_utils"
unit-flag: core-utils
workers: 2
reruns: 1
timeout-minutes: 20
@ -69,11 +70,8 @@ jobs:
- shard: enterprise-routing
artifact-name: enterprise-routing
test-path: >-
tests/test_litellm/google_genai
tests/test_litellm/router_utils
tests/test_litellm/router_strategy
fork-flag: enterprise-routing
test-path: ""
unit-flag: enterprise-routing
workers: 2
reruns: 2
timeout-minutes: 20
@ -81,7 +79,8 @@ jobs:
- shard: integrations
artifact-name: integrations
test-path: "tests/test_litellm/integrations"
test-path: ""
unit-flag: integrations
workers: 2
reruns: 3
timeout-minutes: 20
@ -90,6 +89,7 @@ jobs:
- shard: Vertex AI
artifact-name: llm-vertex-ai
test-path: "tests/test_litellm/llms/vertex_ai"
unit-flag: llm-vertex-ai
workers: 1
reruns: 2
timeout-minutes: 20
@ -98,6 +98,7 @@ jobs:
- shard: All Other Providers
artifact-name: llm-other-providers
test-path: "tests/test_litellm/llms --ignore=tests/test_litellm/llms/vertex_ai"
unit-flag: llm-other-providers
workers: 2
reruns: 2
timeout-minutes: 20
@ -106,26 +107,11 @@ jobs:
- shard: misc
artifact-name: misc
test-path: >-
tests/test_litellm/batches
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol
tests/test_litellm/chat_completions
tests/test_litellm/completion_extras
tests/test_litellm/containers
tests/test_litellm/endpoints
tests/test_litellm/files
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/messages
tests/test_litellm/embeddings
tests/test_litellm/ocr
tests/test_litellm/passthrough
tests/test_litellm/rag
tests/test_litellm/rerank_api
tests/test_litellm/rust_bridge
tests/test_litellm/vector_stores
tests/test_litellm/videos
tests/test_litellm/test_*.py
unit-flag: misc
workers: 2
reruns: 2
timeout-minutes: 20
@ -205,7 +191,7 @@ jobs:
tests/test_litellm/proxy/types_utils
tests/test_litellm/proxy/logging_endpoints
tests/test_litellm/proxy/test_*.py
fork-flag: proxy-infra
unit-flag: proxy-infra
workers: 4
reruns: 2
timeout-minutes: 20
@ -214,7 +200,7 @@ jobs:
- shard: caching-local
artifact-name: caching-local
test-path: ""
fork-flag: caching-local
unit-flag: caching-local
workers: 2
reruns: 2
timeout-minutes: 20
@ -223,7 +209,7 @@ jobs:
- shard: proxy-extras
artifact-name: proxy-extras
test-path: ""
fork-flag: proxy-extras
unit-flag: proxy-extras
workers: 2
reruns: 2
timeout-minutes: 20
@ -232,7 +218,7 @@ jobs:
- shard: enterprise-package
artifact-name: enterprise-package
test-path: ""
fork-flag: enterprise-package
unit-flag: enterprise-package
workers: 4
reruns: 2
timeout-minutes: 20
@ -240,10 +226,8 @@ jobs:
- shard: responses-caching-types
artifact-name: responses-caching-types
test-path: >-
tests/test_litellm/responses
tests/test_litellm/caching
tests/test_litellm/types
test-path: ""
unit-flag: responses-caching-types
workers: 2
reruns: 2
timeout-minutes: 20
@ -251,7 +235,7 @@ jobs:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: ${{ matrix.test-path }}
fork-flag: ${{ matrix.fork-flag || '' }}
unit-flag: ${{ matrix.unit-flag || '' }}
workers: ${{ matrix.workers }}
reruns: ${{ matrix.reruns }}
timeout-minutes: ${{ matrix.timeout-minutes }}

12
Makefile
View file

@ -301,7 +301,7 @@ test-rust-extension:
UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
"$$temporary/venv/bin/python" -I -m mypy.stubtest \
--mypy-config-file tests/test_litellm/rust_bridge/stubtest.ini \
--mypy-config-file tests/unit/rust_bridge/stubtest.ini \
litellm.rust_bridge._native && \
LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
@ -314,7 +314,7 @@ test-unit: install-test-deps
# Matrix test targets (matching CI workflow groups)
test-unit-llms: install-test-deps
$(UV_RUN) pytest tests/test_litellm/llms --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/unit/llms --tb=short -vv -n 4 --durations=20
test-unit-proxy-guardrails: install-test-deps
$(UV_RUN) pytest tests/test_litellm/proxy/guardrails tests/test_litellm/proxy/management_endpoints tests/test_litellm/proxy/management_helpers --tb=short -vv -n 4 --durations=20
@ -326,16 +326,16 @@ test-unit-proxy-misc: install-test-deps
$(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20
test-unit-integrations: install-test-deps
$(UV_RUN) pytest tests/test_litellm/integrations --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/unit/integrations --tb=short -vv -n 4 --durations=20
test-unit-core-utils: install-test-deps
$(UV_RUN) pytest tests/test_litellm/litellm_core_utils --tb=short -vv -n 2 --durations=20
$(UV_RUN) pytest tests/unit/litellm_core_utils --tb=short -vv -n 2 --durations=20
test-unit-other: install-test-deps
$(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/test_litellm/vector_stores tests/test_litellm/a2a_protocol tests/test_litellm/anthropic_interface tests/test_litellm/completion_extras tests/test_litellm/containers tests/unit/enterprise tests/test_litellm/experimental_mcp_client tests/test_litellm/google_genai tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/test_litellm/types --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/unit/caching tests/unit/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/unit/router_strategy tests/unit/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20
test-unit-root: install-test-deps
$(UV_RUN) pytest tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/unit/test_*.py tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20
# Proxy unit tests (tests/unit/proxy split alphabetically)
test-proxy-unit-a: install-test-deps

2
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md
View file

@ -2,7 +2,7 @@
Every `litellm_*` metric family the proxy can expose on `/metrics` (136 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/unit/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard
The first eleven rows need only `callbacks: ["prometheus"]`. The last three rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected

350
litellm-rust/Cargo.lock generated
View file

@ -73,6 +73,15 @@ version = "1.0.104"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470"
[[package]]
name = "arbitrary"
version = "1.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1"
dependencies = [
"derive_arbitrary",
]
[[package]]
name = "arc-swap"
version = "1.9.2"
@ -897,6 +906,12 @@ dependencies = [
"hybrid-array",
]
[[package]]
name = "borrow-or-share"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc0b364ead1874514c8c2855ab558056ebfeb775653e7ae45ff72f28f8f3166c"
[[package]]
name = "bstr"
version = "1.13.1"
@ -914,6 +929,12 @@ version = "3.20.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649"
[[package]]
name = "bytecount"
version = "0.6.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "175812e0be2bccb6abe50bb8d566126198344f707e304f45c648fd8f2cc0365e"
[[package]]
name = "byteorder"
version = "1.5.0"
@ -1458,6 +1479,17 @@ dependencies = [
"serde_core",
]
[[package]]
name = "derive_arbitrary"
version = "1.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "derive_builder"
version = "0.20.2"
@ -1540,6 +1572,15 @@ version = "1.16.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e"
[[package]]
name = "email_address"
version = "0.2.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e079f19b08ca6239f47f8ba8509c11cf3ea30095831f7fed61441475edd8c449"
dependencies = [
"serde",
]
[[package]]
name = "equivalent"
version = "1.0.2"
@ -1622,6 +1663,16 @@ version = "2.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223"
[[package]]
name = "filetime"
version = "0.2.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759"
dependencies = [
"cfg-if",
"libc",
]
[[package]]
name = "find-msvc-tools"
version = "0.1.9"
@ -1639,6 +1690,17 @@ dependencies = [
"zlib-rs",
]
[[package]]
name = "fluent-uri"
version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bc74ac4d8359ae70623506d512209619e5cf8f347124910440dbc221714b328e"
dependencies = [
"borrow-or-share",
"ref-cast",
"serde",
]
[[package]]
name = "fnv"
version = "1.0.7"
@ -1660,6 +1722,16 @@ dependencies = [
"percent-encoding",
]
[[package]]
name = "fraction"
version = "0.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e246562084dde8ebbcc943b261c406ce4f68e5032ec28029a251a47d6a295500"
dependencies = [
"num",
"num-bigint 0.4.8",
]
[[package]]
name = "fs_extra"
version = "1.3.0"
@ -1817,9 +1889,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd"
dependencies = [
"cfg-if",
"js-sys",
"libc",
"r-efi 5.3.0",
"wasip2",
"wasm-bindgen",
]
[[package]]
@ -2660,6 +2734,59 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "jsonschema"
version = "0.55.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b68339c3d874e48151d74ffe256d93a58cffa240983cb0967d3cbaea083a44fe"
dependencies = [
"ahash",
"bytecount",
"data-encoding",
"email_address",
"fancy-regex 0.19.2",
"fraction",
"getrandom 0.3.4",
"itoa",
"jsonschema-regex",
"jsonschema-value",
"num-cmp",
"num-traits",
"percent-encoding",
"referencing",
"regex",
"serde",
"serde_json",
"strum",
"unicode-general-category",
"uuid-simd",
]
[[package]]
name = "jsonschema-regex"
version = "0.55.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6307b5b51216ec9b941b52244c74043fa0b1d6b657b56199f57cb1416d3641c5"
dependencies = [
"regex-syntax",
]
[[package]]
name = "jsonschema-value"
version = "0.55.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0230ac05e09c6111e96c147b75c390579f5cbd45654b980c68ac60fe17b3f129"
dependencies = [
"ahash",
"bytecount",
"fraction",
"getrandom 0.3.4",
"num-cmp",
"num-traits",
"serde_json",
"zmij",
]
[[package]]
name = "lazy_static"
version = "1.5.0"
@ -2995,6 +3122,7 @@ dependencies = [
"tokio-tungstenite",
"url",
"veil",
"wiremock",
]
[[package]]
@ -3038,10 +3166,11 @@ dependencies = [
"aws-smithy-types",
"bytes",
"futures-util",
"proptest",
"rstest",
"sse-stream",
"thiserror 2.0.19",
"tokio",
"tokio-util",
]
[[package]]
@ -3126,14 +3255,14 @@ dependencies = [
name = "litellm-model-catalog"
version = "0.1.0"
dependencies = [
"criterion",
"indexmap 2.14.0",
"litellm-model-catalog",
"jsonschema",
"rstest",
"schemars 1.2.2",
"serde",
"serde_json",
"thiserror 2.0.19",
"time",
]
[[package]]
@ -3355,6 +3484,27 @@ dependencies = [
"veil",
]
[[package]]
name = "litellm-testkit"
version = "0.1.0"
dependencies = [
"flate2",
"futures-util",
"reqwest 0.12.28",
"rstest",
"semver",
"serde",
"serde_json",
"sha2 0.10.9",
"tar",
"target-lexicon",
"tempfile",
"thiserror 2.0.19",
"tokio",
"toml",
"zip",
]
[[package]]
name = "litellm-token-counter"
version = "0.1.0"
@ -3504,6 +3654,12 @@ version = "2.8.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98"
[[package]]
name = "micromap"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a86d3146ed3995b5913c414f6664344b9617457320782e64f0bb44afd49d74"
[[package]]
name = "mime"
version = "0.3.17"
@ -3599,6 +3755,20 @@ dependencies = [
"minimal-lexical",
]
[[package]]
name = "num"
version = "0.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23"
dependencies = [
"num-bigint 0.4.8",
"num-complex",
"num-integer",
"num-iter",
"num-rational",
"num-traits",
]
[[package]]
name = "num-bigint"
version = "0.4.8"
@ -3619,6 +3789,12 @@ dependencies = [
"num-traits",
]
[[package]]
name = "num-cmp"
version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "63335b2e2c34fae2fb0aa2cecfd9f0832a1e24b3b32ecec612c3426d46dc8aaa"
[[package]]
name = "num-complex"
version = "0.4.6"
@ -3643,6 +3819,27 @@ dependencies = [
"num-traits",
]
[[package]]
name = "num-iter"
version = "0.1.46"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c92800bd69a1eac91786bcfe9da64a897eb72911b8dc3095decbd07429e8048b"
dependencies = [
"num-integer",
"num-traits",
]
[[package]]
name = "num-rational"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824"
dependencies = [
"num-bigint 0.4.8",
"num-integer",
"num-traits",
]
[[package]]
name = "num-traits"
version = "0.2.19"
@ -4458,6 +4655,23 @@ dependencies = [
"syn 3.0.0",
]
[[package]]
name = "referencing"
version = "0.55.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a196a5b4a8a12f46b6353174df865a05d41a6055aff212ec30877492788618b6"
dependencies = [
"ahash",
"fluent-uri",
"getrandom 0.3.4",
"hashbrown 0.17.1",
"itoa",
"micromap",
"parking_lot",
"percent-encoding",
"serde_json",
]
[[package]]
name = "regex"
version = "1.13.1"
@ -5044,6 +5258,15 @@ dependencies = [
"serde_core",
]
[[package]]
name = "serde_spanned"
version = "1.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26"
dependencies = [
"serde_core",
]
[[package]]
name = "serde_urlencoded"
version = "0.7.1"
@ -5246,19 +5469,6 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "sse-stream"
version = "0.2.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4"
dependencies = [
"bytes",
"futures-util",
"http-body 1.1.0",
"http-body-util",
"pin-project-lite",
]
[[package]]
name = "stable_deref_trait"
version = "1.2.1"
@ -5375,6 +5585,17 @@ version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
[[package]]
name = "tar"
version = "0.4.46"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3f6221d9a6003c78398e3b239969f352578258df48c8eb051caadae0015bc840"
dependencies = [
"filetime",
"libc",
"xattr",
]
[[package]]
name = "target-lexicon"
version = "0.13.5"
@ -5644,6 +5865,30 @@ dependencies = [
"tokio",
]
[[package]]
name = "toml"
version = "0.9.12+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863"
dependencies = [
"indexmap 2.14.0",
"serde_core",
"serde_spanned",
"toml_datetime 0.7.5+spec-1.1.0",
"toml_parser",
"toml_writer",
"winnow 0.7.15",
]
[[package]]
name = "toml_datetime"
version = "0.7.5+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347"
dependencies = [
"serde_core",
]
[[package]]
name = "toml_datetime"
version = "1.1.1+spec-1.1.0"
@ -5660,9 +5905,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b"
dependencies = [
"indexmap 2.14.0",
"toml_datetime",
"toml_datetime 1.1.1+spec-1.1.0",
"toml_parser",
"winnow",
"winnow 1.0.4",
]
[[package]]
@ -5671,9 +5916,15 @@ version = "1.1.3+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56"
dependencies = [
"winnow",
"winnow 1.0.4",
]
[[package]]
name = "toml_writer"
version = "1.1.2+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2"
[[package]]
name = "tonic"
version = "0.14.6"
@ -5941,6 +6192,12 @@ version = "2.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142"
[[package]]
name = "unicode-general-category"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b993bddc193ae5bd0d623b49ec06ac3e9312875fdae725a975c51db1cc1677f"
[[package]]
name = "unicode-ident"
version = "1.0.24"
@ -6021,6 +6278,16 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "uuid-simd"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "23b082222b4f6619906941c17eb2297fff4c2fb96cb60164170522942a200bd8"
dependencies = [
"outref",
"vsimd",
]
[[package]]
name = "valuable"
version = "0.1.1"
@ -6419,6 +6686,12 @@ version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
[[package]]
name = "winnow"
version = "0.7.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945"
[[package]]
name = "winnow"
version = "1.0.4"
@ -6481,6 +6754,16 @@ dependencies = [
"time",
]
[[package]]
name = "xattr"
version = "1.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156"
dependencies = [
"libc",
"rustix",
]
[[package]]
name = "xmlparser"
version = "0.13.6"
@ -6606,6 +6889,23 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "zip"
version = "2.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fabe6324e908f85a1c52063ce7aa26b68dcb7eb6dbc83a2d148403c9bc3eba50"
dependencies = [
"arbitrary",
"crc32fast",
"crossbeam-utils",
"displaydoc",
"flate2",
"indexmap 2.14.0",
"memchr",
"thiserror 2.0.19",
"zopfli",
]
[[package]]
name = "zlib-rs"
version = "0.6.7"
@ -6617,3 +6917,15 @@ name = "zmij"
version = "1.0.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b"
[[package]]
name = "zopfli"
version = "0.8.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249"
dependencies = [
"bumpalo",
"crc32fast",
"log",
"simd-adler32",
]

6
litellm-rust/Cargo.toml
View file

@ -81,6 +81,12 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
base64 = "0.22"
flate2 = "1"
semver = "1"
tar = "0.4"
target-lexicon = "0.13.5"
tempfile = "3"
zip = { version = "2", default-features = false, features = ["deflate"] }
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"

2
litellm-rust/crates/callbacks-legacy-python/src/lib.rs
View file

@ -33,7 +33,7 @@ mod test_support {
use crate::{LegacyLogging, LegacySurface, PublicCall};
/// The parameters of every `callbacks_legacy_python` function, as the real module declares them.
/// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python
/// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python
/// signatures, and [`namespace`] binds every fake call against it.
pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json");

1
litellm-rust/crates/core/Cargo.toml
View file

@ -40,3 +40,4 @@ litellm-auth-gcp.workspace = true
litellm-llms = { workspace = true, features = ["test-support"] }
rstest.workspace = true
rstest_reuse.workspace = true
wiremock = "0.6.5"

26
litellm-rust/crates/core/src/chat_completions/handler.rs
View file

@ -85,3 +85,29 @@ pub(super) async fn outbound_request(
other => other,
})
}
#[cfg(test)]
mod tests {
use super::{Error, as_response_error};
#[test]
fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() {
for original in [
Error::MissingField("usage"),
Error::Unsupported("non-text response content block"),
Error::InvalidRequest("whatever".to_string()),
Error::Auth(litellm_auth::Error::InvalidHeader),
] {
let label = format!("{original:?}");
assert!(
matches!(as_response_error(original), Error::InvalidResponse(_)),
"{label} must not stay retryable once the provider has answered"
);
}
let upstream = Error::Transport(litellm_http::transport::Error::Http {
status: 500,
body: "boom".to_string(),
});
assert_eq!(as_response_error(upstream.clone()), upstream);
}
}

244
litellm-rust/crates/core/src/chat_completions/prepare.rs
View file

@ -736,248 +736,4 @@ mod tests {
.unwrap_or_else(|error| panic!("prepare declined {messages}: {error}"));
}
}
mod round_trip {
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
};
use super::*;
use crate::chat_completions::chat_completions;
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
let header_end = loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break request.len();
}
request.extend_from_slice(&buffer[..n]);
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n")
{
break position + 4;
}
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while request.len().saturating_sub(header_end) < content_length {
let n = socket.read(&mut buffer).await.expect("reads body");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
}
String::from_utf8(request).expect("request is utf8")
}
fn http_response(status: &str, body: &str) -> String {
format!(
"HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
)
}
/// Serve one request from a stub upstream and hand back what it received.
async fn serve_once(
status: &'static str,
body: &'static str,
) -> (String, tokio::task::JoinHandle<String>) {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let port = listener.local_addr().expect("addr").port();
let handle = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts");
let received = read_http_request(&mut socket).await;
socket
.write_all(http_response(status, body).as_bytes())
.await
.expect("writes response");
socket.flush().await.expect("flushes");
received
});
(format!("http://127.0.0.1:{port}/v1/messages"), handle)
}
fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> {
ChatCompletionsRequest {
model: "anthropic/claude-sonnet-4-5",
messages,
optional_params: match params {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
},
api_key: Some("sk-test"),
api_base: Some(api_base),
custom_llm_provider: None,
extra_headers: None,
timeout: Some(std::time::Duration::from_secs(10)),
}
}
const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
#[tokio::test]
async fn round_trip_sends_the_translated_body_and_normalizes_the_response() {
let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await;
let response = chat_completions(call(
&api_base,
json!([
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"}
]),
json!({"max_tokens": 16}),
))
.await
.expect("call succeeds");
let received = handle.await.expect("server task");
let sent: Value = serde_json::from_str(
received
.split_once("\r\n\r\n")
.expect("request has a body")
.1,
)
.expect("body is json");
assert_eq!(
sent["messages"],
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}])
);
assert_eq!(
sent["system"],
json!([{"type": "text", "text": "be terse"}])
);
assert_eq!(sent["max_tokens"], json!(16));
assert!(received.to_lowercase().contains("x-api-key: sk-test"));
assert_eq!(
response.choices[0].message.content.as_deref(),
Some("hello")
);
assert_eq!(response.usage.total_tokens, 15);
}
#[tokio::test]
async fn a_response_it_cannot_normalize_is_reported_as_already_sent() {
// The provider was called and billed, so the host must not retry this
// on its own path. `MissingField` here would read as a pre-send
// decline and be retried; `InvalidResponse` cannot.
const NO_USAGE: &str =
r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#;
let (api_base, handle) = serve_once("200 OK", NO_USAGE).await;
let err = chat_completions(call(
&api_base,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.await
.expect_err("response cannot be normalized");
handle.await.expect("server task");
assert!(
matches!(err, Error::InvalidResponse(_)),
"expected a post-send error, got {err:?}"
);
}
#[tokio::test]
async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() {
const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#;
let (api_base, handle) = serve_once("200 OK", TOOL_USE).await;
let err = chat_completions(call(
&api_base,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.await
.expect_err("response cannot be normalized");
handle.await.expect("server task");
assert!(
matches!(err, Error::InvalidResponse(_)),
"expected a post-send error, got {err:?}"
);
}
#[tokio::test]
async fn an_upstream_error_status_keeps_its_code() {
let (api_base, handle) =
serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await;
let err = chat_completions(call(
&api_base,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.await
.expect_err("upstream rejects");
handle.await.expect("server task");
assert!(
matches!(
err,
Error::Transport(litellm_http::transport::Error::Http { status: 429, .. })
),
"expected a 429, got {err:?}"
);
}
#[tokio::test]
async fn a_connection_that_is_never_established_declines_instead_of_failing() {
// Nothing was sent, so nothing was billed and the host can still serve
// the request. Classing this with the post-send failures would turn a
// recoverable fallback into a user-facing error on exactly the
// deployments whose transport is configured only on the Python client.
let port = {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
listener.local_addr().expect("has an address").port()
// Dropped here, so the port is closed and the connect is refused.
};
let err = chat_completions(call(
&format!("http://127.0.0.1:{port}/v1/messages"),
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.await
.expect_err("nothing is listening");
assert!(
matches!(
err,
Error::Transport(litellm_http::transport::Error::Connect(_))
),
"expected a pre-send connect failure, got {err:?}"
);
}
#[test]
fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() {
use crate::chat_completions::handler::as_response_error;
for original in [
Error::MissingField("usage"),
Error::Unsupported("non-text response content block"),
Error::InvalidRequest("whatever".to_string()),
Error::Auth(litellm_auth::Error::InvalidHeader),
] {
let label = format!("{original:?}");
assert!(
matches!(as_response_error(original), Error::InvalidResponse(_)),
"{label} must not stay retryable once the provider has answered"
);
}
// An upstream status is already unambiguous, so it survives intact.
assert!(matches!(
as_response_error(Error::Transport(litellm_http::transport::Error::Http {
status: 500,
body: "boom".to_string()
})),
Error::Transport(litellm_http::transport::Error::Http { status: 500, .. })
));
}
}
}

145
litellm-rust/crates/core/src/messages/common_utils.rs
View file

@ -29,151 +29,10 @@ pub(super) fn string_headers(
#[cfg(test)]
mod tests {
use std::{sync::Arc, time::Duration};
use futures_util::future::BoxFuture;
use litellm_secrets::{SecretValue, source::SecretSource};
use serde_json::{Value, json};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
};
use serde_json::json;
use super::{messages_provider_config, string_headers, truncate_error_body};
use crate::messages::{
Error,
route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine},
types::MessagesShaping,
};
struct RecordingSecrets {
values: Vec<(&'static str, String)>,
requested: std::sync::Mutex<Vec<String>>,
}
impl SecretSource for RecordingSecrets {
fn get_secret_str<'a>(
&'a self,
name: &'a str,
) -> BoxFuture<'a, Result<Option<SecretValue>, litellm_secrets::Error>> {
Box::pin(async move {
self.requested.lock().unwrap().push(name.to_string());
Ok(self
.values
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| SecretValue::new(value.clone())))
})
}
}
fn secrets_call() -> MessagesCall {
let Value::Object(body) = json!({
"model": "claude-sonnet-4-5",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hi"}]
}) else {
unreachable!("literal object")
};
MessagesCall {
model: "claude-sonnet-4-5".into(),
body,
api_key: None,
api_base: None,
custom_llm_provider: Some("anthropic".into()),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
}
}
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
let header_end = loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break request.len();
}
request.extend_from_slice(&buffer[..n]);
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
break position + 4;
}
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while request.len().saturating_sub(header_end) < content_length {
let n = socket.read(&mut buffer).await.expect("reads body");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
}
String::from_utf8(request).expect("request is utf8")
}
#[tokio::test]
async fn route_reads_the_provider_credential_and_base_from_the_secret_source() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
request
});
let secrets = Arc::new(RecordingSecrets {
values: vec![
("ANTHROPIC_API_KEY", "sk-from-manager".to_string()),
("ANTHROPIC_BASE_URL", format!("http://{addr}")),
],
requested: std::sync::Mutex::new(Vec::new()),
});
let output = litellm_host::run::run(
messages_machine(secrets.clone()),
&LocalMessagesHost::new(secrets_call()),
)
.await
.expect("messages request succeeds");
assert!(matches!(output, MessagesOutput::Message(_)));
let request = server.await.expect("server task completes");
assert!(
request
.to_ascii_lowercase()
.contains("x-api-key: sk-from-manager"),
"{request}"
);
let requested = secrets.requested.lock().unwrap().clone();
assert_eq!(
requested,
messages_provider_config("anthropic")
.unwrap()
.secret_names()
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
);
}
use crate::messages::Error;
#[test]
fn provider_config_resolves_anthropic_and_azure_ai() {

4
litellm-rust/crates/core/src/messages/handler.rs
View file

@ -17,8 +17,10 @@ pub(super) async fn send(
body: &Value,
timeout: Option<Duration>,
) -> Result<reqwest::Response, Error> {
let encoded = serde_json::to_vec(body)
.map_err(|err| Error::InvalidRequest(format!("failed to encode messages body: {err}")))?;
let builder = headers.iter().fold(
http_client().post(url).json(body),
http_client().post(url).body(encoded),
|builder, (key, value)| builder.header(key, value),
);
let builder = match timeout {

38
litellm-rust/crates/core/src/messages/route.rs
View file

@ -6,7 +6,6 @@ use std::{
use bytes::Bytes;
use litellm_auth::SecretValue;
use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
use litellm_host::{
event::{MachineEvent, RawResponse, RequestContext, WireRequest},
host::{Demand, Host},
@ -22,7 +21,6 @@ use serde_json::{Map, Value};
use super::{
Error,
common_utils::messages_provider_config,
handler::{decode_response, network, provider_error, send},
prepare::{prepare_provider_request, resolve_provider},
types::{MessagesRequest, MessagesShaping},
@ -54,6 +52,11 @@ pub enum MessagesOutput {
Streamed,
}
/// The upstream response as the caller sees it at stream hand-off, before any chunk.
pub struct MessagesStreamHead {
pub headers: Vec<(String, String)>,
}
pub struct Messages;
impl Protocol for Messages {
@ -62,7 +65,7 @@ impl Protocol for Messages {
type Projection = MessagesCall;
type Op = Infallible;
type Chunk = Bytes;
type StreamHead = ();
type StreamHead = MessagesStreamHead;
}
impl From<MachineFault> for Error {
@ -77,19 +80,6 @@ impl From<MachineFault> for Error {
pub type MessagesHost = HostChannel<Messages>;
pub type MessagesMachine = CallMachine<Messages>;
/// Whether this route serves the request, decided before any callback runs so a host
/// can still run its own path.
pub fn supports(model: &str, custom_llm_provider: Option<&str>, stream: bool) -> bool {
let provider = get_custom_llm_provider(model, custom_llm_provider)
.map(|resolved| resolved.custom_llm_provider)
.or(custom_llm_provider);
match provider {
Some(ANTHROPIC_MESSAGES_PROVIDER) => true,
Some(provider) => !stream && messages_provider_config(provider).is_some(),
None => false,
}
}
/// The in-process host for a request already in hand. It answers projection once and
/// observes nothing.
pub struct LocalMessagesHost {
@ -152,8 +142,11 @@ async fn execute(
model: request.model.clone(),
custom_llm_provider: request.provider.clone(),
optional_params: Value::Object(
call.body
.iter()
request
.body
.as_object()
.into_iter()
.flatten()
.filter(|(name, _)| !matches!(name.as_str(), "model" | "messages"))
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
@ -193,7 +186,14 @@ async fn relay(
host: &MessagesHost,
mut response: reqwest::Response,
) -> Result<MessagesOutput, Error> {
if host.open(()).await? == Demand::Detached {
let head = MessagesStreamHead {
headers: response
.headers()
.iter()
.filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string())))
.collect(),
};
if host.open(head).await? == Demand::Detached {
return Ok(MessagesOutput::Streamed);
}
while let Some(chunk) = response.chunk().await.map_err(network)? {

157
litellm-rust/crates/core/src/ocr/document.rs
View file

@ -244,160 +244,3 @@ mod tests {
}
}
}
#[cfg(test)]
mod document_tests {
use litellm_host::event::WireRequest;
use litellm_llms::base_llm::ocr::error::Error;
use rstest::rstest;
use serde_json::{Value, json};
use crate::ocr::route::LocalOcrHost;
use crate::ocr::test_support::{
MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with,
request_body, wire_request_with_document,
};
#[derive(Clone, Copy, Debug)]
enum Route {
Mistral,
AzureAi,
VertexMistral,
AzureCohereParse,
Cohere,
}
impl Route {
fn model(self) -> &'static str {
match self {
Self::Mistral => "mistral/model",
Self::AzureAi => "azure_ai/model",
Self::VertexMistral => "vertex_ai/mistral-ocr-maas",
Self::AzureCohereParse => "azure_ai/cohere-parse",
Self::Cohere => "cohere/model",
}
}
fn document_type(self) -> &'static str {
match self {
Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url",
Self::AzureCohereParse | Self::Cohere => "image_url",
}
}
fn options(self) -> Value {
match self {
Self::Mistral | Self::AzureAi => json!({"pages": [0]}),
Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}),
Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}),
}
}
}
/// What the host does to the wire request in `before_send`.
#[derive(Clone, Copy, Debug)]
enum Host {
Detached,
ReplacesDocument,
}
const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ=";
impl Host {
fn before_send(self, wire: WireRequest) -> WireRequest {
let Value::Object(fields) = wire.body else {
return wire;
};
let body = fields
.into_iter()
.map(|(name, value)| match self {
Self::Detached => (name, value),
Self::ReplacesDocument if name == "document" => {
let document_type = value["type"].clone();
let key = document_type.as_str().unwrap_or_default().to_string();
(name, json!({"type": document_type, key: REPLACED_DOCUMENT}))
}
Self::ReplacesDocument => (name, value),
})
.collect();
WireRequest {
body: Value::Object(body),
..wire
}
}
}
struct Sent {
result: Result<(), Error>,
provider_body: Option<Value>,
}
async fn send(route: Route, host: Host, document_base: &str) -> Sent {
let (base, seen, provider) =
mock_server(vec![MockResponse::json(json!({"pages": []}))]).await;
let document_type = route.document_type();
let document =
json!({"type": document_type, document_type: format!("{document_base}/scan.png")});
let request = wire_request_with_document(route.model(), &base, document, route.options());
let local =
LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire)));
let result = perform_ocr_with(local).await.map(|_| ());
match result {
Ok(()) => provider.await.unwrap(),
Err(_) => provider.abort(),
}
let provider_body = seen
.lock()
.unwrap()
.first()
.map(|request| request_body(request));
Sent {
result,
provider_body,
}
}
fn served_document_uri() -> String {
use base64::Engine;
format!(
"data:image/png;base64,{}",
base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT)
)
}
#[rstest]
#[case::azure_ai(Route::AzureAi)]
#[case::vertex_mistral(Route::VertexMistral)]
#[case::azure_cohere_parse(Route::AzureCohereParse)]
#[tokio::test]
async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) {
let (document_base, _documents) = document_server().await;
let sent = send(route, Host::Detached, &document_base).await;
sent.result.unwrap();
assert_eq!(
sent.provider_body.unwrap()["document"][route.document_type()],
json!(served_document_uri())
);
}
#[rstest]
#[tokio::test]
async fn document_replaced_by_the_host_reaches_the_provider(
#[values(
Route::Mistral,
Route::AzureAi,
Route::VertexMistral,
Route::AzureCohereParse,
Route::Cohere
)]
route: Route,
) {
let (document_base, _documents) = document_server().await;
let sent = send(route, Host::ReplacesDocument, &document_base).await;
sent.result.unwrap();
assert_eq!(
sent.provider_body.unwrap()["document"][route.document_type()],
json!(REPLACED_DOCUMENT)
);
}
}

209
litellm-rust/crates/core/src/ocr/mod.rs
View file

@ -7,212 +7,3 @@ pub mod provider_config;
pub mod route;
pub mod types;
pub mod wire;
#[cfg(test)]
pub(crate) mod test_support {
use std::sync::{Arc, Mutex};
use futures_util::future::BoxFuture;
use litellm_host::event::WireRequest;
use litellm_llms::base_llm::ocr::{
error::Error,
handler::{CallHooks, OcrClient},
transformation::LiteLLMOcrResponse,
};
use serde_json::{Value, json};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
};
use crate::ocr::{
route::{LocalOcrHost, ocr_machine},
types::LiteLLMOcrRequest,
wire::{OcrWireRequest, decode_request},
};
/// Stands in for a host with no hooks registered: the wire request goes out unchanged
/// and response events go nowhere.
pub(crate) struct NoHooks;
impl CallHooks<Error> for NoHooks {
fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result<WireRequest, Error>> {
Box::pin(async move { Ok(wire) })
}
fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> {
Box::pin(async { Ok(()) })
}
}
pub(crate) fn ocr_client() -> OcrClient {
let document_http = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test document client builds");
OcrClient::for_test(reqwest::Client::new(), document_http)
}
pub(crate) async fn perform_ocr(
request: LiteLLMOcrRequest,
) -> Result<LiteLLMOcrResponse, Error> {
crate::ocr::client::perform(&ocr_client(), request).await
}
pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result<LiteLLMOcrResponse, Error> {
litellm_host::run::run(ocr_machine(ocr_client()), &host).await
}
pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest {
wire_request_with_document(
model,
base,
json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}),
options,
)
}
pub(crate) fn wire_request_with_document(
model: &str,
base: &str,
document: Value,
options: Value,
) -> LiteLLMOcrRequest {
decode_request(OcrWireRequest {
model: model.into(),
document,
api_key: Some(litellm_auth::SecretValue::new("test-key")),
api_base: Some(base.into()),
custom_llm_provider: None,
extra_headers: None,
optional_params: options.as_object().unwrap().clone(),
input_sources: Default::default(),
timeout_seconds: Some(2.0),
})
.unwrap()
}
pub(crate) fn resolved_request(
request: LiteLLMOcrRequest,
) -> crate::ocr::types::ResolvedOcrRequest {
request
.map_document(crate::ocr::document::prepare_document)
.unwrap()
}
pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest {
let request = resolved_request(request);
let document = request.document.clone().with_source(source.into());
request.with_document(document.into())
}
pub(crate) fn request_body(request: &str) -> Value {
serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
}
pub(crate) const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document";
/// Serves [`SERVED_DOCUMENT`] as `image/png` to every connection until aborted.
pub(crate) async fn document_server() -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let task = tokio::spawn(async move {
loop {
let (mut socket, _) = listener.accept().await.unwrap();
let mut buffer = [0u8; 4096];
let _ = socket.read(&mut buffer).await.unwrap();
let head = format!(
"HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
SERVED_DOCUMENT.len()
);
socket.write_all(head.as_bytes()).await.unwrap();
socket.write_all(SERVED_DOCUMENT).await.unwrap();
}
});
(base, task)
}
pub(crate) struct MockResponse {
pub status: u16,
pub headers: Vec<(&'static str, String)>,
pub body: Value,
}
impl MockResponse {
pub fn json(body: Value) -> Self {
Self {
status: 200,
headers: vec![],
body,
}
}
}
pub(crate) async fn mock_server(
responses: Vec<MockResponse>,
) -> (String, Arc<Mutex<Vec<String>>>, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let requests = Arc::new(Mutex::new(Vec::new()));
let seen = requests.clone();
let server_base = base.clone();
let task = tokio::spawn(async move {
for response in responses {
let (mut socket, _) = listener.accept().await.unwrap();
let mut bytes = Vec::new();
let mut buffer = [0u8; 4096];
let header_end = loop {
let n = socket.read(&mut buffer).await.unwrap();
assert!(n > 0);
bytes.extend_from_slice(&buffer[..n]);
if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") {
break index + 4;
}
};
let length = String::from_utf8_lossy(&bytes[..header_end])
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().unwrap())
})
.unwrap_or(0);
while bytes.len() < header_end + length {
let n = socket.read(&mut buffer).await.unwrap();
assert!(n > 0);
bytes.extend_from_slice(&buffer[..n]);
}
seen.lock()
.unwrap()
.push(String::from_utf8_lossy(&bytes).into_owned());
let body = serde_json::to_vec(&response.body).unwrap();
let headers = response
.headers
.into_iter()
.map(|(name, value)| {
format!("{name}: {}\r\n", value.replace("{base}", &server_base))
})
.collect::<String>();
let head = format!(
"HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n",
response.status,
body.len(),
headers
);
socket.write_all(head.as_bytes()).await.unwrap();
socket.write_all(&body).await.unwrap();
}
});
(base, requests, task)
}
pub(crate) fn header<'a>(request: &'a str, name: &str) -> Option<&'a str> {
request
.lines()
.take_while(|line| !line.is_empty())
.find_map(|line| {
let (key, value) = line.split_once(':')?;
key.eq_ignore_ascii_case(name).then(|| value.trim())
})
}
}

205
litellm-rust/crates/core/src/ocr/prepare.rs
View file

@ -70,20 +70,203 @@ pub(crate) fn prepare_request(
}
}
#[cfg(test)]
pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedOcrRequest {
prepare_request(
request,
true,
&OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()),
std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment),
)
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use futures_util::future::BoxFuture;
use litellm_core_utils::call_arguments::{CallArguments, compose_body, parse_options};
use serde_json::json;
use litellm_host::event::WireRequest;
use litellm_llms::{
base_llm::ocr::{
error::Error,
handler::{CallHooks, OcrClient},
transformation::{BaseOcrConfig, OcrResponseFormat},
},
cohere::ocr::transformation::CohereParseConfig,
mistral::ocr::transformation::MistralOcrConfig,
vertex_ai::ocr::transformation::VertexAiOcrConfig,
};
use serde_json::{Value, json};
use super::*;
use crate::ocr::{
document::prepare_document,
types::LiteLLMOcrRequest,
wire::{OcrWireRequest, decode_request},
};
/// Stands in for a host with no hooks registered.
struct NoHooks;
impl CallHooks<Error> for NoHooks {
fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result<WireRequest, Error>> {
Box::pin(async move { Ok(wire) })
}
fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> {
Box::pin(async { Ok(()) })
}
}
fn client() -> OcrClient {
OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new())
}
fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest {
decode_request(OcrWireRequest {
model: model.into(),
document,
api_key: Some(litellm_auth::SecretValue::new("test-key")),
api_base: Some(base.into()),
custom_llm_provider: None,
extra_headers: None,
optional_params: options.as_object().unwrap().clone(),
input_sources: Default::default(),
timeout_seconds: Some(2.0),
})
.unwrap()
}
fn prepared(request: LiteLLMOcrRequest) -> PreparedOcrRequest {
prepare_request(
request.map_document(prepare_document).unwrap(),
true,
&client(),
std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment),
)
}
fn image(url: &str) -> Value {
json!({"type": "image_url", "image_url": url})
}
#[tokio::test]
async fn cohere_body_keeps_native_document_fields_and_untyped_overrides() {
let request = request(
"cohere/parse",
"https://example.com",
image("https://example.com/original.png"),
json!({
"output_format": "markdown", "timeout": 30,
"extra_body": {
"output_format": {"future": true},
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
"provider_options": {"nested": [false, 0, null]}}
}
}),
);
let http = CohereParseConfig
.prepare_request(&prepared(request), &client(), &NoHooks)
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(
body,
json!({
"model": "parse", "output_format": {"future": true},
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
"provider_options": {"nested": [false, 0, null]}}
})
);
}
#[tokio::test]
async fn explicit_null_options_use_defaults_before_http() {
let request = request(
"cohere/parse",
"https://example.com",
image("https://example.com/a.png"),
json!({"output_format": null, "req_format": null}),
);
assert_eq!(
request.response_format().unwrap(),
OcrResponseFormat::Litellm
);
let http = CohereParseConfig
.prepare_request(&prepared(request), &client(), &NoHooks)
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(body["output_format"], "markdown");
assert!(body.get("req_format").is_none());
}
#[tokio::test]
async fn direct_and_vertex_mistral_build_the_same_request_and_share_normalization() {
let options = json!({
"pages": [0, 2],
"include_image_base64": true,
"vertex_project": "project-1",
"vertex_location": "us-central1",
"unknown": "preserved"
});
let document =
json!({"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"});
let direct = prepared(request(
"mistral/mistral-ocr-maas",
"https://mistral.test",
document.clone(),
options.clone(),
));
let vertex = prepared(request(
"vertex_ai/mistral-ocr-maas",
"https://vertex.test",
document,
options,
));
let direct_http = MistralOcrConfig
.prepare_request(&direct, &client(), &NoHooks)
.await
.unwrap();
let vertex_http = VertexAiOcrConfig
.prepare_request(&vertex, &client(), &NoHooks)
.await
.unwrap();
assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr");
assert_eq!(
vertex_http.url(),
"https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
for http in [&direct_http, &vertex_http] {
assert_eq!(http.header("authorization").unwrap(), "Bearer test-key");
assert_eq!(http.header("content-type").unwrap(), "application/json");
assert_eq!(http.timeout(), Some(Duration::from_secs(2)));
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(
body,
json!({
"model": "mistral-ocr-maas",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
"pages": [0, 2],
"include_image_base64": true,
"unknown": "preserved"
})
);
}
let payload = serde_json::to_vec(
&json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}),
)
.unwrap();
let direct_response = MistralOcrConfig
.transform_ocr_response(&direct.model, &payload, OcrResponseFormat::Litellm)
.unwrap()
.into_json();
let vertex_response = VertexAiOcrConfig
.transform_ocr_response(&vertex.model, &payload, OcrResponseFormat::Litellm)
.unwrap()
.into_json();
assert_eq!(direct_response, vertex_response);
assert_eq!(direct_response["model"], "mistral-ocr-maas");
assert_eq!(direct_response["object"], "ocr");
assert_eq!(direct_response["extra"], "preserved");
}
#[derive(serde::Deserialize)]
struct KnownParams {

3521
litellm-rust/crates/core/src/ocr/route.rs
View file

File diff suppressed because it is too large Load diff

268
litellm-rust/crates/core/tests/audio_transcription.rs
View file

@ -1,50 +1,250 @@
use std::{
io::{Read, Write},
net::TcpListener,
thread,
use litellm_core::audio_transcription::{
Error, audio_transcription, types::AudioTranscriptionRequest,
};
use rstest::{fixture, rstest};
use serde_json::{Map, Value, json};
use wiremock::ResponseTemplate;
use litellm_core::audio_transcription::{audio_transcription, types::AudioTranscriptionRequest};
use serde_json::{Map, json};
mod support;
use support::*;
#[tokio::test]
async fn bedrock_request_is_signed_and_contains_audio() {
let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
let address = listener.local_addr().expect("address");
let server = thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("connection");
let mut request = Vec::new();
let mut buffer = [0_u8; 16_384];
let count = stream.read(&mut buffer).expect("request");
request.extend_from_slice(&buffer[..count]);
let request = String::from_utf8_lossy(&request);
assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse"));
assert!(request.contains("authorization: AWS4-HMAC-SHA256"));
assert!(request.contains("x-amz-date:"));
assert!(request.contains("\"bytes\":\"AQI=\""));
assert!(request.contains("Transcribe the audio. Respond with only the transcript."));
let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}";
stream.write_all(response).expect("response");
});
const MODEL: &str = "mistral.voxtral-mini-3b-2507";
let optional_params = Map::from_iter([
fn transcript_response(text: &str) -> ResponseTemplate {
json_response(json!({"output": {"message": {"content": [{"text": text}]}}}))
}
fn aws_params(region: &str) -> Map<String, Value> {
Map::from_iter([
("aws_access_key_id".to_string(), json!("access-key")),
("aws_secret_access_key".to_string(), json!("secret-key")),
("aws_region_name".to_string(), json!("us-east-1")),
]);
let api_base = format!("http://{address}");
let response = audio_transcription(AudioTranscriptionRequest {
model: "mistral.voxtral-mini-3b-2507",
("aws_region_name".to_string(), json!(region)),
])
}
#[fixture]
fn request() -> AudioTranscriptionRequest<'static> {
AudioTranscriptionRequest {
model: MODEL,
audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}),
api_key: None,
api_base: Some(&api_base),
api_base: None,
custom_llm_provider: Some("bedrock"),
extra_headers: None,
optional_params,
optional_params: aws_params("us-east-1"),
timeout: None,
}
}
#[rstest]
#[case::us_east_1("us-east-1")]
#[case::eu_west_1("eu-west-1")]
#[tokio::test]
async fn bedrock_converse_request_is_signed_for_the_requested_region(
request: AudioTranscriptionRequest<'static>,
#[case] region: &str,
) {
let upstream = upstream([transcript_response("hello")]).await;
let base = upstream.uri();
let response = audio_transcription(AudioTranscriptionRequest {
api_base: Some(&base),
optional_params: aws_params(region),
..request
})
.await
.expect("transcription");
assert_eq!(response, json!({"text": "hello"}));
server.join().expect("server");
let sent = only_request(&upstream).await;
assert_eq!(sent.method.as_str(), "POST");
assert_eq!(sent.url.path(), format!("/model/{MODEL}/converse"));
let authorization = sent.header("authorization").expect("request is signed");
assert!(
authorization.starts_with("AWS4-HMAC-SHA256 Credential=access-key/"),
"{authorization}"
);
assert!(
authorization.contains(&format!("/{region}/bedrock/aws4_request")),
"{authorization}"
);
assert!(sent.header("x-amz-date").is_some());
assert!(!sent.body_text().contains("secret-key"));
}
#[rstest]
#[tokio::test]
async fn the_provider_can_come_from_the_model_prefix(request: AudioTranscriptionRequest<'static>) {
let upstream = upstream([transcript_response("hello")]).await;
let base = upstream.uri();
let model = format!("bedrock/{MODEL}");
audio_transcription(AudioTranscriptionRequest {
model: &model,
custom_llm_provider: None,
api_base: Some(&base),
..request
})
.await
.expect("transcription");
assert_eq!(
only_request(&upstream).await.url.path(),
format!("/model/{MODEL}/converse")
);
}
#[rstest]
#[tokio::test]
async fn audio_and_transcription_params_reach_the_converse_body(
request: AudioTranscriptionRequest<'static>,
#[values("wav", "mp3", "flac", "ogg")] format: &str,
) {
let upstream = upstream([transcript_response("hello")]).await;
let base = upstream.uri();
let optional_params = aws_params("us-east-1")
.into_iter()
.chain([
("language".to_string(), json!("fr")),
("temperature".to_string(), json!(0.2)),
])
.collect();
audio_transcription(AudioTranscriptionRequest {
audio: json!({"data": "AQI=", "format": format}),
api_base: Some(&base),
optional_params,
..request
})
.await
.expect("transcription");
let body = only_request(&upstream).await.json();
let content = &body["messages"][0]["content"];
assert_eq!(
content[0],
json!({"audio": {"format": format, "source": {"bytes": "AQI="}}})
);
let instruction = content[1]["text"].as_str().expect("instruction text");
assert!(instruction.contains("fr"), "{instruction}");
assert_eq!(body["inferenceConfig"]["temperature"], 0.2);
}
#[rstest]
#[case::unknown_format(json!({"data": "AQI=", "format": "aac"}))]
#[case::missing_data(json!({"format": "wav"}))]
#[case::not_an_object(json!("AQI="))]
#[tokio::test]
async fn invalid_audio_is_rejected_before_sending(
request: AudioTranscriptionRequest<'static>,
#[case] audio: Value,
) {
let upstream = upstream([transcript_response("hello")]).await;
let base = upstream.uri();
let error = audio_transcription(AudioTranscriptionRequest {
audio,
api_base: Some(&base),
..request
})
.await
.expect_err("invalid audio is rejected");
assert!(
matches!(
error,
Error::InvalidRequest(_) | Error::MissingField(_) | Error::InvalidType { .. }
),
"{error:?}"
);
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[case::unknown_provider(MODEL, Some("openai"), "openai")]
#[case::unresolvable_model(
"no-such-model",
None,
"unable to resolve custom_llm_provider for audio transcription request"
)]
#[tokio::test]
async fn unsupported_providers_are_rejected_before_sending(
request: AudioTranscriptionRequest<'static>,
#[case] model: &'static str,
#[case] provider: Option<&'static str>,
#[case] reported: &str,
) {
let error = audio_transcription(AudioTranscriptionRequest {
model,
custom_llm_provider: provider,
api_base: Some(UNREACHABLE_BASE),
..request
})
.await
.expect_err("unsupported provider errors");
assert_eq!(error, Error::InvalidProvider(reported.into()));
}
#[rstest]
#[tokio::test]
async fn a_non_string_extra_header_is_rejected(request: AudioTranscriptionRequest<'static>) {
let error = audio_transcription(AudioTranscriptionRequest {
extra_headers: Some(Map::from_iter([("x-count".to_string(), json!(3))])),
api_base: Some(UNREACHABLE_BASE),
..request
})
.await
.expect_err("a non-string header is rejected");
assert!(matches!(error, Error::Headers(_)), "{error:?}");
}
#[rstest]
#[case::throttled(429)]
#[case::server_error(500)]
#[tokio::test]
async fn an_upstream_error_keeps_its_status_and_body(
request: AudioTranscriptionRequest<'static>,
#[case] status: u16,
) {
let upstream =
upstream([ResponseTemplate::new(status).set_body_string("upstream said no")]).await;
let base = upstream.uri();
let error = audio_transcription(AudioTranscriptionRequest {
api_base: Some(&base),
..request
})
.await
.expect_err("upstream error propagates");
assert_eq!(
error,
Error::Transport(litellm_http::transport::Error::Http {
status,
body: "upstream said no".into()
})
);
}
#[rstest]
#[case::not_json(ResponseTemplate::new(200).set_body_string("not json"))]
#[case::no_output(json_response(json!({"unexpected": true})))]
#[tokio::test]
async fn an_unreadable_success_body_is_an_invalid_response(
request: AudioTranscriptionRequest<'static>,
#[case] response: ResponseTemplate,
) {
let upstream = upstream([response]).await;
let base = upstream.uri();
let error = audio_transcription(AudioTranscriptionRequest {
api_base: Some(&base),
..request
})
.await
.expect_err("an unreadable body fails");
assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}");
}

320
litellm-rust/crates/core/tests/chat_completions.rs Normal file
View file

@ -0,0 +1,320 @@
use std::time::Duration;
use litellm_core::chat_completions::{
Error, chat_completions, chat_completions_decline_reason, types::ChatCompletionsRequest,
};
use litellm_http::transport::Error as TransportError;
use rstest::{fixture, rstest};
use serde_json::{Map, Value, json};
use wiremock::ResponseTemplate;
mod support;
use support::*;
const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
fn object(value: Value) -> Map<String, Value> {
let Value::Object(map) = value else {
panic!("expected a json object, got {value}");
};
map
}
fn anthropic_response(body: &str) -> ResponseTemplate {
ResponseTemplate::new(200).set_body_raw(body, "application/json")
}
fn hi() -> Value {
json!([{"role": "user", "content": "hi"}])
}
#[fixture]
fn request() -> ChatCompletionsRequest<'static> {
ChatCompletionsRequest {
model: "anthropic/claude-sonnet-4-5",
messages: hi(),
optional_params: object(json!({"max_tokens": 16})),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
timeout: Some(Duration::from_secs(10)),
}
}
#[rstest]
#[tokio::test]
async fn anthropic_round_trip_translates_the_conversation_and_normalizes_the_response(
request: ChatCompletionsRequest<'static>,
) {
let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await;
let base = upstream.uri();
let response = chat_completions(ChatCompletionsRequest {
messages: json!([
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"}
]),
api_base: Some(&base),
..request
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert_eq!(sent.url.path(), "/v1/messages");
assert_eq!(sent.header_values("x-api-key"), ["sk-test"]);
let body = sent.json();
assert_eq!(body["model"], "claude-sonnet-4-5");
assert_eq!(
body["messages"],
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}])
);
assert_eq!(
body["system"],
json!([{"type": "text", "text": "be terse"}])
);
assert_eq!(body["max_tokens"], 16);
assert_eq!(
response.choices[0].message.content.as_deref(),
Some("hello")
);
assert_eq!(response.usage.total_tokens, 15);
}
#[rstest]
#[tokio::test]
async fn the_deployment_key_replaces_a_caller_supplied_x_api_key(
request: ChatCompletionsRequest<'static>,
) {
let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await;
let base = upstream.uri();
chat_completions(ChatCompletionsRequest {
api_base: Some(&base),
extra_headers: Some(object(
json!({"x-api-key": "caller-key", "x-trace": "kept"}),
)),
..request
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert_eq!(sent.header_values("x-api-key"), ["sk-test"]);
assert_eq!(sent.header("x-trace"), Some("kept"));
}
#[rstest]
#[tokio::test]
async fn bedrock_round_trip_is_signed_and_normalized(request: ChatCompletionsRequest<'static>) {
let upstream = upstream([json_response(json!({
"output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}
}))])
.await;
let base = upstream.uri();
let response = chat_completions(ChatCompletionsRequest {
model: "bedrock/anthropic.claude-sonnet-4-5",
optional_params: object(json!({
"aws_access_key_id": "access-key",
"aws_secret_access_key": "secret-key",
"aws_region_name": "eu-west-1"
})),
api_key: None,
api_base: Some(&base),
..request
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert_eq!(
sent.url.path(),
"/model/anthropic.claude-sonnet-4-5/converse"
);
let authorization = sent.header("authorization").expect("request is signed");
assert!(
authorization.contains("/eu-west-1/bedrock/aws4_request"),
"{authorization}"
);
assert_eq!(
sent.json()["messages"],
json!([{"role": "user", "content": [{"text": "hi"}]}])
);
assert_eq!(
response.choices[0].message.content.as_deref(),
Some("hello")
);
assert_eq!(response.usage.total_tokens, 15);
}
/// The provider already answered and billed these, so the host must not retry them on
/// its own path: they surface as `InvalidResponse`, never as a pre-send decline.
#[rstest]
#[case::missing_usage(
r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#
)]
#[case::tool_use_block(r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#)]
#[case::not_json("not json")]
#[tokio::test]
async fn a_response_it_cannot_normalize_is_reported_as_already_sent(
request: ChatCompletionsRequest<'static>,
#[case] body: &str,
) {
let upstream = upstream([anthropic_response(body)]).await;
let base = upstream.uri();
let error = chat_completions(ChatCompletionsRequest {
api_base: Some(&base),
..request
})
.await
.expect_err("response cannot be normalized");
assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}");
}
#[rstest]
#[case::rate_limited(429)]
#[case::server_error(500)]
#[tokio::test]
async fn an_upstream_error_status_keeps_its_code_and_body(
request: ChatCompletionsRequest<'static>,
#[case] status: u16,
) {
let upstream = upstream([ResponseTemplate::new(status).set_body_string("slow down")]).await;
let base = upstream.uri();
let error = chat_completions(ChatCompletionsRequest {
api_base: Some(&base),
..request
})
.await
.expect_err("upstream rejects");
assert_eq!(
error,
Error::Transport(TransportError::Http {
status,
body: "slow down".into()
})
);
}
/// Nothing was sent, so nothing was billed and the host can still serve the request.
#[rstest]
#[tokio::test]
async fn a_connection_that_is_never_established_declines_instead_of_failing(
request: ChatCompletionsRequest<'static>,
) {
let error = chat_completions(ChatCompletionsRequest {
api_base: Some(UNREACHABLE_BASE),
..request
})
.await
.expect_err("nothing is listening");
assert!(
matches!(error, Error::Transport(TransportError::Connect(_))),
"{error:?}"
);
}
#[rstest]
#[tokio::test]
async fn a_timeout_after_sending_is_not_a_pre_send_decline(
request: ChatCompletionsRequest<'static>,
) {
let upstream =
upstream([anthropic_response(ANTHROPIC_MESSAGE).set_delay(Duration::from_secs(5))]).await;
let base = upstream.uri();
let error = chat_completions(ChatCompletionsRequest {
api_base: Some(&base),
timeout: Some(Duration::from_millis(100)),
..request
})
.await
.expect_err("the call times out");
assert!(
matches!(error, Error::Transport(TransportError::Network(_))),
"{error:?}"
);
}
#[rstest]
#[case::accepted("anthropic/claude-sonnet-4-5", None, hi(), json!({"max_tokens": 16}), None)]
#[case::accepted_bedrock("bedrock/anthropic.claude-sonnet-4-5", None, hi(), json!({}), None)]
#[case::unknown_provider(
"gpt-4o",
Some("openai"),
hi(),
json!({}),
Some("provider is not on the rust chat completions path")
)]
#[case::unreadable_messages(
"anthropic/claude-sonnet-4-5",
None,
json!("hi"),
json!({}),
Some("unreadable message list")
)]
#[case::empty_messages("anthropic/claude-sonnet-4-5", None, json!([]), json!({}), Some("empty message list"))]
#[case::streaming(
"anthropic/claude-sonnet-4-5",
None,
hi(),
json!({"stream": true}),
Some("streaming")
)]
#[case::unrecognized_param(
"anthropic/claude-sonnet-4-5",
None,
hi(),
json!({"not_a_param": 1}),
Some("unrecognized request parameter")
)]
#[case::opens_on_assistant_turn(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "assistant", "content": "hi"}]),
json!({}),
Some("conversation does not open on a user turn")
)]
fn decline_reason_names_why_the_core_would_not_serve_the_request(
#[case] model: &str,
#[case] provider: Option<&str>,
#[case] messages: Value,
#[case] params: Value,
#[case] reason: Option<&str>,
) {
assert_eq!(
chat_completions_decline_reason(model, provider, messages, &object(params)),
reason
);
}
/// A request the decline check accepts must not be declined by the call itself.
#[rstest]
#[tokio::test]
async fn a_declined_request_fails_the_call_before_sending(
request: ChatCompletionsRequest<'static>,
) {
let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await;
let base = upstream.uri();
let error = chat_completions(ChatCompletionsRequest {
optional_params: object(json!({"stream": true})),
api_base: Some(&base),
..request
})
.await
.expect_err("streaming is declined");
assert_eq!(error, Error::Unsupported("streaming"));
assert!(received(&upstream).await.is_empty());
}

471
litellm-rust/crates/core/tests/messages.rs
View file

@ -1,471 +0,0 @@
use std::{sync::Arc, time::Duration};
use futures_util::future::BoxFuture;
use litellm_core::messages::{
Error, messages,
route::{LocalMessagesHost, MessagesCall, messages_machine},
types::{MessagesRequest, MessagesShaping},
};
use litellm_secrets::{SecretValue, source::SecretSource};
use serde_json::{Map, Value, json};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
};
struct RecordingSecrets {
values: Vec<(&'static str, String)>,
fails: bool,
requested: std::sync::Mutex<Vec<String>>,
}
impl RecordingSecrets {
fn new(values: Vec<(&'static str, String)>, fails: bool) -> Self {
Self {
values,
fails,
requested: std::sync::Mutex::new(Vec::new()),
}
}
}
impl SecretSource for RecordingSecrets {
fn get_secret_str<'a>(
&'a self,
name: &'a str,
) -> BoxFuture<'a, Result<Option<SecretValue>, litellm_secrets::Error>> {
Box::pin(async move {
self.requested.lock().unwrap().push(name.to_string());
if self.fails {
return Err(litellm_secrets::Error::ManagedSecretMissing);
}
Ok(self
.values
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| SecretValue::new(value.clone())))
})
}
}
fn secrets_call() -> MessagesCall {
let Value::Object(body) = json!({
"model": "claude-sonnet-4-5",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hi"}]
}) else {
unreachable!("literal object")
};
MessagesCall {
model: "claude-sonnet-4-5".into(),
body,
api_key: None,
api_base: None,
custom_llm_provider: Some("anthropic".into()),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
}
}
#[tokio::test]
async fn route_surfaces_a_secret_manager_failure_before_the_call() {
let Err(error) = litellm_host::run::run(
messages_machine(Arc::new(RecordingSecrets::new(Vec::new(), true))),
&LocalMessagesHost::new(secrets_call()),
)
.await
else {
panic!("a secret manager failure fails the call");
};
assert!(
matches!(&error, Error::Secret(source) if matches!(source.source_error(), litellm_secrets::Error::ManagedSecretMissing)),
"{error:?}"
);
}
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
let header_end = loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break request.len();
}
request.extend_from_slice(&buffer[..n]);
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
break position + 4;
}
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while request.len().saturating_sub(header_end) < content_length {
let n = socket.read(&mut buffer).await.expect("reads body");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
}
String::from_utf8(request).expect("request is utf8")
}
fn write_response(body: &str) -> String {
format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
)
}
#[tokio::test]
async fn messages_round_trip_builds_azure_request_and_passes_response_through() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#;
socket
.write_all(write_response(response_body).as_bytes())
.await
.expect("writes response");
request
});
let response = messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({
"model": "claude-sonnet-4-5",
"max_tokens": 1024,
"messages": [{
"role": "user",
"content": [{
"type": "text",
"text": "hi",
"cache_control": {"type": "ephemeral", "scope": "global"}
}]
}]
}),
api_key: Some("sk-azure"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
})
.await
.expect("messages request succeeds");
assert_eq!(response.content[0]["text"], "hi");
assert_eq!(response.stop_reason.as_deref(), Some("end_turn"));
let request = server.await.expect("server task completes");
let (head, body) = request.split_once("\r\n\r\n").expect("has body");
assert!(head.starts_with("POST /anthropic/v1/messages "), "{head}");
let head_lower = head.to_ascii_lowercase();
assert!(head_lower.contains("x-api-key: sk-azure"), "{head}");
assert!(
head_lower.contains("anthropic-version: 2023-06-01"),
"{head}"
);
assert!(
head_lower.contains("content-type: application/json"),
"{head}"
);
let sent_body: Value = serde_json::from_str(body).expect("body is json");
assert_eq!(
sent_body["messages"][0]["content"][0]["cache_control"],
json!({"type": "ephemeral"})
);
}
#[tokio::test]
async fn messages_round_trip_builds_native_anthropic_request() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#;
socket
.write_all(write_response(response_body).as_bytes())
.await
.expect("writes response");
request
});
let response = messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({
"model": "claude-sonnet-4-5",
"max_tokens": 1024,
"messages": [{"role": "user", "content": "hi"}]
}),
api_key: Some("sk-ant"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("anthropic"),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
})
.await
.expect("messages request succeeds");
assert_eq!(response.content[0]["text"], "hi");
assert_eq!(response.stop_reason.as_deref(), Some("end_turn"));
let request = server.await.expect("server task completes");
let (head, _) = request.split_once("\r\n\r\n").expect("has body");
assert!(head.starts_with("POST /v1/messages "), "{head}");
let head_lower = head.to_ascii_lowercase();
assert!(head_lower.contains("x-api-key: sk-ant"), "{head}");
assert!(
head_lower.contains("anthropic-version: 2023-06-01"),
"{head}"
);
}
#[tokio::test]
async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body =
r#"{"id":"msg_2","type":"message","role":"assistant","content":[],"model":"m"}"#;
socket
.write_all(write_response(response_body).as_bytes())
.await
.expect("writes response");
request
});
let mut headers = Map::new();
headers.insert(
"x-api-key".to_string(),
Value::String("from-python".to_string()),
);
headers.insert(
"anthropic-beta".to_string(),
Value::String("token-efficient-tools-2025-02-19".to_string()),
);
messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
api_key: Some("rust-fallback-key"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: Some(headers),
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
})
.await
.expect("messages request succeeds");
let request = server.await.expect("server task completes");
let head = request
.split_once("\r\n\r\n")
.expect("has body")
.0
.to_ascii_lowercase();
let api_key_count = head
.lines()
.filter(|line| line.starts_with("x-api-key:"))
.count();
assert_eq!(api_key_count, 1, "{head}");
assert!(head.contains("x-api-key: from-python"), "{head}");
assert!(
head.contains("anthropic-beta: token-efficient-tools-2025-02-19"),
"{head}"
);
assert!(!head.contains("rust-fallback-key"), "{head}");
}
#[tokio::test]
async fn messages_forwards_entra_id_bearer_without_requiring_api_key() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body =
r#"{"id":"msg_3","type":"message","role":"assistant","content":[],"model":"m"}"#;
socket
.write_all(write_response(response_body).as_bytes())
.await
.expect("writes response");
request
});
let mut headers = Map::new();
headers.insert(
"Authorization".to_string(),
Value::String("Bearer entra-token".to_string()),
);
messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
api_key: None,
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: Some(headers),
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
})
.await
.expect("entra id request succeeds without api key");
let request = server.await.expect("server task completes");
let head = request
.split_once("\r\n\r\n")
.expect("has body")
.0
.to_ascii_lowercase();
assert!(head.contains("authorization: bearer entra-token"), "{head}");
assert!(!head.contains("x-api-key"), "{head}");
}
#[tokio::test]
async fn messages_requires_auth_when_no_key_and_no_header() {
let err = messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
api_key: None,
api_base: Some("http://127.0.0.1:1"),
custom_llm_provider: Some("azure_ai"),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_millis(50)),
shaping: MessagesShaping::default(),
})
.await
.expect_err("missing auth errors");
assert!(matches!(err, Error::Auth(_)));
}
#[tokio::test]
async fn messages_ignores_malformed_authorization_and_uses_api_key() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body =
r#"{"id":"msg_4","type":"message","role":"assistant","content":[],"model":"m"}"#;
socket
.write_all(write_response(response_body).as_bytes())
.await
.expect("writes response");
request
});
let mut headers = Map::new();
headers.insert(
"Authorization".to_string(),
Value::String("Bearer ".to_string()),
);
messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
api_key: Some("sk-azure"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: Some(headers),
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
})
.await
.expect("falls back to api key");
let request = server.await.expect("server task completes");
let head = request
.split_once("\r\n\r\n")
.expect("has body")
.0
.to_ascii_lowercase();
assert!(head.contains("x-api-key: sk-azure"), "{head}");
}
#[tokio::test]
async fn messages_maps_provider_error_status_to_http_error() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let _ = read_http_request(&mut socket).await;
let body = "unauthorized";
let response = format!(
"HTTP/1.1 401 Unauthorized\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
});
let err = messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
api_key: Some("sk-azure"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
})
.await
.expect_err("provider error propagates");
assert!(matches!(
err,
Error::Transport(litellm_http::transport::Error::Http { status: 401, .. })
));
}
#[tokio::test]
async fn messages_rejects_unsupported_provider() {
let err = messages(MessagesRequest {
model: "claude-3-5-sonnet",
body: json!({"model": "claude-3-5-sonnet", "max_tokens": 8, "messages": []}),
api_key: Some("sk"),
api_base: Some("http://127.0.0.1:1"),
custom_llm_provider: Some("openai"),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_millis(50)),
shaping: MessagesShaping::default(),
})
.await
.expect_err("unsupported provider errors");
assert!(matches!(err, Error::InvalidProvider(provider) if provider == "openai"));
}

210
litellm-rust/crates/core/tests/messages/host.rs Normal file
View file

@ -0,0 +1,210 @@
use std::{convert::Infallible, sync::Mutex};
use litellm_core::messages::route::Messages;
use litellm_host::{
event::{CallEvent, MachineEvent, RequestContext, WireRequest},
host::Host,
};
use litellm_llms::anthropic::common_utils::AnthropicModelCapabilities;
use rstest::rstest;
use super::*;
type Rewrite = Box<dyn Fn(WireRequest) -> Result<WireRequest, Error> + Send + Sync>;
/// Projects like `LocalMessagesHost`, answers `before_send` through `rewrite`, and keeps
/// every event the driver emits.
struct RecordingHost {
call: LocalMessagesHost,
rewrite: Rewrite,
events: Mutex<Vec<CallEvent>>,
optional_params: Mutex<Vec<Value>>,
}
impl RecordingHost {
fn new(call: MessagesCall, rewrite: Rewrite) -> Self {
Self {
call: LocalMessagesHost::new(call),
rewrite,
events: Mutex::new(Vec::new()),
optional_params: Mutex::new(Vec::new()),
}
}
fn passthrough(call: MessagesCall) -> Self {
Self::new(call, Box::new(Ok))
}
fn raw_responses(&self) -> Vec<String> {
self.events
.lock()
.unwrap()
.iter()
.filter_map(|event| match event {
CallEvent::Machine(MachineEvent::ResponseReceived { raw }) => {
Some(raw.body.clone())
}
_ => None,
})
.collect()
}
}
impl Host<Messages> for RecordingHost {
async fn project(&self) -> Result<MessagesCall, Error> {
self.call.project().await
}
async fn custom_op(&self, op: Infallible) -> Result<(), Error> {
match op {}
}
async fn before_send(
&self,
wire: WireRequest,
context: &RequestContext,
) -> Result<WireRequest, Error> {
self.optional_params
.lock()
.unwrap()
.push(context.optional_params.clone());
(self.rewrite)(wire)
}
async fn emit(&self, event: &CallEvent) -> Result<(), Error> {
self.events.lock().unwrap().push(event.clone());
Ok(())
}
}
async fn run_through(host: &RecordingHost) -> Result<MessagesOutput, Error> {
litellm_host::run::run(messages_machine(Arc::new(RecordingSecrets::empty())), host).await
}
fn authenticated(call: MessagesCall, api_base: String) -> MessagesCall {
MessagesCall {
api_key: Some("sk-ant".into()),
api_base: Some(api_base),
..call
}
}
#[rstest]
#[tokio::test]
async fn what_before_send_returns_is_what_the_provider_receives(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let host = RecordingHost::new(
authenticated(call, upstream.uri()),
Box::new(|wire| {
let mut body = wire.body;
body["system"] = json!("added by the host");
Ok(WireRequest {
headers: wire
.headers
.into_iter()
.chain([("x-host".to_string(), "seen".to_string())])
.collect(),
body,
..wire
})
}),
);
run_through(&host).await.expect("messages call succeeds");
let request = only_request(&upstream).await;
assert_eq!(request.json()["system"], "added by the host");
assert_eq!(request.header("x-host"), Some("seen"));
assert_eq!(request.header("x-api-key"), Some("sk-ant"));
}
#[rstest]
#[tokio::test]
async fn a_before_send_failure_never_sends(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let host = RecordingHost::new(
authenticated(call, upstream.uri()),
Box::new(|_| Err(Error::InvalidRequest("vetoed by the host".into()))),
);
let error = run_through(&host)
.await
.err()
.expect("the host failure fails the call");
assert_eq!(error, Error::InvalidRequest("vetoed by the host".into()));
assert!(received(&upstream).await.is_empty());
assert!(host.raw_responses().is_empty());
}
#[rstest]
#[tokio::test]
async fn the_raw_upstream_text_is_emitted_once_for_a_message(call: MessagesCall) {
let raw = message_body();
let upstream = upstream([json_response(raw.clone())]).await;
let host = RecordingHost::passthrough(authenticated(call, upstream.uri()));
let output = run_through(&host).await.expect("messages call succeeds");
assert!(matches!(output, MessagesOutput::Message(_)));
let [emitted] = <[String; 1]>::try_from(host.raw_responses())
.unwrap_or_else(|raws| panic!("expected one raw response, got {}", raws.len()));
assert_eq!(serde_json::from_str::<Value>(&emitted).unwrap(), raw);
}
#[rstest]
#[case::upstream_error(ResponseTemplate::new(500).set_body_string("boom"))]
#[case::stream(ResponseTemplate::new(200).set_body_raw("event: message_stop\ndata: {}\n\n", "text/event-stream"))]
#[tokio::test]
async fn no_raw_response_is_emitted_for_a_stream_or_a_failure(
call: MessagesCall,
#[case] response: ResponseTemplate,
) {
let upstream = upstream([response]).await;
let mut body = call.body.clone();
body.insert("stream".into(), json!(true));
let host =
RecordingHost::passthrough(authenticated(MessagesCall { body, ..call }, upstream.uri()));
let _ = run_through(&host).await;
assert_eq!(received(&upstream).await.len(), 1);
assert!(host.raw_responses().is_empty());
}
/// Python logs `optional_params` as what it is about to send, so a dropped param must
/// not resurface in callbacks.
#[rstest]
#[tokio::test]
async fn the_request_context_carries_the_shaped_params_without_model_or_messages(
call: MessagesCall,
) {
let upstream = upstream([message_response()]).await;
let body: Map<String, Value> = call
.body
.clone()
.into_iter()
.chain([("temperature".to_string(), json!(0.2))])
.collect();
let host = RecordingHost::passthrough(authenticated(
MessagesCall {
body,
shaping: MessagesShaping {
capabilities: AnthropicModelCapabilities {
supports_sampling_params: false,
..AnthropicModelCapabilities::default()
},
drop_params: true,
..MessagesShaping::default()
},
..call
},
upstream.uri(),
));
run_through(&host).await.expect("messages call succeeds");
let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap())
.unwrap_or_else(|seen| panic!("before_send runs once, saw {}", seen.len()));
assert_eq!(optional_params, json!({"max_tokens": 16}));
}

95
litellm-rust/crates/core/tests/messages/main.rs Normal file
View file

@ -0,0 +1,95 @@
use std::{sync::Arc, time::Duration};
use litellm_core::messages::{
Error,
route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine},
types::MessagesShaping,
};
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use rstest::fixture;
use serde_json::{Map, Value, json};
use wiremock::ResponseTemplate;
#[path = "../support/mod.rs"]
mod support;
use support::*;
mod host;
mod request;
mod response;
mod secrets;
mod stream;
const MODEL: &str = "claude-sonnet-4-5";
fn object(value: Value) -> Map<String, Value> {
let Value::Object(map) = value else {
panic!("expected a json object, got {value}");
};
map
}
fn message_body() -> Value {
json!({
"id": "msg_1",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "hi"}],
"model": MODEL,
"stop_reason": "end_turn",
"usage": {"input_tokens": 1, "output_tokens": 2}
})
}
fn message_response() -> ResponseTemplate {
json_response(message_body())
}
/// A non-streaming call with nothing that would authenticate or route it, so each test
/// states the provider, credentials, and base it depends on.
#[fixture]
fn call() -> MessagesCall {
MessagesCall {
model: MODEL.into(),
body: object(json!({
"model": MODEL,
"max_tokens": 16,
"messages": [{"role": "user", "content": "hi"}]
})),
api_key: None,
api_base: None,
custom_llm_provider: Some("anthropic".into()),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
}
}
fn headers<'a>(pairs: impl IntoIterator<Item = (&'a str, &'a str)>) -> Option<Map<String, Value>> {
Some(
pairs
.into_iter()
.map(|(name, value)| (name.to_string(), Value::from(value)))
.collect(),
)
}
async fn run_with(
secrets: Arc<RecordingSecrets>,
call: MessagesCall,
) -> Result<MessagesOutput, Error> {
litellm_host::run::run(messages_machine(secrets), &LocalMessagesHost::new(call)).await
}
/// Runs the route with a secret source that knows nothing, so no environment leaks in.
async fn run(call: MessagesCall) -> Result<MessagesOutput, Error> {
run_with(Arc::new(RecordingSecrets::empty()), call).await
}
async fn run_message(call: MessagesCall) -> AnthropicMessagesResponse {
match run(call).await.expect("messages call succeeds") {
MessagesOutput::Message(message) => *message,
MessagesOutput::Streamed => panic!("a non-streaming call returned a stream"),
}
}

675
litellm-rust/crates/core/tests/messages/request.rs Normal file
View file

@ -0,0 +1,675 @@
use litellm_llms::anthropic::common_utils::{
ANTHROPIC_ADVISOR_TOOL_TYPE, ANTHROPIC_OAUTH_BETA_HEADER, AnthropicModelCapabilities,
SupportedEffortTiers, beta,
};
use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders};
use rstest::rstest;
use super::*;
#[rstest]
#[case::anthropic_key("anthropic", Some("sk-ant"), &[], ("x-api-key", "sk-ant"), &["authorization"])]
#[case::azure_key("azure_ai", Some("sk-azure"), &[], ("x-api-key", "sk-azure"), &["authorization"])]
#[case::caller_x_api_key_wins(
"azure_ai",
Some("rust-fallback-key"),
&[("x-api-key", "from-python")],
("x-api-key", "from-python"),
&["authorization"]
)]
#[case::entra_bearer_without_key(
"azure_ai",
None,
&[("Authorization", "Bearer entra-token")],
("authorization", "Bearer entra-token"),
&["x-api-key"]
)]
#[case::empty_bearer_falls_back_to_key(
"azure_ai",
Some("sk-azure"),
&[("Authorization", "Bearer ")],
("x-api-key", "sk-azure"),
&[]
)]
#[case::anthropic_forwards_caller_authorization(
"anthropic",
Some("sk-ant"),
&[("Authorization", "Bearer caller")],
("authorization", "Bearer caller"),
&["x-api-key"]
)]
#[case::anthropic_oauth_key_becomes_bearer(
"anthropic",
Some("sk-ant-oat01-token"),
&[],
("authorization", "Bearer sk-ant-oat01-token"),
&["x-api-key"]
)]
#[tokio::test]
async fn credentials_become_exactly_one_auth_header(
call: MessagesCall,
#[case] provider: &str,
#[case] api_key: Option<&str>,
#[case] extra_headers: &[(&str, &str)],
#[case] expected: (&str, &str),
#[case] absent: &[&str],
) {
let upstream = upstream([message_response()]).await;
run_message(MessagesCall {
custom_llm_provider: Some(provider.into()),
api_key: api_key.map(Into::into),
api_base: Some(upstream.uri()),
extra_headers: headers(extra_headers.iter().copied()),
..call
})
.await;
let request = only_request(&upstream).await;
let (name, value) = expected;
assert_eq!(request.header_values(name), [value]);
for name in absent {
assert_eq!(request.header(name), None, "{name} must not be sent");
}
}
#[rstest]
#[case::anthropic("anthropic")]
#[case::azure_ai("azure_ai")]
#[tokio::test]
async fn a_call_without_credentials_fails_before_sending(
call: MessagesCall,
#[case] provider: &str,
) {
let upstream = upstream([message_response()]).await;
let error = run(MessagesCall {
custom_llm_provider: Some(provider.into()),
api_base: Some(upstream.uri()),
..call
})
.await
.err()
.expect("a call without credentials fails");
assert!(
matches!(
error,
Error::Auth(litellm_auth::Error::MissingApiKey { .. })
),
"{error:?}"
);
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[case::anthropic(MODEL, Some("anthropic"), "", "/v1/messages")]
#[case::anthropic_base_with_trailing_slash(MODEL, Some("anthropic"), "/", "/v1/messages")]
#[case::anthropic_base_with_the_messages_path(
MODEL,
Some("anthropic"),
"/v1/messages",
"/v1/messages"
)]
#[case::azure_ai(MODEL, Some("azure_ai"), "", "/anthropic/v1/messages")]
#[case::provider_from_model_prefix("anthropic/claude-sonnet-4-5", None, "", "/v1/messages")]
#[tokio::test]
async fn each_provider_posts_to_its_messages_endpoint(
call: MessagesCall,
#[case] model: &str,
#[case] provider: Option<&str>,
#[case] base_suffix: &str,
#[case] path: &str,
) {
let upstream = upstream([message_response()]).await;
run_message(MessagesCall {
model: model.into(),
custom_llm_provider: provider.map(Into::into),
api_key: Some("sk".into()),
api_base: Some(format!("{}{base_suffix}", upstream.uri())),
..call
})
.await;
let request = only_request(&upstream).await;
assert_eq!(request.method.as_str(), "POST");
assert_eq!(request.url.path(), path);
assert_eq!(request.json()["model"], MODEL);
assert_eq!(request.header_values("anthropic-version"), ["2023-06-01"]);
assert_eq!(request.header_values("content-type"), ["application/json"]);
}
#[rstest]
#[case::unknown_provider(MODEL, Some("openai"), "openai")]
#[case::unresolvable_model(
"no-such-model",
None,
"unable to resolve custom_llm_provider for messages request"
)]
#[tokio::test]
async fn unsupported_providers_are_rejected_before_sending(
call: MessagesCall,
#[case] model: &str,
#[case] provider: Option<&str>,
#[case] reported: &str,
) {
let error = run(MessagesCall {
model: model.into(),
custom_llm_provider: provider.map(Into::into),
api_key: Some("sk".into()),
api_base: Some(UNREACHABLE_BASE.into()),
..call
})
.await
.err()
.expect("unsupported provider errors");
assert_eq!(error, Error::InvalidProvider(reported.into()));
}
#[rstest]
#[tokio::test]
async fn caller_headers_and_provider_scoped_headers_are_forwarded(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let scoped = |provider: &str, value: &str| ProviderSpecificHeader {
custom_llm_provider: provider.into(),
extra_headers: object(json!({"x-scoped": value})),
};
run_message(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
extra_headers: headers([("anthropic-beta", "token-efficient-tools-2025-02-19")]),
provider_specific_header: Some(ProviderSpecificHeaders::Many(vec![
scoped("bedrock", "other-provider"),
scoped("azure_ai, anthropic", "this-provider"),
])),
..call
})
.await;
let request = only_request(&upstream).await;
assert_eq!(
request.header("anthropic-beta"),
Some("token-efficient-tools-2025-02-19")
);
assert_eq!(request.header_values("x-scoped"), ["this-provider"]);
}
#[rstest]
#[tokio::test]
async fn azure_strips_the_cache_control_scope_anthropic_rejects(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
run_message(MessagesCall {
custom_llm_provider: Some("azure_ai".into()),
api_key: Some("sk-azure".into()),
api_base: Some(upstream.uri()),
body: object(json!({
"model": MODEL,
"max_tokens": 16,
"messages": [{
"role": "user",
"content": [{
"type": "text",
"text": "hi",
"cache_control": {"type": "ephemeral", "scope": "global"}
}]
}]
})),
..call
})
.await;
assert_eq!(
only_request(&upstream).await.json()["messages"][0]["content"][0]["cache_control"],
json!({"type": "ephemeral"})
);
}
#[rstest]
#[tokio::test]
async fn additional_drop_params_remove_fields_before_sending(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let mut body = call.body.clone();
body.insert("temperature".into(), json!(0.5));
body.insert("top_k".into(), json!(3));
run_message(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
body,
shaping: MessagesShaping {
additional_drop_params: vec!["temperature".into()],
..MessagesShaping::default()
},
..call
})
.await;
let sent = only_request(&upstream).await.json();
assert_eq!(sent.get("temperature"), None);
assert_eq!(sent["top_k"], 3);
}
fn with_fields(call: MessagesCall, fields: Value) -> MessagesCall {
let body: Map<String, Value> = call.body.into_iter().chain(object(fields)).collect();
MessagesCall { body, ..call }
}
fn sent_betas(request: &wiremock::Request) -> Vec<String> {
let [header] = <[&str; 1]>::try_from(request.header_values("anthropic-beta"))
.unwrap_or_else(|values| panic!("expected one anthropic-beta header, got {values:?}"));
header
.split(',')
.map(str::trim)
.map(str::to_string)
.collect()
}
#[rstest]
#[case::structured_output(json!({"output_format": {"type": "json_schema"}}), &[beta::STRUCTURED_OUTPUT])]
#[case::fast_mode(json!({"speed": "fast"}), &[beta::FAST_MODE_2026_02_01])]
#[case::compaction(json!({"compaction": {"enabled": true}}), &[beta::COMPACT_2026_09_04])]
#[case::context_management_edits(
json!({"context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}}),
&[beta::CONTEXT_MANAGEMENT_2025_06_27]
)]
#[case::per_message_output_config(
json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}),
&[beta::PER_TURN_CONTROL_2026_07_01]
)]
#[case::advisor_tool(
json!({"tools": [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", "model": MODEL}]}),
&[beta::ADVISOR_TOOL_2026_03_01]
)]
#[case::several_features_at_once(
json!({"speed": "fast", "output_format": {"type": "json_schema"}}),
&[beta::STRUCTURED_OUTPUT, beta::FAST_MODE_2026_02_01]
)]
#[tokio::test]
async fn feature_betas_join_the_callers_betas_in_one_sorted_header(
call: MessagesCall,
#[case] fields: Value,
#[case] features: &[&str],
) {
let upstream = upstream([message_response()]).await;
let capabilities = AnthropicModelCapabilities {
supports_speed: true,
..AnthropicModelCapabilities::default()
};
run_message(with_fields(
MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]),
shaping: MessagesShaping {
capabilities,
..MessagesShaping::default()
},
..call
},
fields,
))
.await;
let sent = sent_betas(&only_request(&upstream).await);
let mut expected: Vec<String> = features
.iter()
.map(|feature| feature.to_string())
.chain(["caller-beta-2025-01-01".to_string()])
.collect();
expected.sort();
assert_eq!(sent, expected);
}
#[rstest]
#[tokio::test]
async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
run_message(MessagesCall {
api_key: Some("sk-ant-oat01-token".into()),
api_base: Some(upstream.uri()),
..call
})
.await;
let request = only_request(&upstream).await;
assert_eq!(
request.header("anthropic-dangerous-direct-browser-access"),
Some("true")
);
assert_eq!(sent_betas(&request), [ANTHROPIC_OAUTH_BETA_HEADER]);
assert_eq!(request.header("x-api-key"), None);
}
#[rstest]
#[case::anthropic("anthropic")]
#[case::azure_ai("azure_ai")]
#[tokio::test]
async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[case] provider: &str) {
let upstream = upstream([message_response()]).await;
run_message(MessagesCall {
custom_llm_provider: Some(provider.into()),
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
extra_headers: headers([
("Anthropic-Version", "2024-01-01"),
("Content-Type", "application/json; charset=utf-8"),
]),
..call
})
.await;
let request = only_request(&upstream).await;
assert_eq!(request.header_values("anthropic-version"), ["2024-01-01"]);
assert_eq!(
request.header_values("content-type"),
["application/json; charset=utf-8"]
);
}
fn sampling_removed() -> AnthropicModelCapabilities {
AnthropicModelCapabilities {
supports_sampling_params: false,
..AnthropicModelCapabilities::default()
}
}
#[rstest]
#[case::sampling_params(sampling_removed(), json!({"temperature": 0.2, "top_p": 0.9, "top_k": 5}), &["temperature", "top_p", "top_k"], "temperature=0.2")]
#[case::speed(AnthropicModelCapabilities::default(), json!({"speed": "fast"}), &["speed"], "speed='fast'")]
#[tokio::test]
async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_it(
call: MessagesCall,
#[case] capabilities: AnthropicModelCapabilities,
#[case] fields: Value,
#[case] dropped: &[&str],
#[case] rejected_as: &str,
) {
let upstream = upstream([message_response(), message_response()]).await;
let shaped = |drop_params: bool| {
with_fields(
MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
capabilities,
drop_params,
..MessagesShaping::default()
},
body: call.body.clone(),
custom_llm_provider: call.custom_llm_provider.clone(),
extra_headers: None,
provider_specific_header: None,
model: call.model.clone(),
timeout: call.timeout,
},
fields.clone(),
)
};
let error = run(shaped(false))
.await
.err()
.expect("an unsupported param is rejected without drop_params");
assert!(
matches!(&error, Error::InvalidRequest(message) if message.contains(rejected_as)),
"{error:?}"
);
assert!(received(&upstream).await.is_empty());
run_message(shaped(true)).await;
let sent = only_request(&upstream).await.json();
for name in dropped {
assert_eq!(sent.get(*name), None, "{name} must be dropped");
}
assert_eq!(sent["max_tokens"], 16);
}
#[rstest]
#[case::adaptive_thinking(json!({"type": "adaptive"}), json!({"type": "adaptive", "display": "summarized"}))]
#[case::disabled_thinking(json!({"type": "disabled"}), json!({"type": "disabled"}))]
#[tokio::test]
async fn reasoning_auto_summary_marks_active_thinking_on_the_wire(
call: MessagesCall,
#[case] thinking: Value,
#[case] expected: Value,
) {
let upstream = upstream([message_response()]).await;
run_message(with_fields(
MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
capabilities: AnthropicModelCapabilities {
supports_reasoning: true,
supports_adaptive_thinking: true,
..AnthropicModelCapabilities::default()
},
reasoning_auto_summary: true,
..MessagesShaping::default()
},
..call
},
json!({"thinking": thinking}),
))
.await;
assert_eq!(only_request(&upstream).await.json()["thinking"], expected);
}
#[rstest]
#[case::reasoning_effort_on_an_adaptive_model(
AnthropicModelCapabilities {
supports_reasoning: true,
supports_adaptive_thinking: true,
supports_output_config: true,
effort_tiers: SupportedEffortTiers { high: true, ..SupportedEffortTiers::default() },
..AnthropicModelCapabilities::default()
},
json!({"reasoning_effort": "high"}),
json!({"thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "high"}})
)]
#[case::reasoning_effort_on_a_legacy_model_caps_the_budget_below_max_tokens(
AnthropicModelCapabilities {
supports_reasoning: true,
..AnthropicModelCapabilities::default()
},
json!({"reasoning_effort": "high"}),
json!({"thinking": {"type": "enabled", "budget_tokens": 2999}})
)]
#[case::adaptive_payload_on_a_legacy_model_becomes_a_capped_budget(
AnthropicModelCapabilities {
supports_reasoning: true,
..AnthropicModelCapabilities::default()
},
json!({"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}, "temperature": 0}),
json!({"thinking": {"type": "enabled", "budget_tokens": 2999}})
)]
#[case::adaptive_payload_on_a_model_without_reasoning_is_dropped(
AnthropicModelCapabilities::default(),
json!({"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}),
json!({})
)]
#[tokio::test]
async fn reasoning_is_translated_by_the_model_capabilities(
call: MessagesCall,
#[case] capabilities: AnthropicModelCapabilities,
#[case] fields: Value,
#[case] expected: Value,
) {
let upstream = upstream([message_response()]).await;
run_message(with_fields(
MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
capabilities,
..MessagesShaping::default()
},
..call
},
[("max_tokens".to_string(), json!(3000))]
.into_iter()
.chain(object(fields))
.collect(),
))
.await;
let sent = only_request(&upstream).await.json();
assert_eq!(sent.get("reasoning_effort"), None);
assert_eq!(sent.get("temperature"), None);
let reasoning: Map<String, Value> = ["thinking", "output_config"]
.into_iter()
.filter_map(|name| Some((name.to_string(), sent.get(name)?.clone())))
.collect();
assert_eq!(Value::Object(reasoning), expected);
}
#[rstest]
#[case::empty_text_blocks(
json!([{"role": "assistant", "content": [{"type": "text", "text": " "}, {"type": "text", "text": "kept"}]}]),
json!([{"role": "assistant", "content": [{"type": "text", "text": "kept"}]}])
)]
#[case::provider_specific_fields(
json!([{"role": "assistant", "content": [{"type": "text", "text": "kept", "provider_specific_fields": {"x": 1}}]}]),
json!([{"role": "assistant", "content": [{"type": "text", "text": "kept"}]}])
)]
#[case::unencrypted_web_search_results_become_text(
json!([{"role": "assistant", "content": [{
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_1",
"content": [{"type": "web_search_result", "title": "T", "url": "https://e.x", "page_age": null}]
}]}]),
json!([{"role": "assistant", "content": [{"type": "text", "text": "Web search results:\n\nTitle: T\nURL: https://e.x"}]}])
)]
#[tokio::test]
async fn replayed_history_is_cleaned_before_sending(
call: MessagesCall,
#[case] history: Value,
#[case] expected: Value,
) {
let upstream = upstream([message_response()]).await;
run_message(with_fields(
MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
},
json!({"messages": history}),
))
.await;
assert_eq!(only_request(&upstream).await.json()["messages"], expected);
}
#[rstest]
#[tokio::test]
async fn metadata_is_reduced_to_the_user_id(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
run_message(with_fields(
MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
},
json!({"metadata": {"user_id": "u-1", "trace_id": "internal", "tags": ["a"]}}),
))
.await;
assert_eq!(
only_request(&upstream).await.json()["metadata"],
json!({"user_id": "u-1"})
);
}
#[rstest]
#[case::numeric_user_id(json!({"metadata": {"user_id": 7}}))]
#[case::missing_max_tokens(json!({"max_tokens": null}))]
#[tokio::test]
async fn an_invalid_request_fails_before_sending(call: MessagesCall, #[case] fields: Value) {
let upstream = upstream([message_response()]).await;
let error = run(with_fields(
MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
},
fields,
))
.await
.err()
.expect("the request is rejected");
assert!(error.is_request(), "{error:?}");
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[tokio::test]
async fn azure_folds_system_role_messages_into_the_system_prompt(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
run_message(with_fields(
MessagesCall {
custom_llm_provider: Some("azure_ai".into()),
api_key: Some("sk-azure".into()),
api_base: Some(upstream.uri()),
..call
},
json!({
"system": "top level",
"messages": [
{"role": "system", "content": "from a message"},
{"role": "user", "content": "hi"}
]
}),
))
.await;
let sent = only_request(&upstream).await.json();
assert_eq!(
sent["system"],
json!([
{"type": "text", "text": "top level"},
{"type": "text", "text": "from a message"}
])
);
assert_eq!(sent["messages"], json!([{"role": "user", "content": "hi"}]));
}
#[rstest]
#[case::bare_model(MODEL, MODEL)]
#[case::one_prefix("anthropic/claude-sonnet-4-5", MODEL)]
#[case::doubled_prefix_loses_one_segment(
"anthropic/anthropic/claude-sonnet-4-5",
"anthropic/claude-sonnet-4-5"
)]
#[tokio::test]
async fn the_provider_prefix_is_stripped_exactly_once(
call: MessagesCall,
#[case] model: &str,
#[case] sent_model: &str,
) {
let upstream = upstream([message_response()]).await;
run_message(MessagesCall {
model: model.into(),
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await;
assert_eq!(only_request(&upstream).await.json()["model"], sent_model);
}

221
litellm-rust/crates/core/tests/messages/response.rs Normal file
View file

@ -0,0 +1,221 @@
use litellm_core::messages::{messages, types::MessagesRequest};
use litellm_http::transport::Error as TransportError;
use rstest::rstest;
use super::*;
#[rstest]
#[case::anthropic("anthropic")]
#[case::azure_ai("azure_ai")]
#[tokio::test]
async fn the_provider_message_is_returned(call: MessagesCall, #[case] provider: &str) {
let upstream = upstream([message_response()]).await;
let message = run_message(MessagesCall {
custom_llm_provider: Some(provider.into()),
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await;
assert_eq!(message.id, "msg_1");
assert_eq!(message.content, [json!({"type": "text", "text": "hi"})]);
assert_eq!(message.stop_reason.as_deref(), Some("end_turn"));
}
/// A refusal and fields the route does not model come back exactly as the provider sent
/// them, since the Python side returns the raw message and the router decides what to do.
#[rstest]
#[tokio::test]
async fn the_message_passes_through_losslessly(call: MessagesCall) {
let upstream_body = json!({
"id": "msg_2",
"type": "message",
"role": "assistant",
"model": MODEL,
"content": [
{"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "q"}},
{"type": "text", "text": "no", "citations": [{"type": "web_search_result_location", "url": "https://e.x"}]}
],
"stop_reason": "refusal",
"stop_sequence": null,
"stop_details": {"type": "safeguard", "safeguard_types": ["dangerous_tool_use"]},
"container": {"id": "container_1", "expires_at": "2026-01-01T00:00:00Z"},
"context_management": {"applied_edits": []},
"usage": {"input_tokens": 1, "output_tokens": 2, "server_tool_use": {"web_search_requests": 1}},
"unknown_future_field": {"nested": true}
});
let upstream = upstream([json_response(upstream_body.clone())]).await;
let message = run_message(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await;
assert_eq!(message.stop_reason.as_deref(), Some("refusal"));
assert_eq!(serde_json::to_value(&message).unwrap(), upstream_body);
}
#[rstest]
#[tokio::test]
async fn a_json_error_envelope_is_kept_verbatim(call: MessagesCall) {
let envelope =
json!({"type": "error", "error": {"type": "invalid_request_error", "message": "bad"}});
let upstream = upstream([status_response(400, envelope.clone())]).await;
let error = run(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await
.err()
.expect("upstream error propagates");
let Error::Transport(TransportError::Http { status, body }) = error else {
panic!("{error:?}");
};
assert_eq!(status, 400);
assert_eq!(serde_json::from_str::<Value>(&body).unwrap(), envelope);
}
#[rstest]
#[tokio::test]
async fn a_long_error_body_is_truncated_at_the_documented_cap(call: MessagesCall) {
let long = "x".repeat(600);
let upstream = upstream([ResponseTemplate::new(500).set_body_string(long.clone())]).await;
let error = run(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await
.err()
.expect("upstream error propagates");
assert_eq!(
error,
Error::Transport(TransportError::Http {
status: 500,
body: format!("{}... (truncated)", &long[..256])
})
);
}
#[rstest]
#[case::bad_request(400)]
#[case::unauthorized(401)]
#[case::rate_limited(429)]
#[case::server_error(500)]
#[case::overloaded(529)]
#[tokio::test]
async fn an_upstream_error_keeps_its_status_and_body(call: MessagesCall, #[case] status: u16) {
let upstream =
upstream([ResponseTemplate::new(status).set_body_string("upstream said no")]).await;
let error = run(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await
.err()
.expect("upstream error propagates");
assert_eq!(
error,
Error::Transport(TransportError::Http {
status,
body: "upstream said no".into()
})
);
}
#[rstest]
#[case::not_json(ResponseTemplate::new(200).set_body_string("not json"))]
#[case::not_a_message(json_response(json!({"unexpected": true})))]
#[tokio::test]
async fn an_unreadable_success_body_is_an_invalid_response(
call: MessagesCall,
#[case] response: ResponseTemplate,
) {
let upstream = upstream([response]).await;
let error = run(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await
.err()
.expect("an unreadable body fails");
assert!(error.is_response(), "{error:?}");
}
#[rstest]
#[tokio::test]
async fn a_provider_slower_than_the_timeout_fails_the_call(call: MessagesCall) {
let upstream = upstream([message_response().set_delay(Duration::from_secs(5))]).await;
let error = run(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
timeout: Some(Duration::from_millis(100)),
..call
})
.await
.err()
.expect("the call times out");
assert!(matches!(error, Error::Transport(_)), "{error:?}");
}
fn facade_request(body: Value, api_base: &str) -> MessagesRequest<'_> {
MessagesRequest {
model: MODEL,
body,
api_key: Some("sk-ant"),
api_base: Some(api_base),
custom_llm_provider: Some("anthropic"),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
}
}
#[tokio::test]
async fn the_facade_runs_the_route_in_process() {
let upstream = upstream([message_response()]).await;
let base = upstream.uri();
let message = messages(facade_request(
json!({"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]}),
&base,
))
.await
.expect("messages request succeeds");
assert_eq!(message.id, "msg_1");
assert_eq!(
only_request(&upstream).await.header("x-api-key"),
Some("sk-ant")
);
}
#[tokio::test]
async fn the_facade_rejects_a_body_that_is_not_an_object() {
let error = messages(facade_request(json!([]), UNREACHABLE_BASE))
.await
.expect_err("a non-object body is rejected");
assert_eq!(
error,
Error::InvalidRequest("messages body must be an object".into())
);
}

200
litellm-rust/crates/core/tests/messages/secrets.rs Normal file
View file

@ -0,0 +1,200 @@
use rstest::rstest;
use super::*;
#[rstest]
#[case::anthropic(
"anthropic",
"ANTHROPIC_API_KEY",
"ANTHROPIC_BASE_URL",
"/v1/messages",
&["ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"]
)]
#[case::azure_ai(
"azure_ai",
"AZURE_API_KEY",
"AZURE_API_BASE",
"/anthropic/v1/messages",
&["AZURE_API_KEY", "AZURE_API_BASE"]
)]
#[tokio::test]
async fn the_credential_and_base_come_from_the_secret_source(
call: MessagesCall,
#[case] provider: &str,
#[case] key_name: &str,
#[case] base_name: &str,
#[case] path: &str,
#[case] looked_up: &[&str],
) {
let upstream = upstream([message_response()]).await;
let base = upstream.uri();
let secrets = Arc::new(RecordingSecrets::new([
(key_name, "sk-from-manager"),
(base_name, base.as_str()),
]));
let output = run_with(
secrets.clone(),
MessagesCall {
custom_llm_provider: Some(provider.into()),
..call
},
)
.await
.expect("messages call succeeds");
assert!(matches!(output, MessagesOutput::Message(_)));
let request = only_request(&upstream).await;
assert_eq!(request.url.path(), path);
assert_eq!(request.header("x-api-key"), Some("sk-from-manager"));
assert_eq!(secrets.requested(), looked_up);
}
#[rstest]
#[tokio::test]
async fn call_arguments_win_over_the_secret_source(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let secrets = Arc::new(RecordingSecrets::new([
("ANTHROPIC_API_KEY", "sk-from-manager"),
("ANTHROPIC_BASE_URL", UNREACHABLE_BASE),
]));
run_with(
secrets,
MessagesCall {
api_key: Some("sk-from-call".into()),
api_base: Some(upstream.uri()),
..call
},
)
.await
.expect("messages call succeeds");
assert_eq!(
only_request(&upstream).await.header("x-api-key"),
Some("sk-from-call")
);
}
#[rstest]
#[tokio::test]
async fn a_secret_manager_failure_fails_the_call_before_sending(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let error = run_with(
Arc::new(RecordingSecrets::failing()),
MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
},
)
.await
.err()
.expect("a secret manager failure fails the call");
assert!(
matches!(&error, Error::Secret(source) if matches!(source.source_error(), litellm_secrets::Error::ManagedSecretMissing)),
"{error:?}"
);
assert!(received(&upstream).await.is_empty());
}
#[derive(Clone, Copy)]
enum Base {
Upstream,
Unreachable,
Blank,
Absent,
}
fn base_value(base: Base, upstream: &str) -> Option<String> {
match base {
Base::Upstream => Some(upstream.to_string()),
Base::Unreachable => Some(UNREACHABLE_BASE.to_string()),
Base::Blank => Some(" ".to_string()),
Base::Absent => None,
}
}
#[rstest]
#[case::api_base_beats_base_url(Base::Upstream, Base::Unreachable)]
#[case::blank_api_base_falls_through_to_base_url(Base::Blank, Base::Upstream)]
#[case::base_url_alone(Base::Absent, Base::Upstream)]
#[tokio::test]
async fn the_anthropic_base_env_precedence_picks_the_upstream(
call: MessagesCall,
#[case] api_base: Base,
#[case] base_url: Base,
) {
let upstream = upstream([message_response()]).await;
let uri = upstream.uri();
let values: Vec<(&str, &str)> = [
("ANTHROPIC_API_KEY", Some("sk-env".to_string())),
("ANTHROPIC_API_BASE", base_value(api_base, &uri)),
("ANTHROPIC_BASE_URL", base_value(base_url, &uri)),
]
.iter()
.filter_map(|(name, value)| Some((*name, value.as_deref()?)))
.map(|(name, value)| (name, Box::leak(value.to_string().into_boxed_str()) as &str))
.collect();
run_with(Arc::new(RecordingSecrets::new(values)), call)
.await
.expect("messages call reaches the upstream the precedence picks");
assert_eq!(only_request(&upstream).await.url.path(), "/v1/messages");
}
#[rstest]
#[case::auth_token_alone(
&[("ANTHROPIC_AUTH_TOKEN", "tok")],
("authorization", "Bearer tok"),
"x-api-key"
)]
#[case::api_key_beats_the_auth_token(
&[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "tok")],
("x-api-key", "sk-env"),
"authorization"
)]
#[tokio::test]
async fn the_auth_token_env_is_a_bearer_only_without_a_key(
call: MessagesCall,
#[case] values: &[(&str, &str)],
#[case] expected: (&str, &str),
#[case] absent: &str,
) {
let upstream = upstream([message_response()]).await;
run_with(
Arc::new(RecordingSecrets::new(values.iter().copied())),
MessagesCall {
api_base: Some(upstream.uri()),
..call
},
)
.await
.expect("messages call succeeds");
let request = only_request(&upstream).await;
let (name, value) = expected;
assert_eq!(request.header_values(name), [value]);
assert_eq!(request.header(absent), None);
}
#[rstest]
#[tokio::test]
async fn azure_without_a_base_anywhere_fails_before_sending(call: MessagesCall) {
let error = run_with(
Arc::new(RecordingSecrets::new([("AZURE_API_KEY", "sk-azure")])),
MessagesCall {
custom_llm_provider: Some("azure_ai".into()),
..call
},
)
.await
.err()
.expect("azure needs a base");
assert_eq!(error, Error::Auth(litellm_auth::Error::MissingAzureApiBase));
}

266
litellm-rust/crates/core/tests/messages/stream.rs Normal file
View file

@ -0,0 +1,266 @@
use std::{convert::Infallible, sync::Mutex};
use bytes::Bytes;
use litellm_core::messages::route::{Messages, MessagesStreamHead};
use litellm_host::host::{Demand, Host};
use rstest::rstest;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
};
use super::*;
const UPSTREAM_HEADERS: [(&str, &str); 2] = [
("request-id", "req_upstream_123"),
("anthropic-ratelimit-requests-remaining", "41"),
];
const SSE_BODY: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
enum Seen {
Open(Vec<(String, String)>),
Deliver(Bytes),
}
/// Projects like `LocalMessagesHost`, records every stream op in the order the route
/// performs it, and detaches after `detach_after` ops.
struct RecordingStreamHost {
call: LocalMessagesHost,
detach_after: usize,
seen: Mutex<Vec<Seen>>,
}
impl RecordingStreamHost {
fn new(call: MessagesCall, detach_after: usize) -> Self {
Self {
call: LocalMessagesHost::new(call),
detach_after,
seen: Mutex::new(Vec::new()),
}
}
fn record(&self, op: Seen) -> Demand {
let mut seen = self.seen.lock().unwrap();
seen.push(op);
match seen.len() < self.detach_after {
true => Demand::More,
false => Demand::Detached,
}
}
}
impl Host<Messages> for RecordingStreamHost {
async fn project(&self) -> Result<MessagesCall, Error> {
self.call.project().await
}
async fn custom_op(&self, op: Infallible) -> Result<(), Error> {
match op {}
}
async fn open(&self, head: MessagesStreamHead) -> Result<Demand, Error> {
Ok(self.record(Seen::Open(head.headers)))
}
async fn deliver(&self, chunk: Bytes) -> Result<Demand, Error> {
Ok(self.record(Seen::Deliver(chunk)))
}
}
fn streaming(call: MessagesCall, api_base: String) -> MessagesCall {
let mut body = call.body.clone();
body.insert("stream".into(), json!(true));
MessagesCall {
api_key: Some("sk-ant".into()),
api_base: Some(api_base),
body,
..call
}
}
fn sse_response() -> ResponseTemplate {
UPSTREAM_HEADERS.iter().fold(
ResponseTemplate::new(200).set_body_raw(SSE_BODY, "text/event-stream"),
|response, (name, value)| response.insert_header(*name, *value),
)
}
async fn stream_through(host: &RecordingStreamHost) -> Result<MessagesOutput, Error> {
litellm_host::run::run(messages_machine(Arc::new(RecordingSecrets::empty())), host).await
}
#[rstest]
#[tokio::test]
async fn upstream_headers_are_on_the_stream_head_before_the_first_chunk(call: MessagesCall) {
let upstream = upstream([sse_response()]).await;
let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX);
let outcome = stream_through(&host).await.expect("streamed call succeeds");
assert!(matches!(outcome, MessagesOutput::Streamed));
let seen = host.seen.into_inner().unwrap();
let [Seen::Open(headers), chunks @ ..] = seen.as_slice() else {
panic!("the stream opens before any chunk is delivered");
};
let surfaced: Vec<(&str, &str)> = headers
.iter()
.filter(|(name, _)| {
UPSTREAM_HEADERS
.iter()
.any(|(upstream, _)| upstream == name)
})
.map(|(name, value)| (name.as_str(), value.as_str()))
.collect();
assert_eq!(surfaced, UPSTREAM_HEADERS);
let delivered: Vec<u8> = chunks
.iter()
.flat_map(|step| match step {
Seen::Deliver(chunk) => chunk.to_vec(),
Seen::Open(_) => panic!("the stream opens exactly once"),
})
.collect();
assert_eq!(delivered, SSE_BODY.as_bytes());
}
#[rstest]
#[case::at_open(1)]
#[case::after_the_first_chunk(2)]
#[tokio::test]
async fn a_detached_caller_receives_nothing_more(call: MessagesCall, #[case] detach_after: usize) {
let upstream = upstream([sse_response()]).await;
let host = RecordingStreamHost::new(streaming(call, upstream.uri()), detach_after);
let outcome = stream_through(&host)
.await
.expect("a detached stream still completes");
assert!(matches!(outcome, MessagesOutput::Streamed));
assert_eq!(host.seen.into_inner().unwrap().len(), detach_after);
}
#[rstest]
#[case::text_body(ResponseTemplate::new(429).set_body_string("slow down"), "slow down")]
#[case::json_envelope(
status_response(429, json!({"type": "error", "error": {"type": "rate_limit_error", "message": "slow down"}})),
r#"{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}"#
)]
#[tokio::test]
async fn an_upstream_error_fails_the_call_without_opening_the_stream(
call: MessagesCall,
#[case] response: ResponseTemplate,
#[case] body: &str,
) {
let upstream = upstream([response]).await;
let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX);
let error = stream_through(&host)
.await
.err()
.expect("upstream error propagates");
assert_eq!(
error,
Error::Transport(litellm_http::transport::Error::Http {
status: 429,
body: body.into()
})
);
assert!(host.seen.into_inner().unwrap().is_empty());
}
/// The native route relays bytes as they are. Python's synthetic `api_error` for a stream
/// that never reaches `message_stop` lives in its SSE wrapper, above this route.
#[rstest]
#[tokio::test]
async fn a_stream_that_ends_without_message_stop_is_relayed_as_is(call: MessagesCall) {
const INCOMPLETE: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\n";
let upstream =
upstream([ResponseTemplate::new(200).set_body_raw(INCOMPLETE, "text/event-stream")]).await;
let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX);
stream_through(&host).await.expect("streamed call succeeds");
let delivered: Vec<u8> = host
.seen
.into_inner()
.unwrap()
.iter()
.flat_map(|step| match step {
Seen::Deliver(chunk) => chunk.to_vec(),
Seen::Open(_) => Vec::new(),
})
.collect();
assert_eq!(delivered, INCOMPLETE.as_bytes());
}
/// Serves one SSE chunk and then holds the connection open without ever finishing.
async fn stalling_upstream() -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut request = vec![0; 4096];
let _ = socket.read(&mut request).await;
socket
.write_all(
b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n\
1f\r\nevent: message_start\ndata: {}\n\n\r\n",
)
.await
.unwrap();
std::future::pending::<()>().await;
});
base
}
#[rstest]
#[tokio::test]
async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) {
let base = stalling_upstream().await;
let host = RecordingStreamHost::new(
MessagesCall {
timeout: Some(Duration::from_millis(300)),
..streaming(call, base)
},
usize::MAX,
);
let error = tokio::time::timeout(Duration::from_secs(5), stream_through(&host))
.await
.expect("the stalled stream gives up within the timeout")
.err()
.expect("a stalled body fails the call");
assert!(matches!(error, Error::Transport(_)), "{error:?}");
let seen = host.seen.into_inner().unwrap();
assert!(
matches!(seen.as_slice(), [Seen::Open(_), Seen::Deliver(chunk)] if chunk.as_ref() == b"event: message_start\ndata: {}\n\n"),
"the chunk before the stall reached the caller, saw {} ops",
seen.len()
);
}
#[rstest]
#[tokio::test]
async fn streaming_is_refused_for_providers_that_cannot_stream(call: MessagesCall) {
let upstream = upstream([sse_response()]).await;
let host = RecordingStreamHost::new(
MessagesCall {
custom_llm_provider: Some("azure_ai".into()),
..streaming(call, upstream.uri())
},
usize::MAX,
);
let error = stream_through(&host)
.await
.err()
.expect("azure streaming is refused");
assert_eq!(
error,
Error::Unsupported("streaming messages for this provider")
);
assert!(received(&upstream).await.is_empty());
}

173
litellm-rust/crates/core/tests/ocr/aws_textract.rs Normal file
View file

@ -0,0 +1,173 @@
use std::{collections::BTreeMap, time::SystemTime};
use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post};
use rstest::rstest;
use time::{PrimitiveDateTime, format_description};
use wiremock::Request;
use super::*;
const ACCESS_KEY_ID: &str = "AKIDEXAMPLE";
const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY";
const DETECT: &str = "aws_textract/detect-document-text";
const ANALYZE: &str = "aws_textract/analyze-document";
fn textract_request(model: &str, base: &str) -> LiteLLMOcrRequest {
ocr_request_with_document(
model,
&format!("{base}/"),
json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}),
json!({
"aws_access_key_id": ACCESS_KEY_ID,
"aws_secret_access_key": SECRET_ACCESS_KEY,
"aws_region_name": "eu-west-1"
}),
)
}
fn textract_response() -> ResponseTemplate {
json_response(json!({
"DocumentMetadata": {"Pages": 1},
"Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}]
}))
}
/// Recomputes SigV4 over the request the upstream received, at the time the client claimed.
fn expected_authorization(url: &str, sent: &Request) -> String {
let format =
format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z")
.unwrap();
let signed_at: SystemTime =
PrimitiveDateTime::parse(sent.header("x-amz-date").unwrap(), &format)
.unwrap()
.assume_utc()
.into();
let headers: BTreeMap<String, String> = ["content-type", "x-amz-target"]
.into_iter()
.map(|name| (name.to_string(), sent.header(name).unwrap().to_string()))
.collect();
sign_post(
url,
&sent.body,
&aws_signature_headers(&headers),
"eu-west-1",
"textract",
&Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"),
signed_at,
)
.unwrap()["Authorization"]
.clone()
}
/// The recorded URL names wiremock's host, not the address the client signed for.
fn assert_signed(upstream: &MockServer, sent: &Request) {
let url = format!("{}/", upstream.uri());
assert_eq!(
sent.header("authorization"),
Some(expected_authorization(&url, sent).as_str())
);
}
#[tokio::test]
async fn detect_document_text_is_signed_and_lines_become_the_page() {
let upstream = upstream([textract_response()]).await;
let response = perform_with(LocalOcrHost::new(textract_request(DETECT, &upstream.uri())))
.await
.unwrap();
let sent = only_request(&upstream).await;
assert_eq!(
sent.header("x-amz-target"),
Some("Textract.DetectDocumentText")
);
assert_eq!(
sent.header("content-type"),
Some("application/x-amz-json-1.1")
);
assert_eq!(sent.json(), json!({"Document": {"Bytes": "b3JpZ2luYWw="}}));
assert_signed(&upstream, &sent);
assert_eq!(response.pages[0].markdown, "Invoice 12345");
assert_eq!(response.usage_info.unwrap().pages_processed, Some(1));
}
#[tokio::test]
async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() {
let upstream = upstream([textract_response()]).await;
let host = LocalOcrHost::new(textract_request(DETECT, &upstream.uri())).with_before_send(
|mut wire, _| {
assert!(
!wire
.headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("authorization")),
"the hook ran after signing"
);
wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ=");
Ok(wire)
},
);
perform_with(host).await.unwrap();
let sent = only_request(&upstream).await;
assert_eq!(sent.json(), json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}}));
assert_signed(&upstream, &sent);
}
#[tokio::test]
async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() {
let upstream = upstream([json_response(json!({
"DocumentMetadata": {"Pages": 1},
"Blocks": [
{"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"},
{"Id": "t", "BlockType": "LAYOUT_TITLE",
"Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]}
]
}))])
.await;
let response = perform_with(LocalOcrHost::new(textract_request(
ANALYZE,
&upstream.uri(),
)))
.await
.unwrap();
let sent = only_request(&upstream).await;
assert_eq!(
sent.header("x-amz-target"),
Some("Textract.AnalyzeDocument")
);
assert_eq!(sent.json()["FeatureTypes"], json!(["LAYOUT", "TABLES"]));
assert_signed(&upstream, &sent);
assert_eq!(response.pages[0].markdown, "# Quarterly Report");
}
#[rstest]
#[case::detect(DETECT)]
#[case::analyze(ANALYZE)]
#[tokio::test]
async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit(#[case] model: &str) {
let upstream = upstream([status_response(
400,
json!({
"__type": "UnsupportedDocumentException",
"Message": "Request has unsupported document format"
}),
)])
.await;
let error = perform_with(LocalOcrHost::new(textract_request(model, &upstream.uri())))
.await
.unwrap_err();
let Error::Provider { status, body, .. } = error else {
panic!("expected a provider error, got {error:?}");
};
assert_eq!(status, 400);
assert!(
body.contains("multi-page documents are not supported"),
"{body}"
);
}

270
litellm-rust/crates/core/tests/ocr/azure_ai.rs Normal file
View file

@ -0,0 +1,270 @@
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use litellm_auth::{
ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle,
};
use rstest::rstest;
use super::*;
#[derive(Debug)]
struct CountingToken {
token: fn(usize) -> String,
calls: AtomicUsize,
}
impl CountingToken {
fn new(token: fn(usize) -> String) -> Arc<Self> {
Arc::new(Self {
token,
calls: AtomicUsize::new(0),
})
}
fn calls(&self) -> usize {
self.calls.load(Ordering::SeqCst)
}
}
impl TokenProvider for CountingToken {
fn acquire(&self) -> TokenFuture<'_> {
let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
let token = SecretValue::new((self.token)(call));
Box::pin(async move {
Ok(ResolvedCredential::AccessToken {
token,
expires_on: None,
})
})
}
}
fn numbered_token(call: usize) -> String {
format!("callback-{call}")
}
fn azure_request(
provider: &Arc<CountingToken>,
api_base: Option<&str>,
api_key: Option<&str>,
extra_headers: Value,
optional_params: Value,
) -> LiteLLMOcrRequest {
let wire = serde_json::from_value(json!({
"model": "azure_ai/mistral-ocr-latest",
"document": {"type": "document_url", "document_url": INLINE_PDF},
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": null,
"extra_headers": extra_headers,
"optional_params": optional_params,
"timeout_seconds": 2.0
}))
.unwrap();
let mut request = decode_request(wire).unwrap();
request.azure_ad_token_provider = Some(TokenProviderHandle::new(provider.clone()));
request
}
fn ocr_page() -> ResponseTemplate {
json_response(json!({"pages": [{"index": 0, "markdown": "hello"}]}))
}
#[tokio::test]
async fn mistral_on_azure_sends_the_prepared_bearer_and_the_mistral_body() {
let upstream = upstream([json_response(json!({
"pages": [{"index": 0, "markdown": "hello"}],
"usage_info": {"pages_processed": 1}
}))])
.await;
let request = with_headers(
without_api_key(ocr_request(
"azure_ai/model",
&upstream.uri(),
json!({"include_image_base64": true}),
)),
&[("Authorization", "Bearer python-prepared-token")],
);
let result = perform(request).await.unwrap();
assert_eq!(result.pages[0].markdown, "hello");
let sent = only_request(&upstream).await;
assert_eq!(sent.url.path(), "/providers/mistral/azure/ocr");
assert_eq!(
sent.header("authorization"),
Some("Bearer python-prepared-token")
);
assert_eq!(
sent.json(),
json!({
"model": "model",
"document": {"type": "document_url", "document_url": INLINE_PDF},
"include_image_base64": true
})
);
}
#[tokio::test]
async fn a_static_entra_token_becomes_the_bearer() {
let upstream = upstream([pages_response()]).await;
let request = without_api_key(ocr_request(
"azure_ai/model",
&upstream.uri(),
json!({"azure_ad_token": "rust-owned-token"}),
));
perform(request).await.unwrap();
assert_eq!(
only_request(&upstream).await.header("authorization"),
Some("Bearer rust-owned-token")
);
}
#[tokio::test]
async fn a_guardrail_that_swaps_in_a_remote_document_is_rejected() {
let host = LocalOcrHost::new(ocr_request("azure_ai/model", UNREACHABLE_BASE, json!({})))
.with_before_send(|mut wire, _| {
wire.body["document"] = json!({
"type": "document_url",
"document_url": "https://example.com/not-inline.pdf"
});
Ok(wire)
});
let error = perform_with(host).await.unwrap_err();
assert!(error.to_string().contains("data URI"), "{error}");
}
#[tokio::test]
async fn the_token_provider_is_the_bearer_and_is_acquired_for_each_request() {
let provider = CountingToken::new(numbered_token);
let upstream = upstream([ocr_page(), ocr_page()]).await;
let base = upstream.uri();
for _ in 0..2 {
perform(azure_request(
&provider,
Some(&base),
None,
Value::Null,
json!({}),
))
.await
.unwrap();
}
assert_eq!(provider.calls(), 2);
let authorizations: Vec<String> = received(&upstream)
.await
.iter()
.map(|request| {
request
.header("authorization")
.unwrap_or_default()
.to_string()
})
.collect();
assert_eq!(authorizations, ["Bearer callback-1", "Bearer callback-2"]);
}
#[rstest]
#[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)]
#[case::provider_beats_static_token(
None,
Value::Null,
json!({"azure_ad_token": "static-token"}),
"Bearer callback-1",
1
)]
#[case::header_wins_on_the_wire_but_provider_still_runs(
None,
json!({"Authorization": "Bearer override"}),
json!({}),
"Bearer override",
1
)]
#[tokio::test]
async fn credential_precedence(
#[case] api_key: Option<&str>,
#[case] extra_headers: Value,
#[case] optional_params: Value,
#[case] expected_authorization: &str,
#[case] expected_calls: usize,
) {
let provider = CountingToken::new(numbered_token);
let upstream = upstream([ocr_page()]).await;
perform(azure_request(
&provider,
Some(&upstream.uri()),
api_key,
extra_headers,
optional_params,
))
.await
.unwrap();
assert_eq!(provider.calls(), expected_calls);
assert_eq!(
only_request(&upstream).await.header_values("authorization"),
[expected_authorization]
);
}
#[rstest]
#[case::missing_api_base(
false,
json!({}),
numbered_token,
|error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase {
provider: "Azure AI",
environment_variable: "AZURE_AI_API_BASE",
})),
0
)]
#[case::unsupported_oidc_reference(
true,
json!({"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}),
numbered_token,
|error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)),
0
)]
#[case::empty_provider_token_ignores_static_token(
true,
json!({"azure_ad_token": "static-token"}),
|_| String::new(),
|error: &Error| matches!(error, Error::MissingAzureAiCredentials),
1
)]
#[tokio::test]
async fn credential_failures_send_no_provider_request(
#[case] with_api_base: bool,
#[case] optional_params: Value,
#[case] token: fn(usize) -> String,
#[case] expected: fn(&Error) -> bool,
#[case] expected_calls: usize,
) {
let provider = CountingToken::new(token);
let upstream = upstream([ocr_page()]).await;
let base = upstream.uri();
let error = perform(azure_request(
&provider,
with_api_base.then_some(base.as_str()),
None,
Value::Null,
optional_params,
))
.await
.unwrap_err();
assert!(expected(&error), "unexpected error: {error:?}");
assert_eq!(provider.calls(), expected_calls);
assert!(received(&upstream).await.is_empty());
}

441
litellm-rust/crates/core/tests/ocr/azure_document_intelligence.rs Normal file
View file

@ -0,0 +1,441 @@
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_host::event::{CallEvent, MachineEvent};
use litellm_llms::base_llm::ocr::settings::OcrSettings;
use rstest::rstest;
use super::*;
const MODEL: &str = "azure_ai/doc-intelligence/prebuilt-read";
fn read_request(base: &str, options: Value) -> LiteLLMOcrRequest {
ocr_request(MODEL, base, options)
}
#[tokio::test]
async fn pages_features_and_extra_options_map_to_the_analyze_call() {
let upstream = upstream([json_response(json!({
"status": "succeeded",
"analyzeResult": {"pages": []}
}))])
.await;
let request = read_request(
&upstream.uri(),
json!({
"pages": [2, 0, 0, 1],
"features": ["keyValuePairs", "languages"],
"future_option": {"nested": null},
"extra_body": {"provider_option": false}
}),
)
.with_document(
document(
json!({"type": "document_url", "document_url": "https://example.com/document.pdf"}),
)
.into(),
);
perform(request).await.unwrap();
let sent = only_request(&upstream).await;
assert!(
sent.url.path().ends_with("/prebuilt-read:analyze"),
"{}",
sent.url
);
assert_eq!(sent.query("pages").as_deref(), Some("1,2,3"));
assert_eq!(
sent.query("features").as_deref(),
Some("keyValuePairs,languages")
);
assert_eq!(
sent.json(),
json!({
"urlSource": "https://example.com/document.pdf",
"future_option": {"nested": null},
"provider_option": false
})
);
}
#[rstest]
#[case(json!({"pages": [true]}), Error::Pages("expected only integers or only strings".into()))]
#[case(json!({"pages": [1, "2"]}), Error::Pages("expected only integers or only strings".into()))]
#[case(json!({"pages": [-1]}), Error::Pages("negative page index".into()))]
#[case(json!({"pages": "1&&features=bad"}), Error::Pages("invalid native page range".into()))]
#[case(json!({"features": "languages&pages=1"}), Error::Features)]
#[case(json!({"req_format": "azure"}), Error::RequestFormat)]
#[tokio::test]
async fn invalid_pages_features_and_format_are_rejected_before_sending(
#[case] options: Value,
#[case] expected: Error,
) {
let upstream = upstream([json_response(json!({}))]).await;
let result = match decode_request(wire(
MODEL,
&upstream.uri(),
json!({"type": "document_url", "document_url": "https://example.com/a.pdf"}),
options.clone(),
)) {
Ok(request) => perform(request).await,
Err(error) => Err(error),
};
assert!(
received(&upstream).await.is_empty(),
"sent invalid options: {options}"
);
let error = result.unwrap_err();
assert_eq!(
std::mem::discriminant(&error),
std::mem::discriminant(&expected)
);
assert_eq!(error.http_status_code(), Some(400));
assert_eq!(error.to_string(), expected.to_string());
}
#[rstest]
#[case::no_options(json!({}))]
#[case::litellm_format(json!({"req_format": "litellm"}))]
#[tokio::test]
async fn an_inline_document_is_sent_as_base64_and_only_page_text_is_kept(#[case] options: Value) {
let upstream = upstream([json_response(json!({
"status": "succeeded",
"analyzeResult": {"pages": [{"pageNumber": 1, "lines": [{"content": "hello"}]}]}
}))])
.await;
let response = perform(read_request(&upstream.uri(), options))
.await
.unwrap();
assert_eq!(response.pages.len(), 1);
assert_eq!(response.pages[0].index, 0);
assert_eq!(response.pages[0].markdown, "hello");
assert_eq!(response.provider_native_response, None);
let serialized = response.into_json();
for field in ["content", "tables", "keyValuePairs"] {
assert_eq!(serialized.get(field), Some(&Value::Null), "{field}");
}
let sent = only_request(&upstream).await;
for field in ["pages", "features", "req_format"] {
assert_eq!(sent.query(field), None, "{field}");
}
assert_eq!(sent.json(), json!({"base64Source": "YWJj"}));
}
#[tokio::test]
async fn native_format_normalizes_pages_and_keeps_the_provider_response() {
let operation = json!({
"status": "succeeded",
"operationExtension": 42,
"analyzeResult": {
"content": "A\n\nB",
"tables": [{"cells": []}],
"keyValuePairs": [{"key": {"content": "A"}}],
"pages": [{
"pageNumber": "2",
"width": "8.5",
"height": 11,
"unit": "inch",
"lines": [{"content": "A"}, {"content": null}, {"content": "B"}]
}]
}
});
let upstream = upstream([json_response(operation.clone())]).await;
let result = perform(read_request(
&upstream.uri(),
json!({"req_format": "native"}),
))
.await
.unwrap();
assert_eq!(result.pages[0].index, 1);
assert_eq!(result.pages[0].markdown, "A\n\nB");
assert_eq!(
serde_json::to_value(&result.pages[0].dimensions).unwrap(),
json!({"width": 816, "height": 1056, "dpi": 96})
);
assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1));
let serialized = result.clone().into_json();
assert_eq!(serialized["content"], "A\n\nB");
assert_eq!(serialized["tables"], json!([{"cells": []}]));
assert_eq!(
serialized["keyValuePairs"],
json!([{"key": {"content": "A"}}])
);
assert!(serialized.get("key_value_pairs").is_none());
assert_eq!(
result.provider_native_response.map(Value::Object),
Some(operation)
);
}
#[tokio::test]
async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() {
let upstream = upstream([json_response(json!({
"status": "succeeded",
"analyzeResult": {"pages": [{"pageNumber": 1, "width": 8.5, "height": 11, "unit": "inch"}]}
}))])
.await;
let client = ocr_client().with_settings(OcrSettings {
document_intelligence_api_version: "2099-01-01".into(),
document_intelligence_dpi: 72,
..OcrSettings::default()
});
let result =
litellm_core::ocr::client::perform(&client, read_request(&upstream.uri(), json!({})))
.await
.unwrap();
assert_eq!(
only_request(&upstream)
.await
.query("api-version")
.as_deref(),
Some("2099-01-01")
);
assert_eq!(
serde_json::to_value(&result.pages[0].dimensions).unwrap(),
json!({"width": 612, "height": 792, "dpi": 72})
);
}
#[tokio::test]
async fn an_accepted_response_polls_to_success_with_only_credentials() {
let operation = json!({"status": "succeeded", "analyzeResult": {"pages": []}});
let upstream = MockServer::start().await;
respond_in_order(
&upstream,
[
accepted(&upstream, json!({})),
json_response(json!({"status": "running"})).insert_header("Retry-After", "0"),
json_response(operation.clone()),
],
)
.await;
let request = with_headers(
read_request(&upstream.uri(), json!({"req_format": "native"})),
&[("X-Trace", "initial-only")],
);
let result = perform(request).await.unwrap();
assert_eq!(
result.provider_native_response.map(Value::Object),
Some(operation)
);
let requests = received(&upstream).await;
assert_eq!(requests.len(), 3);
assert_eq!(requests[0].header("x-trace"), Some("initial-only"));
for poll in &requests[1..] {
assert_eq!(poll.method.as_str(), "GET");
assert_eq!(poll.url.path(), "/operation");
assert_eq!(poll.header("x-trace"), None);
assert_eq!(poll.header("ocp-apim-subscription-key"), Some("test-key"));
}
}
#[tokio::test]
async fn polling_forwards_bearer_credentials() {
let upstream = MockServer::start().await;
respond_in_order(
&upstream,
[
accepted(&upstream, json!({})),
json_response(json!({"status": "succeeded"})),
],
)
.await;
let request = with_headers(
without_api_key(read_request(&upstream.uri(), json!({}))),
&[("Authorization", "Bearer token")],
);
perform(request).await.unwrap();
assert_eq!(
received(&upstream).await[1].header("authorization"),
Some("Bearer token")
);
}
#[tokio::test]
async fn response_received_fires_for_the_submission_and_the_completed_poll() {
let upstream = MockServer::start().await;
respond_in_order(
&upstream,
[
accepted(&upstream, json!({"submitted": true})),
json_response(json!({"status": "succeeded"})),
],
)
.await;
let observed = Arc::new(Mutex::new(Vec::new()));
let recorder = observed.clone();
let host =
LocalOcrHost::new(read_request(&upstream.uri(), json!({}))).with_observer(move |event| {
if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event {
recorder.lock().unwrap().push(raw.body.clone());
}
});
perform_with(host).await.unwrap();
assert_eq!(received(&upstream).await.len(), 2);
assert_eq!(
*observed.lock().unwrap(),
[r#"{"submitted":true}"#, r#"{"status":"succeeded"}"#]
);
}
#[tokio::test]
async fn polling_does_not_follow_redirects() {
let upstream = MockServer::start().await;
respond_in_order(
&upstream,
[
accepted(&upstream, json!({})),
ResponseTemplate::new(302)
.insert_header("Location", format!("{}/redirected", upstream.uri())),
json_response(json!({"status": "succeeded"})),
],
)
.await;
let error = perform(read_request(&upstream.uri(), json!({})))
.await
.unwrap_err();
assert!(error.to_string().contains("status 302"), "{error}");
assert_eq!(received(&upstream).await.len(), 2);
}
#[tokio::test]
async fn a_failed_operation_is_an_error() {
let upstream = MockServer::start().await;
respond_in_order(
&upstream,
[
accepted(&upstream, json!({})),
json_response(json!({"status": "failed"})),
],
)
.await;
let error = perform(read_request(&upstream.uri(), json!({})))
.await
.unwrap_err();
assert!(error.to_string().contains("status failed"), "{error}");
}
#[tokio::test]
async fn the_polling_deadline_bounds_the_retry_delay() {
let upstream = MockServer::start().await;
respond_in_order(
&upstream,
[
accepted(&upstream, json!({})),
json_response(json!({"status": "notStarted"})).insert_header("Retry-After", "9999"),
],
)
.await;
let client = ocr_client().with_settings(OcrSettings {
poll_timeout: Duration::from_millis(100),
..OcrSettings::default()
});
let error = tokio::time::timeout(
Duration::from_secs(1),
litellm_core::ocr::client::perform(&client, read_request(&upstream.uri(), json!({}))),
)
.await
.expect("the deadline cuts the retry delay short")
.unwrap_err();
assert!(error.to_string().contains("timed out"), "{error}");
}
#[rstest]
#[case::null_pages(json!({"pages": null}), "pages")]
#[case::null_page(json!({"pages": [null]}), "pages[0]")]
#[case::null_lines(json!({"pages": [{"lines": null}]}), "lines")]
#[case::bad_width(json!({"pages": [{"width": "bad"}]}), "width")]
#[tokio::test]
async fn malformed_provider_pages_report_the_response_path(
#[case] analysis: Value,
#[case] path: &str,
) {
let upstream = upstream([json_response(json!({
"status": "succeeded",
"analyzeResult": analysis
}))])
.await;
let error = perform(read_request(&upstream.uri(), json!({})))
.await
.unwrap_err();
assert!(error.to_string().contains(path), "{error}");
}
#[rstest]
#[case::missing(None)]
#[case::relative(Some("/relative"))]
#[case::cross_origin(Some("http://example.com/operation"))]
#[case::with_userinfo(Some("http://user:password@127.0.0.1/operation"))]
#[tokio::test]
async fn an_unusable_operation_location_is_rejected(#[case] location: Option<&str>) {
let response = location
.into_iter()
.fold(ResponseTemplate::new(202), |response, location| {
response.insert_header("Operation-Location", location)
});
let upstream = upstream([response]).await;
let error = perform(read_request(&upstream.uri(), json!({})))
.await
.unwrap_err();
assert!(error.to_string().contains("operation-location"), "{error}");
assert_eq!(received(&upstream).await.len(), 1);
}
#[tokio::test]
async fn the_model_id_is_percent_encoded() {
let upstream = upstream([json_response(json!({"status": "succeeded"}))]).await;
perform(ocr_request(
"azure_ai/doc-intelligence/a ?#é",
&upstream.uri(),
json!({}),
))
.await
.unwrap();
let sent = only_request(&upstream).await;
assert!(
sent.url.path().ends_with("/a%20%3F%23%C3%A9:analyze"),
"{}",
sent.url
);
}
#[rstest]
#[case::dot("azure_ai/doc-intelligence/.")]
#[case::dot_dot("azure_ai/doc-intelligence/..")]
#[tokio::test]
async fn dot_segment_model_ids_are_rejected(#[case] model: &str) {
let error = perform(ocr_request(model, UNREACHABLE_BASE, json!({})))
.await
.unwrap_err();
assert!(error.to_string().contains("dot segment"), "{error}");
}

42
litellm-rust/crates/core/tests/ocr/cohere.rs Normal file
View file

@ -0,0 +1,42 @@
use rstest::rstest;
use super::*;
#[rstest]
#[case::cohere("cohere/parse-v5.0", "/v2/parse")]
#[case::azure_ai("azure_ai/Cohere-parse-v5.0", "/providers/cohere/v2/parse")]
#[tokio::test]
async fn an_image_goes_to_the_parse_endpoint_with_the_bearer_key(
#[case] model: &str,
#[case] path: &str,
) {
let upstream = upstream([pages_response()]).await;
let request = ocr_request_with_document(
model,
&upstream.uri(),
json!({"type": "image_url", "image_url": "data:image/png;base64,YWJj"}),
json!({}),
);
perform(request).await.unwrap();
let sent = only_request(&upstream).await;
assert_eq!(sent.method.as_str(), "POST");
assert_eq!(sent.url.path(), path);
assert_eq!(sent.header("authorization"), Some("Bearer test-key"));
}
#[rstest]
#[tokio::test]
async fn a_non_image_document_is_rejected_before_sending(
#[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str,
) {
let upstream = upstream([pages_response()]).await;
let error = perform(ocr_request(model, &upstream.uri(), json!({})))
.await
.unwrap_err();
assert!(matches!(error, Error::CohereImageOnly), "{error:?}");
assert!(received(&upstream).await.is_empty());
}

182
litellm-rust/crates/core/tests/ocr/documents.rs Normal file
View file

@ -0,0 +1,182 @@
use base64::Engine;
use litellm_core::ocr::types::OcrDocumentInput;
use litellm_host::event::WireRequest;
use rstest::rstest;
use wiremock::{Mock, matchers::any};
use super::*;
const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document";
const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ=";
#[derive(Clone, Copy, Debug)]
enum Route {
Mistral,
AzureAi,
VertexMistral,
AzureCohereParse,
Cohere,
}
impl Route {
fn model(self) -> &'static str {
match self {
Self::Mistral => "mistral/model",
Self::AzureAi => "azure_ai/model",
Self::VertexMistral => "vertex_ai/mistral-ocr-maas",
Self::AzureCohereParse => "azure_ai/cohere-parse",
Self::Cohere => "cohere/model",
}
}
fn document_type(self) -> &'static str {
match self {
Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url",
Self::AzureCohereParse | Self::Cohere => "image_url",
}
}
fn options(self) -> Value {
match self {
Self::Mistral | Self::AzureAi => json!({"pages": [0]}),
Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}),
Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}),
}
}
}
/// What the host does to the wire request in `before_send`.
#[derive(Clone, Copy, Debug)]
enum Guardrail {
Detached,
ReplacesDocument,
}
impl Guardrail {
fn before_send(self, wire: WireRequest) -> WireRequest {
let Value::Object(fields) = wire.body else {
return wire;
};
let body = fields
.into_iter()
.map(|(name, value)| match self {
Self::ReplacesDocument if name == "document" => {
let document_type = value["type"].clone();
let key = document_type.as_str().unwrap_or_default().to_string();
(name, json!({"type": document_type, key: REPLACED_DOCUMENT}))
}
Self::Detached | Self::ReplacesDocument => (name, value),
})
.collect();
WireRequest {
body: Value::Object(body),
..wire
}
}
}
/// Serves [`SERVED_DOCUMENT`] as `image/png` to every request.
async fn document_server() -> MockServer {
let server = MockServer::start().await;
Mock::given(any())
.respond_with(ResponseTemplate::new(200).set_body_raw(SERVED_DOCUMENT, "image/png"))
.mount(&server)
.await;
server
}
/// Sends a remote document through `route` and returns the document the provider saw.
async fn provider_document(route: Route, guardrail: Guardrail) -> Value {
let documents = document_server().await;
let upstream = upstream([pages_response()]).await;
let document_type = route.document_type();
let request = ocr_request_with_document(
route.model(),
&upstream.uri(),
json!({"type": document_type, document_type: format!("{}/scan.png", documents.uri())}),
route.options(),
);
let host =
LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(guardrail.before_send(wire)));
perform_with(host).await.unwrap();
only_request(&upstream).await.json()["document"][document_type].clone()
}
#[rstest]
#[case::azure_ai(Route::AzureAi)]
#[case::vertex_mistral(Route::VertexMistral)]
#[case::azure_cohere_parse(Route::AzureCohereParse)]
#[tokio::test]
async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) {
let expected = format!(
"data:image/png;base64,{}",
base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT)
);
assert_eq!(
provider_document(route, Guardrail::Detached).await,
expected
);
}
#[rstest]
#[tokio::test]
async fn a_document_replaced_by_the_host_reaches_the_provider(
#[values(
Route::Mistral,
Route::AzureAi,
Route::VertexMistral,
Route::AzureCohereParse,
Route::Cohere
)]
route: Route,
) {
assert_eq!(
provider_document(route, Guardrail::ReplacesDocument).await,
REPLACED_DOCUMENT
);
}
#[tokio::test]
async fn an_empty_byte_document_fails_before_sending() {
let upstream = upstream([pages_response()]).await;
let request = ocr_request("mistral/model", &upstream.uri(), json!({})).with_document(
OcrDocumentInput::Bytes {
bytes: Default::default(),
file_name: None,
mime_type: None,
},
);
let error = perform(request).await.unwrap_err();
assert!(matches!(error, Error::EmptyFile), "{error:?}");
assert!(received(&upstream).await.is_empty());
}
#[tokio::test]
async fn a_missing_path_document_fails_before_sending() {
let upstream = upstream([pages_response()]).await;
let path =
std::env::temp_dir().join(format!("litellm-ocr-missing-{}.png", rand::random::<u64>()));
let request = ocr_request("mistral/model", &upstream.uri(), json!({})).with_document(
OcrDocumentInput::Path {
path: path.clone(),
mime_type: None,
},
);
let error = perform(request).await.unwrap_err();
assert!(
matches!(
&error,
Error::FileRead { path: failed, source }
if *failed == path && source.kind() == std::io::ErrorKind::NotFound
),
"{error:?}"
);
assert!(received(&upstream).await.is_empty());
}

269
litellm-rust/crates/core/tests/ocr/lifecycle.rs Normal file
View file

@ -0,0 +1,269 @@
use std::sync::{Arc, Mutex};
use litellm_core::ocr::{
route::{Ocr, OcrOp, OcrProjection, ocr_machine},
types::OcrDocumentInput,
};
use litellm_host::{
event::{CallEvent, MachineEvent, RequestContext, WireRequest},
host::Host,
};
use rstest::rstest;
use super::*;
pub(crate) fn event_name(event: &CallEvent) -> &'static str {
match event {
CallEvent::Started { .. } => "started",
CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response",
CallEvent::Succeeded { .. } => "success",
CallEvent::Failed { .. } => "failure",
}
}
fn recording_host(
request: LiteLLMOcrRequest,
events: Arc<Mutex<Vec<&'static str>>>,
block: bool,
) -> LocalOcrHost {
let before_send_events = events.clone();
LocalOcrHost::new(request)
.with_before_send(move |wire, _| {
before_send_events.lock().unwrap().push("before_send");
match block {
true => Err(Error::InvalidRequest("blocked".into())),
false => Ok(wire),
}
})
.with_observer(move |event| events.lock().unwrap().push(event_name(event)))
}
#[tokio::test]
async fn hooks_run_in_order_and_one_success_is_emitted() {
let upstream = upstream([pages_response()]).await;
let events = Arc::new(Mutex::new(Vec::new()));
perform_with(recording_host(
ocr_request("mistral/model", &upstream.uri(), json!({})),
events.clone(),
false,
))
.await
.unwrap();
assert_eq!(
*events.lock().unwrap(),
["started", "before_send", "response", "success"]
);
assert_eq!(received(&upstream).await.len(), 1);
}
#[tokio::test]
async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() {
let upstream = upstream([pages_response()]).await;
let events = Arc::new(Mutex::new(Vec::new()));
let error = perform_with(recording_host(
ocr_request("mistral/model", &upstream.uri(), json!({})),
events.clone(),
true,
))
.await
.unwrap_err();
assert!(
matches!(&error, Error::InvalidRequest(message) if message == "blocked"),
"{error:?}"
);
assert_eq!(
*events.lock().unwrap(),
["started", "before_send", "failure"]
);
assert!(received(&upstream).await.is_empty());
}
#[tokio::test]
async fn an_upstream_failure_emits_one_terminal_failure() {
let upstream = upstream([status_response(500, json!({"error": "failed"}))]).await;
let events = Arc::new(Mutex::new(Vec::new()));
let result = perform_with(recording_host(
ocr_request("mistral/model", &upstream.uri(), json!({})),
events.clone(),
false,
))
.await;
assert!(result.is_err());
assert_eq!(
*events.lock().unwrap(),
["started", "before_send", "failure"]
);
assert_eq!(received(&upstream).await.len(), 1);
}
#[tokio::test]
async fn an_invalid_provider_response_is_observed_before_normalization_fails() {
let upstream = upstream([json_response(json!({"pages": "invalid"}))]).await;
let observed = Arc::new(Mutex::new(Vec::new()));
let recorder = observed.clone();
let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({})))
.with_observer(move |event| {
if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event {
recorder.lock().unwrap().push(raw.body.clone());
}
});
let error = perform_with(host).await.unwrap_err();
assert!(matches!(error, Error::ResponseField { .. }), "{error:?}");
assert_eq!(*observed.lock().unwrap(), [r#"{"pages":"invalid"}"#]);
}
#[tokio::test]
async fn headers_returned_by_before_send_are_sent() {
let upstream = upstream([pages_response()]).await;
let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({})))
.with_before_send(|mut wire, _| {
wire.headers
.push(("x-core-callback".into(), "edited".into()));
Ok(wire)
});
perform_with(host).await.unwrap();
assert_eq!(
only_request(&upstream).await.header("x-core-callback"),
Some("edited")
);
}
async fn before_send_context(request: LiteLLMOcrRequest) -> (WireRequest, RequestContext) {
let observed = Arc::new(Mutex::new(None));
let captured = observed.clone();
let host = LocalOcrHost::new(request).with_before_send(move |wire, context| {
*captured.lock().unwrap() = Some((wire.clone(), context.clone()));
Ok(wire)
});
perform_with(host).await.unwrap();
let context = observed.lock().unwrap().take();
context.expect("before_send ran")
}
#[tokio::test]
async fn before_send_sees_the_route_its_params_and_the_body() {
let upstream = upstream([pages_response()]).await;
let (wire, context) = before_send_context(ocr_request(
"mistral/model",
&upstream.uri(),
json!({"pages": [0], "req_format": "native"}),
))
.await;
assert_eq!(context.custom_llm_provider, "mistral");
assert_eq!(context.model, "model");
assert_eq!(context.optional_params["req_format"], "native");
assert!(context.secret_fields.is_empty());
assert_eq!(wire.body["pages"], json!([0]));
}
#[rstest]
#[case::client_secret(json!({"client_secret": "shh", "tenant_id": "t"}), &["client_secret"])]
#[case::no_secrets(json!({"tenant_id": "t"}), &[])]
#[tokio::test]
async fn before_send_names_the_secret_params(#[case] options: Value, #[case] secrets: &[&str]) {
let upstream = upstream([pages_response()]).await;
let request = ocr_request("azure_ai/model", &upstream.uri(), options).with_document(
OcrDocumentInput::Bytes {
bytes: b"abc".as_slice().into(),
file_name: None,
mime_type: Some("application/pdf".into()),
},
);
let (_, context) = before_send_context(request).await;
assert_eq!(context.secret_fields, secrets);
}
/// Hands the route a caller-owned Azure token and rewrites the bearer in `before_send`.
struct CallerTokenHost {
request: Mutex<Option<LiteLLMOcrRequest>>,
trace: Mutex<Vec<String>>,
}
impl Host<Ocr> for CallerTokenHost {
async fn project(&self) -> Result<OcrProjection, Error> {
self.trace.lock().unwrap().push("project".into());
Ok(OcrProjection {
request: self.request.lock().unwrap().take().unwrap(),
caller_token: true,
})
}
async fn custom_op(&self, op: OcrOp) -> Result<(), Error> {
match op {
OcrOp::AcquireAzureAdToken(reply) => {
self.trace.lock().unwrap().push("token".into());
reply.send(litellm_auth::ResolvedCredential::Static(
litellm_auth::SecretValue::new("caller-token"),
));
Ok(())
}
}
}
async fn before_send(
&self,
wire: WireRequest,
_: &RequestContext,
) -> Result<WireRequest, Error> {
let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization");
let authorization = wire
.headers
.iter()
.find(|(name, _)| is_authorization(name))
.map(|(_, value)| value.clone())
.unwrap_or_default();
self.trace
.lock()
.unwrap()
.push(format!("before_send:{authorization}"));
let headers = wire
.headers
.into_iter()
.map(|(name, value)| match is_authorization(&name) {
true => (name, "Bearer edited".to_string()),
false => (name, value),
})
.collect();
Ok(WireRequest { headers, ..wire })
}
}
#[tokio::test]
async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_replace_it() {
let upstream = upstream([pages_response()]).await;
let host = CallerTokenHost {
request: Mutex::new(Some(without_api_key(ocr_request(
"azure_ai/model",
&upstream.uri(),
json!({}),
)))),
trace: Mutex::new(Vec::new()),
};
litellm_host::run::run(ocr_machine(ocr_client()), &host)
.await
.unwrap();
assert_eq!(
*host.trace.lock().unwrap(),
["project", "token", "before_send:Bearer caller-token"]
);
assert_eq!(
only_request(&upstream).await.header_values("authorization"),
["Bearer edited"]
);
}

284
litellm-rust/crates/core/tests/ocr/machine.rs Normal file
View file

@ -0,0 +1,284 @@
use std::{
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
time::Duration,
};
use litellm_core::ocr::{
route::{OcrMachine, OcrOp, OcrProjection},
types::OcrDocumentInput,
};
use litellm_host::{
event::{CallEvent, WireRequest},
host::{Host, HostOp},
machine::{HostFailure, Machine, MachineStep},
};
use litellm_llms::base_llm::ocr::transformation::OcrTransportConfig;
use rstest::rstest;
use tokio::{io::AsyncReadExt, net::TcpListener, sync::Notify};
use super::{lifecycle::event_name, *};
/// Drives the machine by hand, answering every op through `host` except `before_send`,
/// which `intercept` answers so a test can fail or cancel exactly there.
async fn drive_until(
host: &LocalOcrHost,
mut intercept: impl FnMut(WireRequest) -> Result<WireRequest, HostFailure<Error>>,
) -> (
Result<LiteLLMOcrResponse, Error>,
Vec<&'static str>,
OcrMachine,
) {
let mut machine = ocr_machine(ocr_client());
let mut ops = Vec::new();
let outcome = loop {
let op = match machine.resume().await {
Ok(MachineStep::Host(op)) => op,
Ok(MachineStep::Complete(response)) => break Ok(response),
Err(error) => break Err(error),
};
let answer = match op {
HostOp::Project(reply) => {
ops.push("Project");
host.project()
.await
.map(|projection| reply.send(projection))
.map_err(HostFailure::Error)
}
HostOp::Custom(op) => {
ops.push(match op {
OcrOp::AcquireAzureAdToken(_) => "AcquireAzureAdToken",
});
host.custom_op(op).await.map_err(HostFailure::Error)
}
HostOp::BeforeSend { wire, reply, .. } => {
ops.push("BeforeSend");
intercept(*wire).map(|wire| reply.send(wire))
}
HostOp::Emit(event, reply) => {
let event = CallEvent::Machine(event);
ops.push(event_name(&event));
host.emit(&event)
.await
.map(|()| reply.send(()))
.map_err(HostFailure::Error)
}
};
if let Err(failure) = answer {
break machine.interrupt(failure).await;
}
};
(outcome, ops, machine)
}
/// Answers every op until `stop` fires, leaving the machine suspended mid-call.
async fn drive_until_notified(machine: &mut OcrMachine, host: &LocalOcrHost, stop: &Notify) {
tokio::time::timeout(Duration::from_secs(2), async {
loop {
tokio::select! {
_ = stop.notified() => break,
step = machine.resume() => {
match step.unwrap() {
MachineStep::Host(HostOp::Project(reply)) => reply.send(host.project().await.unwrap()),
MachineStep::Host(HostOp::Custom(op)) => host.custom_op(op).await.unwrap(),
MachineStep::Host(HostOp::BeforeSend { wire, reply, .. }) => reply.send(*wire),
MachineStep::Host(HostOp::Emit(_, reply)) => reply.send(()),
MachineStep::Complete(_) => panic!("the stalled call completed"),
}
}
}
}
})
.await
.expect("the call reached the stall point");
}
#[tokio::test]
async fn a_hand_driven_machine_performs_the_same_call() {
let upstream = upstream([json_response(json!({
"pages": [{"index": 0, "markdown": "native"}]
}))])
.await;
let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({})));
let (outcome, ops, mut machine) = drive_until(&host, Ok).await;
assert_eq!(outcome.unwrap().pages[0].markdown, "native");
assert_eq!(received(&upstream).await.len(), 1);
assert_eq!(ops, ["Project", "BeforeSend", "response"]);
assert!(matches!(
machine.resume().await,
Err(Error::InvalidRequest(_))
));
}
#[tokio::test]
async fn a_path_document_is_read_by_core_without_a_host_operation() {
let upstream = upstream([json_response(json!({
"pages": [{"index": 0, "markdown": "path"}]
}))])
.await;
let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::<u64>()));
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("scan.png");
std::fs::write(&path, b"abc").unwrap();
let request = ocr_request("mistral/model", &upstream.uri(), json!({})).with_document(
OcrDocumentInput::Path {
path,
mime_type: None,
},
);
let (response, ops, _) = drive_until(&LocalOcrHost::new(request), Ok).await;
std::fs::remove_dir_all(&dir).unwrap();
assert_eq!(response.unwrap().pages[0].markdown, "path");
assert_eq!(ops, ["Project", "BeforeSend", "response"]);
assert_eq!(
only_request(&upstream).await.json()["document"]["image_url"],
"data:image/png;base64,YWJj"
);
}
#[rstest]
#[case::failed(HostFailure::Error(Error::InvalidRequest("before_send failed".into())), "before_send failed")]
#[case::cancelled(HostFailure::Cancelled(Error::InvalidRequest("cancelled".into())), "cancelled")]
#[tokio::test]
async fn a_before_send_failure_ends_the_call_without_reaching_transport(
#[case] failure: HostFailure<Error>,
#[case] message: &str,
) {
let upstream = upstream([pages_response()]).await;
let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({})));
let failure = Arc::new(std::sync::Mutex::new(Some(failure)));
let (outcome, ops, mut machine) = drive_until(&host, |_| {
Err(failure
.lock()
.unwrap()
.take()
.expect("before_send is asked once"))
})
.await;
assert!(
matches!(&outcome, Err(Error::InvalidRequest(actual)) if actual == message),
"{outcome:?}"
);
assert_eq!(ops, ["Project", "BeforeSend"]);
assert!(machine.resume().await.is_err());
assert!(received(&upstream).await.is_empty());
}
#[tokio::test]
async fn resuming_before_answering_keeps_the_pending_operation() {
let request = ocr_request("mistral/model", UNREACHABLE_BASE, json!({}));
let mut machine = ocr_machine(ocr_client());
let Ok(MachineStep::Host(HostOp::Project(reply))) = machine.resume().await else {
panic!("expected the projection op first");
};
assert!(machine.resume().await.is_err());
reply.send(OcrProjection {
request,
caller_token: false,
});
assert!(matches!(
machine.resume().await,
Ok(MachineStep::Host(HostOp::BeforeSend { .. }))
));
}
#[derive(Debug)]
struct PendingToken {
entered: Arc<Notify>,
dropped: Arc<AtomicBool>,
}
struct TokenFutureDrop(Arc<AtomicBool>);
impl Drop for TokenFutureDrop {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
impl litellm_auth::TokenProvider for PendingToken {
fn acquire(&self) -> litellm_auth::TokenFuture<'_> {
Box::pin(async move {
let _guard = TokenFutureDrop(self.dropped.clone());
self.entered.notify_one();
std::future::pending().await
})
}
}
#[tokio::test]
async fn interrupt_drops_provider_captures_before_returning() {
let entered = Arc::new(Notify::new());
let dropped = Arc::new(AtomicBool::new(false));
let mut request = ocr_request("azure_ai/mistral-ocr", "https://example.invalid", json!({}));
request.transport = OcrTransportConfig {
extra_headers: vec![("authorization".into(), "Bearer test-key".into())],
..request.transport
};
request.azure_ad_token_provider = Some(litellm_auth::TokenProviderHandle::new(Arc::new(
PendingToken {
entered: entered.clone(),
dropped: dropped.clone(),
},
)));
let host = LocalOcrHost::new(request);
let mut machine = ocr_machine(ocr_client());
drive_until_notified(&mut machine, &host, &entered).await;
assert!(!dropped.load(Ordering::SeqCst));
let acknowledgement = machine.interrupt(HostFailure::Cancelled(Error::InvalidRequest(
"cancelled".into(),
)));
assert!(
dropped.load(Ordering::SeqCst),
"interrupt returned while provider captures were still alive"
);
assert!(
matches!(acknowledgement.await, Err(Error::InvalidRequest(message)) if message == "cancelled")
);
}
#[tokio::test]
async fn interrupting_an_in_flight_provider_request_closes_its_connection() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let received = Arc::new(Notify::new());
let server_received = received.clone();
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut request = Vec::new();
let mut buffer = [0u8; 4096];
while !request.windows(4).any(|window| window == b"\r\n\r\n") {
let read = socket.read(&mut buffer).await.unwrap();
request.extend_from_slice(&buffer[..read]);
}
server_received.notify_one();
while socket.read(&mut buffer).await.unwrap() != 0 {}
});
let host = LocalOcrHost::new(ocr_request("mistral/model", &base, json!({})));
let mut machine = ocr_machine(ocr_client());
drive_until_notified(&mut machine, &host, &received).await;
let cancelled = Error::InvalidRequest("cancelled".into());
assert!(
machine
.interrupt(HostFailure::Cancelled(cancelled))
.await
.is_err()
);
tokio::time::timeout(Duration::from_secs(1), server)
.await
.expect("the provider connection stayed open after the interrupt")
.unwrap();
}

125
litellm-rust/crates/core/tests/ocr/main.rs Normal file
View file

@ -0,0 +1,125 @@
use litellm_core::ocr::{
document::prepare_document,
route::{LocalOcrHost, ocr_machine},
types::LiteLLMOcrRequest,
wire::{OcrWireRequest, decode_request},
};
use litellm_llms::base_llm::ocr::{
error::Error,
handler::OcrClient,
transformation::{LiteLLMOcrResponse, OcrDocument},
};
use serde_json::{Map, Value, json};
use wiremock::{MockServer, ResponseTemplate};
#[path = "../support/mod.rs"]
mod support;
use support::*;
mod aws_textract;
mod azure_ai;
mod azure_document_intelligence;
mod cohere;
mod documents;
mod lifecycle;
mod machine;
mod mistral;
mod reducto;
mod vertex_ai;
const INLINE_PDF: &str = "data:application/pdf;base64,YWJj";
fn object(value: Value) -> Map<String, Value> {
let Value::Object(map) = value else {
panic!("expected a json object, got {value}");
};
map
}
fn ocr_client() -> OcrClient {
let document_http = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test document client builds");
OcrClient::for_test(reqwest::Client::new(), document_http)
}
async fn perform(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
litellm_core::ocr::client::perform(&ocr_client(), request).await
}
async fn perform_with(host: LocalOcrHost) -> Result<LiteLLMOcrResponse, Error> {
litellm_host::run::run(ocr_machine(ocr_client()), &host).await
}
fn wire(model: &str, base: &str, document: Value, options: Value) -> OcrWireRequest {
OcrWireRequest {
model: model.into(),
document,
api_key: Some(litellm_auth::SecretValue::new("test-key")),
api_base: Some(base.into()),
custom_llm_provider: None,
extra_headers: None,
optional_params: object(options),
input_sources: Default::default(),
timeout_seconds: Some(2.0),
}
}
/// A request for an inline PDF, authenticated with `test-key`.
fn ocr_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest {
ocr_request_with_document(
model,
base,
json!({"type": "document_url", "document_url": INLINE_PDF}),
options,
)
}
fn ocr_request_with_document(
model: &str,
base: &str,
document: Value,
options: Value,
) -> LiteLLMOcrRequest {
decode_request(wire(model, base, document, options)).expect("request decodes")
}
fn document(value: Value) -> OcrDocument {
serde_json::from_value(value).expect("document parses")
}
/// Points the request's resolved document at `source`, keeping its type.
fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest {
let resolved = request
.map_document(prepare_document)
.expect("document resolves");
let document = resolved.document.clone().with_source(source.into());
resolved.with_document(document.into())
}
fn with_headers(request: LiteLLMOcrRequest, headers: &[(&str, &str)]) -> LiteLLMOcrRequest {
let mut request = request;
request.transport.extra_headers = headers
.iter()
.map(|(name, value)| (name.to_string(), value.to_string()))
.collect();
request
}
fn without_api_key(request: LiteLLMOcrRequest) -> LiteLLMOcrRequest {
let mut request = request;
request.credentials.api_key = None;
request
}
fn pages_response() -> ResponseTemplate {
json_response(json!({"pages": []}))
}
/// An Azure Document Intelligence 202 whose operation lives on `server`.
fn accepted(server: &MockServer, body: Value) -> ResponseTemplate {
ResponseTemplate::new(202)
.insert_header("Operation-Location", format!("{}/operation", server.uri()))
.set_body_json(body)
}

248
litellm-rust/crates/core/tests/ocr/mistral.rs Normal file
View file

@ -0,0 +1,248 @@
use std::sync::Arc;
use litellm_auth_gcp::VertexAuth;
use litellm_http::{
HttpClientPool, HttpSettings, Resolution,
media::{PublicDnsResolver, UrlPolicy},
};
use litellm_llms::{
base_llm::ocr::{
settings::OcrSettings,
transformation::{BaseOcrConfig, OCR_RESPONSE_MAX_BYTES},
},
mistral::ocr::transformation::MistralOcrConfig,
};
use rstest::rstest;
use super::*;
#[tokio::test]
async fn direct_mistral_sends_one_request_with_every_option() {
let upstream = upstream([json_response(json!({
"pages": [{"index": 0, "markdown": "hello", "custom": "preserved"}],
"usage_info": {"pages_processed": 1}
}))])
.await;
let result = perform(ocr_request(
"mistral/model",
&upstream.uri(),
json!({"pages": "0,2-4", "extract_header": true, "unknown": "ignored"}),
))
.await
.unwrap();
assert_eq!(result.pages[0].markdown, "hello");
assert_eq!(result.pages[0].extra_fields["custom"], "preserved");
let sent = only_request(&upstream).await;
assert_eq!(sent.url.path(), "/v1/ocr");
assert_eq!(sent.header("authorization"), Some("Bearer test-key"));
assert_eq!(
sent.json(),
json!({
"model": "model",
"document": {"type": "document_url", "document_url": INLINE_PDF},
"pages": "0,2-4",
"extract_header": true,
"unknown": "ignored"
})
);
}
#[rstest]
#[case::litellm_format(json!({}), false)]
#[case::native_format(json!({"req_format": "native"}), true)]
#[tokio::test]
async fn the_native_response_is_kept_only_when_requested(
#[case] options: Value,
#[case] kept: bool,
) {
let provider_response = json!({
"pages": [{"index": 0, "markdown": "hello"}],
"usage_info": {"pages_processed": 1},
"provider_only": "preserved"
});
let upstream = upstream([json_response(provider_response.clone())]).await;
let response = perform(ocr_request("mistral/model", &upstream.uri(), options))
.await
.unwrap();
assert_eq!(
response.provider_native_response.map(Value::Object),
kept.then_some(provider_response)
);
}
#[rstest]
#[case::mistral("mistral/model", json!({}))]
#[case::vertex(
"vertex_ai/mistral-ocr-latest",
json!({"vertex_project": "test-project", "vertex_location": "us-central1"})
)]
#[tokio::test]
async fn an_upstream_error_keeps_its_status_whole_body_and_headers(
#[case] model: &str,
#[case] options: Value,
) {
let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))});
let expected_body = serde_json::to_string(&payload).unwrap();
let upstream = upstream([status_response(422, payload)
.insert_header("Retry-After", "17")
.insert_header("X-Request-ID", "request-123")
.insert_header("X-Future-Header", "retained")])
.await;
let error = perform(ocr_request(model, &upstream.uri(), options))
.await
.unwrap_err();
assert_eq!(received(&upstream).await.len(), 1);
let Error::Provider {
status,
body,
headers,
} = error
else {
panic!("expected provider error, got {error:?}");
};
assert_eq!(status, 422);
for (name, value) in [
("retry-after", "17"),
("x-request-id", "request-123"),
("x-future-header", "retained"),
] {
assert!(
headers
.iter()
.any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value),
"{name} missing from {headers:?}"
);
}
assert_eq!(body, expected_body);
}
#[rstest]
#[case::mistral_prefix("mistral/model", None, true)]
#[case::unknown_provider("model", Some("unknown"), false)]
fn decoding_accepts_known_providers_and_rejects_unknown_ones(
#[case] model: &str,
#[case] provider: Option<&str>,
#[case] accepted: bool,
) {
let request = OcrWireRequest {
custom_llm_provider: provider.map(Into::into),
..wire(
model,
"https://example.com",
json!({"type": "document_url", "document_url": "https://example.com/doc.pdf"}),
json!({"extract_header": true, "unknown": 42}),
)
};
assert_eq!(decode_request(request).is_ok(), accepted);
}
#[rstest]
#[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")]
#[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")]
#[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")]
#[tokio::test]
async fn missing_credentials_come_from_the_injected_secret_source(
#[case] secrets: &[(&str, &str)],
#[case] expected_key: &str,
) {
let upstream = upstream([pages_response()]).await;
let base = upstream.uri();
let source = Arc::new(RecordingSecrets::new(
secrets
.iter()
.copied()
.chain([("MISTRAL_AZURE_API_BASE", base.as_str())]),
));
let client = ocr_client().with_secrets(source.clone());
let request = decode_request(OcrWireRequest {
api_key: None,
api_base: None,
..wire(
"mistral/model",
&base,
json!({"type": "document_url", "document_url": INLINE_PDF}),
json!({}),
)
})
.unwrap();
litellm_core::ocr::client::perform(&client, request)
.await
.unwrap();
assert_eq!(source.requested(), MistralOcrConfig.secret_names());
assert_eq!(
only_request(&upstream).await.header("authorization"),
Some(format!("Bearer {expected_key}").as_str())
);
}
#[tokio::test]
async fn the_client_uses_the_injected_http_pool_configuration() {
let upstream = upstream([pages_response()]).await;
let settings = HttpSettings {
user_agent: Some("host-owned/1".into()),
..HttpSettings::default()
};
let client = OcrClient::new(
&HttpClientPool::new(Arc::new(PublicDnsResolver)),
&Resolution::from(&settings).config,
UrlPolicy::default(),
VertexAuth::default(),
OcrSettings::default(),
Arc::new(litellm_secrets::source::EnvironmentSecrets::default()),
)
.unwrap();
litellm_core::ocr::client::perform(
&client,
ocr_request("mistral/model", &upstream.uri(), json!({})),
)
.await
.unwrap();
assert_eq!(
only_request(&upstream).await.header("user-agent"),
Some("host-owned/1")
);
}
#[test]
fn a_valid_response_limit_is_consumed_and_not_forwarded() {
let request = ocr_request(
"mistral/model",
UNREACHABLE_BASE,
json!({"max_response_bytes": 123}),
);
assert_eq!(request.transport.max_response_bytes, 123);
assert!(!request.optional_params.contains_key("max_response_bytes"));
}
#[rstest]
#[case::zero(json!(0))]
#[case::negative(json!(-1))]
#[case::boolean(json!(true))]
#[case::string(json!("123"))]
#[case::fraction(json!(1.5))]
#[case::above_the_cap(json!(OCR_RESPONSE_MAX_BYTES + 1))]
#[case::null(Value::Null)]
fn an_invalid_response_limit_is_rejected(#[case] limit: Value) {
let Err(error) = decode_request(wire(
"mistral/model",
UNREACHABLE_BASE,
json!({"type": "document_url", "document_url": INLINE_PDF}),
json!({"max_response_bytes": limit}),
)) else {
panic!("invalid response limit {limit} accepted");
};
assert!(error.to_string().contains("max_response_bytes"), "{error}");
}

321
litellm-rust/crates/core/tests/ocr/reducto.rs Normal file
View file

@ -0,0 +1,321 @@
use std::sync::{Arc, Mutex};
use litellm_host::event::{CallEvent, MachineEvent, WireRequest};
use rstest::rstest;
use super::*;
fn upload_response() -> ResponseTemplate {
json_response(json!({"file_id": "reducto://uploaded.pdf"}))
}
fn chunks_response(chunks: Value) -> ResponseTemplate {
json_response(json!({"result": {"chunks": chunks}}))
}
fn source_field(model: &str) -> &'static str {
match model.ends_with("parse-legacy") {
true => "document_url",
false => "input",
}
}
#[rstest]
#[case::v3(
"reducto/parse-v3",
json!({
"formatting": {"table_output_format": "html"},
"retrieval": {"chunk_mode": "section"},
"settings": {"ocr_system": "standard"},
"future_ocr_option": true,
"extra_body": {"provider_option": "value"}
}),
"reducto://already.pdf",
json!({
"input": "reducto://already.pdf",
"formatting": {"table_output_format": "html"},
"retrieval": {"chunk_mode": "section"},
"settings": {"ocr_system": "standard"},
"future_ocr_option": true,
"provider_option": "value"
})
)]
#[case::legacy(
"reducto/parse-legacy",
json!({
"enhance": {"agentic": [{"type": "table"}]},
"future_ocr_option": true,
"extra_body": {"provider_option": "value"}
}),
"reducto://legacy.pdf",
json!({
"document_url": "reducto://legacy.pdf",
"options": {"enhance": {"agentic": [{"type": "table"}]}},
"future_ocr_option": true,
"provider_option": "value"
})
)]
#[tokio::test]
async fn an_uploaded_document_is_parsed_with_mapped_options(
#[case] model: &str,
#[case] options: Value,
#[case] source: &str,
#[case] expected: Value,
) {
let upstream = upstream([chunks_response(json!([]))]).await;
perform(with_source(
ocr_request(model, &upstream.uri(), options),
source,
))
.await
.unwrap();
let sent = only_request(&upstream).await;
assert_eq!(sent.url.path(), "/parse");
assert_eq!(sent.json(), expected);
}
#[rstest]
#[tokio::test]
async fn an_inline_document_is_uploaded_as_multipart_then_parsed(
#[values("parse-v3", "parse-legacy")] model: &str,
#[values("application/pdf", "image/png")] mime_type: &str,
) {
let upstream = upstream([
upload_response(),
chunks_response(json!([{"content": "hello"}])),
])
.await;
let data_uri = format!("data:{mime_type};base64,YWJj");
let document = match mime_type.starts_with("image/") {
true => json!({"type": "image_url", "image_url": data_uri}),
false => json!({"type": "document_url", "document_url": data_uri}),
};
let request = with_headers(
ocr_request_with_document(
&format!("reducto/{model}"),
&upstream.uri(),
document,
json!({}),
),
&[
("Content-Type", "application/json"),
("X-Trace", "upload-test"),
],
);
let response = perform(request).await.unwrap();
assert_eq!(response.pages[0].markdown, "hello");
let requests = received(&upstream).await;
let [upload, parse] = requests.as_slice() else {
panic!(
"expected an upload and a parse, got {} requests",
requests.len()
);
};
assert_eq!(upload.url.path(), "/upload");
assert!(
upload
.header("content-type")
.is_some_and(|value| value.starts_with("multipart/form-data; boundary=")),
"{:?}",
upload.header("content-type")
);
assert_eq!(upload.header("x-trace"), Some("upload-test"));
let multipart = upload.body_text();
assert!(
multipart.contains(&format!("Content-Type: {mime_type}\r\n")),
"{multipart}"
);
assert!(multipart.contains("\r\n\r\nabc\r\n--"), "{multipart}");
assert_eq!(parse.url.path(), "/parse");
assert_eq!(
parse.json(),
json!({source_field(model): "reducto://uploaded.pdf"})
);
for request in &requests {
assert_eq!(request.header("authorization"), Some("Bearer test-key"));
}
}
#[tokio::test]
async fn response_received_fires_once_for_the_parse_response() {
let upstream = upstream([upload_response(), chunks_response(json!([]))]).await;
let observed = Arc::new(Mutex::new(Vec::new()));
let recorder = observed.clone();
let host = LocalOcrHost::new(ocr_request("reducto/parse-v3", &upstream.uri(), json!({})))
.with_observer(move |event| {
if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event {
recorder.lock().unwrap().push(raw.body.clone());
}
});
perform_with(host).await.unwrap();
assert_eq!(received(&upstream).await.len(), 2);
assert_eq!(*observed.lock().unwrap(), [r#"{"result":{"chunks":[]}}"#]);
}
#[rstest]
#[case::empty_id(json_response(json!({"file_id": ""})))]
#[case::missing_id(json_response(json!({})))]
#[case::null_id(json_response(json!({"file_id": null})))]
#[case::upload_failure(status_response(503, json!({"error": "unavailable"})))]
#[tokio::test]
async fn a_failed_upload_stops_before_parse(#[case] upload: ResponseTemplate) {
let upstream = upstream([upload]).await;
let result = perform(ocr_request("reducto/parse-v3", &upstream.uri(), json!({}))).await;
assert!(result.is_err());
assert_eq!(received(&upstream).await.len(), 1);
}
#[rstest]
#[case::remote_url("https://example.com/a.pdf", Error::ReductoSource)]
#[case::empty_file_id("reducto://", Error::RequestField { path: "document file id".into() })]
#[case::data_uri_without_payload("data:application/pdf;base64", Error::InvalidDataUri)]
#[case::invalid_base64("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)]
#[tokio::test]
async fn invalid_document_sources_are_rejected_before_sending(
#[case] source: &str,
#[case] expected: Error,
) {
let upstream = upstream([json_response(json!({}))]).await;
let result = perform(with_source(
ocr_request("reducto/parse-v3", &upstream.uri(), json!({})),
source,
))
.await;
assert!(
received(&upstream).await.is_empty(),
"sent invalid source: {source}"
);
let error = result.unwrap_err();
assert_eq!(
std::mem::discriminant(&error),
std::mem::discriminant(&expected)
);
assert_eq!(error.http_status_code(), Some(400));
assert_eq!(error.to_string(), expected.to_string());
}
#[tokio::test]
async fn a_forwarded_authorization_wins_and_the_native_response_is_omitted_by_default() {
let upstream = upstream([json_response(
json!({"job_id": "job-1", "result": {"chunks": []}}),
)])
.await;
let request = with_headers(
with_source(
ocr_request("reducto/parse-v3", &upstream.uri(), json!({})),
"reducto://ready.pdf",
),
&[("authorization", "Bearer existing")],
);
let response = perform(request).await.unwrap();
assert_eq!(response.provider_native_response, None);
assert_eq!(
only_request(&upstream).await.header_values("authorization"),
["Bearer existing"]
);
}
#[tokio::test]
async fn native_format_retains_the_provider_response() {
let raw = json!({
"result": {"chunks": [{"content": "native OCR response"}]},
"usage": {"num_pages": 1}
});
let upstream = upstream([json_response(raw.clone())]).await;
let response = perform(with_source(
ocr_request(
"reducto/parse-v3",
&upstream.uri(),
json!({"req_format": "native"}),
),
"reducto://ready.pdf",
))
.await
.unwrap();
assert_eq!(response.pages[0].markdown, "native OCR response");
assert_eq!(
response.provider_native_response.map(Value::Object),
Some(raw)
);
}
#[tokio::test]
async fn an_unknown_model_reaches_parse_and_keeps_its_name() {
let upstream = upstream([chunks_response(
json!([{"content": "future model response"}]),
)])
.await;
let response = perform(with_source(
ocr_request("reducto/future-parse-model", &upstream.uri(), json!({})),
"reducto://ready.pdf",
))
.await
.unwrap();
assert_eq!(response.model, "future-parse-model");
assert_eq!(response.pages[0].markdown, "future model response");
let sent = only_request(&upstream).await;
assert_eq!(sent.url.path(), "/parse");
assert_eq!(sent.json(), json!({"input": "reducto://ready.pdf"}));
}
#[tokio::test]
async fn a_guardrail_can_replace_the_document_before_upload() {
let upstream = upstream([chunks_response(json!([]))]).await;
let host = LocalOcrHost::new(ocr_request("reducto/parse-v3", &upstream.uri(), json!({})))
.with_before_send(|wire, _| {
assert_eq!(wire.body["document_url"], INLINE_PDF);
Ok(WireRequest {
body: json!({"type": "document_url", "document_url": "reducto://guarded.pdf"}),
..wire
})
});
perform_with(host).await.unwrap();
let sent = only_request(&upstream).await;
assert_eq!(sent.url.path(), "/parse");
assert_eq!(sent.json(), json!({"input": "reducto://guarded.pdf"}));
}
#[rstest]
#[tokio::test]
async fn guardrail_headers_reach_both_upload_and_parse(
#[values("reducto/parse-v3", "reducto/parse-legacy")] model: &str,
) {
let upstream = upstream([upload_response(), chunks_response(json!([]))]).await;
let request = with_headers(
ocr_request(model, &upstream.uri(), json!({})),
&[("authorization", "Bearer original")],
);
let host = LocalOcrHost::new(request).with_before_send(|wire, _| {
Ok(WireRequest {
headers: vec![("authorization".into(), "Bearer guarded".into())],
..wire
})
});
perform_with(host).await.unwrap();
let requests = received(&upstream).await;
let paths: Vec<&str> = requests.iter().map(|request| request.url.path()).collect();
assert_eq!(paths, ["/upload", "/parse"]);
for request in &requests {
assert_eq!(request.header_values("authorization"), ["Bearer guarded"]);
}
}

184
litellm-rust/crates/core/tests/ocr/vertex_ai.rs Normal file
View file

@ -0,0 +1,184 @@
use litellm_auth::{InputSource, Sourced};
use litellm_core::ocr::arguments::is_supported_request;
use litellm_llms::base_llm::ocr::settings::OcrSettings;
use rstest::rstest;
use super::*;
#[tokio::test]
async fn mistral_is_served_at_the_resolved_project_and_location() {
let upstream = upstream([json_response(json!({
"pages": [{"index": 0, "markdown": "hello"}],
"usage_info": {"pages_processed": 1}
}))])
.await;
let response = perform(ocr_request(
"vertex_ai/mistral-ocr-maas",
&upstream.uri(),
json!({
"vertex_project": "project-1",
"vertex_location": "europe-west4",
"extract_footer": true
}),
))
.await
.unwrap();
assert_eq!(response.pages[0].markdown, "hello");
let sent = only_request(&upstream).await;
assert_eq!(
sent.url.path(),
"/v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
assert_eq!(sent.header("authorization"), Some("Bearer test-key"));
assert_eq!(
sent.json(),
json!({
"model": "mistral-ocr-maas",
"document": {"type": "document_url", "document_url": INLINE_PDF},
"extract_footer": true
})
);
}
#[tokio::test]
async fn configured_project_and_location_apply_when_the_call_sets_neither() {
let upstream = upstream([pages_response()]).await;
let client = ocr_client().with_settings(OcrSettings {
vertex_project: Some("configured-project".into()),
vertex_location: Some("europe-west4".into()),
..OcrSettings::default()
});
litellm_core::ocr::client::perform(
&client,
ocr_request("vertex_ai/mistral-ocr-maas", &upstream.uri(), json!({})),
)
.await
.unwrap();
assert_eq!(
only_request(&upstream).await.url.path(),
"/v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
}
#[tokio::test]
async fn a_supplied_authorization_is_forwarded_without_a_static_token() {
let upstream = upstream([pages_response()]).await;
let request = with_headers(
without_api_key(ocr_request(
"vertex_ai/model",
&upstream.uri(),
json!({"vertex_project": "project-1"}),
)),
&[("authorization", "Bearer supplied")],
);
perform(request).await.unwrap();
assert_eq!(
only_request(&upstream).await.header_values("authorization"),
["Bearer supplied"]
);
}
#[tokio::test]
async fn invalid_credentials_fail_before_sending() {
let error = perform(ocr_request(
"vertex_ai/model",
UNREACHABLE_BASE,
json!({"vertex_credentials": true}),
))
.await
.unwrap_err();
assert!(error.to_string().contains("vertex_credentials"), "{error}");
}
#[rstest]
#[tokio::test]
async fn a_request_controlled_api_base_is_rejected_before_vertex_auth(
#[values("vertex_ai/mistral-ocr-maas", "vertex_ai/deepseek-ocr-maas")] model: &str,
) {
let mut request = ocr_request(
model,
"https://caller.example",
json!({"vertex_project": "project-1"}),
);
request.credentials.api_base = Some(Sourced::new(
"https://caller.example".into(),
InputSource::Request,
));
let error = perform(request).await.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Vertex AI endpoint"),
"{error}"
);
}
#[tokio::test]
async fn deepseek_is_served_at_the_openai_compatible_endpoint() {
let upstream = upstream([json_response(json!({
"choices": [{"message": {"content": "recognized"}}],
"usage": {"prompt_tokens": 1}
}))])
.await;
let request = with_source(
ocr_request(
"vertex_ai/deepseek-ocr-maas",
&upstream.uri(),
json!({
"vertex_project": "project-1",
"vertex_location": "europe-west4",
"temperature": 0.1,
"future_ocr_option": true,
"extra_body": {"provider_option": "value"}
}),
),
"gs://bucket/document.pdf",
);
let response = perform(request).await.unwrap();
assert_eq!(response.pages[0].markdown, "recognized");
assert_eq!(
response.usage_info.unwrap().extra_fields["prompt_tokens"],
1
);
let sent = only_request(&upstream).await;
assert_eq!(
sent.url.path(),
"/v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions"
);
assert_eq!(sent.header("authorization"), Some("Bearer test-key"));
let body = sent.json();
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
assert_eq!(body["temperature"], 0.1);
assert_eq!(body["future_ocr_option"], true);
assert_eq!(body["provider_option"], "value");
assert!(body.get("vertex_project").is_none());
assert!(body.get("extra_body").is_none());
assert_eq!(
body["messages"][0]["content"][0],
json!({"type": "image_url", "image_url": "gs://bucket/document.pdf"})
);
}
#[rstest]
#[case::deepseek("deepseek-ocr-maas", Some("vertex_ai"), true)]
#[case::mistral("mistral-ocr-maas", Some("vertex_ai"), true)]
#[case::prefixed("vertex_ai/mistral-ocr-maas", None, true)]
#[case::unknown_provider("model", Some("unknown"), false)]
fn supported_requests_follow_the_registered_configs(
#[case] model: &str,
#[case] provider: Option<&str>,
#[case] supported: bool,
) {
assert_eq!(is_supported_request(model, provider), supported);
}

155
litellm-rust/crates/core/tests/support/mod.rs Normal file
View file

@ -0,0 +1,155 @@
//! Shared fixtures for route integration tests: a scripted upstream and a recording
//! secret source.
#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset
use std::sync::Mutex;
use futures_util::future::BoxFuture;
use litellm_secrets::{SecretValue, source::SecretSource};
use serde_json::Value;
use wiremock::{Mock, MockServer, Request, ResponseTemplate, matchers::any};
/// A port nothing listens on, for calls that must fail before any request is sent.
pub const UNREACHABLE_BASE: &str = "http://127.0.0.1:1";
/// Starts an upstream that answers its n-th request with the n-th response and 404s after.
pub async fn upstream(responses: impl IntoIterator<Item = ResponseTemplate>) -> MockServer {
let server = MockServer::start().await;
respond_in_order(&server, responses).await;
server
}
/// Scripts responses on a started server, for responses that need its address.
pub async fn respond_in_order(
server: &MockServer,
responses: impl IntoIterator<Item = ResponseTemplate>,
) {
for response in responses {
Mock::given(any())
.respond_with(response)
.up_to_n_times(1)
.mount(server)
.await;
}
}
pub async fn received(server: &MockServer) -> Vec<Request> {
server
.received_requests()
.await
.expect("request recording is on")
}
pub async fn only_request(server: &MockServer) -> Request {
let [request] = <[Request; 1]>::try_from(received(server).await)
.unwrap_or_else(|requests| panic!("expected one request, got {}", requests.len()));
request
}
pub fn json_response(body: Value) -> ResponseTemplate {
ResponseTemplate::new(200).set_body_json(body)
}
pub fn status_response(status: u16, body: Value) -> ResponseTemplate {
ResponseTemplate::new(status).set_body_json(body)
}
pub trait ReceivedRequest {
fn header(&self, name: &str) -> Option<&str>;
fn header_values(&self, name: &str) -> Vec<&str>;
fn json(&self) -> Value;
fn body_text(&self) -> String;
/// The path and query, as the request line carried them.
fn target(&self) -> String;
fn query(&self, name: &str) -> Option<String>;
}
impl ReceivedRequest for Request {
fn header(&self, name: &str) -> Option<&str> {
self.headers.get(name).and_then(|value| value.to_str().ok())
}
fn header_values(&self, name: &str) -> Vec<&str> {
self.headers
.get_all(name)
.iter()
.filter_map(|value| value.to_str().ok())
.collect()
}
fn json(&self) -> Value {
serde_json::from_slice(&self.body).expect("request body is json")
}
fn body_text(&self) -> String {
String::from_utf8_lossy(&self.body).into_owned()
}
fn target(&self) -> String {
match self.url.query() {
Some(query) => format!("{}?{query}", self.url.path()),
None => self.url.path().to_string(),
}
}
fn query(&self, name: &str) -> Option<String> {
self.url
.query_pairs()
.find_map(|(key, value)| (key == name).then(|| value.into_owned()))
}
}
/// A secret source that answers from a fixed table and records every name it was asked for.
pub struct RecordingSecrets {
values: Vec<(String, String)>,
fails: bool,
requested: Mutex<Vec<String>>,
}
impl RecordingSecrets {
pub fn new<'a>(values: impl IntoIterator<Item = (&'a str, &'a str)>) -> Self {
Self {
values: values
.into_iter()
.map(|(name, value)| (name.to_string(), value.to_string()))
.collect(),
fails: false,
requested: Mutex::new(Vec::new()),
}
}
pub fn empty() -> Self {
Self::new([])
}
pub fn failing() -> Self {
Self {
fails: true,
..Self::empty()
}
}
pub fn requested(&self) -> Vec<String> {
self.requested.lock().unwrap().clone()
}
}
impl SecretSource for RecordingSecrets {
fn get_secret_str<'a>(
&'a self,
name: &'a str,
) -> BoxFuture<'a, Result<Option<SecretValue>, litellm_secrets::Error>> {
Box::pin(async move {
self.requested.lock().unwrap().push(name.to_string());
if self.fails {
return Err(litellm_secrets::Error::ManagedSecretMissing);
}
Ok(self
.values
.iter()
.find(|(key, _)| key == name)
.map(|(_, value)| SecretValue::new(value.clone())))
})
}
}

5
litellm-rust/crates/framer/Cargo.toml
View file

@ -8,16 +8,17 @@ repository.workspace = true
[features]
default = ["aws", "sse"]
aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"]
sse = ["dep:sse-stream"]
sse = []
[dependencies]
aws-smithy-eventstream = { version = "=0.61.4", optional = true }
aws-smithy-types = { version = "1.6.1", optional = true }
bytes = "1"
futures-util.workspace = true
sse-stream = { version = "=0.2.6", optional = true }
thiserror.workspace = true
tokio-util = { version = "0.7", features = ["codec", "io"] }
[dev-dependencies]
proptest.workspace = true
rstest.workspace = true
tokio.workspace = true

95
litellm-rust/crates/framer/src/aws_event_stream.rs
View file

@ -1,66 +1,47 @@
use bytes::{Buf, Bytes, BytesMut};
use futures_util::{Stream, StreamExt};
use aws_smithy_eventstream::frame::{read_message_from, write_message_to};
pub use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
use bytes::BytesMut;
use tokio_util::codec::{Decoder, Encoder};
use aws_smithy_eventstream::frame::read_message_from;
use aws_smithy_types::event_stream::Header;
use crate::{Error, Framer};
use crate::EventStreamError;
const MIN_FRAME_BYTES: usize = 16;
const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
#[derive(Clone, Debug, PartialEq)]
pub struct AwsEventStreamFrame {
pub headers: Vec<Header>,
pub payload: Bytes,
}
#[derive(Clone, Copy, Debug, Default)]
pub struct AwsEventStreamFramer;
pub struct AwsEventStreamCodec;
impl Framer for AwsEventStreamFramer {
type Frame = AwsEventStreamFrame;
impl Decoder for AwsEventStreamCodec {
type Item = Message;
type Error = EventStreamError;
fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<Self::Frame, Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
futures_util::stream::try_unfold(
(Box::pin(input), BytesMut::new()),
|(mut input, mut buffer)| async move {
loop {
if buffer.len() >= 4 {
let length = (&buffer[..4]).get_u32() as usize;
if !(16..=MAX_FRAME_BYTES).contains(&length) {
return Err(Error::InvalidLength(length));
}
if buffer.len() >= length {
let raw = buffer.split_to(length).freeze();
let message = read_message_from(raw)?;
let frame = AwsEventStreamFrame {
headers: message.headers().to_vec(),
payload: message.payload().clone(),
};
return Ok(Some((frame, (input, buffer))));
}
}
match input.next().await {
Some(Ok(mut chunk)) => {
while chunk.has_remaining() {
let bytes = chunk.chunk();
buffer.extend_from_slice(bytes);
let length = bytes.len();
chunk.advance(length);
}
}
Some(Err(error)) => return Err(Error::Body(Box::new(error))),
None if buffer.is_empty() => return Ok(None),
None => return Err(Error::Truncated),
}
}
},
)
.fuse()
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Message>, EventStreamError> {
let Some(prefix) = src.first_chunk::<4>() else {
return Ok(None);
};
let length = u32::from_be_bytes(*prefix) as usize;
if !(MIN_FRAME_BYTES..=MAX_FRAME_BYTES).contains(&length) {
return Err(EventStreamError::InvalidLength(length));
}
if src.len() < length {
return Ok(None);
}
Ok(Some(read_message_from(src.split_to(length).freeze())?))
}
fn decode_eof(&mut self, src: &mut BytesMut) -> Result<Option<Message>, EventStreamError> {
match self.decode(src)? {
Some(message) => Ok(Some(message)),
None if src.is_empty() => Ok(None),
None => Err(EventStreamError::Truncated),
}
}
}
impl Encoder<Message> for AwsEventStreamCodec {
type Error = EventStreamError;
fn encode(&mut self, message: Message, dst: &mut BytesMut) -> Result<(), EventStreamError> {
Ok(write_message_to(&message, dst)?)
}
}

24
litellm-rust/crates/framer/src/error.rs
View file

@ -1,17 +1,21 @@
#[cfg(feature = "sse")]
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[cfg(feature = "sse")]
#[error("SSE framing failed: {0}")]
Sse(#[from] sse_stream::Error),
#[cfg(feature = "aws")]
#[error("AWS EventStream framing failed: {0}")]
Aws(#[from] aws_smithy_eventstream::error::Error),
pub enum SseError {
#[error("body stream failed: {0}")]
Body(#[source] Box<dyn std::error::Error + Send + Sync>),
#[cfg(feature = "aws")]
Body(#[from] std::io::Error),
#[error("SSE field is not UTF-8: {0}")]
InvalidUtf8(#[from] std::str::Utf8Error),
}
#[cfg(feature = "aws")]
#[derive(Debug, thiserror::Error)]
pub enum EventStreamError {
#[error("body stream failed: {0}")]
Body(#[from] std::io::Error),
#[error("invalid AWS EventStream frame length: {0}")]
InvalidLength(usize),
#[cfg(feature = "aws")]
#[error("truncated AWS EventStream frame")]
Truncated,
#[error("malformed AWS EventStream frame: {0}")]
Malformed(#[from] aws_smithy_eventstream::error::Error),
}

21
litellm-rust/crates/framer/src/framed.rs Normal file
View file

@ -0,0 +1,21 @@
use std::io;
use bytes::Buf;
use futures_util::{Stream, StreamExt, TryStreamExt};
use tokio_util::{
codec::{Decoder, FramedRead},
io::StreamReader,
};
pub fn frames<S, B, E, D>(
input: S,
codec: D,
) -> impl Stream<Item = Result<D::Item, D::Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
D: Decoder + Send,
{
FramedRead::new(StreamReader::new(input.map_err(io::Error::other)), codec).fuse()
}

4
litellm-rust/crates/framer/src/lib.rs
View file

@ -1,8 +1,8 @@
mod error;
mod framer;
mod framed;
pub use error::*;
pub use framer::*;
pub use framed::frames;
#[cfg(feature = "aws")]
pub mod aws_event_stream;

189
litellm-rust/crates/framer/src/sse.rs
View file

@ -1,43 +1,170 @@
use futures_util::{Stream, StreamExt};
use std::str;
use crate::{Error, Framer};
use bytes::{Buf, BufMut, BytesMut};
use tokio_util::codec::{Decoder, Encoder};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SseFrame {
use crate::SseError;
const BOM: &[u8] = b"\xEF\xBB\xBF";
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct SseEvent {
pub event: Option<String>,
pub data: Option<String>,
pub data: String,
pub id: Option<String>,
pub retry: Option<u64>,
}
#[derive(Clone, Copy, Debug, Default)]
pub struct SseFramer;
pub struct SseCodec {
past_bom: bool,
}
impl Framer for SseFramer {
type Frame = SseFrame;
impl Decoder for SseCodec {
type Item = SseEvent;
type Error = SseError;
fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<SseFrame, Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: bytes::Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
let frames = Box::pin(sse_stream::SseStream::from_bytes_stream(input));
futures_util::stream::try_unfold(frames, |mut frames| async move {
let Some(frame) = frames.next().await else {
return Ok(None);
};
let frame = frame?;
Ok(Some((
SseFrame {
event: frame.event,
data: frame.data,
id: frame.id,
retry: frame.retry,
},
frames,
)))
})
.fuse()
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
if !self.skip_bom(src) {
return Ok(None);
}
while let Some(end) = block_end(src) {
let block = src.split_to(end);
let pending = lines(&block)
.map(|(line, _)| line)
.take_while(|line| !line.is_empty())
.try_fold(Pending::default(), Pending::apply)?;
if let Some(event) = pending.dispatch() {
return Ok(Some(event));
}
}
Ok(None)
}
fn decode_eof(&mut self, _pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
Ok(None)
}
}
impl SseCodec {
fn skip_bom(&mut self, src: &mut BytesMut) -> bool {
if self.past_bom {
return true;
}
if src.starts_with(BOM) {
src.advance(BOM.len());
} else if BOM.starts_with(src) {
return false;
}
self.past_bom = true;
true
}
}
fn block_end(bytes: &[u8]) -> Option<usize> {
lines(bytes)
.find(|(line, _)| line.is_empty())
.map(|(_, end)| end)
}
fn lines(bytes: &[u8]) -> impl Iterator<Item = (&[u8], usize)> {
let mut cursor: usize = 0;
std::iter::from_fn(move || {
let rest = &bytes[cursor..];
let end = rest.iter().position(|byte| matches!(byte, b'\n' | b'\r'))?;
cursor += end + terminator_len(&rest[end..]);
Some((&rest[..end], cursor))
})
}
fn terminator_len(terminated: &[u8]) -> usize {
match terminated {
[b'\r', b'\n', ..] => 2,
_ => 1,
}
}
#[derive(Default)]
struct Pending {
event: Option<String>,
data: Option<String>,
id: Option<String>,
retry: Option<u64>,
}
impl Pending {
fn apply(self, line: &[u8]) -> Result<Self, SseError> {
let (name, value) = split_field(line);
Ok(match name {
b"event" => Self {
event: Some(str::from_utf8(value)?.to_owned()),
..self
},
b"data" => Self {
data: Some(append_data(self.data, str::from_utf8(value)?)),
..self
},
b"id" if !value.contains(&0) => Self {
id: Some(str::from_utf8(value)?.to_owned()),
..self
},
b"retry" => Self {
retry: parse_retry(value).or(self.retry),
..self
},
_ => self,
})
}
fn dispatch(self) -> Option<SseEvent> {
Some(SseEvent {
event: self.event,
data: self.data?,
id: self.id,
retry: self.retry,
})
}
}
fn split_field(line: &[u8]) -> (&[u8], &[u8]) {
let Some(colon) = line.iter().position(|byte| *byte == b':') else {
return (line, &[]);
};
let value = &line[colon + 1..];
(&line[..colon], value.strip_prefix(b" ").unwrap_or(value))
}
fn append_data(buffer: Option<String>, line: &str) -> String {
match buffer {
Some(existing) => format!("{existing}\n{line}"),
None => line.to_owned(),
}
}
fn parse_retry(value: &[u8]) -> Option<u64> {
if !value.iter().all(u8::is_ascii_digit) {
return None;
}
str::from_utf8(value).ok()?.parse().ok()
}
impl Encoder<SseEvent> for SseCodec {
type Error = SseError;
fn encode(&mut self, event: SseEvent, dst: &mut BytesMut) -> Result<(), SseError> {
if let Some(name) = event.event {
dst.put_slice(format!("event: {name}\n").as_bytes());
}
for line in event.data.split('\n') {
dst.put_slice(format!("data: {line}\n").as_bytes());
}
if let Some(id) = event.id {
dst.put_slice(format!("id: {id}\n").as_bytes());
}
if let Some(retry) = event.retry {
dst.put_slice(format!("retry: {retry}\n").as_bytes());
}
dst.put_u8(b'\n');
Ok(())
}
}

209
litellm-rust/crates/framer/tests/aws_event_stream.rs
View file

@ -4,89 +4,174 @@ mod support;
use std::io;
use futures_util::TryStreamExt;
use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
use litellm_framing::{Error, Framer};
use bytes::Bytes;
use futures_util::{StreamExt, TryStreamExt, stream};
use litellm_framing::{
EventStreamError,
aws_event_stream::{AwsEventStreamCodec, Header, HeaderValue, Message},
frames,
};
use proptest::prelude::*;
use rstest::{fixture, rstest};
use support::{body_cause, cut_at, encode_all, every, input, runtime};
use support::encode;
async fn collect_aws(bytes: &[u8], chunk_size: usize) -> Result<Vec<AwsEventStreamFrame>, Error> {
AwsEventStreamFramer
.frame(futures_util::stream::iter(
bytes.chunks(chunk_size).map(Ok::<_, io::Error>),
))
async fn collect(pieces: Vec<Bytes>) -> Result<Vec<Message>, EventStreamError> {
frames(input(pieces), AwsEventStreamCodec)
.try_collect()
.await
}
#[fixture]
fn two_frames() -> Vec<u8> {
[encode(b"\xff\x00"), encode(b"second")].concat()
fn message(payload: &[u8]) -> Message {
Message::new(Bytes::copy_from_slice(payload))
.add_header(Header::new(
":event-type",
HeaderValue::String("payload".into()),
))
.add_header(Header::new("sequence", HeaderValue::Int32(7)))
}
#[fixture]
fn payload_frame() -> Vec<u8> {
encode(b"payload")
encode_all(AwsEventStreamCodec, [message(b"payload")])
}
fn header_value() -> impl Strategy<Value = HeaderValue> {
prop_oneof![
"[a-z]{0,8}".prop_map(|text| HeaderValue::String(text.into())),
any::<i32>().prop_map(HeaderValue::Int32),
any::<bool>().prop_map(HeaderValue::Bool),
proptest::collection::vec(any::<u8>(), 0..8)
.prop_map(|bytes| HeaderValue::ByteArray(bytes.into())),
]
}
fn arbitrary_message() -> impl Strategy<Value = Message> {
(
proptest::collection::vec(("[a-z:-]{1,12}", header_value()), 0..3),
proptest::collection::vec(any::<u8>(), 0..32),
)
.prop_map(|(headers, payload)| {
headers.into_iter().fold(
Message::new(Bytes::from(payload)),
|message, (name, value)| message.add_header(Header::new(name, value)),
)
})
}
proptest! {
#[test]
fn any_messages_survive_a_round_trip_through_any_cuts(
messages in proptest::collection::vec(arbitrary_message(), 1..4),
cuts in proptest::collection::vec(0_usize..512, 0..4),
) {
let wire = encode_all(AwsEventStreamCodec, messages.clone());
let decoded = runtime().block_on(collect(cut_at(&wire, cuts))).unwrap();
prop_assert_eq!(decoded, messages);
}
}
#[rstest]
#[case(1)]
#[case(3)]
#[case(12)]
#[case(usize::MAX)]
#[case::prelude_crc(8)]
#[case::message_crc(usize::MAX)]
#[tokio::test]
async fn fragmented_and_coalesced_frames_preserve_typed_headers_and_binary_payloads(
two_frames: Vec<u8>,
#[case] chunk_size: usize,
) {
let chunk_size = chunk_size.min(two_frames.len());
let frames = collect_aws(&two_frames, chunk_size).await.unwrap();
assert_eq!(frames.len(), 2);
assert_eq!(frames[0].payload, &b"\xff\x00"[..]);
assert_eq!(frames[1].payload, "second");
assert_eq!(
frames[0].headers[0].value().as_string().unwrap().as_str(),
"payload"
);
assert_eq!(frames[0].headers[1].value().as_int32(), Ok(7));
}
#[rstest]
#[case(8)]
#[case(0)]
#[tokio::test]
async fn rejects_corrupt_crcs(payload_frame: Vec<u8>, #[case] index: usize) {
let corrupt_index = if index == 0 {
payload_frame.len() - 1
} else {
index
};
async fn a_corrupt_crc_is_malformed(payload_frame: Vec<u8>, #[case] index: usize) {
let mut corrupt = payload_frame;
corrupt[corrupt_index] ^= 1;
assert!(matches!(collect_aws(&corrupt, 3).await, Err(Error::Aws(_))));
}
#[rstest]
#[case(0_u32)]
#[case(15)]
#[case(u32::MAX)]
#[tokio::test]
async fn rejects_invalid_lengths(#[case] length: u32) {
let flipped = index.min(corrupt.len() - 1);
corrupt[flipped] ^= 1;
assert!(matches!(
collect_aws(&length.to_be_bytes(), 1).await,
Err(Error::InvalidLength(_))
collect(every(&corrupt, 3)).await,
Err(EventStreamError::Malformed(_))
));
}
#[rstest]
#[case(1)]
#[case(3)]
#[case(5)]
#[case::zero(0)]
#[case::below_minimum(15)]
#[case::above_maximum(16 * 1024 * 1024 + 1)]
#[case::u32_max(u32::MAX)]
#[tokio::test]
async fn rejects_truncation(payload_frame: Vec<u8>, #[case] end: usize) {
async fn a_length_outside_the_frame_bounds_fails_before_buffering(#[case] length: u32) {
assert!(matches!(
collect_aws(&payload_frame[..end], 1).await,
Err(Error::Truncated)
collect(every(&length.to_be_bytes(), 1)).await,
Err(EventStreamError::InvalidLength(seen)) if seen == length as usize
));
}
#[rstest]
#[case::before_the_length(1)]
#[case::inside_the_prelude(5)]
#[case::one_byte_short(usize::MAX)]
#[tokio::test]
async fn eof_inside_a_frame_is_truncation(payload_frame: Vec<u8>, #[case] end: usize) {
let end = end.min(payload_frame.len() - 1);
assert!(matches!(
collect(every(&payload_frame[..end], 1)).await,
Err(EventStreamError::Truncated)
));
}
const FRAME_OVERHEAD_BYTES: usize = 16;
const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
#[tokio::test]
async fn a_frame_at_exactly_the_maximum_length_decodes() {
let largest = Message::new(vec![0xAB; MAX_FRAME_BYTES - FRAME_OVERHEAD_BYTES]);
let wire = encode_all(AwsEventStreamCodec, [largest.clone()]);
assert_eq!(wire.len(), MAX_FRAME_BYTES);
assert_eq!(collect(every(&wire, 1 << 20)).await.unwrap(), vec![largest]);
}
#[tokio::test]
async fn a_frame_one_byte_over_the_maximum_length_is_rejected_by_its_prelude() {
let oversized = Message::new(vec![0xAB; MAX_FRAME_BYTES - FRAME_OVERHEAD_BYTES + 1]);
let wire = encode_all(AwsEventStreamCodec, [oversized]);
assert!(matches!(
collect(every(&wire[..4], 1)).await,
Err(EventStreamError::InvalidLength(length)) if length == MAX_FRAME_BYTES + 1
));
}
#[tokio::test]
async fn an_empty_body_yields_nothing() {
assert_eq!(collect(vec![]).await.unwrap(), vec![]);
}
#[tokio::test]
async fn a_complete_frame_precedes_a_truncated_following_frame() {
let wire = encode_all(AwsEventStreamCodec, [message(b"first"), message(b"second")]);
let mut messages = Box::pin(frames(
input(every(&wire[..wire.len() - 1], 3)),
AwsEventStreamCodec,
));
assert_eq!(messages.next().await.unwrap().unwrap(), message(b"first"));
assert!(matches!(
messages.next().await,
Some(Err(EventStreamError::Truncated))
));
assert!(messages.next().await.is_none());
}
#[tokio::test]
async fn a_body_error_after_a_complete_frame_preserves_its_cause() {
let first = encode_all(AwsEventStreamCodec, [message(b"first")]);
let mut messages = Box::pin(frames(
stream::iter([
Ok(cut_at(&first, [5])[0].clone()),
Ok(cut_at(&first, [5])[1].clone()),
Ok(Bytes::from_static(b"\0\0\0")),
Err(io::Error::new(io::ErrorKind::ConnectionReset, "reset")),
]),
AwsEventStreamCodec,
));
assert_eq!(messages.next().await.unwrap().unwrap(), message(b"first"));
let Some(Err(EventStreamError::Body(body))) = messages.next().await else {
panic!("the body error surfaces");
};
assert_eq!(
body_cause::<io::Error>(&body).unwrap().kind(),
io::ErrorKind::ConnectionReset
);
assert!(messages.next().await.is_none());
}

74
litellm-rust/crates/framer/tests/chaining.rs
View file

@ -2,28 +2,64 @@
mod support;
use std::io;
use bytes::Bytes;
use futures_util::{StreamExt, TryStreamExt};
use litellm_framing::{
EventStreamError, SseError,
aws_event_stream::{AwsEventStreamCodec, Message},
frames,
sse::{SseCodec, SseEvent},
};
use proptest::prelude::*;
use support::{body_cause, cut_at, encode_all, every, input, runtime};
use futures_util::TryStreamExt;
use litellm_framing::Framer;
use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
use litellm_framing::sse::SseFramer;
fn delta(data: &str) -> SseEvent {
SseEvent {
event: Some("delta".into()),
data: data.into(),
id: Some("7".into()),
retry: None,
}
}
use support::encode;
fn envelopes(payloads: Vec<Bytes>) -> Vec<u8> {
encode_all(AwsEventStreamCodec, payloads.into_iter().map(Message::new))
}
proptest! {
#[test]
fn an_sse_event_cut_anywhere_across_envelopes_is_reassembled(cut in 0_usize..64, chunk in 1_usize..8) {
let sse = encode_all(SseCodec::default(), [delta("hello")]);
let wire = envelopes(cut_at(&sse, [cut.min(sse.len())]));
let events = runtime().block_on(async {
let payloads = frames(input(every(&wire, chunk)), AwsEventStreamCodec)
.map_ok(|message| message.payload().clone());
frames(payloads, SseCodec::default()).try_collect::<Vec<_>>().await
})
.unwrap();
prop_assert_eq!(events, vec![delta("hello")]);
}
}
#[tokio::test]
async fn hosting_payloads_feed_the_same_sse_framer_across_envelope_boundaries() {
let bytes = [encode(b"event: delta\ndata: hel"), encode(b"lo\nid: 7\n\n")].concat();
let envelopes = AwsEventStreamFramer.frame(futures_util::stream::iter(
bytes.chunks(3).map(Ok::<_, io::Error>),
async fn a_truncated_envelope_after_an_sse_event_keeps_the_event_and_its_cause() {
let complete = encode_all(SseCodec::default(), [delta("complete")]);
let incomplete = encode_all(SseCodec::default(), [delta("incomplete")]);
let wire = envelopes(vec![complete.into(), incomplete.into()]);
let payloads = frames(
input(every(&wire[..wire.len() - 1], 3)),
AwsEventStreamCodec,
)
.map_ok(|message| message.payload().clone());
let mut events = Box::pin(frames(payloads, SseCodec::default()));
assert_eq!(events.next().await.unwrap().unwrap(), delta("complete"));
let Some(Err(SseError::Body(body))) = events.next().await else {
panic!("the envelope error surfaces through the SSE layer");
};
assert!(matches!(
body_cause::<EventStreamError>(&body),
Some(EventStreamError::Truncated)
));
let frames = SseFramer
.frame(envelopes.map_ok(|frame: AwsEventStreamFrame| frame.payload))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(frames.len(), 1);
assert_eq!(frames[0].event.as_deref(), Some("delta"));
assert_eq!(frames[0].data.as_deref(), Some("hello"));
assert_eq!(frames[0].id.as_deref(), Some("7"));
assert!(events.next().await.is_none());
}

188
litellm-rust/crates/framer/tests/sse.rs
View file

@ -1,67 +1,169 @@
#![cfg(feature = "sse")]
mod support;
use std::io;
use futures_util::{StreamExt, TryStreamExt};
use litellm_framing::sse::{SseFrame, SseFramer};
use litellm_framing::{Error, Framer};
use bytes::Bytes;
use futures_util::{StreamExt, TryStreamExt, stream};
use litellm_framing::{
SseError, frames,
sse::{SseCodec, SseEvent},
};
use proptest::prelude::*;
use rstest::rstest;
use support::{body_cause, cut_at, encode_all, every, input, runtime};
async fn collect_sse(chunks: &[&[u8]]) -> Result<Vec<SseFrame>, Error> {
SseFramer
.frame(futures_util::stream::iter(
chunks.iter().copied().map(Ok::<_, io::Error>),
))
async fn collect(pieces: Vec<Bytes>) -> Result<Vec<SseEvent>, SseError> {
frames(input(pieces), SseCodec::default())
.try_collect()
.await
}
fn event(name: Option<&str>, data: &str) -> SseEvent {
SseEvent {
event: name.map(str::to_owned),
data: data.to_owned(),
id: None,
retry: None,
}
}
fn sse_event() -> impl Strategy<Value = SseEvent> {
(
proptest::option::of("[^\r\n\0]{0,8}"),
"[^\r\0]{0,16}",
proptest::option::of("[^\r\n\0]{0,8}"),
proptest::option::of(any::<u64>()),
)
.prop_map(|(event, data, id, retry)| SseEvent {
event,
data,
id,
retry,
})
}
fn terminators() -> impl Strategy<Value = &'static [u8]> {
prop_oneof![Just(&b"\n"[..]), Just(&b"\r\n"[..]), Just(&b"\r"[..])]
}
proptest! {
#[test]
fn any_events_survive_a_round_trip_through_any_terminator_and_any_cuts(
events in proptest::collection::vec(sse_event(), 1..4),
terminator in terminators(),
cuts in proptest::collection::vec(0_usize..256, 0..4),
bom in any::<bool>(),
) {
let lf_wire = encode_all(SseCodec::default(), events.clone());
let body: Vec<u8> = lf_wire
.iter()
.flat_map(|byte| if *byte == b'\n' { terminator.to_vec() } else { vec![*byte] })
.collect();
let wire = if bom { [&b"\xEF\xBB\xBF"[..], &body].concat() } else { body };
let decoded = runtime().block_on(collect(cut_at(&wire, cuts))).unwrap();
prop_assert_eq!(decoded, events);
}
}
#[rstest]
#[case(
&[&b":ping\r\nevent: delta\r\nid: 7\r\nretry: 10\r\ndata: \xe2"[..], &b"\x82"[..], &b"\xac\r"[..], &b"\ndata: next\r\n\r"[..], &b"\ndata: [DONE]\n\n"[..]],
vec![
SseFrame {
event: Some("delta".into()),
data: Some("€\nnext".into()),
id: Some("7".into()),
retry: Some(10),
},
SseFrame {
event: None,
data: Some("[DONE]".into()),
id: None,
retry: None,
},
]
)]
#[case::comment(b":ping\ndata: x\n\n")]
#[case::unknown_field(b"vendor: 1\ndata: x\n\n")]
#[case::field_without_colon(b"garbage\ndata: x\n\n")]
#[case::retry_with_non_digits(b"retry: soon\ndata: x\n\n")]
#[case::retry_with_a_sign(b"retry: +5\ndata: x\n\n")]
#[case::retry_without_a_value(b"retry:\ndata: x\n\n")]
#[case::id_with_nul(b"id: a\0b\ndata: x\n\n")]
#[tokio::test]
async fn fragmented_utf8_crlf_and_multiline_data_retain_metadata_and_sentinel(
#[case] chunks: &[&[u8]],
#[case] expected: Vec<SseFrame>,
async fn lines_the_spec_ignores_do_not_change_the_event(#[case] wire: &[u8]) {
assert_eq!(
collect(every(wire, 1)).await.unwrap(),
vec![event(None, "x")]
);
}
#[rstest]
#[case::no_data_at_all(b"event: ping\nid: 1\n\ndata: x\n\n", vec![event(None, "x")])]
#[case::empty_data_field(b"data:\n\n", vec![event(None, "")])]
#[case::one_leading_space_stripped(b"data: x\n\n", vec![event(None, " x")])]
#[case::multiline_data(b"data: a\ndata: b\ndata:\n\n", vec![event(None, "a\nb\n")])]
#[case::last_event_name_wins(b"event: a\nevent: b\ndata: x\n\n", vec![event(Some("b"), "x")])]
#[case::last_retry_wins(b"retry: 1\nretry: 2\ndata: x\n\n", vec![SseEvent { retry: Some(2), ..event(None, "x") }])]
#[case::split_utf8_across_lines_is_not_joined(b"data: \xe2\x82\xac\ndata: \xe2\x82\xac\n\n", vec![event(None, "€\n€")])]
#[tokio::test]
async fn dispatch_follows_the_data_buffer(#[case] wire: &[u8], #[case] expected: Vec<SseEvent>) {
assert_eq!(collect(every(wire, 1)).await.unwrap(), expected);
}
#[rstest]
#[case::unterminated_single(b"data: partial\n", vec![])]
#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
#[case::lone_cr_terminates_at_eof(b"data: x\r\r", vec![event(None, "x")])]
#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
#[tokio::test]
async fn eof_dispatches_only_terminated_events(
#[case] wire: &[u8],
#[case] expected: Vec<SseEvent>,
) {
assert_eq!(collect_sse(chunks).await.unwrap(), expected);
assert_eq!(
collect(vec![Bytes::copy_from_slice(wire)]).await.unwrap(),
expected
);
}
#[rstest]
#[case::inside_the_first_line(vec![&b"data: a\r"[..], &b"\ndata: b\r\n\r\n"[..]])]
#[case::inside_the_blank_line(vec![&b"data: a\r\ndata: b\r\n\r"[..], &b"\n"[..]])]
#[tokio::test]
async fn a_crlf_split_across_chunks_is_one_terminator(#[case] pieces: Vec<&[u8]>) {
let pieces = pieces.into_iter().map(Bytes::copy_from_slice).collect();
assert_eq!(collect(pieces).await.unwrap(), vec![event(None, "a\nb")]);
}
#[tokio::test]
async fn eof_does_not_dispatch_an_unterminated_frame() {
assert!(collect_sse(&[b"data: partial\n"]).await.unwrap().is_empty());
async fn a_bom_is_stripped_only_at_the_start_of_the_stream() {
let wire = b"\xEF\xBB\xBFdata: a\n\n\xEF\xBB\xBFdata: b\ndata: c\n\n";
let decoded = collect(every(wire, 2)).await.unwrap();
assert_eq!(decoded, vec![event(None, "a"), event(None, "c")]);
}
#[tokio::test]
async fn invalid_utf8_in_a_field_fails_after_earlier_events_and_terminates() {
let mut events = Box::pin(frames(
input(every(b"data: ok\n\ndata: \xff\n\n", 3)),
SseCodec::default(),
));
assert_eq!(events.next().await.unwrap().unwrap(), event(None, "ok"));
assert!(matches!(
events.next().await,
Some(Err(SseError::InvalidUtf8(_)))
));
assert!(events.next().await.is_none());
}
#[rstest]
#[case(io::ErrorKind::ConnectionReset)]
#[case(io::ErrorKind::UnexpectedEof)]
#[tokio::test]
async fn framing_errors_terminate_and_preserve_input_error_causes(#[case] kind: io::ErrorKind) {
let mut frames = Box::pin(SseFramer.frame(futures_util::stream::iter([
Err(io::Error::new(kind, "reset")),
Ok(&b"data: later\n\n"[..]),
])));
let error = frames.next().await.unwrap().unwrap_err();
assert!(matches!(
error,
Error::Sse(sse_stream::Error::Body(ref cause))
if cause.downcast_ref::<io::Error>().unwrap().kind() == kind
async fn a_body_error_keeps_earlier_events_and_its_cause_then_terminates(
#[case] kind: io::ErrorKind,
) {
let mut events = Box::pin(frames(
stream::iter([
Ok(&b"data: first\n\ndata: partial"[..]),
Err(io::Error::new(kind, "reset")),
Ok(&b"\n\n"[..]),
]),
SseCodec::default(),
));
assert!(frames.next().await.is_none());
assert!(frames.next().await.is_none());
assert_eq!(events.next().await.unwrap().unwrap(), event(None, "first"));
let Some(Err(SseError::Body(body))) = events.next().await else {
panic!("the body error surfaces");
};
assert_eq!(body_cause::<io::Error>(&body).unwrap().kind(), kind);
assert!(events.next().await.is_none());
assert!(events.next().await.is_none());
}

68
litellm-rust/crates/framer/tests/support/mod.rs
View file

@ -1,15 +1,57 @@
use aws_smithy_eventstream::frame::write_message_to;
use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
use bytes::Bytes;
#![allow(dead_code)]
pub fn encode(payload: &'static [u8]) -> Vec<u8> {
let message = Message::new(Bytes::from_static(payload))
.add_header(Header::new(
":event-type",
HeaderValue::String("payload".into()),
))
.add_header(Header::new("sequence", HeaderValue::Int32(7)));
let mut bytes = Vec::new();
write_message_to(&message, &mut bytes).unwrap();
bytes
use std::{error::Error, io};
use bytes::{Bytes, BytesMut};
use futures_util::{Stream, stream};
use tokio_util::codec::Encoder;
pub fn encode_all<C, I>(mut codec: C, items: impl IntoIterator<Item = I>) -> Vec<u8>
where
C: Encoder<I>,
C::Error: std::fmt::Debug,
{
let mut wire = BytesMut::new();
for item in items {
codec.encode(item, &mut wire).unwrap();
}
wire.to_vec()
}
pub fn cut_at(bytes: &[u8], offsets: impl IntoIterator<Item = usize>) -> Vec<Bytes> {
let mut sorted: Vec<usize> = offsets
.into_iter()
.filter(|offset| *offset <= bytes.len())
.collect();
sorted.sort_unstable();
sorted.dedup();
let bounds = std::iter::once(0)
.chain(sorted)
.chain(std::iter::once(bytes.len()))
.collect::<Vec<_>>();
bounds
.windows(2)
.map(|pair| Bytes::copy_from_slice(&bytes[pair[0]..pair[1]]))
.collect()
}
pub fn every(bytes: &[u8], size: usize) -> Vec<Bytes> {
bytes
.chunks(size.max(1))
.map(Bytes::copy_from_slice)
.collect()
}
pub fn input(pieces: Vec<Bytes>) -> impl Stream<Item = Result<Bytes, io::Error>> + Send {
stream::iter(pieces.into_iter().map(Ok))
}
pub fn body_cause<T: Error + 'static>(body: &io::Error) -> Option<&T> {
body.get_ref()?.downcast_ref::<T>()
}
pub fn runtime() -> tokio::runtime::Runtime {
tokio::runtime::Builder::new_current_thread()
.build()
.unwrap()
}

7
litellm-rust/crates/host-python/src/adapter.rs
View file

@ -134,6 +134,13 @@ pub trait ProtocolHost: Send + Sync {
response: <Self::Protocol as Protocol>::Response,
) -> PyResult<Py<PyAny>>;
/// What the stream carries at hand-off, as the caller's stream receives it.
fn head(
&mut self,
py: Python<'_>,
head: <Self::Protocol as Protocol>::StreamHead,
) -> PyResult<Py<PyAny>>;
/// One streamed chunk as the caller receives it.
fn chunk(
&mut self,

187
litellm-rust/crates/host-python/src/driver.rs
View file

@ -134,10 +134,10 @@ where
}
match driver.resume(None)? {
ExecutionStep::Return(value) => Ok(value),
ExecutionStep::Open => py
ExecutionStep::Open(head) => py
.import("litellm.rust_bridge.lifecycle")?
.getattr("SyncStream")?
.call1((Py::new(py, Execution::suspended(driver))?,))
.call1((Py::new(py, Execution::suspended(driver))?, head))
.map(Bound::unbind),
ExecutionStep::Await(_) | ExecutionStep::Yield(_) => {
Err(PyRuntimeError::new_err("sync call suspended"))
@ -312,7 +312,7 @@ where
Ok(_) => return Err(missing_state()),
Err(error) => Err(error),
},
HostOp::Open(_, reply) => return self.opened(py, reply).map(Next::Return),
HostOp::Open(head, reply) => return self.opened(py, head, reply).map(Next::Return),
HostOp::Deliver(chunk, reply) => {
return self.delivered(py, chunk, reply).map(Next::Return);
}
@ -340,12 +340,21 @@ where
}
}
fn opened(&mut self, py: Python<'_>, reply: Reply<Demand>) -> PyResult<ExecutionStep> {
fn opened(
&mut self,
py: Python<'_>,
head: <ProtocolOf<H> as Protocol>::StreamHead,
reply: Reply<Demand>,
) -> PyResult<ExecutionStep> {
self.stage = Stage::Streaming;
let head = match self.host.head(py, head) {
Ok(head) => head,
Err(error) => return self.interrupt(py, error),
};
match self.adapter.opened(py) {
Ok(()) => {
self.pending = Some(Pending::Consumer(reply));
Ok(ExecutionStep::Open)
Ok(ExecutionStep::Open(head))
}
Err(error) => self.interrupt(py, error),
}
@ -699,6 +708,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
.map(|answer| reply.send(answer))
}
fn head(&mut self, _: Python<'_>, head: std::convert::Infallible) -> PyResult<Py<PyAny>> {
match head {}
}
fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult<Py<PyAny>> {
match chunk {}
}
@ -945,6 +958,163 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
});
}
struct Streaming;
impl Protocol for Streaming {
type Response = ();
type Error = Error;
type Projection = ();
type Op = std::convert::Infallible;
type Chunk = &'static str;
type StreamHead = Vec<(&'static str, &'static str)>;
}
struct StreamingHost;
impl ProtocolHost for StreamingHost {
type Protocol = Streaming;
type Failure = Classified;
fn project(
&mut self,
_: Python<'_>,
_: &Bound<'_, PyDict>,
) -> Result<(), InvokeError<Error>> {
Ok(())
}
fn invoke(
&mut self,
_: Python<'_>,
op: std::convert::Infallible,
) -> Result<(), InvokeError<Error>> {
match op {}
}
fn head(
&mut self,
py: Python<'_>,
head: Vec<(&'static str, &'static str)>,
) -> PyResult<Py<PyAny>> {
let headers = PyDict::new(py);
for (name, value) in head {
headers.set_item(name, value)?;
}
let hidden = PyDict::new(py);
hidden.set_item("additional_headers", headers)?;
Ok(hidden.into_any().unbind())
}
fn chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult<Py<PyAny>> {
Ok(pyo3::types::PyString::new(py, chunk).into_any().unbind())
}
fn complete(&mut self, py: Python<'_>, (): ()) -> PyResult<Py<PyAny>> {
Ok(py.None())
}
fn classify(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
Ok(Classified(error.0))
}
fn host_error(error: &PyErr) -> Error {
Error(error.to_string())
}
fn close(&mut self, _: Python<'_>) {}
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
Ok(())
}
}
fn streaming_machine() -> CallMachine<Streaming> {
CallMachine::new(|host| {
Box::pin(async move {
host.project().await?;
if host.open(vec![("request-id", "req_1")]).await? == Demand::Detached {
return Ok(());
}
for chunk in ["first", "second"] {
if host.deliver(chunk).await? == Demand::Detached {
break;
}
}
Ok(())
})
})
}
/// Drives a `Stream` (async) or `SyncStream` to completion from a sync test.
fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec<String> {
if !asynchronous {
return stream
.try_iter()
.unwrap()
.map(|chunk| chunk.unwrap().extract().unwrap())
.collect();
}
std::iter::from_fn(|| {
let stop = stream
.call_method0("__anext__")
.unwrap()
.call_method1("send", (py.None(),))
.unwrap_err();
if stop.is_instance_of::<pyo3::exceptions::PyStopAsyncIteration>(py) {
return None;
}
assert!(stop.is_instance_of::<pyo3::exceptions::PyStopIteration>(py));
Some(stop.value(py).getattr("value").unwrap().extract().unwrap())
})
.collect()
}
#[test]
fn a_stream_carries_its_head_as_hidden_params_before_the_first_chunk() {
let _guard = PYTHON_GLOBALS
.lock()
.unwrap_or_else(|error| error.into_inner());
crate::initialize_python();
Python::attach(|py| {
install_lifecycle_module(py);
for asynchronous in [false, true] {
let log = Log::default();
let adapter = SyntheticAdapter {
log: Log(log.0.clone()),
script: AdapterScript::Plain,
};
let handed = run_call(
py,
streaming_machine(),
StreamingHost,
Box::new(adapter),
PyDict::new(py).unbind(),
asynchronous,
)
.unwrap();
let stream = if asynchronous {
let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err();
stop.value(py).getattr("value").unwrap()
} else {
handed.into_bound(py)
};
let hidden: std::collections::HashMap<
String,
std::collections::HashMap<String, String>,
> = stream.getattr("_hidden_params").unwrap().extract().unwrap();
assert_eq!(
hidden["additional_headers"],
std::collections::HashMap::from([(
"request-id".to_string(),
"req_1".to_string()
)])
);
assert_eq!(log.entries(), ["started", "begin", "opened"]);
assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]);
}
});
}
fn failing_machine() -> CallMachine<Synthetic> {
CallMachine::new(|host| {
Box::pin(async move {
@ -1202,6 +1372,13 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
) -> Result<(), InvokeError<Error>> {
Err(missing_state().into())
}
fn head(
&mut self,
_: Python<'_>,
head: std::convert::Infallible,
) -> PyResult<Py<PyAny>> {
match head {}
}
fn chunk(
&mut self,
_: Python<'_>,

8
litellm-rust/crates/host-python/src/handle.rs
View file

@ -8,9 +8,9 @@ use pyo3::prelude::*;
pub enum ExecutionStep {
Return(Py<PyAny>),
Await(Py<PyAny>),
/// The call streams: the caller gets a stream over this execution, which stays
/// suspended until the stream asks for a chunk.
Open,
/// The call streams: the caller gets a stream over this execution carrying this head,
/// and the execution stays suspended until the stream asks for a chunk.
Open(Py<PyAny>),
Yield(Py<PyAny>),
}
@ -75,7 +75,7 @@ impl Execution {
let step = body.resume(result)?;
let (tag, value, suspended) = match step {
ExecutionStep::Await(value) => ("Await", value, true),
ExecutionStep::Open => ("Open", py.None(), true),
ExecutionStep::Open(head) => ("Open", head, true),
ExecutionStep::Yield(value) => ("Yield", value, true),
ExecutionStep::Return(value) => ("Complete", value, false),
};

37
litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/streaming_iterator.rs
View file

@ -2,9 +2,9 @@ use base64::Engine;
use bytes::Buf;
use futures_util::{Stream, StreamExt};
use litellm_framing::{
Framer,
aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer},
sse::{SseFrame, SseFramer},
aws_event_stream::{AwsEventStreamCodec, Message},
frames,
sse::{SseCodec, SseEvent},
};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
@ -13,8 +13,6 @@ use serde_json::{Map, Value};
pub enum Error {
#[error("stream framing failed: {0}")]
StreamFraming(String),
#[error("Anthropic SSE frame has no data")]
MissingStreamData,
#[error("Anthropic stream event is invalid: {0}")]
InvalidStreamEvent(String),
#[error("Bedrock event payload is invalid: {0}")]
@ -165,15 +163,14 @@ struct BedrockChunkPayload {
bytes: String,
}
pub fn decode_anthropic_sse_frame(frame: SseFrame) -> Result<AnthropicMessagesStreamEvent, Error> {
let data = frame.data.ok_or(Error::MissingStreamData)?;
serde_json::from_str(&data).map_err(|error| Error::InvalidStreamEvent(error.to_string()))
pub fn decode_anthropic_sse_frame(event: SseEvent) -> Result<AnthropicMessagesStreamEvent, Error> {
serde_json::from_str(&event.data).map_err(|error| Error::InvalidStreamEvent(error.to_string()))
}
pub fn decode_bedrock_anthropic_frame(
frame: AwsEventStreamFrame,
message: Message,
) -> Result<AnthropicMessagesStreamEvent, Error> {
let payload: BedrockChunkPayload = serde_json::from_slice(&frame.payload)
let payload: BedrockChunkPayload = serde_json::from_slice(message.payload())
.map_err(|error| Error::InvalidBedrockPayload(error.to_string()))?;
let event = base64::engine::general_purpose::STANDARD
.decode(payload.bytes)
@ -189,9 +186,8 @@ where
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
SseFramer.frame(input).map(|frame| {
let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?;
decode_anthropic_sse_frame(frame)
frames(input, SseCodec::default()).map(|event| {
decode_anthropic_sse_frame(event.map_err(|error| Error::StreamFraming(error.to_string()))?)
})
}
@ -203,9 +199,10 @@ where
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
AwsEventStreamFramer.frame(input).map(|frame| {
let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?;
decode_bedrock_anthropic_frame(frame)
frames(input, AwsEventStreamCodec).map(|message| {
decode_bedrock_anthropic_frame(
message.map_err(|error| Error::StreamFraming(error.to_string()))?,
)
})
}
@ -247,12 +244,10 @@ mod tests {
#[test]
fn decodes_citations_delta_events() {
let event = decode_anthropic_sse_frame(SseFrame {
let event = decode_anthropic_sse_frame(SseEvent {
event: Some("content_block_delta".into()),
data: Some(
r#"{"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#
.into(),
),
data: r#"{"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#
.into(),
id: None,
retry: None,
})

44
litellm-rust/crates/llms/src/reducto/ocr/transformation.rs
View file

@ -554,6 +554,50 @@ async fn upload_bytes_async(
mod tests {
use super::*;
#[tokio::test]
async fn v3_body_keeps_explicit_null_options_and_drops_unknown_ones() {
use crate::base_llm::ocr::{handler::OcrClient, transformation::OcrRequestContext};
let overrides =
serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true}))
.unwrap();
let params = ReductoParseV3Config
.map_ocr_params(&overrides, "parse-v3")
.unwrap();
let client = OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new());
let connection = OcrConnection::default();
let document = serde_json::from_value(
json!({"type":"document_url","document_url":"reducto://ready.pdf"}),
)
.unwrap();
let body = ReductoParseV3Config
.async_transform_ocr_request(
"parse-v3",
document,
&params,
&[],
OcrRequestContext {
client: &client,
connection: &connection,
},
)
.await
.unwrap();
assert_eq!(
serde_json::to_value(body).unwrap(),
json!({"input":"reducto://ready.pdf", "formatting":null, "settings":{}})
);
let absent = ReductoParseV3Config
.map_ocr_params(
&litellm_core_utils::call_arguments::CallArguments::default(),
"parse-v3",
)
.unwrap();
assert_eq!(serde_json::to_value(absent).unwrap(), json!({}));
}
#[test]
fn options_preserve_null_and_select_the_provider_fields() {
let overrides = serde_json::from_value(json!({

79
litellm-rust/crates/llms/tests/ocr_handler.rs Normal file
View file

@ -0,0 +1,79 @@
use std::time::Duration;
use litellm_llms::base_llm::ocr::{error::Error, handler::read_response_bytes};
use rstest::rstest;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
};
/// Answers one request with raw `response` bytes and then holds the connection open, so a
/// read that waits for the rest of an oversized body hangs instead of passing.
async fn read_bounded(response: String, limit: usize) -> Result<bytes::Bytes, Error> {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut request = [0; 4096];
assert!(socket.read(&mut request).await.unwrap() > 0);
socket.write_all(response.as_bytes()).await.unwrap();
std::future::pending::<()>().await;
});
let response = reqwest::Client::new()
.get(format!("http://{address}"))
.send()
.await
.unwrap();
let result =
tokio::time::timeout(Duration::from_secs(2), read_response_bytes(response, limit)).await;
server.abort();
result.expect("bounded reads must finish without waiting for the rest of an oversized body")
}
#[rstest]
#[case::declared("HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh")]
#[case::chunked(
"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n"
)]
#[tokio::test]
async fn a_body_of_exactly_the_limit_is_read(#[case] response: &str) {
assert_eq!(read_bounded(response.into(), 8).await.unwrap(), "abcdefgh");
}
#[rstest]
#[case::declared("HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n")]
#[case::chunked("HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n")]
#[tokio::test]
async fn a_body_over_the_limit_is_rejected(#[case] response: &str) {
assert!(matches!(
read_bounded(response.into(), 8).await,
Err(Error::TooLarge { limit: 8 })
));
}
#[rstest]
#[case::declared("Content-Length: 1000000")]
#[case::chunked("Transfer-Encoding: chunked")]
#[tokio::test]
async fn an_oversized_error_keeps_its_status_and_a_bounded_body_without_draining(
#[case] headers: &str,
) {
let prefix = "x".repeat(4096);
let body = match headers.starts_with("Transfer") {
true => format!("{:x}\r\n{prefix}\r\n", prefix.len()),
false => prefix.clone(),
};
let error = read_bounded(
format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"),
prefix.len(),
)
.await
.unwrap_err();
let Error::Transport(litellm_http::transport::Error::Http { status, body }) = error else {
panic!("unexpected error: {error}");
};
assert_eq!(status, 429);
assert_eq!(body, prefix);
}

6
litellm-rust/crates/model-catalog/AGENTS.md Normal file
View file

@ -0,0 +1,6 @@
## Validation
For `model_prices_and_context_window.json` validation, we should eventually:
- Remove any schema file like `model_prices_and_context_window.schema.json`
- Stop skipping this crate's tests

8
litellm-rust/crates/model-catalog/Cargo.toml
View file

@ -14,12 +14,8 @@ schemars = { version = "1.0", optional = true }
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
time.workspace = true
[dev-dependencies]
criterion.workspace = true
jsonschema = { version = "0.55.1", default-features = false }
rstest.workspace = true
litellm-model-catalog = { path = ".", features = ["schema"] }
[[bench]]
name = "catalog"
harness = false

25
litellm-rust/crates/model-catalog/README.md
View file

@ -1,25 +0,0 @@
# Model catalog
`litellm-model-catalog` builds an immutable snapshot from caller supplied JSON bytes. It has no network, Python, registration, or refresh behavior. The caller supplies optional source, revision, and ETag provenance. Parse and validation are separate so small synthetic catalogs can use explicit integrity limits
The parser treats `sample_spec` and `fallback_generalizations` as reserved top level metadata. `fallback_rules()` exposes the typed rule array when present; this crate does not execute regex generalizations. Model entries retain all JSON fields except `aliases`, including unknown fields. `field()` returns `None` for an absent key and a JSON null, false, or zero value for a present key. The returned values are borrowed, so callers cannot mutate the snapshot
Each entry also deserializes into `ModelInfo`, a typed mirror of `model_prices_and_context_window.schema.json`'s `modelEntry` definition, reachable via `ModelEntry::info()`. All schema fields are optional on `ModelInfo`, including `litellm_provider` which the schema marks required, so small synthetic catalogs still parse. Unknown fields are not part of `ModelInfo`; they remain on `fields()`. Building with the `schema` feature adds `schemars` derives and exposes `model_entry_json_schema()` for emitting the entry's JSON Schema. Parse and validation failures are reported by the `Error` enum in `error.rs`, while catalog logic lives in `catalog.rs`
The integration tests read the repository's catalog and schema files at test time, assert every entry round-trips through `ModelInfo`, and verify that the generated schema's properties match the repository schema
Aliases point to their canonical entries. An alias that exactly matches any canonical key is skipped; the first canonical entry claiming an alias wins. Invalid alias lists and nonstring names are skipped and reported by `alias_issues()`. Exact lookup wins. For a case insensitive miss, the last key with the same lowercase spelling wins, following Python's lowercase map built after aliases are appended. This uses Rust Unicode lowercasing, which can differ from Python for unusual Unicode model IDs
`validate()` counts canonical entries before alias expansion and excludes both reserved keys. It enforces an explicit minimum and backup shrink ratio, with Python defaults of 50 models and 0.5. Parsing rejects nonobject model entries and known fields with the wrong JSON type, but ignores unknown fields. It does not enforce every constraint in the JSON schema, calculate prices, resolve providers, or check provenance authenticity. The caller decides how to handle validation failures
This snapshot does not represent Python's live mutable `litellm.model_cost`, nested dict and list mutation, or mutation of dicts previously returned by Python APIs. It has no bridge or runtime integration
## Benchmarks
`cargo bench -p litellm-model-catalog --bench catalog` measures parsing plus alias indexing and exact lookup. For a local Python baseline on the same fixture, use:
```sh
python3 -m timeit -s 'import json, pathlib; body = pathlib.Path("../model_prices_and_context_window.json").read_bytes()' 'json.loads(body)'
```
Run these commands from `litellm-rust`. Python's command measures JSON loading only, without alias expansion or snapshot construction. The Rust benchmark does not include future Python object materialization, so these numbers are not an end to end runtime comparison

21
litellm-rust/crates/model-catalog/benches/catalog.rs
View file

@ -1,21 +0,0 @@
use criterion::{Criterion, criterion_group, criterion_main};
use litellm_model_catalog::{Catalog, Provenance};
use std::hint::black_box;
fn benchmarks(c: &mut Criterion) {
let body = include_bytes!("../../../../model_prices_and_context_window.json");
c.bench_function("parse_current_catalog", |b| {
b.iter(|| Catalog::parse(black_box(body), Provenance::default()).unwrap())
});
let catalog = Catalog::parse(body, Provenance::default()).unwrap();
let key = catalog
.model_names()
.next()
.expect("catalog must have a benchmark key");
c.bench_function("lookup_catalog_key", |b| {
b.iter(|| black_box(&catalog).lookup(black_box(key)))
});
}
criterion_group!(benches, benchmarks);
criterion_main!(benches);

80
litellm-rust/crates/model-catalog/src/capabilities.rs Normal file
View file

@ -0,0 +1,80 @@
use serde::{Deserialize, Serialize};
/// Primary API surface / task type of the model.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum Mode {
AudioSpeech,
AudioTranscription,
Chat,
Completion,
Embedding,
Evaluation,
Guardrail,
ImageEdit,
ImageGeneration,
Moderation,
Ocr,
Realtime,
Rerank,
Responses,
Search,
VectorStore,
VideoGeneration,
}
/// Reasoning effort level accepted or applied by the model.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum ReasoningEffort {
None,
Minimal,
Low,
Medium,
High,
Xhigh,
Max,
}
/// Gemini audio generation API the model is served through.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum VertexAiAudioApi {
LyriaPredict,
LyriaInteractions,
}
/// Audio container format the model can return.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum AudioFormat {
Mp3,
Wav,
}
/// Input modality the model accepts.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum InputModality {
Text,
Image,
Audio,
Video,
}
/// Output modality the model can produce.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum OutputModality {
Text,
Image,
Audio,
Video,
Code,
}

40
litellm-rust/crates/model-catalog/src/catalog.rs
View file

@ -1,5 +1,6 @@
use crate::error::Error;
use crate::model_info::{FallbackGeneralizations, FallbackRule, ModelInfo};
use crate::fallback::{FallbackGeneralizations, FallbackRule};
use crate::model_info::ModelInfo;
use indexmap::IndexMap;
use serde::Deserialize;
use serde_json::{Map, Value};
@ -14,19 +15,9 @@ pub struct Provenance {
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct IntegrityLimits {
pub backup_model_count: usize,
pub reference_model_count: usize,
pub min_model_count: usize,
pub min_backup_ratio: f64,
}
impl IntegrityLimits {
pub fn python_defaults(backup_model_count: usize) -> Self {
Self {
backup_model_count,
min_model_count: 50,
min_backup_ratio: 0.5,
}
}
pub min_reference_ratio: f64,
}
#[derive(Clone, Debug, PartialEq, Eq)]
@ -99,16 +90,11 @@ impl Catalog {
}
_ => {}
}
let Value::Object(ref object) = value else {
let Value::Object(mut fields) = value else {
return Err(Error::EntryNotObject { model: name });
};
let info = ModelInfo::deserialize(object)?;
let Value::Object(mut fields) = value else {
unreachable!("value checked is_object above")
};
if let Some(aliases) = fields.remove("aliases")
&& !aliases.is_null()
{
let info = ModelInfo::deserialize(&fields)?;
if let Some(aliases) = fields.remove("aliases") {
match aliases {
Value::Array(names) => alias_lists.push((name.clone(), names)),
_ => alias_issues.push(AliasIssue::InvalidList {
@ -161,7 +147,9 @@ impl Catalog {
}
pub fn validate(&self, limits: IntegrityLimits) -> Result<(), Error> {
if !limits.min_backup_ratio.is_finite() || !(0.0..=1.0).contains(&limits.min_backup_ratio) {
if !limits.min_reference_ratio.is_finite()
|| !(0.0..=1.0).contains(&limits.min_reference_ratio)
{
return Err(Error::InvalidRatio);
}
let actual = self.entries.len();
@ -171,13 +159,13 @@ impl Catalog {
minimum: limits.min_model_count,
});
}
if limits.backup_model_count > 0
&& (actual as f64) < (limits.backup_model_count as f64) * limits.min_backup_ratio
if limits.reference_model_count > 0
&& (actual as f64) < (limits.reference_model_count as f64) * limits.min_reference_ratio
{
return Err(Error::Shrunk {
actual,
backup: limits.backup_model_count,
ratio: limits.min_backup_ratio,
reference: limits.reference_model_count,
ratio: limits.min_reference_ratio,
});
}
Ok(())

11
litellm-rust/crates/model-catalog/src/error.rs
View file

@ -1,6 +1,5 @@
use thiserror::Error;
/// Failures from parsing or validating a catalog snapshot.
#[derive(Debug, Error)]
pub enum Error {
/// The body is not valid JSON, or a model entry fails typed deserialization.
@ -15,14 +14,14 @@ pub enum Error {
/// Canonical entry count is under the configured minimum.
#[error("catalog has {actual} models, below minimum {minimum}")]
BelowMinimum { actual: usize, minimum: usize },
/// Canonical entry count is under the configured backup shrink ratio.
#[error("catalog has {actual} models, below {ratio} of backup count {backup}")]
/// Canonical entry count is under the configured reference ratio.
#[error("catalog has {actual} models, below {ratio} of reference count {reference}")]
Shrunk {
actual: usize,
backup: usize,
reference: usize,
ratio: f64,
},
/// The configured minimum backup ratio is not finite or outside `[0, 1]`.
#[error("minimum backup ratio must be finite and between zero and one")]
/// The configured minimum reference ratio is not finite or outside `[0, 1]`.
#[error("minimum reference ratio must be finite and between zero and one")]
InvalidRatio,
}

23
litellm-rust/crates/model-catalog/src/fallback.rs Normal file
View file

@ -0,0 +1,23 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::BTreeMap;
/// One regex rule generalizing unknown model ids to known families.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct FallbackRule {
pub name: String,
pub pattern: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(flatten)]
pub extra: BTreeMap<String, Value>,
}
/// Regex rules that generalize unknown model ids to known families; not a model entry.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(deny_unknown_fields)]
pub struct FallbackGeneralizations {
pub rules: Vec<FallbackRule>,
}

24
litellm-rust/crates/model-catalog/src/lib.rs
View file

@ -1,16 +1,20 @@
mod capabilities;
mod catalog;
mod error;
mod fallback;
mod model_info;
mod pricing;
mod validation;
pub use capabilities::*;
pub use catalog::*;
pub use error::*;
pub use fallback::*;
pub use model_info::*;
pub use pricing::*;
pub use validation::*;
#[cfg(feature = "schema")]
mod schema;
pub use catalog::{AliasIssue, Catalog, IntegrityLimits, ModelEntry, ModelMatch, Provenance};
pub use error::Error;
pub use model_info::{
AudioFormat, FallbackGeneralizations, FallbackRule, InputModality, Mode, ModelInfo,
OffPeakPricing, OffPeakWindow, OutputModality, ReasoningEffort, SearchContextCostPerQuery,
TieredRate, UtcHours, VertexAiAudioApi, WebSearchBillingUnit, Weekday,
};
#[cfg(feature = "schema")]
pub use schema::model_entry_json_schema;
pub use schema::*;

587
litellm-rust/crates/model-catalog/src/model_info.rs
View file

@ -1,673 +1,482 @@
use crate::capabilities::{
AudioFormat, InputModality, Mode, OutputModality, ReasoningEffort, VertexAiAudioApi,
};
use crate::pricing::{OffPeakPricing, SearchContextCostPerQuery, TieredRate, WebSearchBillingUnit};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::BTreeMap;
/// Primary API surface / task type of the model.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum Mode {
AudioSpeech,
AudioTranscription,
Chat,
Completion,
Embedding,
Evaluation,
Guardrail,
ImageEdit,
ImageGeneration,
Moderation,
Ocr,
Realtime,
Rerank,
Responses,
Search,
VectorStore,
VideoGeneration,
}
/// Reasoning effort level accepted or applied by the model.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum ReasoningEffort {
None,
Minimal,
Low,
Medium,
High,
Xhigh,
Max,
}
/// Gemini audio generation API the model is served through.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum VertexAiAudioApi {
LyriaPredict,
LyriaInteractions,
}
/// Whether web search is billed per query or per prompt.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum WebSearchBillingUnit {
PerQuery,
PerPrompt,
}
/// Audio container format the model can return.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum AudioFormat {
Mp3,
Wav,
}
/// Input modality the model accepts.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum InputModality {
Text,
Image,
Audio,
Video,
}
/// Output modality the model can produce.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum OutputModality {
Text,
Image,
Audio,
Video,
Code,
}
/// UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(untagged)]
pub enum UtcHours {
Single(String),
Multiple(Vec<String>),
}
/// ISO-8601 weekday number (1 = Monday .. 7 = Sunday) or English day name.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(untagged)]
pub enum Weekday {
Number(u8),
Name(String),
}
/// One off-peak window entry inside `windows`.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(deny_unknown_fields)]
pub struct OffPeakWindow {
pub hours_utc: UtcHours,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub weekdays: Option<Vec<Weekday>>,
}
/// Rates that replace the same-named base fields inside the stated UTC windows.
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(deny_unknown_fields)]
pub struct OffPeakPricing {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub hours_utc: Option<UtcHours>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub windows: Option<Vec<OffPeakWindow>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub weekday_timezone: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input_cost_per_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output_cost_per_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output_cost_per_reasoning_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost: Option<f64>,
}
/// USD cost per web search query, keyed by search context size.
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(deny_unknown_fields)]
pub struct SearchContextCostPerQuery {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub search_context_size_low: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub search_context_size_medium: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub search_context_size_high: Option<f64>,
}
/// One tier of a context-length or result-count tiered rate.
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(deny_unknown_fields)]
pub struct TieredRate {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub range: Option<[f64; 2]>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_results_range: Option<[f64; 2]>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input_cost_per_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output_cost_per_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output_cost_per_reasoning_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input_cost_per_query: Option<f64>,
}
/// One regex rule generalizing unknown model ids to known families.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct FallbackRule {
pub name: String,
pub pattern: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(flatten)]
pub extra: BTreeMap<String, Value>,
}
/// Regex rules that generalize unknown model ids to known families; not a model entry.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(deny_unknown_fields)]
pub struct FallbackGeneralizations {
pub rules: Vec<FallbackRule>,
}
/// Typed mirror of one catalog model entry.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct ModelInfo {
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub annotation_cost_per_page: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub annotation_cost_per_page_batches: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub audio_transcription_config: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub bedrock_converse_supports_strict_tools: Option<bool>,
/// Highest reasoning effort the Bedrock output_config accepts for this model.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub bedrock_output_config_effort_ceiling: Option<ReasoningEffort>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_audio_token_cost: Option<f64>,
/// USD per token written to the provider's prompt cache.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_1hr: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_1hr_above_200k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_200k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_256k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_272k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_272k_tokens_batches: Option<f64>,
/// Flex service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_272k_tokens_flex: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_272k_tokens_priority: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_32k_tokens: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_batches: Option<f64>,
/// Flex service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_flex: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_priority: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_audio_token_cost: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_image_token_cost: Option<f64>,
/// USD per prompt token served from the provider's prompt cache.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_200k_tokens: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_200k_tokens_priority: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_256k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_272k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_272k_tokens_batches: Option<f64>,
/// Flex service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_272k_tokens_flex: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_272k_tokens_priority: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_32k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_512k_tokens: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_batches: Option<f64>,
/// Flex service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_flex: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_priority: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub citation_cost_per_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub code_interpreter_cost_per_session: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub comment: Option<String>,
/// Reasoning effort the provider applies when the request omits reasoning_effort. Gates whether a non-default temperature or the top_p/logprobs sampling params are accepted, which hold only when the effort resolves to 'none'.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub default_reasoning_effort: Option<ReasoningEffort>,
/// Date the provider deprecates the model, YYYY-MM-DD.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub deprecation_date: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub gemini_audio_only_live: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub gemini_native_audio: Option<bool>,
/// USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub google_maps_grounding_cost_per_query: Option<f64>,
/// USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub guardrail_cost_per_unit: Option<BTreeMap<String, f64>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_audio_per_second: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_audio_per_second_above_128k_tokens: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_audio_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_audio_token_batches: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_audio_token_priority: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_character: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_character_above_128k_tokens: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_image: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_image_above_128k_tokens: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_image_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_image_token_batches: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_pixel: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_query: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_request: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_second: Option<f64>,
/// USD per prompt token.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_200k_tokens: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_200k_tokens_priority: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_256k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_272k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_272k_tokens_batches: Option<f64>,
/// Flex service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_272k_tokens_flex: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_272k_tokens_priority: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_32k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_512k_tokens: Option<f64>,
/// USD per prompt token via the provider's batch API.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_batches: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_cache_hit: Option<f64>,
/// Flex service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_flex: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_priority: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_video_per_second: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_video_per_second_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_video_per_second_above_15s_interval: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_video_per_second_above_8s_interval: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_video_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_video_token_batches: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub input_dbu_cost_per_token: Option<f64>,
/// LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub litellm_provider: Option<String>,
/// Maximum prompt/context tokens the model accepts.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub max_input_tokens: Option<u64>,
/// Maximum tokens the model can generate in one response.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub max_output_tokens: Option<u64>,
/// Legacy field: max output tokens if the provider specifies it, else max input tokens.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<u64>,
/// Free-form notes about the entry (e.g. pricing derivation).
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<BTreeMap<String, Value>>,
/// Primary API surface / task type of the model.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub mode: Option<Mode>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub ocr_cost_per_credit: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub ocr_cost_per_page: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub ocr_cost_per_page_batches: Option<f64>,
/// Rates that replace the same-named base fields while the request falls inside the stated UTC windows.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub off_peak_pricing: Option<OffPeakPricing>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_audio_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_character: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_character_above_128k_tokens: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_image: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_image_1024: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_image_1536: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_image_512: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_image_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_pixel: Option<f64>,
/// USD per reasoning/thinking token, when billed separately.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_reasoning_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_second: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_second_1080p: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_second_2k: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_second_480p: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_second_4k: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_second_720p: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_second_768p: Option<f64>,
/// USD per generated token.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_200k_tokens: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_200k_tokens_priority: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_256k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_272k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_272k_tokens_batches: Option<f64>,
/// Flex service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_272k_tokens_flex: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_272k_tokens_priority: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_32k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_512k_tokens: Option<f64>,
/// USD per generated token via the provider's batch API.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_batches: Option<f64>,
/// Flex service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_flex: Option<f64>,
/// Priority service-tier rate for the same-named base field.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_priority: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_video_per_second: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_video_token: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_dbu_cost_per_token: Option<f64>,
/// Embedding dimension for embedding models.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub output_vector_size: Option<u64>,
/// Smallest prefix the provider will actually cache; absent means the provider default applies.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt_cache_min_tokens: Option<u64>,
/// Provider-internal routing hints (e.g. bedrock_invocation_schema).
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub provider_specific_entry: Option<BTreeMap<String, Value>>,
/// Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub reasoning_effort_levels: Option<Vec<ReasoningEffort>>,
/// Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub regional_endpoint_uplift_multiplier: Option<f64>,
/// Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%).
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub regional_processing_uplift_multiplier_eu: Option<f64>,
/// Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub regional_processing_uplift_multiplier_us: Option<f64>,
/// Provider default requests-per-minute limit.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub rpm: Option<u64>,
/// USD cost per web search query, keyed by search context size.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub search_context_cost_per_query: Option<SearchContextCostPerQuery>,
/// URL of the provider pricing/model page this entry was taken from.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub source: Option<String>,
/// Audio container formats the model can return.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supported_audio_formats: Option<Vec<AudioFormat>>,
/// OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supported_endpoints: Option<Vec<String>>,
/// Input modalities the model accepts.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supported_modalities: Option<Vec<InputModality>>,
/// Output modalities the model can produce.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supported_output_modalities: Option<Vec<OutputModality>>,
/// Cloud regions the model is available in ('global' or region ids).
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supported_regions: Option<Vec<String>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_adaptive_thinking: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_anthropic_compaction: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_anthropic_thinking_payload: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_assistant_prefill: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_audio_input: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_audio_output: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_computer_use: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_embedding_image_input: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_fast_mode: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_forced_tool_use: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_function_calling: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_image_input: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_image_size: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_legacy_thinking: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_low_reasoning_effort: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_max_reasoning_effort: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_mid_conversation_system: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_minimal_reasoning_effort: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_multimodal: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_native_streaming: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_native_structured_output: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_none_reasoning_effort: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_nova_canvas_image_edit: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_output_config: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_parallel_function_calling: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_parallel_tool_use_config: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_pdf_input: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_prompt_cache_breakpoint: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_prompt_caching: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_reasoning: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_response_schema: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_sampling_params: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_speed: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_system_messages: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_thinking_cache_preservation: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_tool_choice: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_tool_search: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_url_context: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_video_input: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_vision: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_web_search: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub supports_xhigh_reasoning_effort: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub thinking_always_on: Option<bool>,
/// Context-length or result-count tiered rates; each tier's costs apply within its range.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub tiered_pricing: Option<Vec<TieredRate>>,
/// Provider default tokens-per-minute limit.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub tpm: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub use_openai_responses_path: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub uses_embed_content: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub vertex_ai_audio_api: Option<VertexAiAudioApi>,
/// Whether web search is billed per query or per prompt.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(skip_serializing_if = "Option::is_none")]
pub web_search_billing_unit: Option<WebSearchBillingUnit>,
}

97
litellm-rust/crates/model-catalog/src/pricing.rs Normal file
View file

@ -0,0 +1,97 @@
use serde::{Deserialize, Serialize};
/// Whether web search is billed per query or per prompt.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum WebSearchBillingUnit {
PerQuery,
PerPrompt,
}
/// UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(untagged)]
pub enum UtcHours {
Single(String),
Multiple(Vec<String>),
}
/// ISO-8601 weekday number (1 = Monday .. 7 = Sunday) or English day name.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(untagged)]
pub enum Weekday {
Number(u8),
Name(String),
}
/// One off-peak window entry inside `windows`.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(deny_unknown_fields)]
pub struct OffPeakWindow {
pub hours_utc: UtcHours,
#[serde(skip_serializing_if = "Option::is_none")]
pub weekdays: Option<Vec<Weekday>>,
}
/// Rates that replace the same-named base fields inside the stated UTC windows.
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(deny_unknown_fields)]
pub struct OffPeakPricing {
#[serde(skip_serializing_if = "Option::is_none")]
pub hours_utc: Option<UtcHours>,
#[serde(skip_serializing_if = "Option::is_none")]
pub windows: Option<Vec<OffPeakWindow>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub weekday_timezone: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_reasoning_token: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost: Option<f64>,
}
/// USD cost per web search query, keyed by search context size.
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(deny_unknown_fields)]
pub struct SearchContextCostPerQuery {
#[serde(skip_serializing_if = "Option::is_none")]
pub search_context_size_low: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub search_context_size_medium: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub search_context_size_high: Option<f64>,
}
/// One tier of a context-length or result-count tiered rate.
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[serde(deny_unknown_fields)]
pub struct TieredRate {
#[serde(skip_serializing_if = "Option::is_none")]
pub range: Option<[f64; 2]>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_results_range: Option<[f64; 2]>,
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_reasoning_token: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_query: Option<f64>,
}

140
litellm-rust/crates/model-catalog/src/schema.rs
View file

@ -1,7 +1,137 @@
use crate::model_info::ModelInfo;
use schemars::Schema;
use serde_json::{Map, Value, json};
/// JSON Schema for one catalog model entry, mirroring
/// `model_prices_and_context_window.schema.json`'s `modelEntry` definition.
pub fn model_entry_json_schema() -> schemars::Schema {
schemars::schema_for!(ModelInfo)
/// JSON Schema for one model entry, including registry validation constraints.
pub fn model_entry_json_schema() -> Schema {
let mut schema = serde_json::to_value(schemars::schema_for!(crate::ModelInfo))
.expect("derived model schema serializes");
remove_nullable_optional_fields(&mut schema);
decorate_model_entry(&mut schema);
Schema::from(
schema
.as_object()
.expect("derived schema is an object")
.clone(),
)
}
/// JSON Schema for the complete model prices registry document.
pub fn registry_json_schema() -> Schema {
let mut entry = model_entry_json_schema().as_value().clone();
let mut definitions = take_definitions(&mut entry);
entry.as_object_mut().unwrap().remove("$schema");
definitions.insert("modelEntry".into(), entry);
let mut fallback = serde_json::to_value(schemars::schema_for!(crate::FallbackGeneralizations))
.expect("derived fallback schema serializes");
remove_nullable_optional_fields(&mut fallback);
definitions.extend(take_definitions(&mut fallback));
fallback.as_object_mut().unwrap().remove("$schema");
let root = json!({
"$schema": "https://json-schema.org/draft/2020-12/schema",
"title": "LiteLLM model prices and context window registry",
"type": "object",
"properties": {
"sample_spec": {"type": "object"},
"fallback_generalizations": fallback
},
"additionalProperties": {"$ref": "#/$defs/modelEntry"},
"$defs": definitions
});
Schema::from(root.as_object().unwrap().clone())
}
fn take_definitions(schema: &mut Value) -> Map<String, Value> {
schema
.as_object_mut()
.unwrap()
.remove("$defs")
.and_then(|value| value.as_object().cloned())
.unwrap_or_default()
}
fn remove_nullable_optional_fields(value: &mut Value) {
match value {
Value::Array(values) => values.iter_mut().for_each(remove_nullable_optional_fields),
Value::Object(map) => {
map.values_mut().for_each(remove_nullable_optional_fields);
if let Some(Value::Array(types)) = map.get_mut("type") {
types.retain(|value| value != "null");
if types.len() == 1 {
let only = types[0].clone();
map.insert("type".into(), only);
}
}
if let Some(Value::Array(branches)) = map.get_mut("anyOf") {
branches.retain(|branch| branch.get("type") != Some(&Value::String("null".into())));
if branches.len() == 1 {
let only = branches[0]
.as_object()
.expect("schema branch is an object")
.clone();
map.remove("anyOf");
map.extend(only);
}
}
}
_ => {}
}
}
fn decorate_model_entry(schema: &mut Value) {
let object = schema.as_object_mut().unwrap();
object.insert("required".into(), json!(["litellm_provider"]));
object.insert("additionalProperties".into(), Value::Bool(true));
let properties = object
.get_mut("properties")
.unwrap()
.as_object_mut()
.unwrap();
properties.insert(
"aliases".into(),
json!({"type": "array", "items": {"type": "string"}}),
);
properties.get_mut("deprecation_date").unwrap()["format"] = json!("date");
properties.get_mut("deprecation_date").unwrap()["pattern"] =
json!(r"^\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\d|3[01])$");
properties.iter_mut().for_each(|(name, property)| {
if name.contains("cost") {
property["minimum"] = json!(0);
} else if name.contains("uplift_multiplier") {
property["minimum"] = json!(1);
}
});
properties.get_mut("guardrail_cost_per_unit").unwrap()["additionalProperties"]["minimum"] =
json!(0);
let definitions = object.get_mut("$defs").unwrap().as_object_mut().unwrap();
for definition in ["OffPeakPricing", "TieredRate", "SearchContextCostPerQuery"] {
let properties = definitions[definition]["properties"]
.as_object_mut()
.unwrap();
properties.iter_mut().for_each(|(name, property)| {
if name.contains("cost") || definition == "SearchContextCostPerQuery" {
property["minimum"] = json!(0);
}
});
}
definitions["OffPeakPricing"]["anyOf"] = json!([
{"required": ["hours_utc"]},
{"required": ["windows"]}
]);
definitions["OffPeakPricing"]["properties"]["windows"]["minItems"] = json!(1);
definitions["OffPeakWindow"]["properties"]["weekdays"]["minItems"] = json!(1);
definitions["TieredRate"]["properties"]["range"]["items"]["minimum"] = json!(0);
definitions["TieredRate"]["properties"]["max_results_range"]["items"]["minimum"] = json!(0);
definitions["Weekday"]["anyOf"][0]["minimum"] = json!(1);
definitions["Weekday"]["anyOf"][0]["maximum"] = json!(7);
definitions["Weekday"]["anyOf"][1]["pattern"] = json!(
r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$"
);
let window_pattern = json!(r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$");
definitions["UtcHours"]["anyOf"][0]["pattern"] = window_pattern.clone();
definitions["UtcHours"]["anyOf"][1]["items"]["pattern"] = window_pattern;
definitions["UtcHours"]["anyOf"][1]["minItems"] = json!(1);
}

218
litellm-rust/crates/model-catalog/src/validation.rs Normal file
View file

@ -0,0 +1,218 @@
use std::collections::BTreeSet;
use serde_json::{Map, Value};
use thiserror::Error;
use crate::{AliasIssue, Catalog, ModelInfo, UtcHours, Weekday};
/// A registry entry violates the checked-in catalog contract.
#[derive(Debug, Error)]
pub enum RegistryValidationError {
#[error("{reason}")]
Entry { model: String, reason: String },
#[error("alias issue: {0:?}")]
Alias(AliasIssue),
}
/// Validate one registry entry without restricting the tolerant catalog reader.
pub fn validate_model_entry(model: &str, value: &Value) -> Result<(), RegistryValidationError> {
validate_entry_inner(model, value).map_err(|reason| RegistryValidationError::Entry {
model: model.to_owned(),
reason,
})
}
/// Check every model and alias in a parsed catalog against registry rules.
pub fn validate_registry(catalog: &Catalog) -> Result<(), RegistryValidationError> {
if let Some(issue) = catalog.alias_issues().first() {
return Err(RegistryValidationError::Alias(issue.clone()));
}
catalog.model_names().try_for_each(|name| {
let entry = catalog.lookup(name).expect("catalog name must resolve");
validate_model_entry(name, &Value::Object(entry.entry.fields().clone()))
})
}
fn json_eq(left: &Value, right: &Value) -> bool {
match (left, right) {
(Value::Number(left), Value::Number(right)) => left.as_f64() == right.as_f64(),
(Value::Array(left), Value::Array(right)) => {
left.len() == right.len() && left.iter().zip(right).all(|(a, b)| json_eq(a, b))
}
(Value::Object(left), Value::Object(right)) => {
left.len() == right.len()
&& left
.iter()
.all(|(key, value)| right.get(key).is_some_and(|other| json_eq(value, other)))
}
_ => left == right,
}
}
fn keys(value: &Map<String, Value>) -> BTreeSet<String> {
value.keys().cloned().collect()
}
fn symmetric_difference(left: &BTreeSet<String>, right: &BTreeSet<String>) -> BTreeSet<String> {
left.symmetric_difference(right).cloned().collect()
}
fn validate_entry_inner(model_name: &str, value: &Value) -> Result<(), String> {
let object = value
.as_object()
.ok_or_else(|| format!("{model_name} must be an object"))?;
if let Some(aliases) = object.get("aliases") {
let names = aliases
.as_array()
.ok_or_else(|| format!("{model_name}.aliases must be an array"))?;
if names.iter().any(|name| !name.is_string()) {
return Err(format!("{model_name}.aliases must contain strings"));
}
}
let info: ModelInfo =
serde_json::from_value(value.clone()).map_err(|error| format!("{model_name}: {error}"))?;
if info.litellm_provider.is_none() {
return Err(format!("{model_name}.litellm_provider is required"));
}
validate_dates_and_windows(model_name, &info)?;
let serialized = serde_json::to_value(info).map_err(|error| error.to_string())?;
let mut expected = object.clone();
expected.remove("aliases");
if !json_eq(&Value::Object(expected.clone()), &serialized) {
let actual = serialized
.as_object()
.expect("ModelInfo serializes as an object");
return Err(format!(
"{model_name} has an unknown field, null, or changed value: {:?}",
symmetric_difference(&keys(&expected), &keys(actual))
));
}
check_prices(model_name, value)
}
fn validate_dates_and_windows(model_name: &str, info: &ModelInfo) -> Result<(), String> {
if let Some(date) = &info.deprecation_date {
let format = time::format_description::parse_borrowed::<2>("[year]-[month]-[day]").unwrap();
time::Date::parse(date, &format)
.map_err(|error| format!("{model_name}.deprecation_date: {error}"))?;
}
let Some(pricing) = &info.off_peak_pricing else {
return Ok(());
};
if pricing.hours_utc.is_none() && pricing.windows.is_none() {
return Err(format!(
"{model_name}.off_peak_pricing needs hours or windows"
));
}
if let Some(hours) = &pricing.hours_utc {
validate_hours(hours)?;
}
if let Some(windows) = &pricing.windows {
if windows.is_empty() {
return Err(format!("{model_name}.off_peak_pricing.windows is empty"));
}
windows.iter().try_for_each(|window| {
validate_hours(&window.hours_utc)?;
if let Some(days) = &window.weekdays
&& (days.is_empty() || days.iter().any(|day| !valid_weekday(day)))
{
return Err(format!("{model_name}.off_peak_pricing.weekdays is invalid"));
}
Ok(())
})?;
}
Ok(())
}
fn validate_hours(hours: &UtcHours) -> Result<(), String> {
let values = match hours {
UtcHours::Single(value) => std::slice::from_ref(value),
UtcHours::Multiple(values) => values.as_slice(),
};
if values.is_empty() || values.iter().any(|value| !valid_utc_window(value)) {
return Err("off_peak_pricing.hours_utc is invalid".into());
}
Ok(())
}
fn valid_utc_window(value: &str) -> bool {
let Some((start, end)) = value.split_once('-') else {
return false;
};
[start, end].into_iter().all(|clock| {
let Some((hour, minute)) = clock.split_once(':') else {
return false;
};
hour.len() == 2
&& minute.len() == 2
&& hour.parse::<u8>().is_ok_and(|hour| hour < 24)
&& minute.parse::<u8>().is_ok_and(|minute| minute < 60)
})
}
fn valid_weekday(day: &Weekday) -> bool {
match day {
Weekday::Number(number) => (1..=7).contains(number),
Weekday::Name(name) => matches!(
name.to_ascii_lowercase().as_str(),
"mon"
| "monday"
| "tue"
| "tues"
| "tuesday"
| "wed"
| "wednesday"
| "thu"
| "thur"
| "thurs"
| "thursday"
| "fri"
| "friday"
| "sat"
| "saturday"
| "sun"
| "sunday"
),
}
}
fn check_prices(path: &str, value: &Value) -> Result<(), String> {
let Some(object) = value.as_object() else {
return Ok(());
};
object.iter().try_for_each(|(key, field)| {
let field_path = format!("{path}.{key}");
if (key.contains("cost")
|| path.ends_with(".guardrail_cost_per_unit")
|| path.ends_with(".search_context_cost_per_query"))
&& let Some(number) = field.as_f64()
&& number < 0.0
{
return Err(format!("{field_path} must be nonnegative"));
}
if key.contains("uplift_multiplier")
&& let Some(number) = field.as_f64()
&& number < 1.0
{
return Err(format!("{field_path} must be at least one"));
}
if matches!(key.as_str(), "range" | "max_results_range")
&& field.as_array().is_some_and(|values| {
values
.iter()
.any(|value| value.as_f64().is_some_and(|n| n < 0.0))
})
{
return Err(format!("{field_path} must be nonnegative"));
}
if matches!(key.as_str(), "metadata" | "provider_specific_entry") {
return Ok(());
}
match field.as_array() {
Some(items) => items.iter().enumerate().try_for_each(|(index, item)| {
check_prices(&format!("{field_path}[{index}]"), item)
}),
None => check_prices(&field_path, field),
}
})
}

82
litellm-rust/crates/model-catalog/tests/catalog.rs
View file

@ -43,6 +43,7 @@ fn fixture_catalog() -> Catalog {
}
#[rstest]
#[ignore]
fn preserves_fields_and_metadata(fixture_catalog: Catalog) {
let catalog = fixture_catalog;
let entry = catalog.lookup("SHORT").unwrap();
@ -68,6 +69,7 @@ fn preserves_fields_and_metadata(fixture_catalog: Catalog) {
}
#[rstest]
#[ignore]
fn snapshot_does_not_borrow_source() {
let mut source = ALPHA_FIXTURE.to_vec();
let catalog = Catalog::parse(&source, Provenance::default()).unwrap();
@ -84,7 +86,8 @@ fn snapshot_does_not_borrow_source() {
#[case("shared", "Second")]
#[case("FIRST", "First")]
#[case("sHaReD", "Second")]
fn alias_collisions_and_case_fallback_follow_python_order(
#[ignore]
fn alias_collisions_and_case_fallback_follow_entry_order(
#[case] lookup: &str,
#[case] expected: &str,
) {
@ -117,6 +120,36 @@ fn alias_collisions_and_case_fallback_follow_python_order(
);
}
#[test]
#[ignore]
fn json_entry_order_controls_alias_ownership_and_case_fallback() {
let forward = Catalog::parse(
br#"{
"Alpha":{"aliases":["shared"]},
"Beta":{"aliases":["shared"]},
"Foo":{},
"fOO":{}
}"#,
Provenance::default(),
)
.unwrap();
let reversed = Catalog::parse(
br#"{
"fOO":{},
"Foo":{},
"Beta":{"aliases":["shared"]},
"Alpha":{"aliases":["shared"]}
}"#,
Provenance::default(),
)
.unwrap();
assert_eq!(forward.lookup("shared").unwrap().canonical_key, "Alpha");
assert_eq!(reversed.lookup("shared").unwrap().canonical_key, "Beta");
assert_eq!(forward.lookup("foo").unwrap().canonical_key, "fOO");
assert_eq!(reversed.lookup("foo").unwrap().canonical_key, "Foo");
}
#[derive(Debug)]
enum ValidationOutcome {
Ok,
@ -128,36 +161,37 @@ enum ValidationOutcome {
#[rstest]
#[case(
IntegrityLimits {
backup_model_count: 2,
reference_model_count: 2,
min_model_count: 1,
min_backup_ratio: 0.5,
min_reference_ratio: 0.5,
},
ValidationOutcome::Ok
)]
#[case(
IntegrityLimits {
backup_model_count: 3,
reference_model_count: 3,
min_model_count: 1,
min_backup_ratio: 0.5,
min_reference_ratio: 0.5,
},
ValidationOutcome::Shrunk
)]
#[case(
IntegrityLimits {
backup_model_count: 0,
reference_model_count: 0,
min_model_count: 2,
min_backup_ratio: 0.5,
min_reference_ratio: 0.5,
},
ValidationOutcome::BelowMinimum
)]
#[case(
IntegrityLimits {
backup_model_count: 0,
reference_model_count: 0,
min_model_count: 0,
min_backup_ratio: f64::NAN,
min_reference_ratio: f64::NAN,
},
ValidationOutcome::InvalidRatio
)]
#[ignore]
fn integrity_uses_canonical_count_and_strict_shrink_boundary(
#[case] limits: IntegrityLimits,
#[case] expected: ValidationOutcome,
@ -195,6 +229,7 @@ enum MalformedOutcome {
br#"{"fallback_generalizations":{},"a":{}}"#,
MalformedOutcome::Json
)]
#[ignore]
fn malformed_input_and_aliases_have_typed_outcomes(
#[case] body: &[u8],
#[case] expected: MalformedOutcome,
@ -210,6 +245,7 @@ fn malformed_input_and_aliases_have_typed_outcomes(
}
#[rstest]
#[ignore]
fn invalid_aliases_are_reported_not_fatal() {
let catalog = Catalog::parse(
br#"{"a":{"aliases":"bad"},"b":{"aliases":[9,"ok"]}}"#,
@ -228,7 +264,8 @@ fn invalid_aliases_are_reported_not_fatal() {
}
#[rstest]
fn parses_current_and_packaged_catalogs_without_pinning_counts(
#[ignore]
fn parses_current_and_packaged_catalogs_against_independent_baseline(
current_catalog: Catalog,
backup_catalog: Catalog,
) {
@ -236,18 +273,17 @@ fn parses_current_and_packaged_catalogs_without_pinning_counts(
assert!(backup_catalog.model_count() > 0);
assert!(current_catalog.sample_spec().is_some());
assert!(backup_catalog.sample_spec().is_some());
assert!(
current_catalog
.validate(IntegrityLimits::python_defaults(
backup_catalog.model_count()
))
.is_ok()
);
for name in current_catalog.model_names() {
// Snapshot from 2026-09-23; the backup file mirrors the current file and cannot detect shrinkage.
const REFERENCE_MODEL_COUNT: usize = 4303;
current_catalog
.validate(IntegrityLimits {
reference_model_count: REFERENCE_MODEL_COUNT,
min_model_count: 50,
min_reference_ratio: 0.9,
})
.unwrap();
assert!(current_catalog.model_names().all(|name| {
let entry = current_catalog.lookup(name).unwrap().entry;
assert_eq!(
entry.info().litellm_provider.is_some(),
entry.field("litellm_provider").is_some()
);
}
entry.info().litellm_provider.is_some() == entry.field("litellm_provider").is_some()
}));
}

76
litellm-rust/crates/model-catalog/tests/registry_validation.rs Normal file
View file

@ -0,0 +1,76 @@
use std::path::{Path, PathBuf};
use litellm_model_catalog::{
Catalog, FallbackGeneralizations, Provenance, validate_model_entry, validate_registry,
};
use rstest::{fixture, rstest};
use serde_json::{Map, Value};
#[fixture]
fn repo_root() -> PathBuf {
Path::new(env!("CARGO_MANIFEST_DIR")).join("../../..")
}
#[rstest]
#[case("model_prices_and_context_window.json")]
#[case("litellm/model_prices_and_context_window_backup.json")]
#[ignore]
fn checked_in_registry_passes_strict_validation(repo_root: PathBuf, #[case] filename: &str) {
let body = std::fs::read(repo_root.join(filename)).unwrap();
let catalog = Catalog::parse(&body, Provenance::default()).unwrap();
validate_registry(&catalog).unwrap();
}
#[rstest]
#[ignore]
fn fallback_generalizations_are_typed(repo_root: PathBuf) {
let body = std::fs::read(repo_root.join("model_prices_and_context_window.json")).unwrap();
let document: Map<String, Value> = serde_json::from_slice(&body).unwrap();
let Some(raw_rules) = document.get("fallback_generalizations") else {
return;
};
let _: FallbackGeneralizations = serde_json::from_value(raw_rules.clone()).unwrap();
let catalog = Catalog::parse(&body, Provenance::default()).unwrap();
assert!(
catalog
.fallback_rules()
.is_some_and(|rules| !rules.is_empty())
);
}
#[rstest]
#[case::missing_provider(serde_json::json!({"mode": "chat"}), "litellm_provider")]
#[case::unknown_field(serde_json::json!({"litellm_provider": "test", "typo": true}), "unknown field")]
#[case::negative_price(serde_json::json!({"litellm_provider": "test", "input_cost_per_token": -1}), "nonnegative")]
#[case::negative_nested_price(serde_json::json!({"litellm_provider": "test", "guardrail_cost_per_unit": {"unit": -1}}), "nonnegative")]
#[case::invalid_mode(serde_json::json!({"litellm_provider": "test", "mode": "invalid"}), "unknown variant")]
#[case::invalid_date(serde_json::json!({"litellm_provider": "test", "deprecation_date": "2026-02-31"}), "deprecation_date")]
#[case::invalid_hours(serde_json::json!({"litellm_provider": "test", "off_peak_pricing": {"hours_utc": "25:00-01:00"}}), "hours_utc")]
#[case::empty_windows(serde_json::json!({"litellm_provider": "test", "off_peak_pricing": {"windows": []}}), "windows is empty")]
#[case::invalid_weekday(serde_json::json!({"litellm_provider": "test", "off_peak_pricing": {"windows": [{"hours_utc": "00:00-01:00", "weekdays": [0]}]}}), "weekdays is invalid")]
#[case::invalid_aliases(serde_json::json!({"litellm_provider": "test", "aliases": ["good", 7]}), "aliases must contain strings")]
#[case::null_aliases(serde_json::json!({"litellm_provider": "test", "aliases": null}), "aliases must be an array")]
#[ignore]
fn registry_validation_rejects_malformed_entries(#[case] entry: Value, #[case] expected: &str) {
assert!(
validate_model_entry("test", &entry)
.unwrap_err()
.to_string()
.contains(expected)
);
}
#[test]
#[ignore]
fn checked_in_catalog_and_backup_match() {
let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../..");
let current = std::fs::read(root.join("model_prices_and_context_window.json")).unwrap();
let backup =
std::fs::read(root.join("litellm/model_prices_and_context_window_backup.json")).unwrap();
assert_eq!(current, backup);
let catalog = Catalog::parse(&current, Provenance::default()).unwrap();
assert!(
catalog.alias_issues().is_empty(),
"invalid registry aliases"
);
}

121
litellm-rust/crates/model-catalog/tests/schema.rs Normal file
View file

@ -0,0 +1,121 @@
#![cfg(feature = "schema")]
use std::collections::BTreeSet;
use std::path::Path;
use litellm_model_catalog::{model_entry_json_schema, registry_json_schema};
use rstest::rstest;
use serde_json::{Value, json};
fn schema() -> Value {
serde_json::to_value(model_entry_json_schema()).expect("generated schema serializes")
}
fn registry_validator() -> jsonschema::Validator {
let schema = serde_json::to_value(registry_json_schema()).unwrap();
jsonschema::options()
.should_validate_formats(true)
.build(&schema)
.expect("generated registry schema is valid")
}
#[rstest]
#[case("model_prices_and_context_window.json")]
#[case("litellm/model_prices_and_context_window_backup.json")]
#[ignore]
fn generated_registry_schema_validates_checked_in_catalog(#[case] path: &str) {
let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../..");
let catalog: Value = serde_json::from_slice(&std::fs::read(root.join(path)).unwrap()).unwrap();
let validator = registry_validator();
let errors: Vec<_> = validator
.iter_errors(&catalog)
.map(|error| error.to_string())
.collect();
assert!(errors.is_empty(), "{path}: {errors:?}");
}
#[rstest]
#[case(json!({"example": {"litellm_provider": "test"}}))]
#[case(json!({"example": {"litellm_provider": "test", "future_field": true}}))]
#[case(json!({"sample_spec": {"litellm_provider": "placeholder"}}))]
#[ignore]
fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value) {
assert!(registry_validator().is_valid(&document));
}
#[rstest]
#[case::missing_provider(json!({"mode": "chat"}))]
#[case::negative_cost(json!({"litellm_provider": "test", "input_cost_per_token": -1}))]
#[case::negative_guardrail_cost(json!({"litellm_provider": "test", "guardrail_cost_per_unit": {"unit": -1}}))]
#[case::negative_search_cost(json!({"litellm_provider": "test", "search_context_cost_per_query": {"search_context_size_low": -1}}))]
#[case::negative_tier_cost(json!({"litellm_provider": "test", "tiered_pricing": [{"input_cost_per_token": -1}]}))]
#[case::negative_tier_range(json!({"litellm_provider": "test", "tiered_pricing": [{"range": [-1, 2]}]}))]
#[case::low_uplift(json!({"litellm_provider": "test", "regional_endpoint_uplift_multiplier": 0.5}))]
#[case::nullable_cost(json!({"litellm_provider": "test", "input_cost_per_token": null}))]
#[case::invalid_mode(json!({"litellm_provider": "test", "mode": "telepathy"}))]
#[case::invalid_date(json!({"litellm_provider": "test", "deprecation_date": "2026-02-31"}))]
#[case::invalid_hours(json!({"litellm_provider": "test", "off_peak_pricing": {"hours_utc": "25:00-01:00"}}))]
#[case::empty_windows(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": []}}))]
#[case::invalid_weekday(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": [{"hours_utc": "00:00-01:00", "weekdays": [0]}]}}))]
#[case::invalid_aliases(json!({"litellm_provider": "test", "aliases": "wrong"}))]
#[case::non_object_model(json!(4))]
#[ignore]
fn generated_registry_schema_rejects_invalid_entries(#[case] entry: Value) {
assert!(!registry_validator().is_valid(&json!({"example": entry})));
}
#[rstest]
#[case("model_prices_and_context_window.json")]
#[case("litellm/model_prices_and_context_window_backup.json")]
#[ignore]
fn generated_schema_covers_catalog_fields(#[case] path: &str) {
let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../..");
let catalog: Value = serde_json::from_slice(&std::fs::read(root.join(path)).unwrap()).unwrap();
let schema = schema();
let properties = schema["properties"]
.as_object()
.expect("ModelInfo schema has properties");
let fields: BTreeSet<&str> = catalog
.as_object()
.expect("catalog is an object")
.iter()
.filter(|(name, _)| *name != "sample_spec" && *name != "fallback_generalizations")
.flat_map(|(_, entry)| entry.as_object().expect("model entry is an object").keys())
.map(String::as_str)
.filter(|name| *name != "aliases")
.collect();
let missing: Vec<_> = fields
.into_iter()
.filter(|name| !properties.contains_key(*name))
.collect();
assert!(
missing.is_empty(),
"{path}: fields missing from schema: {missing:?}"
);
}
#[rstest]
#[case("Mode", "chat")]
#[case("ReasoningEffort", "high")]
#[case("InputModality", "image")]
#[ignore]
fn generated_schema_includes_enum_values(#[case] definition: &str, #[case] value: &str) {
let schema = schema();
let variants = schema["$defs"][definition]["enum"]
.as_array()
.expect("enum definition has variants");
assert!(variants.iter().any(|variant| variant == value));
}
#[test]
#[ignore]
fn generated_schema_includes_nested_pricing_types() {
let schema = schema();
let definitions = schema["$defs"].as_object().expect("schema has definitions");
assert!(definitions.contains_key("OffPeakPricing"));
assert!(definitions.contains_key("TieredRate"));
assert!(definitions.contains_key("UtcHours"));
}

121
litellm-rust/crates/model-catalog/tests/spec_parity.rs
View file

@ -1,121 +0,0 @@
use std::collections::{BTreeSet, HashSet};
use std::path::{Path, PathBuf};
use indexmap::IndexMap;
use litellm_model_catalog::{
Catalog, FallbackGeneralizations, ModelInfo, Provenance, model_entry_json_schema,
};
use rstest::{fixture, rstest};
use serde_json::{Map, Value};
#[fixture]
fn repo_root() -> PathBuf {
Path::new(env!("CARGO_MANIFEST_DIR")).join("../../..")
}
fn json_eq(left: &Value, right: &Value) -> bool {
match (left, right) {
(Value::Number(left), Value::Number(right)) => left.as_f64() == right.as_f64(),
(Value::Array(left), Value::Array(right)) => {
left.len() == right.len() && left.iter().zip(right).all(|(a, b)| json_eq(a, b))
}
(Value::Object(left), Value::Object(right)) => {
left.len() == right.len()
&& left
.iter()
.all(|(key, value)| right.get(key).is_some_and(|other| json_eq(value, other)))
}
_ => left == right,
}
}
fn keys(value: &Map<String, Value>) -> BTreeSet<String> {
value.keys().cloned().collect()
}
fn symmetric_difference(left: &BTreeSet<String>, right: &BTreeSet<String>) -> BTreeSet<String> {
left.symmetric_difference(right).cloned().collect()
}
#[rstest]
#[case("model_prices_and_context_window.json")]
#[case("litellm/model_prices_and_context_window_backup.json")]
fn every_entry_round_trips_through_model_info(repo_root: PathBuf, #[case] filename: &str) {
let body = std::fs::read(repo_root.join(filename)).unwrap();
let document: IndexMap<String, Value> = serde_json::from_slice(&body).unwrap();
for (model_name, value) in document {
if matches!(
model_name.as_str(),
"sample_spec" | "fallback_generalizations"
) {
continue;
}
let object = value
.as_object()
.unwrap_or_else(|| panic!("{model_name} is not an object"));
let info: ModelInfo = serde_json::from_value(value.clone())
.unwrap_or_else(|error| panic!("{model_name} does not deserialize: {error}"));
let serialized = serde_json::to_value(info).unwrap();
let serialized_object = serialized
.as_object()
.unwrap_or_else(|| panic!("{model_name} did not serialize as an object"));
let mut expected = object.clone();
expected.remove("aliases");
let expected_keys = keys(&expected);
let serialized_keys = keys(serialized_object);
assert_eq!(
expected_keys,
serialized_keys,
"{model_name} key difference: {:?}",
symmetric_difference(&expected_keys, &serialized_keys)
);
assert!(
json_eq(&Value::Object(expected), &serialized),
"{model_name} changed during ModelInfo round-trip"
);
}
}
#[rstest]
fn fallback_generalizations_are_typed(repo_root: PathBuf) {
let body = std::fs::read(repo_root.join("model_prices_and_context_window.json")).unwrap();
let document: Map<String, Value> = serde_json::from_slice(&body).unwrap();
let Some(raw_rules) = document.get("fallback_generalizations") else {
return;
};
let _: FallbackGeneralizations = serde_json::from_value(raw_rules.clone()).unwrap();
let catalog = Catalog::parse(&body, Provenance::default()).unwrap();
assert!(
catalog
.fallback_rules()
.is_some_and(|rules| !rules.is_empty())
);
}
#[rstest]
fn generated_schema_properties_match_repo_schema(repo_root: PathBuf) {
let body =
std::fs::read(repo_root.join("model_prices_and_context_window.schema.json")).unwrap();
let document: Value = serde_json::from_slice(&body).unwrap();
let repo_entry_properties = document["$defs"]["modelEntry"]["properties"]
.as_object()
.unwrap();
let generated = serde_json::to_value(model_entry_json_schema()).unwrap();
let generated_properties = generated["properties"].as_object().unwrap();
let expected = keys(repo_entry_properties);
let actual = keys(generated_properties);
assert_eq!(
expected,
actual,
"modelEntry property difference: {:?}",
symmetric_difference(&expected, &actual)
);
let repo_root_properties = document["properties"].as_object().unwrap();
let actual_root: HashSet<String> = repo_root_properties.keys().cloned().collect();
let expected_root: HashSet<String> = ["sample_spec", "fallback_generalizations"]
.into_iter()
.map(str::to_owned)
.collect();
assert_eq!(actual_root, expected_root);
}

9
litellm-rust/crates/python-bridge/src/routes/messages/host.rs
View file

@ -3,7 +3,7 @@ use std::convert::Infallible;
use bytes::Bytes;
use litellm_core::messages::{
Error,
route::{Messages, MessagesCall, MessagesOutput},
route::{Messages, MessagesCall, MessagesOutput, MessagesStreamHead},
types::MessagesShaping,
};
use litellm_host_python::{InvokeError, ProtocolHost, from_py, lookup, to_py};
@ -238,6 +238,13 @@ impl ProtocolHost for MessagesPythonHost {
}
}
fn head(&mut self, py: Python<'_>, head: MessagesStreamHead) -> PyResult<Py<PyAny>> {
py.import(ROUTE_HOST_MODULE)?
.getattr("stream_hidden_params")?
.call1((to_py(py, &head.headers)?,))
.map(Bound::unbind)
}
fn chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult<Py<PyAny>> {
Ok(PyBytes::new(py, &chunk).into_any().unbind())
}

15
litellm-rust/crates/python-bridge/src/routes/messages/mod.rs
View file

@ -4,14 +4,12 @@ use host::MessagesPythonHost;
use litellm_callbacks_legacy_python::{
LegacySurface, PassThroughStream, PublicCall, run_legacy_call,
};
use litellm_core::messages::route::{messages_machine, supports};
use litellm_core::messages::route::messages_machine;
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use crate::errors::RustBridgeDeclined;
const SURFACE: LegacySurface = LegacySurface {
call_type: "anthropic_messages",
input_description: "Messages",
@ -28,17 +26,6 @@ fn run_messages(
kwargs: Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
let model: String = request.getattr("model")?.extract()?;
let provider: Option<String> = request.getattr("custom_llm_provider")?.extract()?;
let stream = request
.getattr("stream")?
.extract::<Option<bool>>()?
.unwrap_or(false);
if !supports(&model, provider.as_deref(), stream) {
return Err(RustBridgeDeclined::new_err(
"the Rust Messages route does not serve this provider",
));
}
let secrets = crate::secrets::source(py)?;
run_legacy_call(
py,

4
litellm-rust/crates/python-bridge/src/routes/ocr/host.rs
View file

@ -117,6 +117,10 @@ impl ProtocolHost for OcrPythonHost {
.map(Bound::unbind)
}
fn head(&mut self, _: Python<'_>, head: std::convert::Infallible) -> PyResult<Py<PyAny>> {
match head {}
}
fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult<Py<PyAny>> {
match chunk {}
}

2
litellm-rust/crates/secrets-hashicorp/tests/secret_manager/support.rs
View file

@ -118,7 +118,7 @@ pub(super) struct ParityCase {
pub(super) fn parity_cases() -> Vec<ParityCase> {
serde_json::from_str(include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../../../tests/test_litellm/secret_managers/hashicorp_vault_parity.json"
"/../../../tests/unit/secret_managers/hashicorp_vault_parity.json"
)))
.unwrap()
}

20
litellm-rust/crates/secrets/PARITY.md
View file

@ -32,7 +32,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo
`_SecretManagerRuntime` is a private implementation detail, not a replacement SDK class. Its async methods return Futures; public `async def` methods retain lazy coroutine creation and `asyncio.create_task` support. Passing the same names and arguments is insufficient to claim parity until the remaining return-value, error, cache and configuration differences above are closed
## [tests/test_litellm/secret_managers/test_aws_secret_manager_replication.py](../../../tests/test_litellm/secret_managers/test_aws_secret_manager_replication.py)
## [tests/unit/secret_managers/test_aws_secret_manager_replication.py](../../../tests/unit/secret_managers/test_aws_secret_manager_replication.py)
| Python test | Rust coverage or boundary |
| --- | --- |
@ -48,7 +48,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo
| `test_replicate_secret_http_error_raises` | [direct_replication_returns_response_or_service_error](../secrets-aws/tests/secret_manager/writes.rs) |
| `test_replicate_secret_timeout_raises` | [write_and_replication_timeouts_remain_errors](../secrets-aws/tests/secret_manager/writes.rs) |
## [tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py](../../../tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py)
## [tests/unit/secret_managers/test_aws_secret_manager_rotation.py](../../../tests/unit/secret_managers/test_aws_secret_manager_rotation.py)
| Python test | Rust coverage or boundary |
| --- | --- |
@ -59,7 +59,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo
| `test_write_secret_to_name_inside_recovery_window_restores_and_stores_new_value` | [recovery_window_alias_is_restored_updated_and_tagged](../secrets-aws/tests/secret_manager/writes.rs) |
| `test_write_secret_to_live_existing_name_still_fails_without_overwriting` | [create_failure_does_not_overwrite_an_alias_without_a_deletion_date](../secrets-aws/tests/secret_manager/writes.rs) |
## [tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py](../../../tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py)
## [tests/unit/secret_managers/test_aws_secret_manager_v2.py](../../../tests/unit/secret_managers/test_aws_secret_manager_v2.py)
| Python test | Rust coverage or boundary |
| --- | --- |
@ -70,14 +70,14 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo
| `test_prepare_request_explicit_bedrock_runtime_endpoint_param_still_wins` | [endpoint_overrides_replace_the_service_and_override_the_region](../secrets-aws/tests/secret_manager/configuration.rs) |
| `test_prepare_request_env_bedrock_runtime_endpoint_still_wins` | [endpoint_overrides_replace_the_service_and_override_the_region](../secrets-aws/tests/secret_manager/configuration.rs) |
## [tests/test_litellm/secret_managers/test_base_secret_manager.py](../../../tests/test_litellm/secret_managers/test_base_secret_manager.py)
## [tests/unit/secret_managers/test_base_secret_manager.py](../../../tests/unit/secret_managers/test_base_secret_manager.py)
| Python test | Rust coverage or boundary |
| --- | --- |
| `test_raise_if_unsafe_secret_name_rejects_traversal_and_line_breaks` | [names_reject_path_traversal_and_control_characters](../secrets-types/tests/rotation.rs) |
| `test_raise_if_unsafe_secret_name_allows_legitimate_aliases` | [names_allow_safe_values](../secrets-types/tests/rotation.rs) |
## [tests/test_litellm/secret_managers/test_custom_secret_manager.py](../../../tests/test_litellm/secret_managers/test_custom_secret_manager.py)
## [tests/unit/secret_managers/test_custom_secret_manager.py](../../../tests/unit/secret_managers/test_custom_secret_manager.py)
| Python test | Rust coverage or boundary |
| --- | --- |
@ -89,7 +89,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo
| `test_custom_secret_manager_integration_with_litellm` | [manager_strings_are_coerced_like_literal_eval](../secrets/tests/resolution.rs) |
| `test_minimal_custom_secret_manager` | Exercises the Python example subclass itself or Python default methods. Caller-authored Python implementations remain Python callbacks; resolver integration is covered by `manager_strings_are_coerced_like_literal_eval` |
## [tests/test_litellm/secret_managers/test_cyberark_secret_manager.py](../../../tests/test_litellm/secret_managers/test_cyberark_secret_manager.py)
## [tests/unit/secret_managers/test_cyberark_secret_manager.py](../../../tests/unit/secret_managers/test_cyberark_secret_manager.py)
| Python test | Rust coverage or boundary |
| --- | --- |
@ -97,7 +97,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo
| `test_async_write_matches_parity_fixture` | [writes_match_python_parity_fixture](../secrets-cyberark/tests/secret_manager/writes.rs) |
| `test_missing_credentials_raise_value_error` | [new_validates_credentials_before_license_and_configuration](../secrets-cyberark/tests/secret_manager/configuration.rs) |
## [tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py](../../../tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py)
## [tests/unit/secret_managers/test_get_azure_ad_token_provider.py](../../../tests/unit/secret_managers/test_get_azure_ad_token_provider.py)
| Python test | Rust coverage or boundary |
| --- | --- |
@ -115,7 +115,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo
| `test_get_azure_ad_token_provider_prefers_workload_identity_over_managed_identity` | Credential-selection contract belongs to `litellm-auth-azure`, not a secrets crate |
| `test_get_azure_ad_token_provider_defaults_to_default_azure_credential` | Credential-selection contract belongs to `litellm-auth-azure`, not a secrets crate |
## [tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py](../../../tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py)
## [tests/unit/secret_managers/test_hashicorp_secret_manager.py](../../../tests/unit/secret_managers/test_hashicorp_secret_manager.py)
| Python test | Rust coverage or boundary |
| --- | --- |
@ -130,13 +130,13 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo
| `test_tls_login_uses_login_namespace` | [tls_login_posts_the_role_and_uses_the_client_identity](../secrets-hashicorp/tests/secret_manager/configuration.rs) |
| `test_configuration_matches_native_parity_fixture` | [configuration_matches_python_parity_fixture](../secrets-hashicorp/tests/secret_manager/configuration.rs) |
## [tests/test_litellm/secret_managers/test_secret_manager_handler.py](../../../tests/test_litellm/secret_managers/test_secret_manager_handler.py)
## [tests/unit/secret_managers/test_secret_manager_handler.py](../../../tests/unit/secret_managers/test_secret_manager_handler.py)
| Python test | Rust coverage or boundary |
| --- | --- |
| `test_azure_key_vault_matches_rust_parity_fixture` | [parity_fixture_matches_python_backend_contract](../secrets-azure/tests/key_vault.rs) |
## [tests/test_litellm/secret_managers/test_secret_managers_main.py](../../../tests/test_litellm/secret_managers/test_secret_managers_main.py)
## [tests/unit/secret_managers/test_secret_managers_main.py](../../../tests/unit/secret_managers/test_secret_managers_main.py)
| Python test | Rust coverage or boundary |
| --- | --- |

2
litellm-rust/crates/secrets/README.md
View file

@ -30,7 +30,7 @@ The HashiCorp Vault backend is enabled with the `hashicorp` feature and reads KV
Native backends consistently distinguish absence from failure instead of swallowing provider errors. Python-compatible resolution maps these results back to the Python handler contract before applying fallback
`hosted_keys` excludes a name for every backend. Python's handler recognizes Azure `SecretClient` and Google `KeyManagementServiceClient` instances before the `local` branch, allowing excluded names to reach those providers. Rust treats that as a routing bug. `test_rust_hosted_keys_exclude_azure_sdk_clients_too` in `tests/test_litellm/rust_bridge/ocr/test_secrets.py` pins this behavior
`hosted_keys` excludes a name for every backend. Python's handler recognizes Azure `SecretClient` and Google `KeyManagementServiceClient` instances before the `local` branch, allowing excluded names to reach those providers. Rust treats that as a routing bug. `test_rust_hosted_keys_exclude_azure_sdk_clients_too` in `tests/unit/rust_bridge/ocr/test_secrets.py` pins this behavior
Google rejects malformed base64 and mismatched CRC32C values instead of accepting corrupted payloads. Python currently ignores the checksum and uses permissive base64 decoding. Rust follows [RFC 4648](https://www.rfc-editor.org/rfc/rfc4648#section-3.3) and [Google's integrity guidance](https://docs.cloud.google.com/secret-manager/docs/data-integrity); `failed_or_missing_reads_are_not_cached` covers rejection and recovery

32
litellm-rust/crates/testkit/Cargo.toml Normal file
View file

@ -0,0 +1,32 @@
[package]
name = "litellm-testkit"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
publish = false
[dependencies]
flate2.workspace = true
reqwest.workspace = true
serde.workspace = true
semver.workspace = true
serde_json.workspace = true
sha2.workspace = true
tar.workspace = true
target-lexicon.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["fs", "process"] }
zip.workspace = true
[dev-dependencies]
flate2.workspace = true
rstest.workspace = true
sha2.workspace = true
tar.workspace = true
target-lexicon.workspace = true
futures-util.workspace = true
tempfile.workspace = true
tokio.workspace = true
toml = "0.9"
zip.workspace = true

181
litellm-rust/crates/testkit/src/agent/claude.rs Normal file
View file

@ -0,0 +1,181 @@
use std::collections::BTreeMap;
use std::path::Path;
use semver::Version;
use serde::Deserialize;
use super::{
Configure, Drive, Install, LaunchSpec, Outcome, Prompt, Settings, Usage, Wire, env, json_lines,
path_string,
};
use crate::install::release::parse;
use crate::install::{Packaging, Release};
use crate::{Error, Fetch, Target};
const RELEASES: &str = "https://downloads.claude.ai/claude-code-releases";
pub struct ClaudeCode;
#[derive(Deserialize)]
struct Manifest {
platforms: BTreeMap<String, Platform>,
}
#[derive(Deserialize)]
struct Platform {
checksum: String,
}
impl Install for ClaudeCode {
fn binary(&self) -> &'static str {
"claude"
}
async fn release(
&self,
fetch: &impl Fetch,
version: &Version,
target: Target,
) -> Result<Release, Error> {
let manifest_url = format!("{RELEASES}/{version}/manifest.json");
let manifest: Manifest = parse(&manifest_url, &fetch.get(&manifest_url).await?)?;
let key = format!(
"{}-{}{}",
target.os_name(),
target.arch_name(),
target.musl_suffix()
);
let platform = manifest
.platforms
.get(&key)
.ok_or_else(|| Error::AssetNotFound(key.clone()))?;
Ok(Release {
url: format!("{RELEASES}/{version}/{key}/claude"),
asset: key,
sha256: platform.checksum.clone(),
packaging: Packaging::Bare,
})
}
}
impl Configure for ClaudeCode {
fn configure(
&self,
_version: &Version,
settings: &Settings,
home: &Path,
) -> Result<LaunchSpec, Error> {
if settings.wire != Wire::Messages {
return Err(Error::UnsupportedWire {
agent: "claude",
wire: settings.wire,
});
}
Ok(LaunchSpec {
env: env([
("HOME", path_string(home)),
("CLAUDE_CONFIG_DIR", path_string(&home.join(".claude"))),
("ANTHROPIC_BASE_URL", settings.base_url.clone()),
("ANTHROPIC_AUTH_TOKEN", settings.api_key.clone()),
("ANTHROPIC_MODEL", settings.model.clone()),
("DISABLE_AUTOUPDATER", "1".to_owned()),
]),
files: BTreeMap::new(),
})
}
}
#[derive(Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum Event {
Assistant {
message: AssistantMessage,
},
Result(Finished),
#[serde(other)]
Other,
}
#[derive(Deserialize)]
struct AssistantMessage {
content: Vec<Block>,
}
#[derive(Deserialize)]
struct Block {
#[serde(rename = "type")]
kind: String,
name: Option<String>,
}
#[derive(Deserialize)]
struct Finished {
is_error: bool,
result: Option<String>,
usage: Option<TokenUsage>,
}
#[derive(Deserialize)]
struct TokenUsage {
input_tokens: u64,
output_tokens: u64,
}
impl Drive for ClaudeCode {
fn args(&self, _version: &Version, settings: &Settings, prompt: &Prompt) -> Vec<String> {
let base = [
"-p",
&prompt.text,
"--output-format",
"stream-json",
"--verbose",
"--model",
&settings.model,
];
let tools = ["--allowedTools", "Bash,Read,Write,Edit"];
base.into_iter()
.chain(tools.into_iter().filter(|_| prompt.allow_tools))
.map(str::to_owned)
.collect()
}
fn parse(&self, _version: &Version, stdout: &str) -> Outcome {
let events: Vec<Event> = json_lines(stdout).collect();
let tool_calls = events
.iter()
.filter_map(|event| match event {
Event::Assistant { message } => Some(&message.content),
_ => None,
})
.flatten()
.filter(|block| block.kind == "tool_use")
.filter_map(|block| block.name.clone())
.collect();
let finished = events.into_iter().find_map(|event| match event {
Event::Result(finished) => Some(finished),
_ => None,
});
let Some(finished) = finished else {
return Outcome {
tool_calls,
..Outcome::default()
};
};
let result = finished.result.unwrap_or_default();
let (text, errors) = if finished.is_error {
(String::new(), vec![result])
} else {
(result, Vec::new())
};
Outcome {
text,
tool_calls,
usage: finished.usage.map_or_else(Usage::default, |usage| Usage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
}),
errors,
exit_code: None,
}
}
}

174
litellm-rust/crates/testkit/src/agent/codex.rs Normal file
View file

@ -0,0 +1,174 @@
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use semver::Version;
use serde::Deserialize;
use super::{
Configure, Drive, Install, LaunchSpec, Outcome, Prompt, Settings, Usage, Wire, env, json_lines,
path_string, quoted, v1,
};
use crate::install::release::github_release;
use crate::install::{Packaging, Release};
use crate::target::{Arch, Os};
use crate::{Error, Fetch, Target};
const RELEASES: &str = "https://api.github.com/repos/openai/codex/releases/tags";
pub struct Codex;
fn triple(target: Target) -> String {
let arch = match target.arch {
Arch::Aarch64 => "aarch64",
Arch::X86_64 => "x86_64",
};
match target.os {
Os::Macos => format!("{arch}-apple-darwin"),
Os::Linux => format!("{arch}-unknown-linux-musl"),
}
}
impl Install for Codex {
fn binary(&self) -> &'static str {
"codex"
}
async fn release(
&self,
fetch: &impl Fetch,
version: &Version,
target: Target,
) -> Result<Release, Error> {
let triple = triple(target);
github_release(
fetch,
RELEASES,
&format!("rust-v{version}"),
&format!("codex-{triple}.tar.gz"),
Packaging::TarGz {
member: format!("codex-{triple}"),
},
)
.await
}
}
impl Configure for Codex {
fn configure(
&self,
_version: &Version,
settings: &Settings,
home: &Path,
) -> Result<LaunchSpec, Error> {
if settings.wire != Wire::Responses {
return Err(Error::UnsupportedWire {
agent: "codex",
wire: settings.wire,
});
}
let config = format!(
"model = {model}\nmodel_provider = \"litellm\"\n\n[model_providers.litellm]\nname = \"LiteLLM\"\nbase_url = {base_url}\nenv_key = \"LITELLM_API_KEY\"\nwire_api = \"responses\"\n",
model = quoted(&settings.model),
base_url = quoted(&v1(settings)),
);
Ok(LaunchSpec {
env: env([
("HOME", path_string(home)),
("CODEX_HOME", path_string(&home.join(".codex"))),
("LITELLM_API_KEY", settings.api_key.clone()),
]),
files: BTreeMap::from([(PathBuf::from(".codex/config.toml"), config)]),
})
}
}
#[derive(Deserialize)]
enum EventKind {
#[serde(rename = "item.completed")]
ItemCompleted,
#[serde(rename = "turn.completed")]
TurnCompleted,
#[serde(rename = "turn.failed")]
TurnFailed,
#[serde(other)]
Other,
}
#[derive(Deserialize)]
struct Event {
#[serde(rename = "type")]
kind: EventKind,
item: Option<Item>,
usage: Option<TokenUsage>,
error: Option<Failure>,
}
#[derive(Deserialize)]
struct Item {
#[serde(rename = "type")]
kind: String,
text: Option<String>,
}
#[derive(Deserialize)]
struct TokenUsage {
input_tokens: u64,
output_tokens: u64,
}
#[derive(Deserialize)]
struct Failure {
message: String,
}
const NON_TOOL_ITEMS: [&str; 3] = ["agent_message", "reasoning", "error"];
impl Drive for Codex {
fn args(&self, _version: &Version, _settings: &Settings, prompt: &Prompt) -> Vec<String> {
let sandbox = ["--sandbox", "workspace-write"];
["exec", "--json", "--skip-git-repo-check"]
.into_iter()
.chain(sandbox.into_iter().filter(|_| prompt.allow_tools))
.chain([prompt.text.as_str()])
.map(str::to_owned)
.collect()
}
fn parse(&self, _version: &Version, stdout: &str) -> Outcome {
let events: Vec<Event> = json_lines(stdout).collect();
let items: Vec<&Item> = events
.iter()
.filter(|event| matches!(event.kind, EventKind::ItemCompleted))
.filter_map(|event| event.item.as_ref())
.collect();
Outcome {
text: items
.iter()
.rev()
.find(|item| item.kind == "agent_message")
.and_then(|item| item.text.clone())
.unwrap_or_default(),
tool_calls: items
.iter()
.filter(|item| !NON_TOOL_ITEMS.contains(&item.kind.as_str()))
.map(|item| item.kind.clone())
.collect(),
usage: events
.iter()
.filter(|event| matches!(event.kind, EventKind::TurnCompleted))
.filter_map(|event| event.usage.as_ref())
.map(|usage| Usage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
})
.fold(Usage::default(), |total, turn| total + turn),
errors: events
.iter()
.filter(|event| matches!(event.kind, EventKind::TurnFailed))
.filter_map(|event| event.error.as_ref())
.map(|failure| failure.message.clone())
.collect(),
exit_code: None,
}
}
}

69
litellm-rust/crates/testkit/src/agent/configure.rs Normal file
View file

@ -0,0 +1,69 @@
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use semver::Version;
use crate::Error;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Wire {
ChatCompletions,
Messages,
Responses,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Settings {
pub base_url: String,
pub api_key: String,
pub model: String,
pub wire: Wire,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LaunchSpec {
pub env: BTreeMap<String, String>,
pub files: BTreeMap<PathBuf, String>,
}
impl LaunchSpec {
pub fn write_files(&self, home: &Path) -> std::io::Result<()> {
self.files.iter().try_for_each(|(relative, contents)| {
let path = home.join(relative);
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, contents)
})
}
}
pub trait Configure {
fn configure(
&self,
version: &Version,
settings: &Settings,
home: &Path,
) -> Result<LaunchSpec, Error>;
}
pub(crate) fn env(
pairs: impl IntoIterator<Item = (&'static str, String)>,
) -> BTreeMap<String, String> {
pairs
.into_iter()
.map(|(key, value)| (key.to_owned(), value))
.collect()
}
pub(crate) fn path_string(path: &Path) -> String {
path.to_string_lossy().into_owned()
}
pub(crate) fn quoted(value: &str) -> String {
serde_json::Value::from(value).to_string()
}
pub(crate) fn v1(settings: &Settings) -> String {
format!("{}/v1", settings.base_url.trim_end_matches('/'))
}

57
litellm-rust/crates/testkit/src/agent/drive.rs Normal file
View file

@ -0,0 +1,57 @@
use std::ops::Add;
use semver::Version;
use crate::Settings;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Prompt {
pub text: String,
pub allow_tools: bool,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Usage {
pub input_tokens: u64,
pub output_tokens: u64,
}
impl Add for Usage {
type Output = Self;
fn add(self, other: Self) -> Self {
Self {
input_tokens: self.input_tokens + other.input_tokens,
output_tokens: self.output_tokens + other.output_tokens,
}
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Outcome {
pub text: String,
pub tool_calls: Vec<String>,
pub usage: Usage,
pub errors: Vec<String>,
pub exit_code: Option<i32>,
}
impl Outcome {
pub fn succeeded(&self) -> bool {
self.exit_code == Some(0) && self.errors.is_empty()
}
}
pub trait Drive {
fn args(&self, version: &Version, settings: &Settings, prompt: &Prompt) -> Vec<String>;
fn parse(&self, version: &Version, stdout: &str) -> Outcome;
}
pub(crate) fn json_lines<'a, T: serde::de::DeserializeOwned + 'a>(
stdout: &'a str,
) -> impl Iterator<Item = T> + 'a {
stdout
.lines()
.filter_map(|line| serde_json::from_str(line).ok())
}

17
litellm-rust/crates/testkit/src/agent/install.rs Normal file
View file

@ -0,0 +1,17 @@
use std::future::Future;
use semver::Version;
use crate::install::Release;
use crate::{Error, Fetch, Target};
pub trait Install: Sync {
fn binary(&self) -> &'static str;
fn release(
&self,
fetch: &impl Fetch,
version: &Version,
target: Target,
) -> impl Future<Output = Result<Release, Error>> + Send;
}

20
litellm-rust/crates/testkit/src/agent/mod.rs Normal file
View file

@ -0,0 +1,20 @@
mod claude;
mod codex;
mod configure;
mod drive;
mod install;
mod opencode;
pub use claude::ClaudeCode;
pub use codex::Codex;
pub use configure::{Configure, LaunchSpec, Settings, Wire};
pub use drive::{Drive, Outcome, Prompt, Usage};
pub use install::Install;
pub use opencode::Opencode;
pub(crate) use configure::{env, path_string, quoted, v1};
pub(crate) use drive::json_lines;
pub trait Agent: Install + Configure + Drive {}
impl<T: Install + Configure + Drive> Agent for T {}

187
litellm-rust/crates/testkit/src/agent/opencode.rs Normal file
View file

@ -0,0 +1,187 @@
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use semver::Version;
use serde::Deserialize;
use super::{
Configure, Drive, Install, LaunchSpec, Outcome, Prompt, Settings, Usage, Wire, env, json_lines,
path_string, v1,
};
use crate::install::release::github_release;
use crate::install::{Packaging, Release};
use crate::target::Os;
use crate::{Error, Fetch, Target};
const RELEASES: &str = "https://api.github.com/repos/sst/opencode/releases/tags";
pub struct Opencode;
impl Install for Opencode {
fn binary(&self) -> &'static str {
"opencode"
}
async fn release(
&self,
fetch: &impl Fetch,
version: &Version,
target: Target,
) -> Result<Release, Error> {
let stem = format!(
"opencode-{}-{}{}",
target.os_name(),
target.arch_name(),
target.musl_suffix()
);
let member = "opencode".to_owned();
let (asset, packaging) = match target.os {
Os::Macos => (format!("{stem}.zip"), Packaging::Zip { member }),
Os::Linux => (format!("{stem}.tar.gz"), Packaging::TarGz { member }),
};
github_release(fetch, RELEASES, &format!("v{version}"), &asset, packaging).await
}
}
impl Configure for Opencode {
fn configure(
&self,
_version: &Version,
settings: &Settings,
home: &Path,
) -> Result<LaunchSpec, Error> {
let npm = match settings.wire {
Wire::ChatCompletions => "@ai-sdk/openai-compatible",
Wire::Responses => "@ai-sdk/openai",
Wire::Messages => "@ai-sdk/anthropic",
};
let config = serde_json::json!({
"$schema": "https://opencode.ai/config.json",
"model": format!("litellm/{}", settings.model),
"provider": {
"litellm": {
"npm": npm,
"name": "LiteLLM",
"options": { "baseURL": v1(settings), "apiKey": settings.api_key },
"models": { settings.model.clone(): { "name": settings.model } },
}
},
});
Ok(LaunchSpec {
env: env([
("HOME", path_string(home)),
("XDG_CONFIG_HOME", path_string(&home.join(".config"))),
("XDG_DATA_HOME", path_string(&home.join(".local/share"))),
("OPENCODE_DISABLE_AUTOUPDATE", "true".to_owned()),
]),
files: BTreeMap::from([(
PathBuf::from(".config/opencode/opencode.json"),
config.to_string(),
)]),
})
}
}
#[derive(Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum Event {
Text {
part: TextPart,
},
ToolUse {
part: ToolPart,
},
StepFinish {
part: StepFinish,
},
Error {
error: Failure,
},
#[serde(other)]
Other,
}
#[derive(Deserialize)]
struct TextPart {
text: String,
}
#[derive(Deserialize)]
struct ToolPart {
tool: String,
}
#[derive(Deserialize)]
struct StepFinish {
tokens: Tokens,
}
#[derive(Deserialize)]
struct Tokens {
input: u64,
output: u64,
}
#[derive(Deserialize)]
struct Failure {
name: String,
data: Option<FailureData>,
}
#[derive(Deserialize)]
struct FailureData {
message: Option<String>,
}
impl Drive for Opencode {
fn args(&self, _version: &Version, _settings: &Settings, prompt: &Prompt) -> Vec<String> {
["run", "--format", "json", &prompt.text]
.map(str::to_owned)
.to_vec()
}
fn parse(&self, _version: &Version, stdout: &str) -> Outcome {
let events: Vec<Event> = json_lines(stdout).collect();
Outcome {
text: events
.iter()
.rev()
.find_map(|event| match event {
Event::Text { part } => Some(part.text.clone()),
_ => None,
})
.unwrap_or_default(),
tool_calls: events
.iter()
.filter_map(|event| match event {
Event::ToolUse { part } => Some(part.tool.clone()),
_ => None,
})
.collect(),
usage: events
.iter()
.filter_map(|event| match event {
Event::StepFinish { part } => Some(Usage {
input_tokens: part.tokens.input,
output_tokens: part.tokens.output,
}),
_ => None,
})
.fold(Usage::default(), |total, step| total + step),
errors: events
.iter()
.filter_map(|event| match event {
Event::Error { error } => Some(
error
.data
.as_ref()
.and_then(|data| data.message.clone())
.unwrap_or_else(|| error.name.clone()),
),
_ => None,
})
.collect(),
exit_code: None,
}
}
}

56
litellm-rust/crates/testkit/src/error.rs Normal file
View file

@ -0,0 +1,56 @@
use std::io;
use std::path::PathBuf;
use thiserror::Error;
use crate::Wire;
#[derive(Debug, Error)]
pub enum Error {
#[error("unsupported target {0}")]
UnsupportedTarget(String),
#[error("{0} is not a plain x.y.z release version")]
InvalidVersion(String),
#[error("request to {url} failed")]
Request {
url: String,
#[source]
source: reqwest::Error,
},
#[error("{url} answered with status {status}")]
Status { url: String, status: u16 },
#[error("release metadata at {url} is malformed")]
Metadata {
url: String,
#[source]
source: serde_json::Error,
},
#[error("release has no asset named {0}")]
AssetNotFound(String),
#[error("release publishes no sha256 for {0}")]
MissingChecksum(String),
#[error("sha256 mismatch for {asset}: expected {expected}, got {actual}")]
ChecksumMismatch {
asset: String,
expected: String,
actual: String,
},
#[error("archive does not contain {0}")]
ArchiveMemberNotFound(String),
#[error("archive is unreadable")]
Archive(#[source] io::Error),
#[error("zip archive is unreadable")]
Zip(#[from] zip::result::ZipError),
#[error("{binary} reports version '{reported}', expected {expected}")]
VersionMismatch {
binary: PathBuf,
expected: String,
reported: String,
},
#[error("{agent} cannot talk to the gateway over {wire:?}")]
UnsupportedWire { agent: &'static str, wire: Wire },
#[error("agent did not finish within {0:?}")]
Timeout(std::time::Duration),
#[error("io failure")]
Io(#[from] io::Error),
}

52
litellm-rust/crates/testkit/src/install/archive.rs Normal file
View file

@ -0,0 +1,52 @@
use std::io::{Cursor, Read};
use flate2::read::GzDecoder;
use sha2::{Digest, Sha256};
use super::release::Packaging;
use crate::Error;
pub(crate) fn verify_sha256(asset: &str, expected: &str, bytes: &[u8]) -> Result<(), Error> {
let actual = format!("{:x}", Sha256::digest(bytes));
if actual.eq_ignore_ascii_case(expected) {
return Ok(());
}
Err(Error::ChecksumMismatch {
asset: asset.to_owned(),
expected: expected.to_owned(),
actual,
})
}
pub(crate) fn extract_binary(packaging: &Packaging, bytes: &[u8]) -> Result<Vec<u8>, Error> {
match packaging {
Packaging::Bare => Ok(bytes.to_vec()),
Packaging::TarGz { member } => extract_tar_gz(member, bytes),
Packaging::Zip { member } => extract_zip(member, bytes),
}
}
fn extract_tar_gz(member: &str, bytes: &[u8]) -> Result<Vec<u8>, Error> {
let mut archive = tar::Archive::new(GzDecoder::new(bytes));
for entry in archive.entries().map_err(Error::Archive)? {
let mut entry = entry.map_err(Error::Archive)?;
let path = entry.path().map_err(Error::Archive)?;
if path.file_name().is_some_and(|name| name == member) {
let mut binary = Vec::new();
entry.read_to_end(&mut binary).map_err(Error::Archive)?;
return Ok(binary);
}
}
Err(Error::ArchiveMemberNotFound(member.to_owned()))
}
fn extract_zip(member: &str, bytes: &[u8]) -> Result<Vec<u8>, Error> {
let mut archive = zip::ZipArchive::new(Cursor::new(bytes))?;
let mut file = archive.by_name(member).map_err(|error| match error {
zip::result::ZipError::FileNotFound => Error::ArchiveMemberNotFound(member.to_owned()),
other => Error::Zip(other),
})?;
let mut binary = Vec::new();
file.read_to_end(&mut binary).map_err(Error::Archive)?;
Ok(binary)
}

55
litellm-rust/crates/testkit/src/install/fetch.rs Normal file
View file

@ -0,0 +1,55 @@
use std::future::Future;
use crate::Error;
pub trait Fetch: Sync {
fn get(&self, url: &str) -> impl Future<Output = Result<Vec<u8>, Error>> + Send;
}
pub struct HttpFetch {
client: reqwest::Client,
github_token: Option<String>,
}
impl HttpFetch {
pub fn new(github_token: Option<String>) -> Self {
Self {
client: reqwest::Client::new(),
github_token,
}
}
pub fn from_env() -> Self {
Self::new(std::env::var("GITHUB_TOKEN").ok())
}
}
impl Fetch for HttpFetch {
async fn get(&self, url: &str) -> Result<Vec<u8>, Error> {
let request = self
.client
.get(url)
.header("user-agent", "litellm-testkit")
.header("accept", "application/json, application/octet-stream");
let request = match (
&self.github_token,
url.starts_with("https://api.github.com/"),
) {
(Some(token), true) => request.bearer_auth(token),
_ => request,
};
let request_error = |source| Error::Request {
url: url.to_owned(),
source,
};
let response = request.send().await.map_err(request_error)?;
let status = response.status();
if !status.is_success() {
return Err(Error::Status {
url: url.to_owned(),
status: status.as_u16(),
});
}
Ok(response.bytes().await.map_err(request_error)?.to_vec())
}
}

118
litellm-rust/crates/testkit/src/install/mod.rs Normal file
View file

@ -0,0 +1,118 @@
mod archive;
mod fetch;
pub(crate) mod release;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::atomic::{AtomicU64, Ordering};
use semver::Version;
use tokio::fs;
use tokio::process::Command;
use crate::{Error, Install, Target};
use archive::{extract_binary, verify_sha256};
static STAGING_COUNTER: AtomicU64 = AtomicU64::new(0);
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Installed {
pub version: Version,
pub binary: PathBuf,
}
pub struct Installer<F> {
fetch: F,
cache_root: PathBuf,
target: Target,
}
impl<F: Fetch> Installer<F> {
pub fn new(fetch: F, cache_root: impl Into<PathBuf>, target: Target) -> Self {
Self {
fetch,
cache_root: cache_root.into(),
target,
}
}
pub async fn install(
&self,
agent: &impl Install,
version: &Version,
) -> Result<Installed, Error> {
validate_release(version)?;
let dir = self
.cache_root
.join(agent.binary())
.join(version.to_string());
let binary = dir.join(agent.binary());
let installed = Installed {
version: version.clone(),
binary: binary.clone(),
};
if fs::try_exists(&binary).await? && probe_version(&binary, version).await.is_ok() {
return Ok(installed);
}
let release = agent.release(&self.fetch, version, self.target).await?;
let archive = self.fetch.get(&release.url).await?;
verify_sha256(&release.asset, &release.sha256, &archive)?;
let contents = extract_binary(&release.packaging, &archive)?;
fs::create_dir_all(&dir).await?;
let staging = dir.join(format!(
".{}.{}.{}.partial",
agent.binary(),
std::process::id(),
STAGING_COUNTER.fetch_add(1, Ordering::Relaxed)
));
fs::write(&staging, contents).await?;
fs::set_permissions(&staging, std::fs::Permissions::from_mode(0o755)).await?;
fs::rename(&staging, &binary).await?;
match probe_version(&binary, version).await {
Ok(()) => Ok(installed),
Err(error) => {
fs::remove_file(&binary).await?;
Err(error)
}
}
}
}
fn validate_release(version: &Version) -> Result<(), Error> {
if version.pre.is_empty() && version.build.is_empty() {
return Ok(());
}
Err(Error::InvalidVersion(version.to_string()))
}
async fn probe_version(binary: &Path, expected: &Version) -> Result<(), Error> {
let home = std::env::temp_dir();
let output = Command::new(binary)
.arg("--version")
.env_clear()
.env("HOME", home)
.env("DISABLE_AUTOUPDATER", "1")
.stdin(Stdio::null())
.output()
.await?;
let stdout = String::from_utf8_lossy(&output.stdout);
if stdout
.split_whitespace()
.filter_map(|token| Version::parse(token).ok())
.any(|reported| &reported == expected)
{
return Ok(());
}
Err(Error::VersionMismatch {
binary: binary.to_owned(),
expected: expected.to_string(),
reported: stdout.trim().to_owned(),
})
}
pub use fetch::{Fetch, HttpFetch};
pub use release::{Packaging, Release};

65
litellm-rust/crates/testkit/src/install/release.rs Normal file
View file

@ -0,0 +1,65 @@
use serde::Deserialize;
use crate::{Error, Fetch};
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Packaging {
Bare,
TarGz { member: String },
Zip { member: String },
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Release {
pub asset: String,
pub url: String,
pub sha256: String,
pub packaging: Packaging,
}
#[derive(Deserialize)]
struct GithubRelease {
assets: Vec<GithubAsset>,
}
#[derive(Deserialize)]
struct GithubAsset {
name: String,
digest: Option<String>,
browser_download_url: String,
}
pub(crate) async fn github_release(
fetch: &impl Fetch,
releases_url: &str,
tag: &str,
asset_name: &str,
packaging: Packaging,
) -> Result<Release, Error> {
let url = format!("{releases_url}/{tag}");
let release: GithubRelease = parse(&url, &fetch.get(&url).await?)?;
let asset = release
.assets
.into_iter()
.find(|asset| asset.name == asset_name)
.ok_or_else(|| Error::AssetNotFound(asset_name.to_owned()))?;
let sha256 = asset
.digest
.as_deref()
.and_then(|digest| digest.strip_prefix("sha256:"))
.ok_or_else(|| Error::MissingChecksum(asset_name.to_owned()))?
.to_owned();
Ok(Release {
asset: asset.name,
url: asset.browser_download_url,
sha256,
packaging,
})
}
pub(crate) fn parse<T: for<'de> Deserialize<'de>>(url: &str, body: &[u8]) -> Result<T, Error> {
serde_json::from_slice(body).map_err(|source| Error::Metadata {
url: url.to_owned(),
source,
})
}

15
litellm-rust/crates/testkit/src/lib.rs Normal file
View file

@ -0,0 +1,15 @@
mod agent;
mod error;
mod install;
mod session;
mod target;
pub use agent::{
Agent, ClaudeCode, Codex, Configure, Drive, Install, LaunchSpec, Opencode, Outcome, Prompt,
Settings, Usage, Wire,
};
pub use error::Error;
pub use install::{Fetch, HttpFetch, Installed, Installer, Packaging, Release};
pub use semver::Version;
pub use session::Session;
pub use target::{Arch, Os, Target};

76
litellm-rust/crates/testkit/src/session.rs Normal file
View file

@ -0,0 +1,76 @@
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::process::Stdio;
use std::time::Duration;
use semver::Version;
use tokio::process::Command;
use tokio::time::timeout;
use crate::{Configure, Drive, Error, Installed, Outcome, Prompt, Settings};
const STDERR_LIMIT_CHARS: usize = 2000;
pub struct Session {
binary: PathBuf,
home: PathBuf,
version: Version,
settings: Settings,
env: BTreeMap<String, String>,
}
impl Session {
pub fn prepare(
agent: &impl Configure,
installed: &Installed,
settings: Settings,
home: impl Into<PathBuf>,
) -> Result<Self, Error> {
let home = home.into();
let spec = agent.configure(&installed.version, &settings, &home)?;
spec.write_files(&home)?;
Ok(Self {
binary: installed.binary.clone(),
home,
version: installed.version.clone(),
settings,
env: spec.env,
})
}
pub async fn run(
&self,
agent: &impl Drive,
prompt: &Prompt,
limit: Duration,
) -> Result<Outcome, Error> {
let child = Command::new(&self.binary)
.args(agent.args(&self.version, &self.settings, prompt))
.env_clear()
.env("PATH", "/usr/bin:/bin")
.envs(&self.env)
.current_dir(&self.home)
.stdin(Stdio::null())
.kill_on_drop(true)
.output();
let output = timeout(limit, child)
.await
.map_err(|_| Error::Timeout(limit))??;
let parsed = agent.parse(&self.version, &String::from_utf8_lossy(&output.stdout));
let failed_silently = !output.status.success() && parsed.errors.is_empty();
Ok(Outcome {
errors: if failed_silently {
vec![
String::from_utf8_lossy(&output.stderr)
.chars()
.take(STDERR_LIMIT_CHARS)
.collect(),
]
} else {
parsed.errors
},
exit_code: output.status.code(),
..parsed
})
}
}

69
litellm-rust/crates/testkit/src/target.rs Normal file
View file

@ -0,0 +1,69 @@
use target_lexicon::{Architecture, Environment, OperatingSystem, Triple};
use crate::Error;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Os {
Macos,
Linux,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Arch {
Aarch64,
X86_64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Target {
pub os: Os,
pub arch: Arch,
pub musl: bool,
}
impl Target {
pub fn host() -> Result<Self, Error> {
Self::try_from(&Triple::host())
}
pub(crate) const fn os_name(self) -> &'static str {
match self.os {
Os::Macos => "darwin",
Os::Linux => "linux",
}
}
pub(crate) const fn arch_name(self) -> &'static str {
match self.arch {
Arch::Aarch64 => "arm64",
Arch::X86_64 => "x64",
}
}
pub(crate) const fn musl_suffix(self) -> &'static str {
if self.musl { "-musl" } else { "" }
}
}
impl TryFrom<&Triple> for Target {
type Error = Error;
fn try_from(triple: &Triple) -> Result<Self, Error> {
let unsupported = || Error::UnsupportedTarget(triple.to_string());
let os = match triple.operating_system {
OperatingSystem::Darwin(_) | OperatingSystem::MacOSX(_) => Os::Macos,
OperatingSystem::Linux => Os::Linux,
_ => return Err(unsupported()),
};
let arch = match triple.architecture {
Architecture::Aarch64(_) => Arch::Aarch64,
Architecture::X86_64 => Arch::X86_64,
_ => return Err(unsupported()),
};
Ok(Self {
os,
arch,
musl: triple.environment == Environment::Musl,
})
}
}

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