scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_ocr_custom_pricing

Browse source
This commit is contained in:
mateo-berri 2026-09-12 15:42:59 -07:00
commit d7158ba795
749 changed files with 66232 additions and 13145 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
.github/ci-coverage-allowlist.yml vendored
View file

@ -106,11 +106,6 @@ dockerfiles:
and lint workflows already exercise that output, so building the image adds no signal about it
paths:
- ui/Dockerfile
- reason: >-
The Rust gateway ships as its own chart and package with a separate release pipeline, so its
image is not part of this repo's Python image set
paths:
- litellm-rust/crates/ai-gateway/Dockerfile
- reason: >-
An example image under cookbook/ that is documentation rather than a shipped artifact
paths:

2
.github/e2e-stack/down.sh vendored
View file

@ -10,7 +10,7 @@ for pid_file in "${STACK_DIR}"/pids/*.pid; do
rm -f "${pid_file}"
done
for container in e2e-nginx e2e-valkey e2e-jaeger e2e-postgres; do
for container in e2e-nginx e2e-keycloak e2e-valkey e2e-jaeger e2e-postgres; do
docker rm -f "${container}" >/dev/null 2>&1
done

1
.github/e2e-stack/select_tests.py vendored
View file

@ -11,6 +11,7 @@ UNSUPPORTED: Final = re.compile(
)
HARNESS: Final = re.compile(
r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"
r"|^tests/e2e/idp_realm\.json$"
r"|^tests/e2e/gateway/"
r"|^\.github/e2e-stack/"
r"|^\.github/workflows/test-e2e-changed\.yml$"

45
.github/e2e-stack/start-idp.sh vendored Normal file
View file

@ -0,0 +1,45 @@
#!/usr/bin/env bash
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
KEYCLOAK_IMAGE="${E2E_KEYCLOAK_IMAGE:-quay.io/keycloak/keycloak@sha256:ff4257d0d64efbe99ed1ddfaf07765cc3c36dc7518bf8324d41961327f441c54}"
KEYCLOAK_PORT="${E2E_KEYCLOAK_PORT:-8081}"
POSTGRES_IMAGE="${E2E_POSTGRES_IMAGE:-postgres:16.6}"
: "${DATABASE_HOST:?}" "${DATABASE_PORT:?}" "${DATABASE_USER:?}" "${DATABASE_PASSWORD:?}" "${DATABASE_NAME:?}"
DB_HOST="${DATABASE_HOST}"
DB_NETWORK_ARGS=(--network bridge)
IDP_NETWORK_ARGS=(-p "127.0.0.1:${KEYCLOAK_PORT}:${KEYCLOAK_PORT}")
if [[ "$(uname)" == "Linux" ]]; then
DB_NETWORK_ARGS=(--network host)
IDP_NETWORK_ARGS=(--network host)
elif [[ "${DB_HOST}" == "127.0.0.1" || "${DB_HOST}" == "localhost" ]]; then
DB_HOST=host.docker.internal
fi
docker run --rm "${DB_NETWORK_ARGS[@]}" -e "PGPASSWORD=${DATABASE_PASSWORD}" \
"${POSTGRES_IMAGE}" psql -h "${DB_HOST}" -p "${DATABASE_PORT}" \
-U "${DATABASE_USER}" -d "${DATABASE_NAME}" -v ON_ERROR_STOP=1 \
-c 'CREATE SCHEMA IF NOT EXISTS keycloak' >/dev/null
docker rm -f e2e-keycloak >/dev/null 2>&1 || true
docker run -d --name e2e-keycloak "${IDP_NETWORK_ARGS[@]}" --memory 1536m \
-v "${REPO_ROOT}/tests/e2e/idp_realm.json:/opt/keycloak/data/import/realm.json:ro" \
-e KC_DB=postgres -e "KC_DB_URL_HOST=${DB_HOST}" -e "KC_DB_URL_PORT=${DATABASE_PORT}" \
-e "KC_DB_URL_DATABASE=${DATABASE_NAME}" -e KC_DB_SCHEMA=keycloak \
-e "KC_DB_USERNAME=${DATABASE_USER}" -e "KC_DB_PASSWORD=${DATABASE_PASSWORD}" \
-e KC_DB_POOL_INITIAL_SIZE=2 -e KC_DB_POOL_MIN_SIZE=2 -e KC_DB_POOL_MAX_SIZE=10 \
-e "KC_HTTP_PORT=${KEYCLOAK_PORT}" -e KC_BOOTSTRAP_ADMIN_USERNAME=admin \
-e KC_BOOTSTRAP_ADMIN_PASSWORD=e2e-ephemeral-idp-not-a-secret \
"${KEYCLOAK_IMAGE}" start-dev --import-realm >/dev/null
deadline=$((SECONDS + ${E2E_KEYCLOAK_STARTUP_TIMEOUT:-300}))
until curl -fsS --connect-timeout 2 --max-time 3 \
"http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e/.well-known/openid-configuration" >/dev/null 2>&1; do
if ((SECONDS >= deadline)); then
echo 'e2e-stack: timed out waiting for the Keycloak realm' >&2
exit 1
fi
sleep 2
done
echo 'e2e-stack: Keycloak realm is up'

9
.github/e2e-stack/up.sh vendored
View file

@ -25,6 +25,7 @@ DATABASE_PASSWORD="${E2E_DATABASE_PASSWORD:-dbpassword9090}"
DATABASE_NAME="${E2E_DATABASE_NAME:-litellm}"
JAEGER_OTLP_PORT="${E2E_JAEGER_OTLP_PORT:-4318}"
JAEGER_QUERY_PORT="${E2E_JAEGER_QUERY_PORT:-16686}"
KEYCLOAK_PORT="${E2E_KEYCLOAK_PORT:-8081}"
MASTER_KEY="${LITELLM_MASTER_KEY:-sk-e2e-$(openssl rand -hex 16)}"
@ -124,6 +125,9 @@ SERVER_ENV=(
"OTEL_EXPORTER_OTLP_ENDPOINT=http://127.0.0.1:${JAEGER_OTLP_PORT}"
"SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem"
"PYTHONPATH=${REPO_ROOT}"
"JWT_PUBLIC_KEY_URL=http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e/protocol/openid-connect/certs"
"JWT_ISSUER=http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e"
"JWT_AUDIENCE=litellm-e2e"
)
if [[ -n "${VERTEXAI_CREDENTIALS:-}" ]]; then
printf '%s' "${VERTEXAI_CREDENTIALS}" > "${STACK_DIR}/vertex-adc.json"
@ -132,6 +136,8 @@ fi
cd "${REPO_ROOT}"
env "${SERVER_ENV[@]}" "E2E_KEYCLOAK_PORT=${KEYCLOAK_PORT}" bash .github/e2e-stack/start-idp.sh
log "running migrations"
env "${SERVER_ENV[@]}" uv run --no-sync python migrations/run.py >"${LOGS_DIR}/migrations.log" 2>&1
@ -200,6 +206,9 @@ LITELLM_MASTER_KEY=${MASTER_KEY}
REDIS_HOST=127.0.0.1
REDIS_PORT=${REDIS_PORT}
E2E_OTEL_QUERY_URL=http://127.0.0.1:${JAEGER_QUERY_PORT}
E2E_KEYCLOAK_URL=http://127.0.0.1:${KEYCLOAK_PORT}
E2E_KEYCLOAK_ADMIN_USER=admin
E2E_KEYCLOAK_ADMIN_PASSWORD=e2e-ephemeral-idp-not-a-secret
SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem
DATABASE_URL=postgresql://${DATABASE_USER}:${DATABASE_PASSWORD}@${DATABASE_HOST}:${DATABASE_PORT}/${DATABASE_NAME}
EOF

2
.github/pull_request_template.md vendored
View file

@ -127,6 +127,8 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
- Low: anything else worth noting: naming, cleanup, an edge case nobody hits
Nest bullets as deep as helps: hierarchy beats one long line when it makes things clearer to a
human reader
If you assumed something instead of testing it, e.g. "only reproduces with X on" or "no
user-observable behavior difference", list it here too with what breaks if it is wrong
Leave this section empty if there are none -->
## QA runbook

4
.github/scripts/verify_linux_native_wheel.py vendored
View file

@ -205,7 +205,7 @@ def main(
native_module: Final = load_native_module(native_path)
native_module_loads: Final = native_module is not None
panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
native_size_limit: Final = 20_000_000
native_size_limit: Final = 25_000_000
native_size_within_limit: Final = native_member.file_size <= native_size_limit
validations: Final = (
(f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG),
@ -222,7 +222,7 @@ def main(
("Python extension entry point is present", extension_entry_point_present),
("Native module loads", native_module_loads),
("Production module omits the panic test hook", panic_test_hook_absent),
("Native extension does not exceed 20 MB", native_size_within_limit),
("Native extension does not exceed 25 MB", native_size_within_limit),
("Wheel contents are valid", not unexpected_members),
)

73
.github/workflows/ai-gateway-image.yml vendored Normal file
View file

@ -0,0 +1,73 @@
name: ai-gateway image
on:
push:
paths:
- "litellm-rust/**"
- "litellm/**"
- "enterprise/**"
- "litellm-proxy-extras/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/workflows/ai-gateway-image.yml"
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "litellm-rust/**"
- "litellm/**"
- "enterprise/**"
- "litellm-proxy-extras/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/workflows/ai-gateway-image.yml"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
ai-gateway-image:
name: ai-gateway release image
runs-on: ubuntu-latest
timeout-minutes: 60
permissions:
contents: read
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Build the release image
run: docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway:${{ github.sha }} .
- name: Start the gateway and wait for readiness
env:
IMAGE: litellm-ai-gateway:${{ github.sha }}
run: |
docker run -d --name ai-gateway -p 4001:4001 \
-e LITELLM_MASTER_KEY=sk-ci-not-a-real-key \
-e OPENAI_API_KEY=sk-ci-not-a-real-key \
"$IMAGE"
for _ in $(seq 1 60); do
if curl -fsS http://127.0.0.1:4001/health/readiness; then
echo "gateway is serving readiness"
exit 0
fi
sleep 2
done
echo "gateway never became ready" >&2
docker logs ai-gateway >&2
exit 1
- name: Assert the gateway loaded the baked config
run: |
docker logs ai-gateway 2>&1 | tee gateway.log
grep 'via python config reader' gateway.log
- name: Stop the gateway
if: always()
run: docker rm -f ai-gateway || true

2
.github/workflows/test-code-quality.yml vendored
View file

@ -81,7 +81,7 @@ jobs:
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.py
- name: test_e2e_changed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py
- name: router_code_coverage
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py

5
.github/workflows/test-e2e-changed.yml vendored
View file

@ -27,6 +27,8 @@ jobs:
sparse-checkout: |
.github/e2e-stack
tests/e2e/access_control
tests/e2e/management/test_jwt_management_e2e.py
tests/e2e/other/test_jwt_auth_e2e.py
persist-credentials: false
ref: ${{ github.sha }}
@ -45,7 +47,8 @@ jobs:
--jq '.[] | select(.status != "removed") | .filename')"
gh api "repos/${REPO}/pulls/${PR_NUMBER}" --jq '.head.sha' | grep -Fxq "${HEAD_SHA}"
tests="$(printf '%s\n' "${files}" \
| python3 .github/e2e-stack/select_tests.py tests/e2e/access_control/test_*.py)"
| python3 .github/e2e-stack/select_tests.py tests/e2e/access_control/test_*.py \
tests/e2e/management/test_jwt_management_e2e.py tests/e2e/other/test_jwt_auth_e2e.py)"
echo "tests=${tests}" >> "${GITHUB_OUTPUT}"
if [ -n "${tests}" ]; then
echo "any=true" >> "${GITHUB_OUTPUT}"

103
.github/workflows/test-e2e-redis-chaos.yml vendored Normal file
View file

@ -0,0 +1,103 @@
name: "Redis Chaos E2E"
on:
workflow_dispatch:
workflow_call:
inputs:
ref:
description: "Commit SHA or ref to test. Defaults to the ref the workflow was triggered on"
required: false
type: string
permissions:
contents: read
jobs:
redis-chaos-e2e:
runs-on: ubuntu-latest-16-cores
timeout-minutes: 30
services:
postgres:
image: postgres:16.6@sha256:557fea37a744d5f4c8faab304b0a90858b53ab119735a88c131fd19dab802f36
env:
POSTGRES_USER: llmproxy
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U llmproxy"
--health-interval 5s
--health-timeout 5s
--health-retries 10
valkey:
image: valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa
ports:
- 6379:6379
options: >-
--health-cmd "valkey-cli ping"
--health-interval 5s
--health-timeout 5s
--health-retries 10
env:
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
LITELLM_MASTER_KEY: sk-redis-chaos-e2e
LITELLM_LOG: WARNING
JSON_LOGS: "true"
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
ref: ${{ inputs.ref || github.sha }}
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
uses: ./.github/actions/cache-cargo-build
- name: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --group e2e-dev --extra proxy
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Start a multi-worker proxy on the chaos config
run: |
nohup uv run --no-sync litellm --config tests/e2e/gateway/redis_chaos_ci_config.yml --port 4000 --num_workers 4 > proxy.log 2>&1 &
echo "E2E_PROXY_PID=$!" >> "$GITHUB_ENV"
echo "E2E_PROXY_LOG=$(pwd)/proxy.log" >> "$GITHUB_ENV"
for _ in $(seq 1 90); do
if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
exit 0
fi
sleep 2
done
echo "proxy never became live"
tail -n 100 proxy.log
exit 1
- name: Run the Redis chaos load test
env:
E2E_REDIS_CHAOS: "1"
LITELLM_PROXY_URL: http://localhost:4000
REDIS_HOST: 127.0.0.1
REDIS_PORT: "6379"
run: |
uv run --no-sync pytest tests/e2e/load/test_redis_chaos_e2e.py -v --tb=short -rA -s
- name: Show proxy log on failure
if: failure()
run: tail -n 300 proxy.log

3
.gitignore vendored
View file

@ -147,3 +147,6 @@ crash.*.log
ui/litellm-dashboard/out/
litellm.log
.coverage-rust
coverage-rust.xml

2
README.md
View file

@ -262,6 +262,8 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
}
```
For MCP OAuth, an upstream may advertise dynamic client registration but refuse requests with HTTP 401 or 403. If the provider requires a pre-registered OAuth app, configure its `credentials.client_id` and, when required, `credentials.client_secret` on the MCP server. This skips dynamic registration in the gateway sign-in flow. The provider must approve the app for MCP access; reaching its authorization page does not establish that login or tool calls will succeed
[**Docs: MCP Gateway**](https://docs.litellm.ai/docs/mcp)
</details>

136
enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py
View file

@ -11,10 +11,14 @@ import sys
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import functools
import configparser
import contextlib
import itertools
import re
import tempfile
from collections.abc import Generator, Iterator, Sequence
from contextvars import ContextVar
from typing import TYPE_CHECKING, ClassVar, Literal, Optional
from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
@ -433,12 +437,101 @@ _default_detect_secrets_config = {
"name": "ZendeskSecretKeyDetector",
"path": _custom_plugins_path + "/zendesk_secret_key.py",
},
{
"name": "CredentialKeywordDetector",
"path": _custom_plugins_path + "/credential_keyword.py",
},
{"name": "Base64HighEntropyString", "limit": 4.5},
{"name": "HexHighEntropyString", "limit": 3.0},
],
}
_CONFIG_SECTION: Final = "litellm-prompt"
_ASSIGNMENT_LINE: Final = re.compile(r"[^\s\[#;:=][^:=]*[:=]")
_SHELL_ASSIGNMENT: Final = re.compile(r"(?P<key>[^\s\[#;:=](?:[^:=]*[^\s:=])?)=(?P<value>\S+)")
_SHELL_OPERATORS: Final = ";&|"
_SHELL_TRAILER: Final = re.compile(r"\\|#.*|-*\w[\w.-]*=\S*")
_SCAN_SUFFIX: Final = ".py"
@contextlib.contextmanager
def _temp_file(text: str) -> Generator[str, None, None]:
temp_file: Final = tempfile.NamedTemporaryFile(suffix=_SCAN_SUFFIX, delete=False)
try:
temp_file.write(text.encode("utf-8"))
temp_file.close()
yield temp_file.name
finally:
temp_file.close()
os.remove(temp_file.name)
def _scan_lines(lines: Sequence[str]) -> frozenset[tuple[str, str]]:
from detect_secrets import SecretsCollection
secrets: Final = SecretsCollection()
with _temp_file("\n".join(lines)) as path:
secrets.scan_file(path)
return frozenset(
(found_secret.secret_value, found_secret.type)
for file in secrets.files
for found_secret in secrets[file]
if found_secret.secret_value is not None
)
def _classify_line(state: tuple[bool, str | None], numbered: tuple[int, str]) -> tuple[bool, str | None]:
open_option: Final = state[0]
number, line = numbered
stripped: Final = line.strip()
if not stripped or stripped[0] in "#;":
return open_option, None
shell_assignment: Final = _SHELL_ASSIGNMENT.match(stripped)
if shell_assignment is not None:
return True, f"{shell_assignment['key']}_{number}={shell_assignment['value']}"
assignment: Final = _ASSIGNMENT_LINE.match(stripped)
if assignment is not None:
return True, f"{assignment.group()[:-1].strip()}_{number}{stripped[assignment.end() - 1 :]}"
if line[0].isspace() and open_option:
return True, line
return False, None
def _parseable_lines(text: str) -> Iterator[str]:
states: Final = itertools.accumulate(enumerate(text.splitlines()), _classify_line, initial=(False, None))
return (line for _, line in states if line is not None)
def _lone_value(line: str) -> str | None:
tokens: Final = line.split()
if not tokens or '"' in tokens[0]:
return None
value: Final = tokens[0].rstrip(_SHELL_OPERATORS)
if len(tokens) == 1 or value != tokens[0] or _SHELL_TRAILER.fullmatch(tokens[1]) is not None:
return value
return None
def _quoted_assignments(text: str) -> tuple[str, ...]:
parser: Final = configparser.ConfigParser(interpolation=None)
parser.optionxform = str # pyright: ignore[reportAttributeAccessIssue] # configparser types optionxform as a method
parser.read_string(f"[{_CONFIG_SECTION}]\n" + "\n".join(_parseable_lines(text)))
return tuple(
f'{key} = "{value}"'
for section in parser
for key, values in parser.items(section)
for line in values.splitlines()
if (value := _lone_value(line)) is not None
)
class _ENTERPRISE_SecretDetection(CustomGuardrail):
# Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail
# path skips should_run_check and never sees data["prompt"]).
@ -449,35 +542,21 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
super().__init__(**kwargs)
def scan_message_for_secrets(self, message_content: str):
from detect_secrets import SecretsCollection
from detect_secrets.settings import transient_settings
temp_file = tempfile.NamedTemporaryFile(delete=False)
temp_file.write(message_content.encode("utf-8"))
temp_file.close()
secrets = SecretsCollection()
detect_secrets_config = (
self.user_defined_detect_secrets_config or _default_detect_secrets_config
)
with transient_settings(detect_secrets_config):
secrets.scan_file(temp_file.name)
os.remove(temp_file.name)
found: Final = _scan_lines(
(*message_content.splitlines(), *_quoted_assignments(message_content))
)
return [
{"type": found_secret.type, "value": found_secret.secret_value}
for file in sorted(secrets.files)
for found_secret in sorted(
secrets[file],
key=lambda secret: (
-len(secret.secret_value or ""),
secret.type,
secret.secret_value or "",
),
{"type": secret_type, "value": value}
for value, secret_type in sorted(
found, key=lambda pair: (-len(pair[0]), pair[1], pair[0])
)
if found_secret.secret_value is not None
]
def redact_text(self, text: str, source: str = "message") -> str:
@ -490,15 +569,16 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
if counts is not None:
for secret in detected_secrets:
counts[secret["type"]] = counts.get(secret["type"], 0) + 1
secret_types = [secret["type"] for secret in detected_secrets]
secret_types: Final = sorted(
dict.fromkeys(secret["type"] for secret in detected_secrets)
)
verbose_proxy_logger.warning(
f"Detected and redacted secrets in {source}: {secret_types}"
"Detected and redacted secrets in %s: %s", source, secret_types
)
return functools.reduce(
lambda redacted, secret: redacted.replace(secret["value"], "[REDACTED]"),
detected_secrets,
text,
pattern: Final = re.compile(
"|".join(re.escape(secret["value"]) for secret in detected_secrets)
)
return pattern.sub("[REDACTED]", text)
async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool:
if user_api_key_dict.permissions is not None:

63
enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/credential_keyword.py Normal file
View file

@ -0,0 +1,63 @@
import re
from collections.abc import Generator, Mapping
from string import punctuation
from typing import Final
from detect_secrets.plugins.keyword import (
QUOTES_REQUIRED_DENYLIST_REGEX_TO_GROUP,
KeywordDetector,
)
_CREDENTIAL_VALUE: Final = re.compile(r"[^\s()\[\]]+")
_ENVIRONMENT_REFERENCE: Final = re.compile(r"os\.environ/\w+", re.IGNORECASE)
_ENVIRONMENT_VARIABLE_NAME: Final = re.compile(r"[A-Z][A-Z0-9]*(?:_[A-Z0-9]+)+")
_LOWERCASE_WORD_SEQUENCE: Final = re.compile(r"[a-z]+(?:[-._/][a-z]+)+")
_ISO_8601_TIMESTAMP: Final = re.compile(
r"\d{4}-\d{2}-\d{2}(?:T\d{2}:\d{2}(?::\d{2}(?:\.\d+)?)?(?:Z|[+-]\d{2}:?\d{2})?)?"
)
_URL_WITHOUT_USERINFO_OR_QUERY: Final = re.compile(r"[A-Za-z][A-Za-z0-9+.-]*://[^\s@?]*")
_BENIGN_VALUES: Final = (
_ENVIRONMENT_REFERENCE,
_ENVIRONMENT_VARIABLE_NAME,
_LOWERCASE_WORD_SEQUENCE,
_ISO_8601_TIMESTAMP,
_URL_WITHOUT_USERINFO_OR_QUERY,
)
class CredentialKeywordDetector(KeywordDetector): # pyright: ignore[reportUntypedBaseClass] # detect_secrets ships no type information
secret_type = "Credential Keyword"
def __init__(self, minimum_length: int = 12, keyword_exclude: str | None = None) -> None:
if (
not isinstance(minimum_length, int) # pyright: ignore[reportUnnecessaryIsInstance] # the value comes from an operator's YAML
or minimum_length < 1
):
raise ValueError(f"minimum_length must be a positive integer, got {minimum_length!r}")
super().__init__(keyword_exclude=keyword_exclude)
self.minimum_length = minimum_length
def _is_credential(self, value: str) -> bool:
core: Final = value.strip(punctuation)
return (
len(value) >= self.minimum_length
and _CREDENTIAL_VALUE.fullmatch(value) is not None
and all(benign.fullmatch(core) is None for benign in _BENIGN_VALUES)
)
def analyze_string(
self,
string: str,
denylist_regex_to_group: Mapping[re.Pattern[str], int] | None = None,
) -> Generator[str, None, None]:
if self.keyword_exclude is not None and self.keyword_exclude.search(string):
return
regex_to_group: Final = (
QUOTES_REQUIRED_DENYLIST_REGEX_TO_GROUP if denylist_regex_to_group is None else denylist_regex_to_group
)
yield from (
match.group(group)
for regex, group in regex_to_group.items()
for match in regex.finditer(string)
if self._is_credential(match.group(group))
)

5
enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py
View file

@ -780,7 +780,10 @@ async def update_project(
# Handle budget updates
budget_fields = LiteLLM_BudgetTable.model_fields.keys()
budget_updates = {k: v for k, v in update_data.items() if k in budget_fields}
budget_updates = {
**{k: v for k, v in update_data.items() if k in budget_fields},
**({"max_budget": None} if "max_budget" in data.model_fields_set and data.max_budget is None else {}),
}
if budget_updates and existing_project.budget_id:
# Update existing budget

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.66"
version = "0.1.67"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.66"
version = "0.1.67"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

4
gateway/routes/allowlist.py
View file

@ -3,7 +3,8 @@
The gateway exposes the LLM data-plane surface: chat/completions, embeddings,
audio, batches, files, fine-tuning, rerank, ocr, rag, video, search, image,
responses, vector stores, passthrough providers, realtime websockets, MCP
tool-call endpoints, and operational endpoints (/health, /metrics).
tool-call endpoints, and operational endpoints (/health, /metrics, and the
/debug/memory/summary read of the serving worker's RSS).
Any path not listed here is dropped from the gateway process so management/UI
endpoints don't ride on the same pods.
@ -121,6 +122,7 @@ GATEWAY_EXACT_PATHS: frozenset[str] = frozenset(
"/docs/oauth2-redirect",
"/redoc",
"/test",
"/debug/memory/summary",
}
)

8
helm/litellm/templates/_helpers.tpl
View file

@ -257,6 +257,14 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
- name: DATABASE_SCHEMA
value: {{ .schema | quote }}
{{- end }}
{{- if .sslMode }}
- name: DATABASE_SSLMODE
value: {{ .sslMode | quote }}
{{- end }}
{{- if .sslRootCert }}
- name: DATABASE_SSLROOTCERT
value: {{ .sslRootCert | quote }}
{{- end }}
{{- if and .useIAMAuth .useAzureEntraAuth }}
{{- fail "database.writer.useIAMAuth and database.writer.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" }}
{{- end }}

11
helm/litellm/templates/ingress.yaml
View file

@ -89,7 +89,7 @@
at "/" Prefix would swallow the whole backend management API) instead of
adding to it.
*/}}
{{- $builtinPathKeys := list "/test|Exact" "/|Prefix" -}}
{{- $builtinPathKeys := list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix" -}}
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
@ -129,6 +129,8 @@ spec:
# --- Gateway data plane ---
# Exact /test only (see the $gatewayPrefixes comment above);
# /test/* MCP management endpoints fall to the backend catch-all.
# Exact /debug/memory/summary reads a serving worker's RSS (the e2e memory
# gate); the rest of /debug/* stays on the backend.
- path: /test
pathType: Exact
backend:
@ -136,6 +138,13 @@ spec:
name: {{ $gatewayName }}
port:
number: {{ $gatewayPort }}
- path: /debug/memory/summary
pathType: Exact
backend:
service:
name: {{ $gatewayName }}
port:
number: {{ $gatewayPort }}
{{- range $gatewayPrefixes }}
{{- $pathType := include "litellm.ingress.pathType" (dict "controller" $controller "path" . "pathType" "Prefix") }}
{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" . $pathType) }}

77
helm/litellm/tests/database_auth_tests.yaml
View file

@ -4,6 +4,7 @@ templates:
- gateway/configmap.yaml
- backend/deployment.yaml
- backend/configmap.yaml
- migrations-job.yaml
values:
- ./values/required.yaml
tests:
@ -67,6 +68,82 @@ tests:
value: "true"
any: true
- it: emits no TLS env by default
template: gateway/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLMODE
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLROOTCERT
any: true
- it: writer sslMode and sslRootCert reach gateway and backend as DATABASE_SSLMODE and DATABASE_SSLROOTCERT
templates:
- gateway/deployment.yaml
- backend/deployment.yaml
set:
database.writer.useIAMAuth: true
database.writer.sslMode: verify-full
database.writer.sslRootCert: /etc/ssl/certs/ca-certificates.crt
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLMODE
value: verify-full
any: true
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLROOTCERT
value: /etc/ssl/certs/ca-certificates.crt
any: true
- it: writer sslMode and sslRootCert reach the collector sidecar and the migrations job, which dial Postgres themselves
set:
gateway.collector.enabled: true
database.connectionPool.enabled: true
database.writer.sslMode: verify-full
database.writer.sslRootCert: /etc/ssl/certs/ca-certificates.crt
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: collector
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_SSLMODE
value: verify-full
any: true
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_SSLROOTCERT
value: /etc/ssl/certs/ca-certificates.crt
any: true
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLMODE
value: verify-full
any: true
template: migrations-job.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLROOTCERT
value: /etc/ssl/certs/ca-certificates.crt
any: true
template: migrations-job.yaml
- it: writer rejects both token sources at once
template: gateway/deployment.yaml
set:

10
helm/litellm/tests/ingress_controller_tests.yaml
View file

@ -97,6 +97,16 @@ tests:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- contains:
path: spec.rules[0].http.paths
content:
path: /debug/memory/summary
pathType: Exact
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- equal:
path: spec.rules[0].http.paths[-1]
value:

11
helm/litellm/tests/ingress_extra_paths_tests.yaml
View file

@ -288,6 +288,17 @@ tests:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path /test with pathType Exact is already routed by this chart, and a duplicate would take it over rather than add to it"
- it: rejects an entry that would take over the exact /debug/memory/summary route
set:
ingress.enabled: true
ingress.extraPaths:
- path: /debug/memory/summary
pathType: Exact
service: backend
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path /debug/memory/summary with pathType Exact is already routed by this chart, and a duplicate would take it over rather than add to it"
- it: allows a built-in path under a different pathType, which is a distinct rule
set:
ingress.enabled: true

5
helm/litellm/values.yaml
View file

@ -208,6 +208,11 @@ database:
name: litellm-writer-secret
usernameKey: username
passwordKey: password
# libpq sslmode / sslrootcert applied to the writer and reader URLs (Prisma and the
# in-container PgBouncer); e.g. verify-full with /etc/ssl/certs/ca-certificates.crt for AWS RDS.
# sslRootCert on its own implies sslMode verify-full
sslMode: ""
sslRootCert: ""
# Optional read-replica routing. When `reader.host` is set, the proxy routes
# reads (find_*, count, group_by, query_raw/_first) to this endpoint while

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_add_autorouter_session_baseline_models/migration.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "baseline_models" JSONB NOT NULL DEFAULT '{}';

27
litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py
View file

@ -37,11 +37,13 @@ raised it above the deploy default keeps that larger budget for deploy unless
the deploy override says otherwise.
"""
import importlib.util
import math
import os
import shutil
import signal
import subprocess
import sys
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
@ -64,6 +66,7 @@ DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
BOOTSTRAP_ARG = "--version"
PRISMA_CONSOLE_SCRIPT = "prisma"
@dataclass(frozen=True)
@ -184,6 +187,28 @@ def _kill_process_group(process: "subprocess.Popen[str]") -> None:
return
def prisma_cli_available() -> bool:
"""Whether some way of running the Prisma CLI exists: the console script on PATH or the importable package."""
if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None:
return True
return importlib.util.find_spec(PRISMA_CONSOLE_SCRIPT) is not None
def resolve_prisma_argv(argv: Sequence[str]) -> tuple[str, ...]:
"""Route a bare ``prisma`` command through ``python -m prisma`` when the console script is not on PATH.
The console script and ``python -m prisma`` are the same entry point, but
only the module form survives an interpreter whose ``bin`` directory is
missing from PATH, which is how the proxy gets started under launchers and
init systems. Any other executable name is left untouched.
"""
if not argv or argv[0] != PRISMA_CONSOLE_SCRIPT:
return tuple(argv)
if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None:
return tuple(argv)
return (sys.executable, "-m", PRISMA_CONSOLE_SCRIPT, *argv[1:])
def run_prisma(
argv: Sequence[str],
*,
@ -200,7 +225,7 @@ def run_prisma(
text unless ``stdout``/``stderr`` say otherwise.
"""
with subprocess.Popen(
argv,
resolve_prisma_argv(argv),
env=env,
stdout=stdout,
stderr=stderr,

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1514,6 +1514,7 @@ model LiteLLM_AutoRouterSession {
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")
baseline_models Json @default("{}")
@@id([api_key, session_id, router_name])
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.96"
version = "0.4.97"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.96"
version = "0.4.97"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

29
litellm-rust/ADDING_A_PROVIDER.md
View file

@ -1,29 +0,0 @@
# Adding a provider / route to litellm-rust
Everything for a route lives in `crates/core/src/<route>/`; `crates/core/src/messages` is the reference. A host (the axum gateway, the Python bridge) only calls the route's entrypoint.
1. **Entrypoint** — `mod.rs`: `pub async fn <route>(request) -> CoreResult<Response>`, the Rust equivalent of `litellm.<route>()`, plus a `<route>_stream` variant when the route streams. It is the only thing a host touches.
2. **Transform contract** — `transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) with types in `types.rs`.
3. **Provider config** — `crates/core/src/providers/<provider>/<route>/transformation.rs`: implement that trait as a `const <PROVIDER>_<ROUTE>_CONFIG`, mirroring the Python provider tree. Add parity unit tests.
4. **Prepare + handler** — `prepare.rs` resolves provider/model, credentials, auth headers, and URL, then transforms the request; `handler.rs` performs the provider call through the shared client in `client.rs` and transforms the response.
## Coding standards
Before writing new logic, look for an existing base to extend. When a change is
“the same behavior for one more provider/endpoint/integration”, the codebase
almost always already has a shared abstraction for it (for example, provider
`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared
helpers in `litellm_core_utils/`, typed request/response models, or factory
functions). Find it first with a search, then add the new variant by inheriting
from or composing that base, overriding only what genuinely differs (model
name, parameter mapping, or auth).
Never copy an existing implementation and edit it in place, and never hand-roll
a parallel version of logic a base already provides. If you catch yourself
writing a second copy of a pattern that exists twice already, stop and extract a
base instead: put the shared shape in one place and make both call sites thin
variants of it. The test for a good abstraction is that adding the next provider
is a few declarative lines, not a new file of duplicated flow. Only diverge from
the base when behavior is genuinely different, and say so explicitly in the PR.
**Calling:** hosts invoke the core entrypoint — the Python bridge and the `ai-gateway` route service both call `litellm_core::messages::messages`. Never add a provider handler to `ai-gateway`. Register new modules in `lib.rs` / `mod.rs`, then run the commands under "Checks" in [CLAUDE.md](CLAUDE.md).

45
litellm-rust/AGENTS.md
View file

@ -1,45 +0,0 @@
# AGENTS.md
litellm-rust has six crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
## Crates
| Crate | Role |
|-------|------|
| litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. |
| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. |
| litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. |
| litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. |
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers, `litellm-token-counter`, and `litellm-python-interop`. The token counter and interop foundations depend on no LiteLLM domain crate.
## Where a route lives
A top-level LiteLLM call is a module under `crates/core/src/<route>/`, shaped like `messages`:
```
core/src/messages/
mod.rs # pub async fn messages(..) -> CoreResult<..> (+ messages_stream for SSE)
types.rs # request/response types, MessagesRequest
transformation.rs # the provider template trait
prepare.rs # provider resolution, auth headers, URL
handler.rs # the provider call
client.rs # the shared reqwest client
```
Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched.
Adding a crate: default to a module. A new crate requires a real trigger: separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these.
Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional.
## Style
All Rust in `litellm-rust/` follows the official Rust Style Guide:
https://doc.rust-lang.org/style-guide/
`rustfmt` implements its formatting by default, so run `cargo fmt` before committing; CI gates every PR on `cargo fmt --check`. Do not hand-format against rustfmt or add a `rustfmt.toml` that diverges from the default style.
Beyond formatting, follow the guide's naming and idiom conventions rustfmt cannot auto-apply: `snake_case` items/functions/modules, `UpperCamelCase` types/traits/variants, `SCREAMING_SNAKE_CASE` constants/statics (acronyms as one word, e.g. `HttpClient`), and the import grouping and item ordering it prescribes. See CLAUDE.md for the detailed version.

189
litellm-rust/CLAUDE.md
View file

@ -1,189 +0,0 @@
# CLAUDE.md
This file defines the rules for Rust work in LiteLLM.
## Provider Coding Standards
Before writing new logic, look for an existing base to extend. When a change is
“the same behavior for one more provider/endpoint/integration”, the codebase
almost always already has a shared abstraction for it (for example, provider
`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared
helpers in `litellm_core_utils/`, typed request/response models, or factory
functions). Find it first with a search, then add the new variant by inheriting
from or composing that base, overriding only what genuinely differs (model
name, parameter mapping, or auth).
Never copy an existing implementation and edit it in place, and never hand-roll
a parallel version of logic a base already provides. If you catch yourself
writing a second copy of a pattern that exists twice already, stop and extract a
base instead: put the shared shape in one place and make both call sites thin
variants of it. The test for a good abstraction is that adding the next provider
is a few declarative lines, not a new file of duplicated flow. Only diverge from
the base when behavior is genuinely different, and say so explicitly in the PR.
## Crates (see AGENTS.md)
`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call.
`litellm-config` is the config-loading boundary and returns resolved core types.
`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and
`litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop`
holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate
is a layer or shared foundation, not a route; add modules, not crates.
## Core Boundary
`litellm-core` owns the whole call. The Rust equivalent of `litellm.messages()`
is `litellm_core::messages::messages(request).await`: you call it, it does the
provider call, and you get a typed non-streaming response back.
Route-level Rust structure mirrors LiteLLM's Python responsibilities:
- `core/src/<route>/` owns the route end to end: the public entrypoint fn named
after the route in `mod.rs`, the request/response types (`types.rs`), the
provider template trait (`transformation.rs`), the provider/auth/URL
resolution (`prepare.rs`), the HTTP client (`client.rs`), and the handler that
performs the call (`handler.rs`). `core/src/messages` is the reference.
- `core/src/providers/<provider>/<route>/transformation.rs` owns the
provider-specific transform. For Anthropic Messages, this means
`core/src/providers/anthropic/messages/transformation.rs`.
- Handlers live in `core`, never in a host. `ai-gateway` must not contain a
route handler that talks to a provider; its axum route reads the HTTP request,
picks a deployment, and calls the `core` entrypoint. `python-bridge` marshals
Python objects and calls the same entrypoint.
Streaming keeps the same shape: the route entrypoint has a `<route>_stream`
variant in `core` that returns the upstream response so a host can splice it to
its own caller; the host still owns no provider logic.
Call-hook and lifecycle instrumentation, including phase timing, usage
accumulation, and callback payload construction, always lives in `core`.
Hosts feed observed events into core and dispatch the completed payloads through
their I/O logger; hosts must not own callback orchestration.
Allowed in `core`:
- The public entrypoint for a top-level LiteLLM call
- Request/response transforms and stream chunk normalization
- Provider resolution, auth header construction, and URL building
- The provider HTTP call itself, through a shared reused client with connect and
request timeouts
- Shared data types and validation errors
- Deterministic token/cost helper logic
Not allowed in `core`:
- Serving HTTP: axum routes, extractors, and transport concerns stay in the host
- Filesystem access
- Database access
- Config file reading and rollout state
- Logging callbacks, spend writes, or custom callbacks
- Global mutable runtime state
Env reads in `core` are limited to credential fallback inside a route's
`prepare.rs` (the `env_lookup` closure), mirroring what the Python SDK does when
no key is passed. Everything else config-shaped is resolved by the host and
passed in.
Routes still hosted in `ai-gateway` (`ocr`, `audio_transcription`, `realtime`)
predate this rule and are being moved into `core` route modules; do not add new
ones there, and prefer moving one when you touch it.
Python owns rollout state and fallback while Rust is being introduced. Rust
paths must be off by default until parity tests prove equivalence with Python.
A new provider/route may instead be implemented rust-only with no Python
reference; then the Python interface is a thin dispatch that calls Rust with no
fallback, and you state the rust-only choice explicitly in the PR. Either way
the Python side stays minimal (it only marshals inputs and calls the Rust
interface), never add a per-route feature flag, and never push provider
dispatch into `litellm/main.py`; put it in a thin dispatch class under
`litellm/llms/<provider>/<route>/`.
## Production Bar
Rust code in this workspace is held to a strict parity and robustness bar from
the first PR:
- Correctness parity is proven with tests. Do not rely on README claims or
manual inspection for a port that mirrors Python behavior.
- Every provider transform must have unit tests for supported-parameter
filtering, request body shape, response normalization, missing/null fields,
and bad-input errors.
- When Rust is exposed through Python, add Python tests that prove disabled,
enabled, and unavailable-bridge fallback behavior.
- Avoid panics on user/provider input. Return typed errors and let the host map
them to Python exceptions or HTTP responses.
- OCR handles documents that often contain personal data. Do not log document
contents, base64 payloads, provider response bodies, or secrets.
- Error messages must be useful but data-minimized. Truncate or sanitize any
upstream body before it crosses a host boundary.
- Treat empty or whitespace-only credentials, URLs, and config values as absent
at the host/config resolution layer.
- Preserve Python output shape intentionally. If a field is always serialized as
`null` for Python parity, leave a short comment explaining that parity choice.
## Network I/O Rules
These rules apply to every module that executes network I/O, whether it is a
`core` route handler or a host such as `ai-gateway`:
- Set connect and full-request timeouts. No unbounded waits.
- Reuse HTTP clients; do not construct clients per request.
- Prefer rustls TLS for portable Python wheels and Linux images unless there is
a documented reason not to.
- Add request IDs and structured tracing at the host layer, without logging OCR
document contents or secrets.
- Do not echo raw upstream response bodies to callers. Sanitize and bound them.
- Avoid `expect`/`unwrap` in server startup and request paths unless the panic is
impossible by construction and documented.
## Rust Style Guide
All Rust in `litellm-rust/` follows the official Rust Style Guide:
https://doc.rust-lang.org/style-guide/
`rustfmt` implements the guide's formatting rules by default, so the mechanical
side is enforced for you: run `cargo fmt` before committing and CI gates every
PR on `cargo fmt --check` (see Checks). Do not hand-format against rustfmt or add
a `rustfmt.toml` that diverges from the default style; the default style *is* the
guide.
The guide also covers conventions rustfmt cannot auto-apply; follow these too:
- Naming: `snake_case` for items, functions, and modules; `UpperCamelCase` for
types, traits, and enum variants; `SCREAMING_SNAKE_CASE` for constants and
statics; acronyms count as one word (`HttpClient`, not `HTTPClient`).
- Ordering and grouping the guide prescribes: imports grouped std / external /
crate-local, derives before other attributes, and consistent item order.
- Idioms the guide recommends over the formatter fighting you (e.g. prefer
restructuring an over-long expression rather than forcing an awkward wrap).
## Constants
Magic numbers and fixed strings go in a crate-level `constants.rs`, never
hardcoded inline — the Rust mirror of Python's `litellm/constants.py`.
- Each crate that needs them has `src/constants.rs` (declared `mod constants;`);
import from it (`use crate::constants::...`). Don't scatter `const` values at
the top of feature modules.
- An env-overridable tunable still lives in `constants.rs` as its `DEFAULT_*`
value; the env read (with fallback to that default) happens at the host/config
resolution layer, not in `core`/`providers`.
- Exception: a value that is purely local to one function and has no meaning
elsewhere may stay inline, but prefer `constants.rs` when in doubt.
## Checks
Run these before pushing Rust changes. The same checks run in GitHub Actions
for changes under `litellm-rust/`.
```bash
cd litellm-rust
cargo fmt --check
cargo clippy --workspace --all-targets -- -D warnings
cargo clippy -p litellm-core --all-targets --features bedrock-auth -- -D warnings
# the ai-gateway binary + server code is behind the `server` feature
cargo clippy -p litellm-ai-gateway --all-targets --all-features -- -D warnings
cargo test --workspace
cargo test -p litellm-core --features bedrock-auth
# the `auth`, `routes`, `state` and `realtime` tests only exist under `server`
cargo test -p litellm-ai-gateway --features server
```
When a Rust path is exposed through Python, add Python parity tests that compare
the existing Python output with the Rust-backed output.

760
litellm-rust/Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

7
litellm-rust/Cargo.toml
View file

@ -16,6 +16,7 @@ license = "MIT"
repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
bytes = "1"
tracing = "0.1"
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
litellm-core = { path = "crates/core" }
@ -41,8 +42,14 @@ 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"
gcp_auth = "0.12.7"
azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
criterion = "0.8.2"
veil = "0.3.0"
[profile.release]
opt-level = 3

56
litellm-rust/README.md
View file

@ -1,56 +0,0 @@
# LiteLLM Rust
This workspace contains the staged Rust implementation for LiteLLM.
`litellm-core` is the LiteLLM SDK in Rust: one entrypoint per top-level call
that makes the LLM call and hands back a typed response, the same shape as
`litellm.messages()` in Python.
```rust
let response = litellm_core::messages::messages(MessagesRequest {
model: "claude-sonnet-4-5",
body,
api_key: Some(key),
..
})
.await?;
```
Python continues to own configuration, retries, routing policy, logging,
callbacks, spend tracking, and customer plugins until each Rust path has parity
coverage and production evidence.
## Crates
| Crate | Role |
|-------|------|
| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. |
| litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. |
| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop.
## Layout
```text
crates/
core/ The SDK: route modules + provider transforms.
src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client
src/providers/anthropic/messages/transformation.rs
config/ Config loading and resolved deployments.
ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints.
python-interop/ Domain-neutral PyO3 conversion and GIL primitives.
python-bridge/ PyO3 API adapter for Python LiteLLM.
```
The folder shape follows the Python provider tree:
`core/src/providers/<provider>/<route>/transformation.rs`. The bridge exposes one
function per top-level route, mirroring the core entrypoints.
## Checks
Run the commands under "Checks" in [CLAUDE.md](CLAUDE.md) before pushing Rust
changes. That list is the single source of truth and matches what GitHub Actions
runs for changes under `litellm-rust/`.

53
litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md
View file

@ -1,53 +0,0 @@
# Provider coding standards (litellm-rust)
Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` (`core/src/messages`, `ANTHROPIC_MESSAGES_CONFIG`) is the reference: a route is a `core` module with a public entrypoint that makes the call and returns a typed response.
## Provider resolution
1. Always resolve the provider/model first with `get_custom_llm_provider` (`core/src/routing_utils/provider.rs`). Nothing downstream may branch on a raw model string.
2. Model/provider is resolved once, in `prepare.rs`, and passed down as typed fields. Don't re-resolve or re-parse it in transforms or handlers.
## Transforms and the base config
3. Every route defines a base config trait with `transform_request` + `transform_response` (+ `complete_url`, `supported_params`), living in `core/src/<route>/transformation.rs` (e.g. `AnthropicMessagesProviderConfig`, mirroring `OcrProviderConfig`).
4. Each provider implements that trait as a `const <PROVIDER>_<ROUTE>_CONFIG` in `core/src/providers/<provider>/<route>/transformation.rs`, mirroring the Python provider tree.
5. Individual configs implement only the request/response transforms. Shared behavior (param filtering, defaults) stays as trait default methods so future providers inherit existing logic instead of reimplementing it.
6. Prefer composition: a provider that extends another reuses the base trait's defaults or wraps another config; don't copy transform bodies between providers.
## Boundaries
7. Layers never cross: `core` = the call itself (entrypoint, types, transforms, provider resolution, auth headers, provider HTTP, lifecycle hooks); `ai-gateway` = serving HTTP/WS (routing, extractors, auth of *our* callers, streaming to the client); `python-bridge` = thin PyO3 adapter. Hosts call the core entrypoint; they never build a provider request.
8. Generic/route files contain zero provider-specific branches. A provider is one module under `core/src/providers/<provider>/<route>/`; a route is a module, never a new crate.
9. Route entry point stays thin: `core::<route>::<route>()` -> `prepare_*` -> handler (or `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing). Axum handlers validate and delegate to a service that calls the entrypoint; no business logic in them.
10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Config-shaped env reads happen at the host/config layer with the `DEFAULT_*` fallback defined in `constants.rs`; the only env read in `core` is the credential fallback in a route's `prepare.rs`.
## Types and errors
11. Typed contracts only: no bare `serde_json::Value` / `String` / `Vec<String>` as a transform input or output. Parse wire bytes into typed structs/enums at the host edge; a `type` discriminator is a typed field, not a raw string.
12. Model failures as values: return typed `CoreError`, don't panic. No `unwrap`/`expect`/`panic!` on user or provider input.
13. No mutation: build values in one shot (comprehensions/iterators, `collect`), prefer immutable bindings and owned typed structs over seeding-and-mutating.
14. Early returns over deep nesting; small focused files over god modules.
15. Preserve Python output shape intentionally. If a field is always serialized as `null` for parity, keep it and pin it with a test.
## Safety and data minimization
16. Never log request/response bodies, base64 payloads, document contents, or secrets. Truncate and bound any upstream body before it crosses a host boundary.
17. Treat empty/whitespace credentials, URLs, and config values as absent at the host resolution layer.
18. Network I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS.
## Tests and rollout
19. Every provider transform ships tests for: supported-param filtering, request body shape, response normalization, missing/null fields, bad input, and `*_match_python` fixture parity.
20. Lifecycle/hook tests cover hook order, success + failure callback payloads, pre-call guardrail blocking before any provider I/O, during-call body mutation, and provider-error mapping.
21. When a route has a Python reference implementation, the Rust path stays off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. A new provider/route may instead be implemented rust-only with no Python reference; then the Python interface is a thin dispatch to Rust with no fallback, and tests cover the rust-backed path plus the unavailable-bridge error. State the rust-only choice explicitly in the PR.
## Python bridge (SDK side)
22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust.
23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms/<provider>/<route>/` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method.
24. Do not add new feature flags unless explicitly requested. Reuse the existing LiteLLM Rust rollout mechanism (`litellm.rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_<ROUTE>`.
## Checks before push
25. Run, and keep green, the commands under "Checks" in `litellm-rust/CLAUDE.md`.
That list is the single source of truth and matches what GitHub Actions runs.

54
litellm-rust/crates/ai-gateway/AGENTS.md
View file

@ -1,54 +0,0 @@
# ai-gateway — folder architecture
The Axum server that fronts the Rust gateway. It owns transport + config + auth
only; deployment selection lives in `core::router`, and the LLM call itself
(transforms, auth headers, provider HTTP) lives behind a `core` route entrypoint
such as `litellm_core::messages::messages`. No provider handler lives here.
```
src/
main.rs # entrypoint: build AppState (router + master key), bind, serve
state.rs # AppState — shared Arc<Router> + master_key
auth/ # authentication as an axum extractor — added to handler args
mod.rs # RequireMasterKey: FromRequestParts, single master key (LITELLM_MASTER_KEY)
routes/ # one module per route, all matching the same template
AGENTS.md # ← the route template (read this before adding a route)
mod.rs # app(): merges every module's router()
health.rs # simple route (one file): router() + liveness/readiness
realtime/ # route with logic → axum surface + a no-axum service:
mod.rs # router() + handler + WS<->events adapter (the axum surface)
service.rs # business logic (select deployment, call provider) — no axum, testable
```
## Rules
- **Routes follow one template.** Each route module exposes
`pub fn router() -> Router<AppState>`; `routes/mod.rs` only merges them. Simple
routes are one file; non-trivial routes are a folder (`handler`/`service`/
`transport`). See `routes/AGENTS.md`.
- **Auth is an extractor.** Add `crate::auth::RequireMasterKey` to a handler's
args; it runs during extraction. Never re-implement the check per route.
- **Handlers are thin.** A handler validates and delegates to its `service`. No
business logic, no provider calls, no transforms in handlers.
- **Services call `core`, they don't reimplement it.** A `service` picks the
deployment and calls the `core` route entrypoint. Provider resolution, auth
headers, URL building, and the HTTP call are `core`'s job; a service that
builds a provider request itself is a bug (`routes/messages/service.rs` is
the reference).
- **State is shared and cheap to clone.** Long-lived handles live behind `Arc` in
`state.rs`; read env/config only in `main.rs` when building state.
## Auth (interim)
A single **master key** (`LITELLM_MASTER_KEY`), enforced by the
`auth::RequireMasterKey` extractor: any caller presenting it as
`Authorization: Bearer <key>` may invoke the gateway. Fails closed (500) when
unset; constant-time compare. The server binds `127.0.0.1` by default (`HOST` to
override). Full per-key auth + budgets/rate-limits are delegated to the Python
proxy in a later phase. Health routes don't add the extractor (unauthenticated).
## Python interop
Python-backed loading lives in `litellm-config` and is **load-time only**. The
gateway's `python-config` feature forwards to that crate. The realtime data path
never takes the GIL.

14
litellm-rust/crates/ai-gateway/ARCHITECTURE.md
View file

@ -1,14 +0,0 @@
# ai-gateway architecture
The Rust ai-gateway does LLM inference (realtime WebSocket). Spend tracking is an
API callback: it POSTs each finished session to the LiteLLM proxy, which records
spend and runs the usual callbacks.
```mermaid
flowchart LR
C[client] <--> G[Rust ai-gateway<br/>LLM inference]
G <--> O[OpenAI realtime]
G -. spend tracking callback .-> P[litellm proxy]
F[litellm-config<br/>load-time only] --> G
F -. Python backend .-> P
```

5
litellm-rust/crates/ai-gateway/Cargo.toml
View file

@ -13,6 +13,11 @@ name = "litellm-ai-gateway"
path = "src/main.rs"
required-features = ["server"]
[[bin]]
name = "trace-parity-gateway"
path = "src/bin/trace_parity_gateway.rs"
required-features = ["trace-parity"]
[dependencies]
tracing.workspace = true
litellm-core = { workspace = true, features = ["bedrock-auth"] }

41
litellm-rust/crates/ai-gateway/Dockerfile
View file

@ -14,15 +14,20 @@
# ---- Chef -------------------------------------------------------------------
# cargo-chef caches the dependency build so only the gateway crate recompiles on
# a source-only change. python3-dev is present in every rust stage because the
# `python-config` feature links libpython via pyo3 (even in the cook step).
FROM rust:1.90-slim-bookworm AS chef
# `python-config` feature links libpython via pyo3 (even in the cook step), and
# python3-pip builds the litellm wheel in the builder stage.
FROM rust:1.98-slim-bookworm AS chef
ENV PYO3_PYTHON=python3.11
# rustup reads rust-toolchain.toml from any parent of the working directory, so
# copying it in is what keeps every cargo call below on the repo's pinned
# channel rather than on whatever the base image happens to ship.
COPY rust-toolchain.toml /build/rust-toolchain.toml
WORKDIR /build/litellm-rust
RUN apt-get update \
&& apt-get install -y --no-install-recommends \
python3 python3-dev pkg-config libssl-dev clang \
python3 python3-dev python3-pip pkg-config libssl-dev clang \
&& rm -rf /var/lib/apt/lists/* \
&& cargo install cargo-chef --locked --version 0.1.77
WORKDIR /build/litellm-rust
# ---- Planner ----------------------------------------------------------------
# Produce the dependency recipe from the rust workspace manifests + Cargo.lock.
@ -43,6 +48,19 @@ RUN cargo chef cook --locked --release \
COPY litellm-rust/ .
RUN cargo build --locked --release -p litellm-ai-gateway --bin litellm-ai-gateway --features server,python-config
# The root pyproject builds with maturin against litellm-rust/crates/python-bridge,
# so the wheel is built here, next to the crate sources and the cargo toolchain,
# and the runtime stage installs the artifact instead of compiling anything.
# litellm[proxy] pins litellm-enterprise and litellm-proxy-extras to the versions
# in this repo, and those hit PyPI hours after every version bump merges, so both
# wheels are built from the repo too instead of being resolved from PyPI.
COPY pyproject.toml README.md LICENSE /build/
COPY litellm/ /build/litellm/
COPY enterprise/ /build/enterprise/
COPY litellm-proxy-extras/ /build/litellm-proxy-extras/
RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /build/dist \
/build /build/enterprise /build/litellm-proxy-extras
# ---- Runtime ----------------------------------------------------------------
# python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3
# 3.11 ABI so the embedded interpreter links and imports cleanly.
@ -56,11 +74,16 @@ RUN apt-get update \
WORKDIR /app
# Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so
# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. Copy the
# package + packaging metadata, then pip install the proxy extra.
COPY pyproject.toml README.md LICENSE ./
COPY litellm/ ./litellm/
RUN pip install --no-cache-dir ".[proxy]"
# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. The two
# sibling wheels come from the builder as well, so the pins in litellm[proxy]
# resolve against them and never wait on a PyPI publish.
COPY --from=builder /build/dist/*.whl /tmp/wheels/
RUN wheel="$(ls /tmp/wheels/litellm-*.whl)" \
&& pip install --no-cache-dir \
/tmp/wheels/litellm_enterprise-*.whl \
/tmp/wheels/litellm_proxy_extras-*.whl \
"${wheel}[proxy]" \
&& rm -rf /tmp/wheels
# The compiled gateway binary (pure-Rust realtime hot path; Python is load-time
# only).

11
litellm-rust/crates/ai-gateway/Dockerfile.dockerignore
View file

@ -9,19 +9,28 @@
# Strategy: ignore everything, then re-include only what the build needs:
# - litellm/ (pip install . needs the full package + proxy reader)
# - litellm-rust/ (the rust workspace; Cargo.lock + crate sources)
# - pyproject.toml / README.md / LICENSE (packaging metadata for pip install)
# - enterprise/ (litellm/proxy/enterprise symlinks into it; maturin walks it)
# - litellm-proxy-extras/ (built into a wheel alongside enterprise/ for litellm[proxy])
# - pyproject.toml / README.md / LICENSE (packaging metadata for the wheel build)
# - rust-toolchain.toml (the pinned channel every cargo call in the build uses)
*
# --- re-include the build inputs ---
!litellm/
!litellm-rust/
!enterprise/
!litellm-proxy-extras/
!pyproject.toml
!rust-toolchain.toml
!README.md
!LICENSE
# --- prune heavy / irrelevant subpaths back out of the re-included trees ---
# Rust build artifacts (huge; regenerated in the builder).
**/target/
# Committed python distribution artifacts; the wheel build does not read them.
enterprise/dist/
litellm-proxy-extras/dist/
# Python caches and compiled bytecode.
**/__pycache__/
**/*.pyc

55
litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md
View file

@ -1,55 +0,0 @@
# Realtime gateway benchmark — pool on/off
Measures what the gateway adds over talking to OpenAI's realtime WebSocket
directly, and what the pre-warmed connection pool removes. See
`../../src/routes/realtime/README.md` for how the pool works.
## Results
5000 calls / 500 concurrency, gateway at 10 instances, pool ON
(`REALTIME_POOL_SIZE=64`), upstream OpenAI `gpt-realtime`. Each leg run twice.
Times in **ms**. Phases per connection: **dial** = TCP+TLS+WS upgrade,
**session** = upgrade → `session.created` (the phase the pool removes),
**1st-audio** = `response.create` → first audio delta (OpenAI inference),
**total** = full wall-clock.
| metric | Direct OpenAI | Gateway (pool ON) | Overhead (ms) | vs OpenAI |
| ------------------ | ------------- | ----------------- | ------------- | ---------- |
| success rate (%) | 99.8 | 99.8 | — | — |
| dial p50 (ms) | 276 | 158 | −118 | **faster** |
| session p50 (ms) | 7 | 0 | −7 | **faster** |
| 1st-audio p50 (ms) | 440 | 664 | +224 | slower¹ |
| total p50 (ms) | 816 | 1010 | +194 | slower¹ |
| total p95 (ms) | 2152 | 1970 | −182 | **faster** |
| total p99 (ms) | 2692 | 2610 | −82 | **faster** |
The gateway is **faster than direct on 4 of 6 metrics**. The warm pool makes the
**session phase sub-millisecond** at the median — ~76% of connects hit the pool,
~70% had session < 1 ms. ¹ The two "slower" rows are not gateway overhead:
`1st-audio` is OpenAI's own inference time (the gateway only relays it), which ran
slower during the gateway legs and drags `total p50` with it.
**Pool OFF** (control, `REALTIME_POOL_SIZE=0`): session p50 was **367 ms** — the
fresh-dial overhead the pool removes.
## Reproduce
The load generator lives in a separate repo:
**https://github.com/ishaan-berri/litellm-realtime-bench**
```bash
git clone https://github.com/ishaan-berri/litellm-realtime-bench
cd litellm-realtime-bench && go build -o wsbench .
# Direct to OpenAI (baseline)
./wsbench -host api.openai.com -key "$OPENAI_API_KEY" -m gpt-realtime -n 5000 -c 500 -t 60
# Through the gateway — run once with pool ON, once with REALTIME_POOL_SIZE=0
./wsbench -host <gateway-host> -key "$LITELLM_MASTER_KEY" -m gpt-realtime -n 5000 -c 500 -t 60
```
Run the gateway with the env stand-in (`OPENAI_REALTIME_MODEL=gpt-realtime`,
`OPENAI_API_KEY`, `LITELLM_MASTER_KEY`, `REALTIME_POOL_SIZE`, `HOST=0.0.0.0`). At
500 concurrency over N instances, size the pool to `≈ 500 / N` per instance (64 was
used here for 10 instances). The bench repo's README covers running 500-concurrency
legs from a hosted multi-vCPU runner. **Never commit keys — pass them via `-key`.**

8
litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs
View file

@ -269,11 +269,15 @@ fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
fn core_error_kind(error: &Error) -> &'static str {
match error {
Error::Auth(_) | Error::MissingApiKey { .. } => "AuthError",
Error::Auth(_)
| Error::MissingApiKey { .. }
| Error::MissingAzureAiCredentials
| Error::MissingAzureDocumentIntelligenceCredentials
| Error::MissingReductoApiKey => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",
Error::MissingField(_) => "MissingField",
Error::MissingField(_) | Error::MissingDocumentUrl => "MissingField",
Error::Http { .. } => "HttpError",
Error::InvalidResponse(_) => "InvalidResponse",
Error::Network(_) => "NetworkError",

40
litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs Normal file
View file

@ -0,0 +1,40 @@
use std::io::Read;
use serde::Deserialize;
use serde_json::Value;
#[derive(Deserialize)]
struct Input {
model_alias: String,
provider_model: String,
api_base: String,
body: Value,
}
#[tokio::main]
async fn main() {
let mut input = String::new();
if let Err(error) = std::io::stdin().read_to_string(&mut input) {
fail(error);
}
let input: Input = match serde_json::from_str(&input) {
Ok(input) => input,
Err(error) => fail(error),
};
let result = litellm_ai_gateway::trace_parity::traced_messages_request(
input.model_alias,
input.provider_model,
input.api_base,
input.body,
)
.await;
match serde_json::to_string(&result) {
Ok(result) => println!("{result}"),
Err(error) => fail(error),
}
}
fn fail(error: impl std::fmt::Display) -> ! {
eprintln!("{error}");
std::process::exit(1)
}

6
litellm-rust/crates/ai-gateway/src/io/realtime.rs
View file

@ -15,6 +15,8 @@ use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::AuthError;
use litellm_core::auth::error::MissingCredential;
use litellm_core::error::Error;
use litellm_core::realtime::transformation::RealtimeProviderConfig;
use litellm_core::realtime::types::RealtimeEvent;
@ -32,8 +34,6 @@ use crate::io::tls::connect_upstream;
/// Environment variable holding the OpenAI API key (last-resort fallback).
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
/// Default **idle** timeout: if neither side sends a frame for this long, the
/// session is reaped. It resets on any activity, so it does not cap a healthy
/// (continuously streaming) session — it only frees a stalled one (e.g. a
@ -59,7 +59,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
.ok()
.filter(|key| !key.trim().is_empty())
})
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
.ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiRealtimeApiKey)))
}
/// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`.

97
litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
View file

@ -1,115 +1,28 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::AuthError;
use litellm_core::Error;
use litellm_core::auth::error::MissingCredential;
use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
use litellm_core::responses::types::ResponsesWsEvent;
use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
use crate::io::tls::connect_upstream;
use litellm_core::responses::websocket::{ResponsesUpstreamWs, connect_upstream};
use crate::constants::{
DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
};
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
pub type ResponsesUpstreamWs = WebSocketStream<MaybeTlsStream<TcpStream>>;
type UpstreamTx = SplitSink<ResponsesUpstreamWs, Message>;
type UpstreamRx = SplitStream<ResponsesUpstreamWs>;
#[derive(Clone)]
pub struct ResponsesWebSocketConnection {
socket: Arc<Mutex<Option<ResponsesUpstreamWs>>>,
}
impl ResponsesWebSocketConnection {
pub async fn connect_url(
url: &str,
headers: &HashMap<String, String>,
timeout: Option<Duration>,
) -> Result<Self, Error> {
let mut request = url
.into_client_request()
.map_err(|error| Error::Network(error.to_string()))?;
for (name, value) in headers {
let header_name = name
.parse::<HeaderName>()
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let header_value = HeaderValue::from_str(value)
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
request.headers_mut().insert(header_name, header_value);
}
let connect = connect_upstream(request);
let result = match timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
Error::Network("Responses WebSocket connection timed out".to_string())
})?,
None => connect.await,
};
let (socket, _) = result.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
},
other => Error::Network(other.to_string()),
})?;
Ok(Self {
socket: Arc::new(Mutex::new(Some(socket))),
})
}
pub async fn send_text(&self, text: String) -> Result<(), Error> {
let mut socket = self.socket.lock().await;
let Some(socket) = socket.as_mut() else {
return Err(Error::Network("Responses WebSocket is closed".to_string()));
};
socket
.send(Message::Text(text))
.await
.map_err(|error| Error::Network(error.to_string()))
}
pub async fn recv_text(&self) -> Result<Option<String>, Error> {
let mut socket_guard = self.socket.lock().await;
let Some(socket) = socket_guard.as_mut() else {
return Ok(None);
};
match socket.next().await {
Some(Ok(Message::Text(text))) => Ok(Some(text)),
Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
.map(Some)
.map_err(|error| Error::InvalidResponse(error.to_string())),
Some(Ok(Message::Close(_))) | None => Ok(None),
Some(Ok(_)) => Ok(None),
Some(Err(error)) => Err(Error::Network(error.to_string())),
}
}
pub async fn close(&self) -> Result<(), Error> {
let mut socket = self.socket.lock().await;
if let Some(socket) = socket.as_mut() {
socket
.close(None)
.await
.map_err(|error| Error::Network(error.to_string()))?;
}
*socket = None;
Ok(())
}
}
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
api_key
.map(str::trim)
@ -120,7 +33,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
.ok()
.filter(|value| !value.trim().is_empty())
})
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
.ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiResponsesApiKey)))
}
async fn dial_upstream(

529
litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs
View file

@ -1,529 +0,0 @@
use std::net::IpAddr;
use std::time::{Duration, Instant};
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use litellm_core::error::Error;
use litellm_core::ocr::transformation::OcrProviderConfig;
use reqwest::Url;
use serde_json::{Map, Value};
use litellm_core::providers::azure_ai::ocr::transformation::{
AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG,
};
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
use litellm_core::providers::reducto::ocr::transformation as reducto;
use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai;
use litellm_core::providers::vertex_ai::ocr::transformation::{
VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
};
use crate::client::http_client;
const ERROR_BODY_MAX_CHARS: usize = 256;
const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0;
const MAX_SAFE_FETCH_REDIRECTS: usize = 10;
pub(super) fn truncate_error_body(body: &str) -> String {
if body.chars().count() <= ERROR_BODY_MAX_CHARS {
return body.to_string();
}
let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect();
format!("{truncated}... (truncated)")
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) fn ocr_provider_config(
provider: &str,
model: &str,
) -> Option<&'static dyn OcrProviderConfig> {
match provider {
"mistral" => Some(&MISTRAL_OCR_CONFIG),
"reducto" => reducto::config_for_model(model),
"azure_ai" if is_azure_document_intelligence_model(model) => {
Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG)
}
"azure_ai" => Some(&AZURE_AI_OCR_CONFIG),
"vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG),
"vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG),
_ => None,
}
}
fn is_azure_document_intelligence_model(model: &str) -> bool {
let model = model.to_ascii_lowercase();
model.contains("doc-intelligence") || model.contains("documentintelligence")
}
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> Result<Vec<(String, String)>, Error> {
extra_headers
.unwrap_or_default()
.into_iter()
.map(|(key, value)| {
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
Error::InvalidRequest(format!(
"OCR extra_headers.{key} must be a string, got {}",
litellm_core::error::json_type_name(&value)
))
})
})
.collect()
}
fn document_url_field(document: &Value) -> Result<Option<(&str, &str)>, Error> {
let Some(object) = document.as_object() else {
return Ok(None);
};
let Some(doc_type) = object.get("type").and_then(Value::as_str) else {
return Ok(None);
};
let field = match doc_type {
"document_url" => "document_url",
"image_url" => "image_url",
_ => return Ok(None),
};
let Some(url) = object.get(field).and_then(Value::as_str) else {
return Ok(None);
};
Ok(Some((field, url)))
}
fn is_url_requiring_fetch(url: &str) -> bool {
!url.starts_with("data:") && (url.starts_with("http://") || url.starts_with("https://"))
}
fn max_document_download_bytes() -> u64 {
let max_size_mb = std::env::var("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB")
.ok()
.and_then(|value| value.parse::<f64>().ok())
.unwrap_or(DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB);
(max_size_mb.max(0.0) * 1024.0 * 1024.0) as u64
}
fn is_blocked_ip(ip: IpAddr) -> bool {
match ip {
IpAddr::V4(ip) => {
ip.is_private()
|| ip.is_loopback()
|| ip.is_link_local()
|| ip.is_broadcast()
|| ip.is_multicast()
|| ip.is_unspecified()
}
IpAddr::V6(ip) => {
let first_segment = ip.segments()[0];
let is_unique_local = (first_segment & 0xfe00) == 0xfc00;
let is_link_local = (first_segment & 0xffc0) == 0xfe80;
ip.is_loopback()
|| ip.is_unspecified()
|| ip.is_multicast()
|| is_unique_local
|| is_link_local
|| ip
.to_ipv4_mapped()
.or_else(|| ip.to_ipv4())
.map(|v4| is_blocked_ip(IpAddr::V4(v4)))
.unwrap_or(false)
}
}
}
fn blocked_url_error(url: &Url) -> Error {
Error::InvalidRequest(format!(
"OCR document URL rejected by SSRF protection: {url}"
))
}
async fn validate_safe_fetch_url(url: &Url) -> Result<(), Error> {
if !matches!(url.scheme(), "http" | "https") {
return Err(blocked_url_error(url));
}
let host = url.host_str().ok_or_else(|| blocked_url_error(url))?;
if let Ok(ip) = host.parse::<IpAddr>() {
if is_blocked_ip(ip) {
return Err(blocked_url_error(url));
}
return Ok(());
}
let port = url
.port_or_known_default()
.ok_or_else(|| blocked_url_error(url))?;
let addresses = tokio::net::lookup_host((host, port))
.await
.map_err(|err| Error::Network(err.to_string()))?;
let mut saw_address = false;
for address in addresses {
saw_address = true;
if is_blocked_ip(address.ip()) {
return Err(blocked_url_error(url));
}
}
if !saw_address {
return Err(blocked_url_error(url));
}
Ok(())
}
fn redirect_location(response: &reqwest::Response, url: &Url) -> Result<Url, Error> {
let location = response
.headers()
.get(reqwest::header::LOCATION)
.and_then(|value| value.to_str().ok())
.ok_or_else(|| {
Error::InvalidResponse("OCR document redirect missing Location header".to_string())
})?;
url.join(location)
.map_err(|err| Error::InvalidResponse(format!("invalid OCR document redirect: {err}")))
}
async fn safe_get_document_url(url: &str) -> Result<(Url, reqwest::Response), Error> {
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|err| Error::Network(err.to_string()))?;
let mut current_url = Url::parse(url)
.map_err(|err| Error::InvalidRequest(format!("invalid OCR document URL: {err}")))?;
for _ in 0..MAX_SAFE_FETCH_REDIRECTS {
validate_safe_fetch_url(&current_url).await?;
let response = client
.get(current_url.clone())
.send()
.await
.map_err(|err| Error::Network(err.to_string()))?;
if !response.status().is_redirection() {
return Ok((current_url, response));
}
current_url = redirect_location(&response, &current_url)?;
}
Err(Error::InvalidRequest(
"Too many redirects while fetching OCR document URL".to_string(),
))
}
fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Result<(), Error> {
if max_bytes == 0 {
return Err(Error::InvalidRequest(format!(
"OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}"
)));
}
if content_length > max_bytes {
let size_mb = content_length as f64 / (1024.0 * 1024.0);
let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0);
return Err(Error::InvalidRequest(format!(
"OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}"
)));
}
Ok(())
}
async fn read_response_with_limit(
mut response: reqwest::Response,
url: &Url,
) -> Result<Vec<u8>, Error> {
let max_bytes = max_document_download_bytes();
if let Some(content_length) = response.content_length() {
enforce_download_size(content_length, max_bytes, url)?;
} else {
enforce_download_size(0, max_bytes, url)?;
}
let mut bytes = Vec::new();
let mut bytes_downloaded: u64 = 0;
while let Some(chunk) = response
.chunk()
.await
.map_err(|err| Error::Network(err.to_string()))?
{
bytes_downloaded += chunk.len() as u64;
enforce_download_size(bytes_downloaded, max_bytes, url)?;
bytes.extend_from_slice(&chunk);
}
Ok(bytes)
}
pub(super) async fn convert_document_url_to_data_uri(document: Value) -> Result<Value, Error> {
let Some((field, url)) = document_url_field(&document)? else {
return Ok(document);
};
if !is_url_requiring_fetch(url) {
return Ok(document);
}
let (final_url, response) = safe_get_document_url(url).await?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&body),
});
}
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("application/octet-stream")
.to_string();
let bytes = read_response_with_limit(response, &final_url).await?;
let data_uri = format!(
"data:{content_type};base64,{}",
BASE64_STANDARD.encode(bytes)
);
let mut transformed = document
.as_object()
.cloned()
.ok_or_else(|| Error::InvalidRequest("OCR document must be an object".to_string()))?;
transformed.insert(field.to_string(), Value::String(data_uri));
Ok(Value::Object(transformed))
}
fn same_origin(left: &str, right: &str) -> bool {
let Ok(left) = reqwest::Url::parse(left) else {
return false;
};
let Ok(right) = reqwest::Url::parse(right) else {
return false;
};
left.scheme() == right.scheme()
&& left.host_str() == right.host_str()
&& left.port_or_known_default() == right.port_or_known_default()
}
fn retry_after_secs(response: &reqwest::Response) -> u64 {
response
.headers()
.get(reqwest::header::RETRY_AFTER)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<u64>().ok())
.unwrap_or(2)
}
fn operation_status(response_json: &Value) -> Result<&str, Error> {
let status = response_json
.get("status")
.and_then(Value::as_str)
.ok_or(Error::MissingField("status"))?;
match status {
"succeeded" => Ok("succeeded"),
"running" | "notStarted" => Ok("running"),
"failed" => {
let message = response_json
.get("error")
.and_then(|error| error.get("message"))
.and_then(Value::as_str)
.unwrap_or("Unknown error");
Err(Error::InvalidResponse(format!(
"Azure Document Intelligence analysis failed: {message}"
)))
}
other => Err(Error::InvalidResponse(format!(
"Unknown operation status: {other}"
))),
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) async fn poll_document_intelligence(
operation_url: &str,
original_url: &str,
headers: &[(String, String)],
timeout: Option<Duration>,
) -> Result<Value, Error> {
if !same_origin(operation_url, original_url) {
return Err(Error::InvalidResponse(
"Azure Document Intelligence: rejected cross-origin polling URL".to_string(),
));
}
let start = Instant::now();
let timeout = timeout.unwrap_or(Duration::from_secs(
AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS,
));
loop {
if start.elapsed() > timeout {
return Err(Error::Network(format!(
"Azure Document Intelligence operation polling timed out after {} seconds",
timeout.as_secs()
)));
}
let mut request_builder = http_client().get(operation_url);
for (key, value) in headers {
if key.eq_ignore_ascii_case("ocp-apim-subscription-key") {
request_builder = request_builder.header(key, value);
}
}
let response = request_builder
.send()
.await
.map_err(|err| Error::Network(err.to_string()))?;
let retry_after = retry_after_secs(&response);
let status = response.status();
let text = response
.text()
.await
.map_err(|err| Error::Network(err.to_string()))?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response_json: Value = serde_json::from_str(&text).map_err(|err| {
Error::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}"))
})?;
if operation_status(&response_json)? == "succeeded" {
return Ok(response_json);
}
tokio::time::sleep(Duration::from_secs(retry_after)).await;
}
}
#[cfg(test)]
mod tests {
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::json;
use super::*;
#[test]
fn blocks_private_and_metadata_ips() {
assert!(is_blocked_ip("127.0.0.1".parse().unwrap()));
assert!(is_blocked_ip("10.0.0.1".parse().unwrap()));
assert!(is_blocked_ip("169.254.169.254".parse().unwrap()));
assert!(is_blocked_ip("::1".parse().unwrap()));
assert!(is_blocked_ip("fd00::1".parse().unwrap()));
assert!(is_blocked_ip("fe80::1".parse().unwrap()));
assert!(is_blocked_ip("::ffff:169.254.169.254".parse().unwrap()));
assert!(is_blocked_ip("::ffff:10.0.0.1".parse().unwrap()));
assert!(!is_blocked_ip("8.8.8.8".parse().unwrap()));
assert!(!is_blocked_ip("::ffff:8.8.8.8".parse().unwrap()));
}
#[tokio::test]
async fn convert_document_url_rejects_loopback_fetch() {
let error = convert_document_url_to_data_uri(json!({
"type": "image_url",
"image_url": "http://127.0.0.1/image.png"
}))
.await
.unwrap_err();
assert!(matches!(
error,
Error::InvalidRequest(message)
if message.contains("SSRF protection")
));
}
#[tokio::test]
async fn convert_document_url_leaves_data_uri_untouched() {
let document = json!({
"type": "image_url",
"image_url": "data:image/png;base64,abcd"
});
let transformed = convert_document_url_to_data_uri(document.clone())
.await
.unwrap();
assert_eq!(transformed, document);
}
#[test]
fn truncate_error_body_passes_short_strings_through() {
let body = "Unauthorized";
assert_eq!(truncate_error_body(body), "Unauthorized");
}
#[test]
fn truncate_error_body_caps_long_payloads() {
let body = "x".repeat(306);
let truncated = truncate_error_body(&body);
assert!(truncated.ends_with("... (truncated)"));
let prefix_chars = truncated
.strip_suffix("... (truncated)")
.expect("truncated marker present")
.chars()
.count();
assert_eq!(prefix_chars, 256);
}
#[test]
fn truncate_error_body_does_not_split_multibyte_chars() {
let body = "é".repeat(266);
let truncated = truncate_error_body(&body);
assert!(truncated.is_char_boundary(truncated.len()));
}
#[test]
fn ocr_dispatch_supports_migrated_providers() {
assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
assert!(
ocr_provider_config("azure_ai", "pixtral-12b-2409")
.expect("azure ai config resolves")
.requires_data_uri_document()
);
assert_eq!(
ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
.expect("document intelligence config resolves")
.response_handling(),
OcrResponseHandling::AzureDocumentIntelligencePoll
);
assert!(
ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
.expect("vertex deepseek config resolves")
.supported_ocr_params()
.contains(&"temperature")
);
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
}
#[test]
fn string_headers_accepts_string_values() {
let headers = json!({
"x-trace-id": "trace-1"
})
.as_object()
.unwrap()
.clone();
assert_eq!(
string_headers(Some(headers)).expect("string headers accepted"),
vec![("x-trace-id".to_string(), "trace-1".to_string())]
);
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({
"x-retry-count": 3
})
.as_object()
.unwrap()
.clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert_eq!(
err,
Error::InvalidRequest(
"OCR extra_headers.x-retry-count must be a string, got number".to_string()
)
);
}
}

84
litellm-rust/crates/ai-gateway/src/ocr/handler.rs
View file

@ -1,84 +0,0 @@
use litellm_core::error::Error;
use litellm_core::http_utils::http_request;
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::Value;
use super::common_utils::{poll_document_intelligence, truncate_error_body};
use super::hooks::OcrLifecycleHooks;
use super::types::PreparedOcrRequest;
use crate::client::http_client;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) async fn execute_ocr_provider_call(
request: PreparedOcrRequest,
hooks: &OcrLifecycleHooks,
) -> Result<Value, Error> {
let request = hooks.prepare_provider_request(request).await?;
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder)
.await
.map_err(|err| Error::Network(err.to_string()))?;
let status = response.status();
if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
&& status.as_u16() == 202
{
let operation_url = response
.headers()
.get("operation-location")
.and_then(|value| value.to_str().ok())
.map(str::to_string)
.ok_or_else(|| {
Error::InvalidResponse(
"Azure Document Intelligence returned 202 but no Operation-Location header found"
.to_string(),
)
})?;
let response_json = poll_document_intelligence(
&operation_url,
&request.url,
&request.upstream_headers,
request.timeout,
)
.await?;
return Ok(request
.config
.transform_ocr_response_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json());
}
let text = response
.text()
.await
.map_err(|err| Error::Network(err.to_string()))?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response_json: Value = serde_json::from_str(&text)
.map_err(|err| Error::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
Ok(request
.config
.transform_ocr_response_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json())
}

401
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
View file

@ -1,401 +0,0 @@
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::Error;
use litellm_core::providers::reducto::ocr::transformation::{
build_upload_request, extract_document_source, extract_upload_file_id,
};
use serde_json::{Map, Value, json};
use std::future::Future;
use std::pin::Pin;
use super::common_utils::{convert_document_url_to_data_uri, string_headers, truncate_error_body};
use super::types::{PreparedOcrRequest, ProviderOcrRequest};
use crate::client::http_client;
use crate::integrations::custom_guardrail::{
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
};
use crate::integrations::custom_logger::{
CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails,
};
use crate::integrations::types::{
RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload,
};
pub(crate) struct OcrLifecycleHooks {
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
}
type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
impl OcrLifecycleHooks {
pub(crate) fn new(
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
) -> Self {
Self {
logger_runner,
guardrail_runner,
request_metadata,
}
}
async fn run_pre_call_guardrails(
&self,
request: PreparedOcrRequest,
) -> Result<PreparedOcrRequest, Error> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
let context = guardrail_context(&self.request_metadata);
let guardrail_request = GuardrailRequest::new(json!({
"model": request.model,
"custom_llm_provider": request.custom_llm_provider,
"document": request.document,
"optional_params": request.optional_params,
}));
let (guardrail_request, _) = self
.guardrail_runner
.run_pre_call(&context, guardrail_request)
.await
.map_err(guardrail_error_to_core_error)?;
let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?;
let optional_params = match &request.config {
Ok(config) => config.map_ocr_params(&optional_params),
Err(_) => optional_params,
};
Ok(PreparedOcrRequest {
document,
optional_params,
..request
})
}
pub(crate) async fn prepare_provider_request(
&self,
request: PreparedOcrRequest,
) -> Result<ProviderOcrRequest, Error> {
let config = request.config?;
let env_lookup = |key: &str| std::env::var(key).ok();
let upstream_headers = config.validate_environment(
string_headers(request.extra_headers)?,
request.api_key.as_deref(),
&env_lookup,
)?;
let url = config.complete_url(
request.api_base.as_deref(),
&request.model,
&request.optional_params,
&env_lookup,
)?;
let model = request.model.clone();
let custom_llm_provider = request.custom_llm_provider.clone();
let is_reducto = custom_llm_provider == "reducto";
let document = if is_reducto {
let guarded_document = self
.run_during_call_guardrails(&model, &custom_llm_provider, &url, request.document)
.await?;
upload_reducto_document(
&guarded_document,
request.api_base.as_deref(),
request.timeout,
&upstream_headers,
)
.await?
} else if config.requires_data_uri_document() {
convert_document_url_to_data_uri(request.document).await?
} else {
request.document
};
let optional_params = request.optional_params;
let body = config
.transform_ocr_request(&request.model, document, optional_params.clone())?
.data;
let body = if is_reducto {
body
} else {
self.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
.await?
};
Ok(ProviderOcrRequest {
model,
config,
url,
body,
optional_params,
upstream_headers,
timeout: request.timeout,
})
}
async fn run_during_call_guardrails(
&self,
model: &str,
custom_llm_provider: &str,
url: &str,
body: Value,
) -> Result<Value, Error> {
if self.guardrail_runner.is_empty() {
return Ok(body);
}
let context = guardrail_context(&self.request_metadata);
let guardrail_request = GuardrailRequest::new(json!({
"model": model,
"custom_llm_provider": custom_llm_provider,
"url": url,
"body": body,
}));
let (guardrail_request, _) = self
.guardrail_runner
.run_during_call(&context, guardrail_request)
.await
.map_err(guardrail_error_to_core_error)?;
parse_ocr_during_call_guardrail_request(guardrail_request)
}
fn standard_logging_payload(
&self,
context: &CallLifecycleContext,
timing: &CallLifecycleTiming,
) -> StandardLoggingPayload {
StandardLoggingPayload {
id: context.litellm_call_id.clone(),
litellm_call_id: context.litellm_call_id.clone(),
call_type: context.call_type.clone(),
model: context.model.clone(),
custom_llm_provider: context.custom_llm_provider.clone(),
response_cost: 0.0,
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
start_time: timing.start_time,
end_time: timing.end_time,
stream: false,
metadata: StandardLoggingMetadata {
user_api_key_hash: self.request_metadata.user_api_key_hash.clone(),
user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(),
user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(),
..Default::default()
},
messages: None,
}
}
}
async fn upload_reducto_document(
document: &Value,
api_base: Option<&str>,
timeout: Option<std::time::Duration>,
upstream_headers: &[(String, String)],
) -> Result<Value, Error> {
let source = extract_document_source(document)?;
let Some(authorization) = upstream_headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
.map(|(_, value)| value.as_str())
else {
return Err(Error::Auth(
"Reducto upload requires an Authorization header".to_string(),
));
};
let Some(upload) = build_upload_request(source, authorization, api_base) else {
return Ok(document.clone());
};
let part = reqwest::multipart::Part::bytes(upload.bytes)
.file_name(upload.file_name)
.mime_str(&upload.mime_type)
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let form = reqwest::multipart::Form::new().part("file", part);
let mut request_builder = http_client().post(upload.url).multipart(form);
for (name, value) in upstream_headers {
if !name.eq_ignore_ascii_case("content-type")
&& !name.eq_ignore_ascii_case("content-length")
{
request_builder = request_builder.header(name, value);
}
}
if let Some(timeout) = timeout {
request_builder = request_builder.timeout(timeout);
}
let response = request_builder
.send()
.await
.map_err(|error| Error::Network(error.to_string()))?;
let status = response.status();
let body = response
.text()
.await
.map_err(|error| Error::Network(error.to_string()))?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&body),
});
}
let response_json: Value = serde_json::from_str(&body).map_err(|error| {
Error::InvalidResponse(format!("invalid Reducto upload response JSON: {error}"))
})?;
let file_id = extract_upload_file_id(&response_json)?;
Ok(json!({"type": "document_url", "document_url": file_id}))
}
impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLifecycleHooks {
type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
type SuccessFuture<'a> = OcrLogFuture<'a>;
type FailureFuture<'a> = OcrLogFuture<'a>;
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: PreparedOcrRequest,
) -> Self::PreCallFuture<'a> {
Box::pin(async move { self.run_pre_call_guardrails(request).await })
}
fn async_during_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: PreparedOcrRequest,
) -> Self::DuringCallFuture<'a> {
Box::pin(async move { Ok(request) })
}
#[tracing::instrument(
name = "success_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_success_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
response: &'a Value,
timing: &'a CallLifecycleTiming,
) -> Self::SuccessFuture<'a> {
Box::pin(async move {
if self.logger_runner.is_empty() {
return;
}
let response_obj = CallbackValue::new("ocr", response.clone());
self.logger_runner
.async_log_success_event(
&ModelCallDetails::from_standard_logging_payload(
self.standard_logging_payload(context, timing),
),
&response_obj,
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
#[tracing::instrument(
name = "failure_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a Error,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
if self.logger_runner.is_empty() {
return;
}
let logging_error = LoggingError {
message: error.to_string(),
kind: core_error_kind(error).to_string(),
};
let response_obj = CallbackValue::new(
"error",
json!({
"message": logging_error.message,
"kind": logging_error.kind,
}),
);
self.logger_runner
.async_log_failure_event(
&ModelCallDetails::from_standard_logging_payload(
self.standard_logging_payload(context, timing),
)
.with_failure_error(logging_error),
Some(&response_obj),
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
}
fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
GuardrailContext {
call_type: CallType::Ocr,
selected_guardrails: Vec::new(),
metadata: std::collections::HashMap::new(),
user_api_key_hash: metadata.user_api_key_hash.clone(),
user_api_key_user_id: metadata.user_api_key_user_id.clone(),
user_api_key_team_id: metadata.user_api_key_team_id.clone(),
trace_parent: None,
}
}
fn parse_ocr_pre_call_guardrail_request(
request: GuardrailRequest,
) -> Result<(Value, Map<String, Value>), Error> {
let Value::Object(mut data) = request.data else {
return Err(Error::InvalidRequest(
"OCR pre_call guardrail must return an object".to_string(),
));
};
let document = data.remove("document").ok_or_else(|| {
Error::InvalidRequest("OCR pre_call guardrail removed document".to_string())
})?;
let optional_params = match data.remove("optional_params") {
Some(Value::Object(params)) => params,
Some(_) => {
return Err(Error::InvalidRequest(
"OCR pre_call guardrail optional_params must be an object".to_string(),
));
}
None => Map::new(),
};
Ok((document, optional_params))
}
fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> Result<Value, Error> {
let Value::Object(mut data) = request.data else {
return Err(Error::InvalidRequest(
"OCR during_call guardrail must return an object".to_string(),
));
};
data.remove("body")
.ok_or_else(|| Error::InvalidRequest("OCR during_call guardrail removed body".to_string()))
}
fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
Error::InvalidRequest(format!("{}: {}", error.kind, error.message))
}
fn core_error_kind(error: &Error) -> &'static str {
match error {
Error::Auth(_) | Error::MissingApiKey { .. } => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",
Error::MissingField(_) => "MissingField",
Error::Http { .. } => "HttpError",
Error::InvalidResponse(_) => "InvalidResponse",
Error::Network(_) => "NetworkError",
Error::Connect(_) => "ConnectError",
Error::Routing(_) => "RoutingError",
Error::Unsupported(_) => "UnsupportedRequest",
}
}

231
litellm-rust/crates/ai-gateway/src/ocr/mod.rs
View file

@ -1,174 +1,127 @@
use litellm_core::Error;
use litellm_core::call_lifecycle::CallLifecycle;
use litellm_core::ocr::{
OcrClient,
wire::{OcrWireRequest, decode_request},
};
use serde_json::Value;
mod common_utils;
mod handler;
mod hooks;
mod prepare;
mod types;
pub use types::OcrRequest;
use handler::execute_ocr_provider_call;
use prepare::{PreparedOcrCall, prepare_ocr_call};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
CallLifecycle::default()
.run_request(request, &hooks, |request| {
execute_ocr_provider_call(request, &hooks)
})
core_ocr(request).await
}
async fn core_ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
validate_host_hooks(&request)?;
let client = OcrClient::new(crate::client::http_client().clone())?;
let core_request = decode_request(OcrWireRequest {
model: request.model.to_string(),
document: request.document,
api_key: request.api_key.map(str::to_string),
api_base: request.api_base.map(str::to_string),
custom_llm_provider: request.custom_llm_provider.map(str::to_string),
extra_headers: request.extra_headers,
optional_params: request.optional_params,
input_sources: Default::default(),
timeout_seconds: request.timeout.map(|timeout| timeout.as_secs_f64()),
})?;
client
.perform(core_request)
.await
.map(|response| response.into_json())
}
fn validate_host_hooks(request: &OcrRequest<'_>) -> Result<(), Error> {
if !request.guardrails.is_empty() {
return Err(Error::Unsupported(
"OCR host guardrails are not wired to the core path",
));
}
if !request.callbacks.is_empty() {
return Err(Error::Unsupported(
"OCR host callbacks are not wired to the core path",
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use litellm_core::ocr::wire::is_supported_request;
use serde_json::{Map, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::{OcrRequest, ocr};
use crate::integrations::types::RequestMetadata;
use super::{OcrRequest, validate_host_hooks};
use crate::integrations::custom_guardrail::{CustomGuardrail, GuardrailEventHook};
use crate::integrations::custom_logger::CustomLogger;
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]);
struct TestGuardrail;
impl CustomGuardrail for TestGuardrail {
fn guardrail_name(&self) -> &str {
"test"
}
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
&[]
}
String::from_utf8(request).expect("request is utf8")
}
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
struct TestLogger;
impl CustomLogger for TestLogger {}
fn request() -> OcrRequest<'static> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
model: "model",
document: json!({"type":"image_url","image_url":"data:image/png;base64,YQ=="}),
api_key: None,
api_base: None,
custom_llm_provider: None,
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
request_metadata: Default::default(),
litellm_call_id: None,
}
}
#[tokio::test]
async fn reducto_file_upload_then_parse_maps_response() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let server = tokio::spawn(async move {
let (mut upload_socket, _) = listener.accept().await.expect("accepts upload request");
let upload_request = read_http_request(&mut upload_socket).await;
let upload_body = r#"{"file_id":"reducto://uploaded.pdf"}"#;
let upload_response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
upload_body.len(),
upload_body
);
upload_socket
.write_all(upload_response.as_bytes())
.await
.expect("writes upload response");
#[test]
fn core_activation_includes_migrated_providers() {
assert!(is_supported_request("model", Some("mistral")));
assert!(is_supported_request("pixtral-12b", Some("azure_ai")));
assert!(is_supported_request(
"doc-intelligence/prebuilt-layout",
Some("azure_ai")
));
assert!(is_supported_request("parse-v3", Some("reducto")));
assert!(is_supported_request("mistral-ocr", Some("vertex_ai")));
assert!(is_supported_request("deepseek-ocr", Some("vertex_ai")));
}
let (mut parse_socket, _) = listener.accept().await.expect("accepts parse request");
let parse_request = read_http_request(&mut parse_socket).await;
let parse_body = r#"{"job_id":"job_123","usage":{"num_pages":3,"credits":3},"result":{"chunks":[{"content":"Page 1 block A","blocks":[{"content":"Page 1 block A","bbox":{"page":1},"kind":"text"}]},{"content":"Page 2 block A","blocks":[{"content":"Page 2 block A","bbox":{"page":2},"kind":"table"}]},{"content":"Page 1 block B","blocks":[{"content":"Page 1 block B","bbox":{"page":1},"kind":"text"}]},{"content":"Page 3 block A","blocks":[{"content":"Page 3 block A","bbox":{"page":3},"kind":"figure"}]}]}}"#;
let parse_response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
parse_body.len(),
parse_body
);
parse_socket
.write_all(parse_response.as_bytes())
.await
.expect("writes parse response");
(upload_request, parse_request)
});
let api_base = format!("http://{address}");
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.api_key = None;
request.extra_headers = Some(Map::from_iter([
("Authorization".to_string(), json!("Bearer test-key")),
("x-trace-id".to_string(), json!("trace-1")),
]));
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
request.optional_params = Map::from_iter([
(
"formatting".to_string(),
json!({"table_output_format": "html"}),
),
("retrieval".to_string(), json!({"chunk_mode": "section"})),
("settings".to_string(), json!({"ocr_system": "standard"})),
]);
#[test]
fn core_path_rejects_unwired_guardrails() {
let request = OcrRequest {
guardrails: vec![Arc::new(TestGuardrail)],
..request()
};
let error = validate_host_hooks(&request).unwrap_err();
assert!(error.to_string().contains("guardrails are not wired"));
}
let response = ocr(request).await.expect("Reducto OCR succeeds");
assert_eq!(response["pages"].as_array().map(Vec::len), Some(3));
assert_eq!(
response["pages"][0]["markdown"],
"Page 1 block A\n\nPage 1 block B"
);
assert_eq!(response["pages"][1]["markdown"], "Page 2 block A");
assert_eq!(response["pages"][2]["markdown"], "Page 3 block A");
assert_eq!(response["usage_info"]["pages_processed"], 3);
assert_eq!(response["usage_info"]["credits"], 3);
assert_eq!(response["provider_native_response"]["job_id"], "job_123");
let (upload_request, parse_request) = server.await.expect("server task completes");
assert!(
upload_request
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
assert!(upload_request.contains("application/pdf"));
assert!(upload_request.contains("%PDF-1.4"));
assert!(upload_request.contains("x-trace-id: trace-1"));
assert!(
parse_request
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
assert!(parse_request.contains(r#""input":"reducto://uploaded.pdf""#));
assert!(parse_request.contains(r#""table_output_format":"html""#));
assert!(parse_request.contains(r#""chunk_mode":"section""#));
assert!(parse_request.contains(r#""ocr_system":"standard""#));
#[test]
fn core_path_rejects_unwired_callbacks() {
let request = OcrRequest {
callbacks: vec![Arc::new(TestLogger)],
..request()
};
let error = validate_host_hooks(&request).unwrap_err();
assert!(error.to_string().contains("callbacks are not wired"));
}
}

163
litellm-rust/crates/ai-gateway/src/ocr/prepare.rs
View file

@ -1,163 +0,0 @@
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use serde_json::{Map, Value};
use super::common_utils::ocr_provider_config;
use super::hooks::OcrLifecycleHooks;
use super::types::{OcrRequest, PreparedOcrRequest};
use crate::integrations::custom_guardrail::CustomGuardrailRunner;
use crate::integrations::custom_logger::CustomLoggerRunner;
pub(crate) struct PreparedOcrCall {
pub(crate) request: PreparedOcrRequest,
pub(crate) hooks: OcrLifecycleHooks,
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
let call_id = request
.litellm_call_id
.map(str::to_string)
.unwrap_or_else(new_ocr_call_id);
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
.unwrap_or(CustomLlmProvider {
model: request.model,
custom_llm_provider: "mistral",
});
let model = provider_info.model.to_string();
let custom_llm_provider = provider_info.custom_llm_provider.to_string();
let config = ocr_provider_config(&custom_llm_provider, &model)
.ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()))
.and_then(|config| {
validate_request_format(config, &request.optional_params, &custom_llm_provider)?;
Ok(config)
});
let optional_params = match &config {
Ok(config) => {
let supported = config.supported_ocr_params();
let mut mapped = config.map_ocr_params(
&request
.optional_params
.iter()
.filter(|(name, _)| supported.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
);
for name in [
"vertex_project",
"vertex_ai_project",
"vertex_location",
"vertex_ai_location",
] {
if let Some(value) = request.optional_params.get(name) {
mapped.insert(name.to_string(), value.clone());
}
}
mapped
}
Err(_) => request.optional_params,
};
PreparedOcrCall {
request: PreparedOcrRequest {
config,
model,
custom_llm_provider,
litellm_call_id: call_id,
document: request.document,
api_key: request.api_key.map(str::to_string),
api_base: request.api_base.map(str::to_string),
extra_headers: request.extra_headers,
optional_params,
timeout: request.timeout,
},
hooks: OcrLifecycleHooks::new(
CustomLoggerRunner::new(request.callbacks),
CustomGuardrailRunner::new(request.guardrails),
request.request_metadata,
),
}
}
fn validate_request_format(
config: &'static dyn litellm_core::ocr::transformation::OcrProviderConfig,
optional_params: &Map<String, Value>,
provider: &str,
) -> Result<(), litellm_core::Error> {
let Some(format) = optional_params.get("req_format") else {
return Ok(());
};
match format.as_str() {
Some("litellm") => Ok(()),
Some("native") if config.supported_ocr_params().contains(&"req_format") => Ok(()),
Some("native") => Err(litellm_core::Error::InvalidRequest(format!(
"`req_format=native` is not supported for provider {provider}"
))),
_ => Err(litellm_core::Error::InvalidRequest(format!(
"Invalid `req_format`: {format}. Expected `litellm` or `native`"
))),
}
}
fn new_ocr_call_id() -> String {
static COUNTER: AtomicU64 = AtomicU64::new(1);
let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_nanos())
.unwrap_or(0);
format!("ocr-{timestamp}-{sequence}")
}
#[cfg(test)]
mod tests {
use litellm_core::error::Error;
use serde_json::{Map, json};
use super::{OcrRequest, prepare_ocr_call};
use crate::integrations::types::RequestMetadata;
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
}
}
fn request_with_format(format: &str) -> OcrRequest<'_> {
let mut request = base_ocr_request("mistral/mistral-ocr-latest");
request.optional_params = Map::from_iter([("req_format".to_string(), json!(format))]);
request
}
#[test]
fn native_format_rejected_for_provider_without_support_as_bad_request() {
let prepared = prepare_ocr_call(request_with_format("native"));
assert!(
matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("not supported for provider"))
);
}
#[test]
fn unknown_format_rejected_for_provider_without_support_as_bad_request() {
let prepared = prepare_ocr_call(request_with_format("raw"));
assert!(
matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("Invalid `req_format`"))
);
}
}

36
litellm-rust/crates/ai-gateway/src/ocr/types.rs
View file

@ -1,8 +1,6 @@
use std::sync::Arc;
use std::time::Duration;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
use litellm_core::ocr::transformation::OcrProviderConfig;
use serde_json::{Map, Value};
use crate::integrations::custom_guardrail::CustomGuardrail;
@ -23,37 +21,3 @@ pub struct OcrRequest<'a> {
pub request_metadata: RequestMetadata,
pub litellm_call_id: Option<&'a str>,
}
pub(crate) struct PreparedOcrRequest {
pub(crate) config: Result<&'static dyn OcrProviderConfig, litellm_core::Error>,
pub(crate) model: String,
pub(crate) custom_llm_provider: String,
pub(crate) litellm_call_id: String,
pub(crate) document: Value,
pub(crate) api_key: Option<String>,
pub(crate) api_base: Option<String>,
pub(crate) extra_headers: Option<Map<String, Value>>,
pub(crate) optional_params: Map<String, Value>,
pub(crate) timeout: Option<Duration>,
}
impl CallLifecycleRequest for PreparedOcrRequest {
fn lifecycle_context(&self) -> CallLifecycleContext {
CallLifecycleContext::new(
"ocr",
self.model.clone(),
self.custom_llm_provider.clone(),
self.litellm_call_id.clone(),
)
}
}
pub(crate) struct ProviderOcrRequest {
pub(crate) model: String,
pub(crate) config: &'static dyn OcrProviderConfig,
pub(crate) url: String,
pub(crate) body: Value,
pub(crate) optional_params: Map<String, Value>,
pub(crate) upstream_headers: Vec<(String, String)>,
pub(crate) timeout: Option<Duration>,
}

8
litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs
View file

@ -105,7 +105,11 @@ impl IntoResponse for MessagesRouteError {
StatusCode::NOT_FOUND,
"no messages deployment is configured for this model".to_string(),
),
Error::Auth(_) => (
Error::Auth(_)
| Error::MissingApiKey { .. }
| Error::MissingAzureAiCredentials
| Error::MissingAzureDocumentIntelligenceCredentials
| Error::MissingReductoApiKey => (
StatusCode::BAD_GATEWAY,
"messages provider authentication failed".to_string(),
),
@ -115,7 +119,7 @@ impl IntoResponse for MessagesRouteError {
| Error::InvalidResponse(_)
| Error::InvalidType { .. }
| Error::MissingField(_)
| Error::MissingApiKey { .. } => (
| Error::MissingDocumentUrl => (
StatusCode::BAD_GATEWAY,
"messages provider request failed".to_string(),
),

33
litellm-rust/crates/ai-gateway/src/trace_parity.rs
View file

@ -10,6 +10,7 @@ use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
use serde::Serialize;
use serde_json::Value;
use tower::ServiceExt;
use tracing::instrument::WithSubscriber;
use crate::io::realtime_pool::RealtimePool;
use crate::routes;
@ -21,6 +22,38 @@ pub struct GatewayResponse {
pub body: Value,
}
#[derive(Debug, Serialize)]
pub struct TracedGatewayResponse {
pub response: Option<GatewayResponse>,
pub error: Option<String>,
pub trace: Vec<litellm_core::observability::FunctionTraceEvent>,
}
pub async fn traced_messages_request(
model_alias: String,
provider_model: String,
api_base: String,
body: Value,
) -> TracedGatewayResponse {
let trace = litellm_core::observability::FunctionTrace::default();
let result = messages_request(model_alias, provider_model, api_base, body)
.with_subscriber(trace.dispatcher())
.await;
let events = trace.events();
match result {
Ok(response) => TracedGatewayResponse {
response: Some(response),
error: None,
trace: events,
},
Err(error) => TracedGatewayResponse {
response: None,
error: Some(error.to_string()),
trace: events,
},
}
}
pub async fn messages_request(
model_alias: String,
provider_model: String,

15
litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs
View file

@ -2,10 +2,10 @@
//! API has to resolve its own crypto provider, in a test binary where nothing
//! has installed a process-wide one, and has to leave it uninstalled.
use std::collections::HashMap;
use std::time::Duration;
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection;
use futures_util::{sink, stream};
use litellm_ai_gateway::io::responses_ws::async_responses_websocket;
use tokio::net::TcpListener;
async fn dead_tls_server() -> u16 {
@ -30,10 +30,15 @@ async fn dead_tls_server() -> u16 {
async fn dialing_wss_returns_an_error_instead_of_panicking() {
let port = dead_tls_server().await;
let result = ResponsesWebSocketConnection::connect_url(
&format!("wss://127.0.0.1:{port}/"),
&HashMap::new(),
let result = async_responses_websocket(
"gpt-5",
Some("test-key"),
Some(&format!("wss://127.0.0.1:{port}/")),
None,
Some(Duration::from_secs(10)),
|_| {},
stream::empty(),
sink::drain(),
)
.await;

641
litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
View file

@ -1,641 +0,0 @@
use std::sync::{Arc, Mutex};
use std::time::Duration;
use litellm_ai_gateway::integrations::custom_guardrail::{
CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,
GuardrailFuture, GuardrailRequest,
};
use litellm_ai_gateway::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
};
use litellm_ai_gateway::integrations::types::RequestMetadata;
use litellm_ai_gateway::ocr::{OcrRequest, ocr};
use litellm_core::error::Error;
#[cfg(feature = "trace-parity")]
use litellm_core::observability::FunctionTrace;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
#[cfg(feature = "trace-parity")]
use tracing::instrument::WithSubscriber;
async fn read_http_headers(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
String::from_utf8(request).expect("request is utf8")
}
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")
}
#[derive(Clone, Debug, PartialEq)]
struct RecordedLogEvent {
hook: &'static str,
model: String,
call_type: String,
user_id: Option<String>,
response_object: Option<String>,
error_kind: Option<String>,
}
#[derive(Default)]
struct RecordingOcrLogger {
events: Mutex<Vec<RecordedLogEvent>>,
}
impl RecordingOcrLogger {
fn events(&self) -> Vec<RecordedLogEvent> {
self.events.lock().unwrap().clone()
}
}
impl CustomLogger for RecordingOcrLogger {
fn async_log_success_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
response_obj: &'a CallbackValue,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push(RecordedLogEvent {
hook: "async_log_success_event",
model: model_call_details.model.clone(),
call_type: model_call_details.call_type.to_string(),
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
response_object: Some(response_obj.object.clone()),
error_kind: None,
});
Ok(())
})
}
fn async_log_failure_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
response_obj: Option<&'a CallbackValue>,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push(RecordedLogEvent {
hook: "async_log_failure_event",
model: model_call_details.model.clone(),
call_type: model_call_details.call_type.to_string(),
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
response_object: response_obj.map(|value| value.object.clone()),
error_kind: model_call_details
.failure_error
.as_ref()
.map(|error| error.kind.clone()),
});
Ok(())
})
}
}
struct RecordingOcrGuardrail {
hooks: Vec<GuardrailEventHook>,
events: Mutex<Vec<&'static str>>,
block_pre_call: bool,
block_during_call: bool,
}
impl RecordingOcrGuardrail {
fn new(hooks: Vec<GuardrailEventHook>) -> Self {
Self {
hooks,
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: false,
}
}
fn blocking_pre_call() -> Self {
Self {
hooks: vec![GuardrailEventHook::PreCall],
events: Mutex::new(Vec::new()),
block_pre_call: true,
block_during_call: false,
}
}
fn blocking_during_call() -> Self {
Self {
hooks: vec![GuardrailEventHook::DuringCall],
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: true,
}
}
fn events(&self) -> Vec<&'static str> {
self.events.lock().unwrap().clone()
}
}
impl CustomGuardrail for RecordingOcrGuardrail {
fn guardrail_name(&self) -> &str {
"recording-ocr-guardrail"
}
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
&self.hooks
}
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a GuardrailContext,
mut request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("async_pre_call_hook");
if self.block_pre_call {
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
"blocked before provider",
)));
}
request.data["document"]["guarded_pre"] = json!(true);
Ok(GuardrailDecision::Mask(request))
})
}
fn async_moderation_hook<'a>(
&'a self,
_context: &'a GuardrailContext,
mut request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("async_moderation_hook");
if self.block_during_call {
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
"blocked before provider",
)));
}
request.data["body"]["guarded_during"] = json!(true);
Ok(GuardrailDecision::Mask(request))
})
}
}
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
}
}
#[tokio::test]
async fn reducto_during_call_guardrail_blocks_before_upload() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let api_base = format!("http://{address}");
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_during_call());
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
request.guardrails = vec![guardrail.clone()];
let error = ocr(request).await.expect_err("guardrail blocks upload");
assert!(matches!(error, Error::InvalidRequest(_)));
assert_eq!(guardrail.events(), vec!["async_moderation_hook"]);
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
assert!(accepted.is_err(), "upload socket should not be touched");
}
#[tokio::test]
async fn reducto_upload_error_body_is_truncated() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts upload request");
let _request = read_http_request(&mut socket).await;
let body = "x".repeat(300);
let response = format!(
"HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes upload response");
});
let api_base = format!("http://{address}");
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
let error = ocr(request).await.expect_err("upload should fail");
assert!(
matches!(error, Error::Http { status: 500, body } if body.chars().count() < 300 && body.ends_with("... (truncated)"))
);
server.await.expect("server task completes");
}
#[tokio::test]
async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let request = read_http_request(&mut socket).await;
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":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 logger = Arc::new(RecordingOcrLogger::default());
let guardrail = Arc::new(RecordingOcrGuardrail::new(vec![
GuardrailEventHook::PreCall,
GuardrailEventHook::DuringCall,
]));
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: vec![logger.clone()],
guardrails: vec![guardrail.clone()],
request_metadata: RequestMetadata {
user_api_key_user_id: Some("user-1".to_string()),
..Default::default()
},
litellm_call_id: Some("ocr-call-1"),
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let response = call.await.expect("ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
assert_eq!(
guardrail.events(),
vec!["async_pre_call_hook", "async_moderation_hook"]
);
assert_eq!(
logger.events(),
vec![RecordedLogEvent {
hook: "async_log_success_event",
model: "mistral-ocr-latest".to_string(),
call_type: "ocr".to_string(),
user_id: Some("user-1".to_string()),
response_object: Some("ocr".to_string()),
error_kind: None,
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["success_callback"]
);
let request = server.await.expect("server task completes");
assert!(request.contains(r#""guarded_pre":true"#), "{request}");
assert!(request.contains(r#""guarded_during":true"#), "{request}");
}
#[tokio::test]
async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let _request = read_http_request(&mut socket).await;
let response_body = "provider failed";
let response = format!(
"HTTP/1.1 500 Internal Server Error\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");
});
let logger = Arc::new(RecordingOcrLogger::default());
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: vec![logger.clone()],
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("ocr-call-2"),
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let err = call.await.expect_err("provider error propagates");
assert!(matches!(err, Error::Http { status: 500, .. }));
server.await.expect("server task completes");
assert_eq!(
logger.events(),
vec![RecordedLogEvent {
hook: "async_log_failure_event",
model: "mistral-ocr-latest".to_string(),
call_type: "ocr".to_string(),
user_id: None,
response_object: Some("error".to_string()),
error_kind: Some("HttpError".to_string()),
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["failure_callback"]
);
}
#[tokio::test]
async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let logger = Arc::new(RecordingOcrLogger::default());
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call());
let err = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_millis(100)),
callbacks: vec![logger.clone()],
guardrails: vec![guardrail.clone()],
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("ocr-call-3"),
})
.await
.expect_err("guardrail blocks request");
assert!(matches!(err, Error::InvalidRequest(_)));
assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]);
assert_eq!(
logger.events(),
vec![RecordedLogEvent {
hook: "async_log_failure_event",
model: "mistral-ocr-latest".to_string(),
call_type: "ocr".to_string(),
user_id: None,
response_object: Some("error".to_string()),
error_kind: Some("InvalidRequest".to_string()),
}]
);
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
assert!(accepted.is_err(), "provider socket should not be touched");
}
#[tokio::test]
async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let request = read_http_headers(&mut socket).await;
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":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 mut headers = Map::new();
headers.insert(
"Authorization".to_string(),
Value::String("Bearer sk-from-python".to_string()),
);
headers.insert(
"x-trace-id".to_string(),
Value::String("trace-1".to_string()),
);
let response = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-for-rust-fallback"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("mistral"),
extra_headers: Some(headers),
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
})
.await
.expect("ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
let request = server.await.expect("server task completes");
let authorization_count = request
.lines()
.filter(|line| line.to_ascii_lowercase().starts_with("authorization:"))
.count();
assert_eq!(authorization_count, 1, "{request}");
assert!(
request.contains("authorization: Bearer sk-from-python")
|| request.contains("Authorization: Bearer sk-from-python"),
"{request}"
);
}
#[tokio::test]
async fn document_intelligence_poll_uses_resolved_subscription_key() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let operation_url = format!("http://{addr}/operations/1");
let server = tokio::spawn(async move {
let (mut post_socket, _) = listener.accept().await.expect("accepts post request");
let post_request = read_http_headers(&mut post_socket).await;
let post_response = format!(
"HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n"
);
post_socket
.write_all(post_response.as_bytes())
.await
.expect("writes post response");
let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request");
let poll_request = read_http_headers(&mut poll_socket).await;
let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#;
let poll_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
);
poll_socket
.write_all(poll_response.as_bytes())
.await
.expect("writes poll response");
(post_request, poll_request)
});
let response = ocr(OcrRequest {
model: "doc-intelligence/prebuilt-read",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("di-key"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
})
.await
.expect("document intelligence request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
let (post_request, poll_request) = server.await.expect("server task completes");
assert!(
post_request
.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: di-key"),
"{post_request}"
);
assert!(
poll_request
.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: di-key"),
"{poll_request}"
);
}

2
litellm-rust/crates/core/AGENTS.md
View file

@ -2,6 +2,6 @@ litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-leve
A route module owns everything the call needs: types, the provider template trait, provider transforms (under `providers/`), provider/auth/URL resolution, and the handler that performs the HTTP call. Handlers belong here, never in a host crate.
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback dispatch. Env reads are limited to credential fallback in a route's `prepare.rs`.
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or host-specific callback execution. Core owns lifecycle sequencing and callback payload construction; hosts execute the selected integrations. Env reads are limited to credential fallback in a route's `prepare.rs`.
Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates.

66
litellm-rust/crates/core/CLAUDE.md
View file

@ -1,66 +0,0 @@
# CLAUDE.md
Rules for `litellm-rust/crates/core`.
## Responsibility
`core` is the LiteLLM SDK in Rust: it makes the LLM call. Every top-level
LiteLLM call has a public entrypoint here, named after the route
(`messages::messages()` is the Rust equivalent of `litellm.messages()`), and
calling it returns a typed non-streaming response.
Allowed:
- The public entrypoint for a route, plus its `<route>_stream` variant when the
route supports streaming.
- Provider resolution, auth header construction, URL building, and the provider
HTTP call (shared reused client, connect + request timeouts).
- Shared request/response structs.
- Typed errors with stable, non-sensitive messages.
- Deterministic validation helpers.
- Serialization helpers that intentionally mirror Python output shape.
- Route templates that match Python base config responsibilities, such as
`messages::transformation::AnthropicMessagesProviderConfig`.
Not allowed:
- Serving HTTP: axum routers, extractors, and other transport concerns.
- Filesystem, database, or cache access.
- Config file reading or rollout state; the host resolves those and passes them
in. Env reads are limited to credential fallback in a route's `prepare.rs`.
- Logging callbacks, tracing spans, spend writes, or customer callbacks.
- Provider-specific branching that belongs in `providers`.
- Panics for user/provider-controlled input.
## Typed Contracts (core rule)
Trait and function boundaries MUST be strongly typed. No stringly-typed JSON
(`&str` / `String` / `Vec<String>` / bare `serde_json::Value`) as a transform
input or output. Parse wire bytes into typed structs/enums at the host edge;
`core` and `providers` operate only on those types (e.g. `RealtimeEvent`,
`RealtimeTransformResult`, `OcrRequestData`). A `type`-style discriminator is a
typed field on a struct, not a raw string threaded through the API.
## Structure
Use route names directly under `src/`: `messages`, `ocr`, future
`chat_completions`, `embeddings`, and similar top-level LiteLLM calls. Do not
invent broad names like `engine` for route contracts.
`src/messages` is the reference shape for a route module:
```
mod.rs pub async fn messages(..) (+ messages_stream)
types.rs request/response types
transformation.rs the provider template trait
prepare.rs provider resolution, auth headers, URL
handler.rs the provider call
client.rs the shared reqwest client
```
## Parity Rules
- Every shared type used by a provider transform needs unit tests for
serialization shape.
- If Python parity requires always emitting a `null` field instead of omitting
it, document that in code and pin it with a test.
- Error enums should preserve enough detail for Python/HTTP hosts to map errors
consistently without exposing document contents or upstream bodies.

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

@ -6,23 +6,33 @@ license.workspace = true
repository.workspace = true
autotests = false
[[test]]
name = "workspace_crate_allowlist"
path = "tests/workspace_crate_allowlist.rs"
[dependencies]
bytes.workspace = true
futures-util.workspace = true
base64.workspace = true
azure_core.workspace = true
azure_identity.workspace = true
data-url = "0.3.2"
gcp_auth.workspace = true
moka.workspace = true
mime_guess = "2.0.5"
rand.workspace = true
reqwest.workspace = true
rustls.workspace = true
rustls-native-certs.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_path_to_error = "0.1"
tokio.workspace = true
strum.workspace = true
subtle.workspace = true
tokio = { workspace = true, features = ["sync"] }
tokio-tungstenite.workspace = true
thiserror.workspace = true
tracing.workspace = true
tracing-subscriber = { workspace = true, optional = true }
sha2.workspace = true
url.workspace = true
veil.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
@ -44,5 +54,4 @@ observability = ["dep:tracing-subscriber"]
[dev-dependencies]
rstest.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
tracing-subscriber.workspace = true

183
litellm-rust/crates/core/src/auth/credential.rs Normal file
View file

@ -0,0 +1,183 @@
use std::future::Future;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
use veil::Redact;
use crate::AuthError;
use super::{ResolvedCredential, SecretValue, TokenProviderHandle};
pub fn credential_index(requested: &str, names: &[String]) -> Option<usize> {
names.iter().position(|name| name == requested)
}
pub fn credential_default_fields<'a>(
supplied: &[String],
credential_fields: &'a [String],
) -> Vec<&'a str> {
credential_fields
.iter()
.filter(|name| !supplied.contains(name))
.map(String::as_str)
.collect()
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialFileRef {
Path(PathBuf),
EnvironmentVariable(String),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialRef {
Explicit(SecretValue),
Env(String),
File(CredentialFileRef),
Request(String),
Host(String),
None,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialLookup {
Found(SecretValue),
Missing,
Declined,
}
pub type CredentialLookupFuture<'a> =
Pin<Box<dyn Future<Output = Result<CredentialLookup, AuthError>> + Send + 'a>>;
pub trait CredentialResolver: std::fmt::Debug + Send + Sync {
fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a>;
}
#[derive(Clone, Redact)]
pub struct CredentialResolverHandle(#[redact(with = "[REDACTED]")] Arc<dyn CredentialResolver>);
impl CredentialResolverHandle {
pub fn new(resolver: Arc<dyn CredentialResolver>) -> Self {
Self(resolver)
}
pub async fn resolve(&self, reference: &CredentialRef) -> Result<CredentialLookup, AuthError> {
self.0.resolve(reference).await
}
}
#[derive(Clone, Debug)]
pub enum CredentialPlan {
Static(CredentialRef),
Caller(TokenProviderHandle),
None,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialPlanResolution {
Resolved(ResolvedCredential),
Unavailable,
}
impl CredentialPlan {
pub async fn resolve(
&self,
resolver: &CredentialResolverHandle,
) -> Result<CredentialPlanResolution, AuthError> {
match self {
Self::Static(CredentialRef::Explicit(secret)) => Ok(
CredentialPlanResolution::Resolved(ResolvedCredential::Static(secret.clone())),
),
Self::Static(CredentialRef::None) | Self::None => {
Ok(CredentialPlanResolution::Unavailable)
}
Self::Static(reference) => match resolver.resolve(reference).await? {
CredentialLookup::Found(secret) => Ok(CredentialPlanResolution::Resolved(
ResolvedCredential::Static(secret),
)),
CredentialLookup::Missing | CredentialLookup::Declined => {
Ok(CredentialPlanResolution::Unavailable)
}
},
Self::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(AuthError::EmptyCallerCredential);
}
Ok(CredentialPlanResolution::Resolved(credential))
}
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::{
CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution,
CredentialRef, CredentialResolver, CredentialResolverHandle,
};
use crate::AuthError;
use crate::auth::SecretValue;
#[derive(Debug)]
struct HostResolver;
impl CredentialResolver for HostResolver {
fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async move {
Ok(match reference {
CredentialRef::Host(name) if name == "rotating-token" => {
CredentialLookup::Found(SecretValue::new("resolved"))
}
_ => CredentialLookup::Declined,
})
})
}
}
#[tokio::test]
async fn static_host_reference_resolves_at_acquisition_time() {
let resolver = CredentialResolverHandle::new(Arc::new(HostResolver));
let plan = CredentialPlan::Static(CredentialRef::Host("rotating-token".to_string()));
let resolved = plan.resolve(&resolver).await.unwrap();
assert!(matches!(resolved, CredentialPlanResolution::Resolved(_)));
}
#[tokio::test]
async fn declined_reference_is_available_for_pre_acquisition_fallback() {
let resolver = CredentialResolverHandle::new(Arc::new(HostResolver));
let plan = CredentialPlan::Static(CredentialRef::Request("api-key".to_string()));
assert_eq!(
plan.resolve(&resolver).await.unwrap(),
CredentialPlanResolution::Unavailable
);
}
#[derive(Debug)]
struct FailingResolver;
impl CredentialResolver for FailingResolver {
fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async { Err(AuthError::UnresolvedOidcReference) })
}
}
#[tokio::test]
async fn acquisition_failure_is_terminal() {
let resolver = CredentialResolverHandle::new(Arc::new(FailingResolver));
let plan = CredentialPlan::Static(CredentialRef::Host("token".to_string()));
let error = plan
.resolve(&resolver)
.await
.expect_err("acquisition errors cannot become fallback");
assert_eq!(error, AuthError::UnresolvedOidcReference);
}
}

128
litellm-rust/crates/core/src/auth/error.rs Normal file
View file

@ -0,0 +1,128 @@
use thiserror::Error;
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AuthError {
#[error("invalid authentication configuration: {0}")]
Configuration(#[from] AuthConfigurationError),
#[error("credential acquisition failed: {0}")]
AzureTokenAcquisition(String),
#[error("credential acquisition failed: Vertex AI credentials: {0}")]
VertexTokenAcquisition(String),
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
CredentialChain(Vec<AuthError>),
#[error("credential caller failed: credential caller returned an empty credential")]
EmptyCallerCredential,
#[error("credential caller failed: Azure AD token provider returned an empty token")]
EmptyAzureToken,
#[error("credential acquisition failed: Azure OIDC reference did not resolve to a value")]
UnresolvedOidcReference,
#[error(
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error(
"Missing {provider} API Base - Set {environment_variable} environment variable or pass api_base parameter"
)]
MissingApiBase {
provider: &'static str,
environment_variable: &'static str,
},
#[error("{0}")]
MissingCredential(#[from] MissingCredential),
#[error("{0}")]
Aws(#[from] AwsAuthError),
#[error("invalid authentication header")]
InvalidHeader,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AuthConfigurationError {
#[error("credential header already exists")]
ExistingCredentialHeader,
#[error("credential plan is not allowed by the provider auth policy")]
DisallowedCredentialPlan,
#[error("credential cannot be empty")]
EmptyCredential,
#[error("invalid Azure credential selector")]
InvalidAzureSelector,
#[error("ClientSecretCredential requires tenant_id, client_id, and client_secret")]
MissingClientSecretFields,
#[error("WorkloadIdentityCredential requires tenant_id")]
MissingWorkloadTenant,
#[error("WorkloadIdentityCredential requires client_id")]
MissingWorkloadClient,
#[error("WorkloadIdentityCredential requires azure_federated_token_file")]
MissingWorkloadTokenFile,
#[error("credential reference requires a host credential resolver")]
MissingHostResolver,
#[error("caller credential plan requires provider-specific inputs")]
MissingCallerInputs,
#[error("credential header {0} already exists")]
DuplicateHeader(&'static str),
#[error("{0} must be a string or null")]
InvalidFieldType(String),
#[error("unsupported OIDC reference")]
UnsupportedOidcReference,
#[error("{0} cannot be empty")]
EmptyReference(String),
#[error("Azure credential initialization failed: {0}")]
AzureCredentialInitialization(String),
#[error("Azure authority must be an HTTPS origin without credentials, query, or fragment")]
InvalidAzureAuthority,
#[error("request-controlled Azure auth inputs cannot be combined with host credentials")]
MixedAzureCredentialSources,
#[error("request-controlled Azure credential references are not allowed")]
RequestAzureCredentialReference,
#[error("host credentials cannot be sent to a request-controlled Azure endpoint")]
RequestAzureCredentialDestination,
#[error("credentials cannot be sent to a request-controlled Vertex AI endpoint")]
RequestVertexCredentialDestination,
#[error(
"request-controlled Vertex credentials must use the canonical Google OAuth token endpoint"
)]
RequestVertexTokenEndpoint,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum MissingCredential {
#[error(
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"
)]
AnthropicApiKey,
#[error("Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable")]
AzureApiKey,
#[error(
"Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. Expected format: https://<resource-name>.services.ai.azure.com/anthropic"
)]
AzureApiBase,
#[error(
"Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"
)]
OpenAiRealtimeApiKey,
#[error(
"Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"
)]
OpenAiResponsesApiKey,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AwsAuthError {
#[error("AWS profile credentials failed: {0}")]
Profile(String),
#[error("AWS default credentials failed: {0}")]
DefaultChain(String),
#[error("AWS role credentials failed: {0}")]
AssumeRole(String),
#[error("AWS web identity credentials failed: {0}")]
WebIdentity(String),
#[error("AWS web identity expiration was invalid: {0}")]
WebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
SigningParameters(String),
#[error("AWS signable request failed: {0}")]
SignableRequest(String),
#[error("AWS request signing failed: {0}")]
Signing(String),
#[error("AWS web identity response had no credentials")]
MissingWebIdentityCredentials,
}

86
litellm-rust/crates/core/src/auth/http.rs Normal file
View file

@ -0,0 +1,86 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlacement {
Bearer,
Header(&'static str),
}
impl CredentialPlacement {
pub fn header_name(self) -> &'static str {
match self {
Self::Bearer => "Authorization",
Self::Header(name) => name,
}
}
}
pub(crate) fn apply_credential(
headers: Vec<(String, String)>,
credential: &str,
placement: CredentialPlacement,
) -> Result<Vec<(String, String)>, AuthError> {
if credential.trim().is_empty() {
return Err(AuthError::Configuration(
AuthConfigurationError::EmptyCredential,
));
}
if headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case(placement.header_name()))
{
return Err(AuthError::Configuration(
AuthConfigurationError::DuplicateHeader(placement.header_name()),
));
}
let value = match placement {
CredentialPlacement::Bearer => format!("Bearer {credential}"),
CredentialPlacement::Header(_) => credential.to_string(),
};
Ok(
std::iter::once((placement.header_name().to_string(), value))
.chain(headers)
.collect(),
)
}
/// How the upstream call is authenticated. API-key strategies are resolved in
/// `prepare`; SigV4 needs the serialized body, so the handler signs it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RequestAuth {
Header { name: &'static str, value: String },
Bearer { token: String },
AwsSigV4 { region: String },
}
#[cfg(test)]
mod tests {
use super::{CredentialPlacement, apply_credential};
#[test]
fn bearer_uses_authorization_header() {
let headers = apply_credential(Vec::new(), "key", CredentialPlacement::Bearer)
.expect("credential applies");
assert_eq!(
headers,
vec![("Authorization".to_string(), "Bearer key".to_string())]
);
}
#[test]
fn named_header_rejects_existing_value() {
let error = apply_credential(
vec![(
"ocp-apim-subscription-key".to_string(),
"caller-key".to_string(),
)],
"configured-key",
CredentialPlacement::Header("Ocp-Apim-Subscription-Key"),
)
.expect_err("provider policy must handle existing credentials");
assert!(error.to_string().contains("already exists"));
}
}

57
litellm-rust/crates/core/src/auth/mod.rs Normal file
View file

@ -0,0 +1,57 @@
mod credential;
pub mod error;
pub(crate) mod vertex;
pub use error::AuthError;
pub(crate) mod http;
mod policy;
mod secret;
mod token;
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum InputSource {
Request,
#[default]
Deployment,
Environment,
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct Sourced<T> {
value: T,
source: InputSource,
}
impl<T> Sourced<T> {
pub fn new(value: T, source: InputSource) -> Self {
Self { value, source }
}
pub fn value(&self) -> &T {
&self.value
}
pub fn source(&self) -> InputSource {
self.source
}
pub fn into_value(self) -> T {
self.value
}
pub fn map<U>(self, map: impl FnOnce(T) -> U) -> Sourced<U> {
Sourced::new(map(self.value), self.source)
}
}
pub use credential::{
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
credential_default_fields, credential_index,
};
pub use http::{CredentialPlacement, RequestAuth};
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
pub use secret::SecretValue;
pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};

114
litellm-rust/crates/core/src/auth/policy.rs Normal file
View file

@ -0,0 +1,114 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
use super::http::apply_credential;
use super::{CredentialPlacement, ResolvedCredential};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlanKind {
Static,
Entra,
Caller,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct CredentialRule {
pub kind: CredentialPlanKind,
pub placement: CredentialPlacement,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExistingHeaderBehavior {
Preserve,
Reject,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ProviderAuthPolicy {
pub rules: &'static [CredentialRule],
pub accepted_existing_headers: &'static [&'static str],
pub existing_header_behavior: ExistingHeaderBehavior,
pub scope: Option<&'static str>,
pub audience: Option<&'static str>,
}
impl ProviderAuthPolicy {
pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, _)| {
self.accepted_existing_headers
.iter()
.any(|accepted| name.eq_ignore_ascii_case(accepted))
})
}
pub fn apply(
&self,
headers: Vec<(String, String)>,
kind: CredentialPlanKind,
credential: &ResolvedCredential,
) -> Result<Vec<(String, String)>, AuthError> {
if self.has_existing_credential(&headers) {
return match self.existing_header_behavior {
ExistingHeaderBehavior::Preserve => Ok(headers),
ExistingHeaderBehavior::Reject => Err(AuthError::Configuration(
AuthConfigurationError::ExistingCredentialHeader,
)),
};
}
let rule =
self.rules
.iter()
.find(|rule| rule.kind == kind)
.ok_or(AuthError::Configuration(
AuthConfigurationError::DisallowedCredentialPlan,
))?;
apply_credential(headers, credential.secret().expose(), rule.placement)
}
}
#[cfg(test)]
mod tests {
use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
use crate::auth::{CredentialPlacement, ResolvedCredential, SecretValue};
const RULES: &[CredentialRule] = &[CredentialRule {
kind: CredentialPlanKind::Static,
placement: CredentialPlacement::Header("x-api-key"),
}];
const POLICY: ProviderAuthPolicy = ProviderAuthPolicy {
rules: RULES,
accepted_existing_headers: &["x-api-key"],
existing_header_behavior: ExistingHeaderBehavior::Preserve,
scope: None,
audience: None,
};
#[test]
fn rules_define_allowed_plans_and_credential_placement() {
let headers = POLICY
.apply(
Vec::new(),
CredentialPlanKind::Static,
&ResolvedCredential::Static(SecretValue::new("secret")),
)
.unwrap();
assert_eq!(
headers,
vec![("x-api-key".to_string(), "secret".to_string())]
);
}
#[test]
fn unsupported_plan_is_rejected() {
let error = POLICY
.apply(
Vec::new(),
CredentialPlanKind::Entra,
&ResolvedCredential::Static(SecretValue::new("secret")),
)
.unwrap_err();
assert!(error.to_string().contains("not allowed"));
}
}

41
litellm-rust/crates/core/src/auth/secret.rs Normal file
View file

@ -0,0 +1,41 @@
use veil::Redact;
#[derive(Redact, Clone)]
pub struct SecretValue(#[redact(with = "[REDACTED]")] String);
impl SecretValue {
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
pub fn expose(&self) -> &str {
&self.0
}
}
impl PartialEq for SecretValue {
fn eq(&self, other: &Self) -> bool {
subtle::ConstantTimeEq::ct_eq(self.0.as_bytes(), other.0.as_bytes()).into()
}
}
impl Eq for SecretValue {}
#[cfg(test)]
mod tests {
use super::SecretValue;
#[test]
fn debug_redacts_plaintext() {
let debug = format!("{:?}", SecretValue::new("credential-value"));
assert!(!debug.contains("credential-value"));
assert!(debug.contains("REDACTED"));
}
#[test]
fn equality_compares_plaintext_values() {
assert_eq!(SecretValue::new("same"), SecretValue::new("same"));
assert_ne!(SecretValue::new("same"), SecretValue::new("different"));
}
}

47
litellm-rust/crates/core/src/auth/token.rs Normal file
View file

@ -0,0 +1,47 @@
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::SystemTime;
use veil::Redact;
use crate::AuthError;
use super::secret::SecretValue;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ResolvedCredential {
Static(SecretValue),
AccessToken {
token: SecretValue,
expires_on: Option<SystemTime>,
},
}
impl ResolvedCredential {
pub fn secret(&self) -> &SecretValue {
match self {
Self::Static(secret) | Self::AccessToken { token: secret, .. } => secret,
}
}
}
pub type TokenFuture<'a> =
Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + 'a>>;
pub trait TokenProvider: std::fmt::Debug + Send + Sync {
fn acquire(&self) -> TokenFuture<'_>;
}
#[derive(Clone, Redact)]
pub struct TokenProviderHandle(#[redact(with = "[REDACTED]")] Arc<dyn TokenProvider>);
impl TokenProviderHandle {
pub fn new(caller: Arc<dyn TokenProvider>) -> Self {
Self(caller)
}
pub async fn acquire(&self) -> Result<ResolvedCredential, AuthError> {
self.0.acquire().await
}
}

592
litellm-rust/crates/core/src/auth/vertex.rs Normal file
View file

@ -0,0 +1,592 @@
use std::collections::BTreeMap;
use std::future::Future;
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use gcp_auth::{CustomServiceAccount, TokenProvider};
use moka::future::Cache;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use crate::auth::error::AuthConfigurationError;
use crate::auth::http::apply_credential;
use crate::auth::{AuthError, CredentialPlacement, InputSource, SecretValue, Sourced};
const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS";
const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY";
const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY";
const VERTEXAI_CREDENTIALS_ENV: &str = "VERTEXAI_CREDENTIALS";
const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT";
const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
#[derive(Clone, Debug, Default)]
pub(crate) struct VertexConfig {
credentials: Option<Sourced<SecretValue>>,
project_id: Option<String>,
location: Option<String>,
}
impl VertexConfig {
pub(crate) fn from_sourced_optional_params(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
) -> Result<Self, AuthError> {
Ok(Self {
credentials: optional_credentials(
params,
sources,
&["vertex_credentials", "vertex_ai_credentials"],
)?,
project_id: optional_string(params, &["vertex_project", "vertex_ai_project"])?,
location: optional_string(params, &["vertex_location", "vertex_ai_location"])?,
})
}
pub(crate) fn project_id(&self) -> Option<&str> {
self.project_id.as_deref()
}
pub(crate) fn location(&self) -> Option<&str> {
self.location.as_deref()
}
}
pub(crate) struct VertexEnvironment {
pub headers: Vec<(String, String)>,
pub project_id: String,
}
struct VertexAccessToken {
token: String,
project_id: String,
}
pub(crate) fn get_vertex_ai_project(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
config
.project_id()
.map(str::to_string)
.or_else(|| non_empty_env(env_lookup, VERTEXAI_PROJECT_ENV))
}
pub(crate) fn get_vertex_ai_location(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
config
.location()
.map(str::to_string)
.or_else(|| non_empty_env(env_lookup, VERTEXAI_LOCATION_ENV))
.or_else(|| non_empty_env(env_lookup, VERTEX_LOCATION_ENV))
}
#[derive(Clone)]
pub(crate) struct VertexAuth {
providers: Cache<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
loader: Arc<dyn VertexProviderLoader>,
}
impl Default for VertexAuth {
fn default() -> Self {
Self::new(Arc::new(GcpProviderLoader))
}
}
impl VertexAuth {
fn new(loader: Arc<dyn VertexProviderLoader>) -> Self {
Self {
providers: Cache::builder().max_capacity(64).build(),
loader,
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) async fn validate_environment(
&self,
headers: Vec<(String, String)>,
api_key: Option<&str>,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<VertexEnvironment, AuthError> {
let has_authorization = headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("Authorization"));
let static_token = api_key
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.or_else(|| non_empty_env(env_lookup, VERTEX_AI_API_KEY_ENV))
.or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV));
let project_id = get_vertex_ai_project(config, env_lookup);
if !has_authorization && static_token.is_none() {
let access = self.get_access_token(config, env_lookup).await?;
return Ok(VertexEnvironment {
headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
project_id: project_id.unwrap_or(access.project_id),
});
}
let project_id = match project_id {
Some(project_id) => project_id,
None => {
self.load_provider(config, env_lookup)
.await?
.project_id()
.await?
}
};
let headers = if has_authorization {
headers
} else {
apply_credential(
headers,
static_token.as_deref().expect("static token was checked"),
CredentialPlacement::Bearer,
)?
};
Ok(VertexEnvironment {
headers,
project_id,
})
}
async fn get_access_token(
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<VertexAccessToken, AuthError> {
let provider = self.load_provider(config, env_lookup).await?;
let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?;
Ok(VertexAccessToken { token, project_id })
}
async fn load_provider(
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Arc<dyn VertexTokenSource>, AuthError> {
let source = credential_source(config, env_lookup);
let key = source.cache_key();
self.providers
.try_get_with(key, self.loader.load(source))
.await
.map_err(|error| (*error).clone())
}
}
trait VertexTokenSource: Send + Sync {
fn project_id(&self) -> VertexAuthFuture<'_, String>;
fn token(&self) -> VertexAuthFuture<'_, String>;
}
trait VertexProviderLoader: Send + Sync {
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>>;
}
type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, AuthError>> + Send + 'a>>;
struct GcpTokenSource(Arc<dyn TokenProvider>);
impl VertexTokenSource for GcpTokenSource {
fn project_id(&self) -> VertexAuthFuture<'_, String> {
Box::pin(async move {
self.0
.project_id()
.await
.map(|project| project.to_string())
.map_err(auth_acquisition_error)
})
}
fn token(&self) -> VertexAuthFuture<'_, String> {
Box::pin(async move {
self.0
.token(&[CLOUD_PLATFORM_SCOPE])
.await
.map(|token| token.as_str().to_string())
.map_err(auth_acquisition_error)
})
}
}
struct GcpProviderLoader;
impl VertexProviderLoader for GcpProviderLoader {
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
Box::pin(async move {
let provider: Arc<dyn TokenProvider> = match source {
CredentialSource::Inline(configured) => Arc::new(
CustomServiceAccount::from_json(validate_request_credentials(
configured.expose(),
)?)
.map_err(auth_acquisition_error)?,
),
CredentialSource::Trusted(configured) => {
let configured = configured.expose();
let service_account = if Path::new(configured).is_file() {
CustomServiceAccount::from_file(configured)
} else {
CustomServiceAccount::from_json(configured)
}
.map_err(auth_acquisition_error)?;
Arc::new(service_account)
}
CredentialSource::ApplicationCredentials(path) => {
Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
}
CredentialSource::Adc => {
gcp_auth::provider().await.map_err(auth_acquisition_error)?
}
};
Ok(Arc::new(GcpTokenSource(provider)) as Arc<dyn VertexTokenSource>)
})
}
}
fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> {
let token_uri = serde_json::from_str::<Value>(configured)
.ok()
.and_then(|credentials| {
credentials
.get("token_uri")
.and_then(Value::as_str)
.map(str::to_string)
});
if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
return Err(AuthConfigurationError::RequestVertexTokenEndpoint.into());
}
Ok(configured)
}
#[derive(Clone, Debug)]
enum CredentialSource {
Inline(SecretValue),
Trusted(SecretValue),
ApplicationCredentials(String),
Adc,
}
impl CredentialSource {
fn cache_key(&self) -> CredentialCacheKey {
match self {
Self::Inline(configured) => {
CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into())
}
Self::Trusted(configured) => {
CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into())
}
Self::ApplicationCredentials(path) => {
CredentialCacheKey::ApplicationCredentials(path.clone())
}
Self::Adc => CredentialCacheKey::Adc,
}
}
}
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
enum CredentialCacheKey {
Inline([u8; 32]),
Trusted([u8; 32]),
ApplicationCredentials(String),
Adc,
}
fn credential_source(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CredentialSource {
if let Some(configured) = config.credentials.clone() {
return match configured.source() {
InputSource::Request => CredentialSource::Inline(configured.into_value()),
InputSource::Deployment | InputSource::Environment => {
CredentialSource::Trusted(configured.into_value())
}
};
}
if let Some(configured) = non_empty_env(env_lookup, VERTEXAI_CREDENTIALS_ENV) {
return CredentialSource::Trusted(SecretValue::new(configured));
}
non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV)
.map(CredentialSource::ApplicationCredentials)
.unwrap_or(CredentialSource::Adc)
}
fn optional_credentials(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
names: &[&str],
) -> Result<Option<Sourced<SecretValue>>, AuthError> {
for name in names {
let source = source_for(sources, name);
match params.get(*name) {
None | Some(Value::Null) => continue,
Some(Value::String(value)) if value.trim().is_empty() => continue,
Some(Value::String(value)) => {
return Ok(Some(Sourced::new(SecretValue::new(value), source)));
}
Some(Value::Object(value)) if value.is_empty() => continue,
Some(Value::Object(value)) => {
return serde_json::to_string(value)
.map(SecretValue::new)
.map(|value| Sourced::new(value, source))
.map(Some)
.map_err(|error| {
AuthError::Configuration(AuthConfigurationError::InvalidFieldType(format!(
"{}: {error}",
names[0]
)))
});
}
Some(_) => {
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
));
}
}
}
Ok(None)
}
fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSource {
sources.get(name).copied().unwrap_or_default()
}
fn optional_string(
params: &Map<String, Value>,
names: &[&str],
) -> Result<Option<String>, AuthError> {
for name in names {
match params.get(*name) {
None | Some(Value::Null) => continue,
Some(Value::String(value)) if value.trim().is_empty() => continue,
Some(Value::String(value)) => return Ok(Some(value.clone())),
Some(_) => {
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
));
}
}
}
Ok(None)
}
fn non_empty_env(env_lookup: &dyn Fn(&str) -> Option<String>, name: &str) -> Option<String> {
env_lookup(name)
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
fn auth_acquisition_error(error: gcp_auth::Error) -> AuthError {
AuthError::VertexTokenAcquisition(error.to_string())
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use serde_json::json;
use super::*;
struct FakeProvider {
calls: Arc<AtomicUsize>,
}
impl VertexTokenSource for FakeProvider {
fn project_id(&self) -> VertexAuthFuture<'_, String> {
self.calls.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Ok("adc-project".into()) })
}
fn token(&self) -> VertexAuthFuture<'_, String> {
self.calls.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Ok("adc-token".into()) })
}
}
struct FakeLoader {
loads: Arc<AtomicUsize>,
provider: Arc<dyn VertexTokenSource>,
}
impl VertexProviderLoader for FakeLoader {
fn load(
&self,
_source: CredentialSource,
) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
let loads = self.loads.clone();
let provider = self.provider.clone();
Box::pin(async move {
loads.fetch_add(1, Ordering::SeqCst);
Ok(provider)
})
}
}
fn config(value: Value) -> VertexConfig {
VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new())
.unwrap()
}
fn auth(calls: Arc<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
let provider: Arc<dyn VertexTokenSource> = Arc::new(FakeProvider { calls });
VertexAuth::new(Arc::new(FakeLoader { loads, provider }))
}
#[test]
fn config_is_typed_and_secrets_are_redacted() {
let config = config(json!({
"vertex_credentials":{"private_key":"secret-key"},
"vertex_project":"project-1",
"vertex_location":"europe-west4"
}));
assert_eq!(config.project_id(), Some("project-1"));
assert_eq!(config.location(), Some("europe-west4"));
assert!(!format!("{config:?}").contains("secret-key"));
assert!(
VertexConfig::from_sourced_optional_params(
json!({"vertex_credentials":true}).as_object().unwrap(),
&BTreeMap::new()
)
.is_err()
);
}
#[test]
fn empty_primary_values_fall_back_to_python_aliases() {
let config = config(json!({
"vertex_credentials": null,
"vertex_ai_credentials": "alias-credentials",
"vertex_project": " ",
"vertex_ai_project": "alias-project",
"vertex_location": null,
"vertex_ai_location": "alias-location"
}));
assert_eq!(
config.credentials.as_ref().unwrap().value().expose(),
"alias-credentials"
);
assert_eq!(config.project_id(), Some("alias-project"));
assert_eq!(config.location(), Some("alias-location"));
}
#[test]
fn project_and_location_prefer_input_then_environment() {
let configured =
config(json!({"vertex_project":"input-project","vertex_location":"input-location"}));
let env = |name: &str| Some(format!("env-{name}"));
assert_eq!(
get_vertex_ai_project(&configured, &env).as_deref(),
Some("input-project")
);
assert_eq!(
get_vertex_ai_location(&configured, &env).as_deref(),
Some("input-location")
);
let empty = VertexConfig::default();
assert_eq!(
get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(),
Some("env-project")
);
assert_eq!(
get_vertex_ai_location(&empty, &|name| (name == VERTEX_LOCATION_ENV)
.then(|| "fallback-location".into()))
.as_deref(),
Some("fallback-location")
);
}
#[test]
fn credential_discovery_prefers_input_then_environment_then_adc() {
let params = json!({"vertex_credentials":"input-json"});
let sources = BTreeMap::from([("vertex_credentials".to_string(), InputSource::Request)]);
let configured =
VertexConfig::from_sourced_optional_params(params.as_object().unwrap(), &sources)
.unwrap();
assert!(
matches!(credential_source(&configured, &|_| Some("environment-value".into())), CredentialSource::Inline(value) if value.expose() == "input-json")
);
let empty = VertexConfig::default();
assert!(
matches!(credential_source(&empty, &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into())), CredentialSource::Trusted(value) if value.expose() == "environment-json")
);
assert!(
matches!(credential_source(&empty, &|name| (name == GOOGLE_APPLICATION_CREDENTIALS_ENV).then(|| "adc.json".into())), CredentialSource::ApplicationCredentials(path) if path == "adc.json")
);
assert!(matches!(
credential_source(&empty, &|_| None),
CredentialSource::Adc
));
assert_ne!(
CredentialSource::Inline(SecretValue::new("same-value")).cache_key(),
CredentialSource::Trusted(SecretValue::new("same-value")).cache_key()
);
}
#[test]
fn request_credentials_require_canonical_token_endpoint() {
assert!(
validate_request_credentials(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#)
.is_ok()
);
assert!(matches!(
validate_request_credentials(r#"{"token_uri":"http://127.0.0.1/token"}"#),
Err(AuthError::Configuration(
AuthConfigurationError::RequestVertexTokenEndpoint
))
));
assert!(matches!(
validate_request_credentials("{}"),
Err(AuthError::Configuration(
AuthConfigurationError::RequestVertexTokenEndpoint
))
));
}
#[tokio::test]
async fn explicit_token_and_header_do_not_acquire_adc() {
let loads = Arc::new(AtomicUsize::new(0));
let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone());
let configured = config(json!({"vertex_project":"project-1"}));
let explicit = auth
.validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None)
.await
.unwrap();
assert_eq!(explicit.headers[0].1, "Bearer access-token");
let existing = auth
.validate_environment(
vec![("authorization".into(), "Bearer existing".into())],
None,
&configured,
&|_| None,
)
.await
.unwrap();
assert_eq!(existing.headers[0].1, "Bearer existing");
assert_eq!(loads.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn provider_is_reused_across_authentication_calls() {
let calls = Arc::new(AtomicUsize::new(0));
let loads = Arc::new(AtomicUsize::new(0));
let auth = auth(calls.clone(), loads.clone());
for _ in 0..2 {
let environment = auth
.validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None)
.await
.unwrap();
assert_eq!(environment.project_id, "adc-project");
assert_eq!(environment.headers[0].1, "Bearer adc-token");
}
assert_eq!(loads.load(Ordering::SeqCst), 1);
assert_eq!(calls.load(Ordering::SeqCst), 4);
}
}

167
litellm-rust/crates/core/src/call_lifecycle/README.md
View file

@ -1,167 +0,0 @@
# Call lifecycle
`litellm_core::call_lifecycle` is the shared execution wrapper for LiteLLM call
types migrated to Rust. It owns lifecycle ordering, phase timing, and trace
observer calls. It must not know about OCR, chat, messages, responses,
completions, provider auth, request transforms, or response normalization.
Call-type modules own their domain behavior. For example, OCR owns document
payloads, OCR provider transforms, safe document fetch, guardrail payload shape,
callback payload shape, and provider HTTP execution.
## Runtime order
Every wrapped call runs in this order:
1. `async_pre_call_hook`
2. `async_during_call_hook`
3. provider call
4. `async_log_success_event` or `async_log_failure_event`
`async_pre_call_hook` receives the initial LiteLLM request shape. It is where
pre-call custom guardrails run.
`async_during_call_hook` converts the initial request into the provider-ready
request. It is where provider config selection, parameter mapping, auth/header
resolution, request transforms, and during-call guardrails belong.
The provider call receives only the provider-ready request. It should execute
I/O and call the provider response transform.
Success and failure callbacks receive `CallLifecycleTiming`. Callback failures
must not replace the original provider or guardrail result.
## Trace contract
The lifecycle runner records:
- full call start and end time
- `pre_call` phase timing
- `during_call` phase timing
- `provider_call` phase timing
- `success_callback` phase timing
- `failure_callback` phase timing
`CallLifecycleObserver` receives phase start and end events. The default
observer is a no-op. Future OTEL support should implement this observer instead
of editing OCR, chat, messages, responses, completions, or provider modules.
## Required shape
Each migrated call type should use this folder shape:
```text
litellm-rust/crates/ai-gateway/src/<call_type>/
mod.rs # thin public entrypoint
types.rs # public request, prepared request, provider request, response types
prepare.rs # model/provider/callback/guardrail setup
hooks.rs # CallLifecycleHooks implementation
handler.rs # provider I/O and response normalization
tests.rs # call-type lifecycle and handler tests
```
Provider transforms can live in `litellm-rust/crates/core/src/providers/...`.
Shared call-type helpers can live beside the call type, but generic lifecycle
code stays in this folder.
## Core API
The prepared request implements `CallLifecycleRequest`:
```rust
impl CallLifecycleRequest for PreparedMessagesRequest {
fn lifecycle_context(&self) -> CallLifecycleContext {
CallLifecycleContext::new(
"messages",
self.model.clone(),
self.custom_llm_provider.clone(),
self.litellm_call_id.clone(),
)
}
}
```
The call-type hooks implement `CallLifecycleHooks`:
```rust
impl CallLifecycleHooks<
PreparedMessagesRequest,
ProviderMessagesRequest,
MessagesResponse,
> for MessagesLifecycleHooks {
fn async_pre_call_hook(...) {
// run pre-call custom guardrails against the LiteLLM request shape
}
fn async_during_call_hook(...) {
// map params, validate env, transform request, run during-call guardrails
}
fn async_log_success_event(...) {
// call async_log_success_event on configured custom loggers
}
fn async_log_failure_event(...) {
// call async_log_failure_event without swallowing the original error
}
}
```
The public entrypoint stays thin:
```rust
pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<MessagesResponse> {
let PreparedMessagesCall { request, hooks } = prepare_messages_call(request)?;
CallLifecycle::default()
.run_request(request, &hooks, execute_messages_provider_call)
.await
}
```
Use `run_request` for new call types. Keep `run` available only for specialized
tests or existing code that already has a `CallLifecycleContext`.
## Adding a new call type
1. Add `<call_type>/types.rs`
Define the public request accepted by the bridge, the prepared request used by
the lifecycle runner, and the provider request consumed by the handler.
2. Implement `CallLifecycleRequest`
Return `call_type`, `model`, `custom_llm_provider`, and `litellm_call_id`.
Do not put provider-specific logic here.
3. Add `<call_type>/prepare.rs`
Resolve model/provider once, generate or preserve `litellm_call_id`, construct
callback and guardrail runners, and return `Prepared<CallType>Call`.
4. Add `<call_type>/hooks.rs`
Implement `CallLifecycleHooks`. Put pre-call guardrail payload construction,
provider config selection, param mapping, request transform, during-call
guardrail payload construction, and callback payload construction here.
5. Add `<call_type>/handler.rs`
Execute the provider request and normalize the provider response. Do not repeat
provider-specific transforms here; call the provider config.
6. Add tests
Cover hook order, success callback payload, failure callback payload, pre-call
guardrail blocking before provider I/O, during-call body mutation, and provider
error mapping.
## Review checklist
- Core lifecycle has no call-type or provider-specific branches
- Public call-type entrypoint only prepares and calls `run_request`
- Provider behavior lives behind provider config/transformation code
- Hook method names map to the Python custom logger and guardrail concepts
- Phase timing is recorded once in lifecycle, not separately per call type
- Callback failures never hide the original provider or guardrail error
- Tests prove the provider socket is not touched when pre-call guardrails block

121
litellm-rust/crates/core/src/call_lifecycle/host.rs Normal file
View file

@ -0,0 +1,121 @@
use std::future::Future;
use std::pin::Pin;
pub enum HostCallStep<O, C> {
Host(O),
Complete(C),
}
pub type HostCallFuture<'a, O, C> =
Pin<Box<dyn Future<Output = Result<HostCallStep<O, C>, crate::Error>> + Send + 'a>>;
pub trait HostCall: Send + Sync {
type Operation: Send + 'static;
type Result: Send + 'static;
type Complete: Send + 'static;
fn resume(
&mut self,
result: Option<Self::Result>,
) -> HostCallFuture<'_, Self::Operation, Self::Complete>;
fn interrupt(
&mut self,
failure: HostFailure,
) -> HostCallFuture<'_, Self::Operation, Self::Complete>;
}
pub enum HostStep<V, S> {
Ready(V),
Suspend(S),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HostPhase {
Setup,
DeploymentPreCall,
Prepare,
Execute,
ConstructResponse,
DeploymentPostCall,
Finalize,
Success,
MapFailure,
DeploymentFailure,
Failure,
AsyncFailure,
Complete,
}
#[derive(Clone, Debug)]
pub enum HostFailure {
Error(crate::Error),
Cancelled(crate::Error),
}
pub struct HostLifecycle {
phase: HostPhase,
asynchronous: bool,
}
impl HostLifecycle {
pub fn new(asynchronous: bool) -> Self {
Self {
phase: HostPhase::Setup,
asynchronous,
}
}
pub fn phase(&self) -> HostPhase {
self.phase
}
pub fn accept(&mut self, result: Result<(), HostFailure>) -> Option<crate::Error> {
if let Err(failure) = result {
if self.phase == HostPhase::DeploymentFailure {
self.phase = HostPhase::Failure;
return None;
}
let error = match failure {
HostFailure::Cancelled(error) => {
self.phase = HostPhase::Complete;
return Some(error);
}
HostFailure::Error(error) => error,
};
match self.phase {
HostPhase::Failure | HostPhase::AsyncFailure => {
self.advance();
return None;
}
HostPhase::Success => self.phase = HostPhase::Complete,
HostPhase::Execute | HostPhase::ConstructResponse => {
self.phase = HostPhase::MapFailure;
}
_ => self.phase = HostPhase::Failure,
}
return Some(error);
}
self.advance();
None
}
fn advance(&mut self) {
self.phase = match self.phase {
HostPhase::Setup if self.asynchronous => HostPhase::DeploymentPreCall,
HostPhase::Setup | HostPhase::DeploymentPreCall => HostPhase::Prepare,
HostPhase::Prepare => HostPhase::Execute,
HostPhase::Execute => HostPhase::ConstructResponse,
HostPhase::ConstructResponse if self.asynchronous => HostPhase::DeploymentPostCall,
HostPhase::ConstructResponse | HostPhase::DeploymentPostCall => HostPhase::Finalize,
HostPhase::Finalize => HostPhase::Success,
HostPhase::MapFailure if self.asynchronous => HostPhase::DeploymentFailure,
HostPhase::MapFailure | HostPhase::DeploymentFailure => HostPhase::Failure,
HostPhase::Failure if self.asynchronous => HostPhase::AsyncFailure,
HostPhase::Failure
| HostPhase::AsyncFailure
| HostPhase::Success
| HostPhase::Complete => HostPhase::Complete,
};
}
}

4
litellm-rust/crates/core/src/call_lifecycle/mod.rs
View file

@ -3,6 +3,10 @@ use std::time::{Instant, SystemTime, UNIX_EPOCH};
use crate::Error;
pub mod host;
#[cfg(test)]
#[path = "../../tests/host_lifecycle.rs"]
mod host_tests;
pub mod types;
pub use types::{

21
litellm-rust/crates/core/src/constants.rs
View file

@ -43,6 +43,27 @@ pub const EMPTY_TEXT_PLACEHOLDER: &str =
"[System: Empty message content sanitised to satisfy protocol]";
pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";
pub(crate) const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10;
pub(crate) const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024;
pub(crate) const OCR_HTTP_TIMEOUT_SECS: u64 = 600;
pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10;
pub const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024;
pub(crate) const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024;
pub(crate) const OCR_MAX_FETCH_REDIRECTS: usize = 10;
pub(crate) const OCR_POLL_TIMEOUT_SECS: u64 = 120;
pub(crate) const OCR_POLL_RETRY_SECS: u64 = 2;
pub(crate) const AZURE_DI_API_VERSION: &str = "2024-11-30";
pub(crate) const AZURE_DI_SUBSCRIPTION_HEADER: &str = "Ocp-Apim-Subscription-Key";
pub(crate) const AZURE_DI_DEFAULT_DPI: i64 = 96;
pub(crate) const AZURE_DI_DEFAULT_WIDTH: f64 = 8.5;
pub(crate) const AZURE_DI_DEFAULT_HEIGHT: f64 = 11.0;
pub(crate) const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
pub(crate) const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
pub(crate) const REDUCTO_ID_PREFIX: &str = "reducto://";
pub(crate) const AZURE_AI_OCR_PATH: &str = "/providers/mistral/azure/ocr";
pub(crate) const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1";
pub(crate) const COHERE_PARSE_API_BASE: &str = "https://api.cohere.com";
pub(crate) const COHERE_API_KEY_ENV: &str = "COHERE_API_KEY";

67
litellm-rust/crates/core/src/error.rs
View file

@ -1,6 +1,6 @@
use thiserror::Error as ThisError;
#[derive(Debug, ThisError, PartialEq, Eq)]
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum Error {
#[error("expected {expected}, got {actual}")]
InvalidType {
@ -9,6 +9,8 @@ pub enum Error {
},
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("Document URL is required")]
MissingDocumentUrl,
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("invalid provider: {0}")]
@ -21,6 +23,18 @@ pub enum Error {
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error(
"invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID"
)]
MissingAzureAiCredentials,
#[error(
"invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID"
)]
MissingAzureDocumentIntelligenceCredentials,
#[error(
"Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()"
)]
MissingReductoApiKey,
#[error("upstream request failed with status {status}: {body}")]
Http { status: u16, body: String },
#[error("upstream network error: {0}")]
@ -40,6 +54,39 @@ pub enum Error {
Unsupported(&'static str),
}
impl Error {
pub const fn http_status_code(&self) -> Option<u16> {
match self {
Self::InvalidRequest(_) => Some(400),
Self::MissingDocumentUrl => Some(500),
Self::Http { status, .. } => Some(*status),
_ => None,
}
}
}
#[derive(Debug, ThisError)]
pub(crate) enum MediaError {
#[error("media URL rejected by network policy")]
BlockedUrl,
#[error("media download is disabled")]
DownloadDisabled,
#[error("media download exceeds the maximum size")]
DownloadTooLarge,
#[error("too many redirects while fetching media")]
TooManyRedirects,
#[error("media redirect is missing a Location header")]
MissingRedirectLocation,
#[error("invalid media redirect")]
InvalidRedirect,
#[error("media download failed with status {0}")]
Http(u16),
#[error("media download timed out")]
Timeout,
#[error("{0}")]
Transport(#[from] TransportError),
}
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum TransportError {
#[error("upstream request failed with status {status}: {body}")]
@ -72,6 +119,7 @@ impl From<crate::ocr::error::OcrRequestError> for Error {
fn from(error: crate::ocr::error::OcrRequestError) -> Self {
match error {
crate::ocr::error::OcrRequestError::MissingField(field) => Self::MissingField(field),
crate::ocr::error::OcrRequestError::MissingDocumentUrl => Self::MissingDocumentUrl,
error => Self::InvalidRequest(error.to_string()),
}
}
@ -93,6 +141,15 @@ impl From<TransportError> for Error {
}
}
impl From<crate::AuthError> for Error {
fn from(error: crate::AuthError) -> Self {
match error {
crate::AuthError::MissingApiKey { provider } => Self::MissingApiKey { provider },
error => Self::Auth(error.to_string()),
}
}
}
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
match value {
serde_json::Value::Null => "null",
@ -108,6 +165,14 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str {
mod transport_tests {
use super::*;
#[test]
fn missing_auth_key_preserves_provider_in_public_error() {
assert_eq!(
Error::from(crate::AuthError::MissingApiKey { provider: "Vertex" }),
Error::MissingApiKey { provider: "Vertex" }
);
}
#[tokio::test]
async fn transport_errors_remove_urls_and_keep_dispatch_context() {
let error = reqwest::Client::builder()

3
litellm-rust/crates/core/src/lib.rs
View file

@ -1,10 +1,12 @@
pub mod audio_transcription;
pub mod auth;
pub mod caching;
pub mod call_lifecycle;
pub mod chat_completions;
pub mod constants;
pub mod error;
pub mod http_utils;
mod media;
pub mod messages;
#[cfg(any(feature = "observability", test))]
pub mod observability;
@ -16,4 +18,5 @@ pub mod router;
pub mod routing_utils;
mod url_utils;
pub use auth::AuthError;
pub use error::Error;

528
litellm-rust/crates/core/src/media.rs Normal file
View file

@ -0,0 +1,528 @@
use std::future::Future;
use std::io;
use std::net::{IpAddr, SocketAddr};
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use reqwest::Url;
use reqwest::dns::{Addrs, Name, Resolve, Resolving};
use crate::constants::MEDIA_CONNECT_TIMEOUT_SECS;
use crate::error::{MediaError, TransportError};
#[derive(Clone)]
pub(crate) struct MediaFetcher {
client: reqwest::Client,
address_resolver: Arc<dyn AddressResolver>,
allow_private_network: bool,
}
type AddressResolution<'a> = Pin<Box<dyn Future<Output = io::Result<Vec<SocketAddr>>> + Send + 'a>>;
trait AddressResolver: Send + Sync {
fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a>;
}
#[derive(Clone, Copy)]
pub(crate) struct DownloadPolicy {
pub(crate) timeout: Duration,
pub(crate) max_bytes: u64,
pub(crate) max_redirects: usize,
}
#[derive(Debug)]
pub(crate) struct DownloadedMedia {
pub(crate) bytes: Vec<u8>,
pub(crate) content_type: String,
}
impl MediaFetcher {
pub(crate) fn new() -> Result<Self, reqwest::Error> {
Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver))
}
fn with_resolvers<R>(
transport_resolver: Arc<R>,
address_resolver: Arc<dyn AddressResolver>,
) -> Result<Self, reqwest::Error>
where
R: Resolve + 'static,
{
let client = reqwest::Client::builder()
.connect_timeout(Duration::from_secs(MEDIA_CONNECT_TIMEOUT_SECS))
.redirect(reqwest::redirect::Policy::none())
.no_proxy()
.dns_resolver(transport_resolver)
.build()?;
Ok(Self {
client,
address_resolver,
allow_private_network: false,
})
}
#[cfg(test)]
pub(crate) fn for_test(client: reqwest::Client) -> Self {
Self {
client,
address_resolver: Arc::new(AllowPrivateResolver),
allow_private_network: true,
}
}
pub(crate) async fn fetch(
&self,
url: Url,
policy: DownloadPolicy,
) -> Result<DownloadedMedia, MediaError> {
if policy.max_bytes == 0 {
return Err(MediaError::DownloadDisabled);
}
tokio::time::timeout(policy.timeout, self.fetch_before_deadline(url, policy))
.await
.map_err(|_| MediaError::Timeout)?
}
async fn fetch_before_deadline(
&self,
mut url: Url,
policy: DownloadPolicy,
) -> Result<DownloadedMedia, MediaError> {
let mut redirects_followed = 0;
loop {
self.validate_url(&url).await?;
let mut response = self
.client
.get(url.clone())
.send()
.await
.map_err(TransportError::from)?;
if response.status().is_redirection() {
if redirects_followed == policy.max_redirects {
return Err(MediaError::TooManyRedirects);
}
let location = response
.headers()
.get(reqwest::header::LOCATION)
.and_then(|value| value.to_str().ok())
.ok_or(MediaError::MissingRedirectLocation)?;
url = url
.join(location)
.map_err(|_| MediaError::InvalidRedirect)?;
redirects_followed += 1;
continue;
}
if !response.status().is_success() {
return Err(MediaError::Http(response.status().as_u16()));
}
enforce_download_size(response.content_length().unwrap_or(0), policy.max_bytes)?;
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("application/octet-stream")
.to_string();
let mut bytes = Vec::new();
while let Some(chunk) = response.chunk().await.map_err(TransportError::from)? {
enforce_download_size(bytes.len() as u64 + chunk.len() as u64, policy.max_bytes)?;
bytes.extend_from_slice(&chunk);
}
return Ok(DownloadedMedia {
bytes,
content_type,
});
}
}
async fn validate_url(&self, url: &Url) -> Result<(), MediaError> {
if !matches!(url.scheme(), "http" | "https")
|| !url.username().is_empty()
|| url.password().is_some()
{
return Err(MediaError::BlockedUrl);
}
let host = url.host_str().ok_or(MediaError::BlockedUrl)?;
if self.allow_private_network {
return Ok(());
}
if let Ok(ip) = host.parse::<IpAddr>() {
return (!is_blocked_ip(ip))
.then_some(())
.ok_or(MediaError::BlockedUrl);
}
let port = url.port_or_known_default().ok_or(MediaError::BlockedUrl)?;
let addresses = self
.address_resolver
.resolve(host, port)
.await
.map_err(|error| TransportError::Network(error.to_string()))?;
validate_addresses(&addresses)
}
}
fn enforce_download_size(length: u64, max_bytes: u64) -> Result<(), MediaError> {
if length > max_bytes {
return Err(MediaError::DownloadTooLarge);
}
Ok(())
}
fn validate_addresses(addresses: &[SocketAddr]) -> Result<(), MediaError> {
if addresses.is_empty() || addresses.iter().any(|address| is_blocked_ip(address.ip())) {
return Err(MediaError::BlockedUrl);
}
Ok(())
}
fn is_blocked_ip(ip: IpAddr) -> bool {
match ip {
IpAddr::V4(ip) => {
let [first, second, third, _] = ip.octets();
first == 0
|| first == 10
|| first == 127
|| (first == 100 && (64..=127).contains(&second))
|| (first == 169 && second == 254)
|| (first == 172 && (16..=31).contains(&second))
|| (first == 192 && second == 0 && (third == 0 || third == 2))
|| (first == 192 && second == 168)
|| (first == 192 && second == 88 && third == 99)
|| (first == 198 && (second == 18 || second == 19))
|| (first == 198 && second == 51 && third == 100)
|| (first == 203 && second == 0 && third == 113)
|| first >= 224
}
IpAddr::V6(ip) => {
let segments = ip.segments();
ip.is_loopback()
|| ip.is_unspecified()
|| ip.is_multicast()
|| (segments[0] & 0xfe00) == 0xfc00
|| (segments[0] & 0xffc0) == 0xfe80
|| (segments[0] & 0xffc0) == 0xfec0
|| (segments[0] == 0x2001 && segments[1] == 0x0db8)
|| ip
.to_ipv4_mapped()
.or_else(|| ip.to_ipv4())
.map(|ipv4| is_blocked_ip(IpAddr::V4(ipv4)))
.unwrap_or(false)
}
}
}
#[derive(Default)]
struct PublicDnsResolver;
struct SystemAddressResolver;
impl AddressResolver for SystemAddressResolver {
fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a> {
Box::pin(async move {
Ok(tokio::net::lookup_host((host, port))
.await?
.collect::<Vec<_>>())
})
}
}
#[cfg(test)]
struct AllowPrivateResolver;
#[cfg(test)]
impl AddressResolver for AllowPrivateResolver {
fn resolve<'a>(&'a self, _host: &'a str, port: u16) -> AddressResolution<'a> {
Box::pin(async move { Ok(vec![SocketAddr::from(([8, 8, 8, 8], port))]) })
}
}
impl Resolve for PublicDnsResolver {
fn resolve(&self, name: Name) -> Resolving {
let host = name.as_str().to_string();
Box::pin(async move {
let addresses = tokio::net::lookup_host((host.as_str(), 0))
.await
.map_err(|error| Box::new(error) as Box<dyn std::error::Error + Send + Sync>)?
.collect::<Vec<_>>();
validate_addresses(&addresses).map_err(|_| {
Box::new(io::Error::other("destination rejected by network policy"))
as Box<dyn std::error::Error + Send + Sync>
})?;
Ok(Box::new(addresses.into_iter()) as Addrs)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
async fn serve(response: &'static [u8]) -> (Url, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has address");
let task = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = [0_u8; 1024];
let bytes_read = socket.read(&mut request).await.expect("reads request");
assert!(bytes_read > 0);
socket.write_all(response).await.expect("writes response");
});
(
Url::parse(&format!("http://{address}/document")).expect("valid test URL"),
task,
)
}
async fn serve_named(
host: &str,
responses: Vec<&'static [u8]>,
) -> (Url, tokio::task::JoinHandle<Vec<String>>, SocketAddr) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has address");
let task = tokio::spawn(async move {
let mut requests = Vec::with_capacity(responses.len());
for response in responses {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = [0_u8; 4096];
let bytes_read = socket.read(&mut request).await.expect("reads request");
requests.push(String::from_utf8_lossy(&request[..bytes_read]).into_owned());
socket.write_all(response).await.expect("writes response");
}
requests
});
(
Url::parse(&format!("http://{host}:{}/document", address.port()))
.expect("valid test URL"),
task,
address,
)
}
struct LoopbackDnsResolver(SocketAddr);
impl Resolve for LoopbackDnsResolver {
fn resolve(&self, _name: Name) -> Resolving {
let address = self.0;
Box::pin(async move { Ok(Box::new(vec![address].into_iter()) as Addrs) })
}
}
struct TestAddressResolver {
blocked_hosts: HashSet<&'static str>,
}
impl AddressResolver for TestAddressResolver {
fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a> {
let blocked = self.blocked_hosts.contains(host);
Box::pin(async move {
let ip = if blocked {
IpAddr::from([127, 0, 0, 1])
} else {
IpAddr::from([8, 8, 8, 8])
};
Ok(vec![SocketAddr::new(ip, port)])
})
}
}
fn policy_checked_fetcher(
address: SocketAddr,
blocked_hosts: HashSet<&'static str>,
) -> MediaFetcher {
MediaFetcher::with_resolvers(
Arc::new(LoopbackDnsResolver(address)),
Arc::new(TestAddressResolver { blocked_hosts }),
)
.expect("test fetcher builds")
}
fn policy(max_bytes: u64, max_redirects: usize) -> DownloadPolicy {
DownloadPolicy {
timeout: Duration::from_secs(1),
max_bytes,
max_redirects,
}
}
#[test]
fn blocks_non_public_addresses() {
for address in [
"0.0.0.1",
"10.0.0.1",
"100.64.0.1",
"127.0.0.1",
"169.254.1.1",
"172.16.0.1",
"192.168.0.1",
"198.18.0.1",
"198.51.100.1",
"203.0.113.1",
"224.0.0.1",
"::1",
"fc00::1",
"fe80::1",
"2001:db8::1",
"::ffff:127.0.0.1",
] {
assert!(is_blocked_ip(address.parse().expect("valid test address")));
}
assert!(!is_blocked_ip(
"8.8.8.8".parse().expect("valid public address")
));
}
#[tokio::test]
async fn fetches_exact_limit_and_normalizes_content_type() {
let (url, server) = serve(
b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc",
)
.await;
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test client builds");
let media = MediaFetcher::for_test(client)
.fetch(url, policy(3, 0))
.await
.expect("download succeeds at exact limit");
server.await.expect("server completes");
assert_eq!(media.bytes, b"abc");
assert_eq!(media.content_type, "application/pdf");
}
#[tokio::test]
async fn rejects_declared_oversize_body() {
let (url, server) = serve(
b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc",
)
.await;
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test client builds");
let error = MediaFetcher::for_test(client)
.fetch(url, policy(2, 0))
.await
.expect_err("oversize body is rejected");
server.await.expect("server completes");
assert!(matches!(error, MediaError::DownloadTooLarge));
}
#[tokio::test]
async fn rejects_streamed_oversize_body() {
let (url, server) = serve(
b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n2\r\nab\r\n2\r\ncd\r\n0\r\n\r\n",
)
.await;
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test client builds");
let error = MediaFetcher::for_test(client)
.fetch(url, policy(3, 0))
.await
.expect_err("stream crossing limit is rejected");
server.await.expect("server completes");
assert!(matches!(error, MediaError::DownloadTooLarge));
}
#[tokio::test]
async fn follows_allowed_redirects_and_revalidates_each_destination() {
let (url, server, address) = serve_named(
"public.test",
vec![
b"HTTP/1.1 302 Found\r\nLocation: /final\r\nContent-Length: 0\r\n\r\n",
b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok",
],
)
.await;
let media = policy_checked_fetcher(address, HashSet::new())
.fetch(url, policy(2, 1))
.await
.expect("redirected fetch succeeds");
let requests = server.await.expect("server completes");
assert_eq!(requests.len(), 2);
assert!(requests[1].starts_with("GET /final "));
assert_eq!(media.bytes, b"ok");
}
#[tokio::test]
async fn blocks_redirected_private_destination_before_second_request() {
let (url, server, address) = serve_named(
"public.test",
vec![b"HTTP/1.1 302 Found\r\nLocation: http://blocked.test/document\r\nContent-Length: 0\r\n\r\n"],
)
.await;
let error = policy_checked_fetcher(address, HashSet::from(["blocked.test"]))
.fetch(url, policy(10, 1))
.await
.expect_err("private redirect is rejected");
let requests = server.await.expect("server completes");
assert_eq!(requests.len(), 1);
assert!(matches!(error, MediaError::BlockedUrl));
}
#[tokio::test]
async fn enforces_total_timeout() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has address");
let server = tokio::spawn(async move {
let (_socket, _) = listener.accept().await.expect("accepts request");
tokio::time::sleep(Duration::from_millis(100)).await;
});
let url = Url::parse(&format!("http://public.test:{}/document", address.port()))
.expect("valid test URL");
let error = policy_checked_fetcher(address, HashSet::new())
.fetch(
url,
DownloadPolicy {
timeout: Duration::from_millis(20),
max_bytes: 10,
max_redirects: 0,
},
)
.await
.expect_err("fetch times out");
server.await.expect("server completes");
assert!(matches!(error, MediaError::Timeout));
}
#[tokio::test]
async fn document_client_does_not_send_ambient_credentials() {
let (url, server, address) = serve_named(
"public.test",
vec![b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"],
)
.await;
policy_checked_fetcher(address, HashSet::new())
.fetch(url, policy(2, 0))
.await
.expect("fetch succeeds");
let requests = server.await.expect("server completes");
assert!(!requests[0].to_ascii_lowercase().contains("authorization:"));
assert!(!requests[0].to_ascii_lowercase().contains("api-key:"));
}
#[tokio::test]
async fn rejects_url_credentials_before_network_access() {
let fetcher = MediaFetcher::new().expect("media fetcher builds");
let url =
Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses");
assert!(matches!(
fetcher.validate_url(&url).await,
Err(MediaError::BlockedUrl)
));
}
}

131
litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs Normal file
View file

@ -0,0 +1,131 @@
use super::super::OcrAdapter;
use crate::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::cohere::{
CohereParams, CohereResponse, transform_request, transform_response, validate_document,
};
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{credential_env, transform_request_body};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use crate::providers::azure_ai::auth::AzureAuthInputs;
use crate::url_utils::ApiUrl;
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
pub(crate) struct AzureCohereAdapter;
impl OcrAdapter for AzureCohereAdapter {
type ProviderResponse = CohereResponse;
const PROVIDER: OcrProvider = OcrProvider::AzureAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let params = super::super::super::wire::decode_request_value::<CohereParams>(
serde_json::Value::Object(request.optional_params.clone()),
"optional_params",
)?;
let mut config = AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
config.azure_ad_token_provider = request.azure_ad_token_provider.clone();
let base = request
.connection
.api_base
.clone()
.or_else(|| credential_env(AZURE_AI_API_BASE_ENV))
.filter(|base| !base.trim().is_empty())
.ok_or_else(|| {
Error::Auth(
"Missing Azure AI API Base - Set AZURE_AI_API_BASE or pass api_base".into(),
)
})?;
let headers =
super::validate_ai_environment(&request.connection, &config, &credential_env).await?;
validate_document(&request.document)?;
let remote = request.document.source().starts_with("http://")
|| request.document.source().starts_with("https://");
let document = inline_remote_document(
client.document_fetcher(),
request.document.clone(),
&request.connection,
)
.await?;
let body = transform_request(&request.model, document, params)?;
transform_request_body(
client,
request,
&complete_url(&base)?,
&headers,
!remote,
body,
|body| {
validate_document(&body.document)?;
validate_inline_document(&body.document)
},
)
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
transform_response(&request.model, response)
}
}
fn complete_url(base: &str) -> Result<String, OcrError> {
let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?;
if !matches!(url.scheme(), "http" | "https") {
return Err(invalid_api_base().into());
}
let path = url.path().trim_end_matches('/').to_string();
if path.ends_with("/v2/parse") {
url.set_path(&path);
return Ok(url.into());
}
url.set_path(path.strip_suffix("/models").unwrap_or(&path));
ApiUrl::parse(url.as_str())
.and_then(|url| url.complete_path(&["providers", "cohere", "v2", "parse"]))
.map(|url| url.into_string())
.map_err(|_| invalid_api_base().into())
}
fn invalid_api_base() -> OcrRequestError {
OcrRequestError::RequestField {
path: "api_base".into(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completes_foundry_urls_without_duplicate_paths_and_preserves_queries() {
for suffix in [
"",
"/models",
"/providers/cohere/v2",
"/providers/cohere/v2/parse",
] {
assert_eq!(
complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(),
"https://example.com/providers/cohere/v2/parse?tenant=a"
);
}
assert_eq!(
complete_url("https://example.com/v2/parse?tenant=a").unwrap(),
"https://example.com/v2/parse?tenant=a"
);
assert!(complete_url("relative/path").is_err());
}
}

215
litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs Normal file
View file

@ -0,0 +1,215 @@
use super::super::OcrAdapter;
use crate::Error;
use crate::auth::{InputSource, Sourced};
use crate::constants::{AZURE_DI_API_VERSION, AZURE_DI_SUBSCRIPTION_HEADER};
use crate::ocr::OcrClient;
use crate::ocr::codecs::document_intelligence::{
self, AzureDocumentIntelligenceOperation, DocumentIntelligenceParams,
};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{credential_env, transform_request_body};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrResponseFormat};
use crate::providers::azure_ai::auth::AzureAuthInputs;
use crate::url_utils::ApiUrl;
mod polling;
const AZURE_DI_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY";
const AZURE_DI_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT";
#[derive(Clone, Debug)]
pub(crate) struct AzureDocumentIntelligenceAdapter;
impl OcrAdapter for AzureDocumentIntelligenceAdapter {
type ProviderResponse = AzureDocumentIntelligenceOperation;
const PROVIDER: OcrProvider = OcrProvider::AzureAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let params = map_ocr_params(request)?;
let mut config = AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
config.azure_ad_token_provider = request.azure_ad_token_provider.clone();
let headers = validate_environment(&request.connection, &config, &credential_env).await?;
let endpoint = nonblank(request.connection.api_base.clone())
.or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV)))
.ok_or_else(|| Error::Auth("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into()))?;
let url = get_complete_url(&endpoint, &request.model, &params)?;
let body = document_intelligence::transform_ocr_request(request.document.clone())?;
transform_request_body(client, request, &url, &headers, false, body, |_| Ok(())).await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
document_intelligence::transform_ocr_response(&request.model, response)
}
async fn read_response(
&self,
client: &OcrClient,
response: reqwest::Response,
url: &str,
headers: &[(String, String)],
request: &LiteLLMOcrRequest,
) -> Result<crate::ocr::wire::DecodedOcrResponse<Self::ProviderResponse>, OcrError> {
polling::read_operation_response(
client.polling_http(),
response,
url,
headers,
&request.connection,
request.response_format()? == OcrResponseFormat::Native,
&request.hooks,
)
.await
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn map_ocr_params(
request: &LiteLLMOcrRequest,
) -> Result<DocumentIntelligenceParams, OcrRequestError> {
let params = document_intelligence::decode_input_params(
request.optional_params.clone(),
"optional_params",
)?;
let crate::ocr::prepare::ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = params;
document_intelligence::map_ocr_params(params)
}
fn get_complete_url(
endpoint: &str,
model: &str,
params: &DocumentIntelligenceParams,
) -> Result<String, OcrError> {
let model = format!("{}:analyze", model_id(model)?);
ApiUrl::parse(endpoint)
.and_then(|url| url.complete_path(&["documentintelligence", "documentModels", &model]))
.map(|url| {
url.append_query_pairs(
[("api-version", AZURE_DI_API_VERSION)]
.into_iter()
.chain(params.pages.iter().map(|pages| ("pages", pages.as_str())))
.chain(
params
.features
.iter()
.map(|features| ("features", features.as_str())),
),
)
.into_string()
})
.map_err(|_| OcrRequestError::RequestField {
path: "api_base".into(),
})
.map_err(OcrError::from)
}
async fn validate_environment(
connection: &OcrConnection,
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization")
|| crate::http_utils::has_header(&connection.extra_headers, AZURE_DI_SUBSCRIPTION_HEADER)
{
super::validate_destination(connection, connection.extra_headers_source)?;
return Ok(connection.extra_headers.clone());
}
let key = nonblank(connection.api_key.clone())
.map(|value| Sourced::new(value, connection.api_key_source))
.or_else(|| {
nonblank(env_lookup(AZURE_DI_API_KEY_ENV))
.map(|value| Sourced::new(value, InputSource::Environment))
});
if let Some(key) = key {
super::validate_destination(connection, key.source())?;
return Ok(
std::iter::once((AZURE_DI_SUBSCRIPTION_HEADER.into(), key.into_value()))
.chain(connection.extra_headers.clone())
.collect(),
);
}
let token = super::resolve_entra(config, env_lookup)
.await?
.ok_or(Error::MissingAzureDocumentIntelligenceCredentials)?;
super::validate_destination(connection, token.source())?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {}", token.value())))
.chain(connection.extra_headers.clone())
.collect(),
)
}
fn model_id(model: &str) -> Result<&str, OcrRequestError> {
let model = model.rsplit('/').next().unwrap_or(model);
if matches!(model, "." | "..") {
return Err(OcrRequestError::DotModel);
}
Ok(model)
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn request_endpoint_cannot_receive_environment_key() {
let connection = OcrConnection {
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let error = validate_environment(&connection, &Default::default(), &|name| {
(name == AZURE_DI_API_KEY_ENV).then(|| "environment-key".into())
})
.await
.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Azure endpoint")
);
}
#[tokio::test]
async fn request_endpoint_accepts_request_owned_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
api_key_source: InputSource::Request,
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let headers = validate_environment(&connection, &Default::default(), &|_| None)
.await
.unwrap();
assert_eq!(
headers[0],
(AZURE_DI_SUBSCRIPTION_HEADER.into(), "request-key".into())
);
}
}

119
litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs Normal file
View file

@ -0,0 +1,119 @@
use std::sync::Arc;
use std::time::Duration;
use reqwest::Url;
use tokio::time::Instant;
use crate::constants::{AZURE_DI_SUBSCRIPTION_HEADER, OCR_POLL_RETRY_SECS};
use crate::ocr::client::read_json_response;
use crate::ocr::codecs::document_intelligence::{
AzureDocumentIntelligenceOperation, OperationStatus,
};
use crate::ocr::error::{OcrError, OcrPollingError, OcrResponseError};
use crate::ocr::hooks::OcrHooks;
use crate::ocr::types::OcrConnection;
use crate::ocr::wire::DecodedOcrResponse;
pub(super) async fn read_operation_response(
http_client: &reqwest::Client,
response: reqwest::Response,
original_url: &str,
headers: &[(String, String)],
connection: &OcrConnection,
native: bool,
hooks: &Arc<dyn OcrHooks>,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, OcrError> {
if response.status() != reqwest::StatusCode::ACCEPTED {
let bytes =
crate::ocr::client::read_response_bytes(response, connection.max_response_bytes)
.await?;
crate::ocr::handler::post_call(hooks, &bytes).await?;
return Ok(crate::ocr::wire::decode_response(&bytes, native)?);
}
let location = response
.headers()
.get("operation-location")
.and_then(|value| value.to_str().ok())
.ok_or(OcrPollingError::PollLocation)?
.to_string();
let original = Url::parse(original_url).map_err(|_| OcrPollingError::PollOrigin)?;
let operation = Url::parse(&location).map_err(|_| OcrPollingError::PollOrigin)?;
if original.origin() != operation.origin()
|| !operation.username().is_empty()
|| operation.password().is_some()
{
return Err(OcrPollingError::PollOrigin.into());
}
let bytes =
crate::ocr::client::read_response_bytes(response, connection.max_response_bytes).await?;
crate::ocr::handler::post_call(hooks, &bytes).await?;
poll_operation(http_client, operation, headers, connection, native, hooks).await
}
async fn poll_operation(
http_client: &reqwest::Client,
url: Url,
headers: &[(String, String)],
connection: &OcrConnection,
native: bool,
hooks: &Arc<dyn OcrHooks>,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, OcrError> {
let deadline = Instant::now()
.checked_add(connection.poll_timeout)
.ok_or(OcrPollingError::PollTimeout)?;
loop {
let remaining = deadline
.checked_duration_since(Instant::now())
.filter(|remaining| !remaining.is_zero())
.ok_or(OcrPollingError::PollTimeout)?;
let builder = http_client
.get(url.clone())
.timeout(remaining.min(connection.timeout));
let builder = crate::http_utils::with_headers(
builder,
headers,
crate::http_utils::HeaderPolicy::Only(&[AZURE_DI_SUBSCRIPTION_HEADER, "authorization"]),
);
let response = tokio::time::timeout_at(deadline, crate::http_utils::http_request(builder))
.await
.map_err(|_| OcrPollingError::PollTimeout)?
.map_err(crate::error::TransportError::from)?;
let retry = response
.headers()
.get(reqwest::header::RETRY_AFTER)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<u64>().ok())
.unwrap_or(OCR_POLL_RETRY_SECS)
.max(1);
let decoded = tokio::time::timeout_at(
deadline,
read_json_response::<AzureDocumentIntelligenceOperation>(
response,
native,
connection.max_response_bytes,
),
)
.await
.map_err(|_| OcrPollingError::PollTimeout)??;
match &decoded.data.status {
Some(OperationStatus::Succeeded) => {
crate::ocr::handler::post_call(hooks, decoded.text.as_bytes()).await?;
return Ok(decoded);
}
Some(OperationStatus::Running | OperationStatus::NotStarted) => {
tokio::time::timeout_at(deadline, tokio::time::sleep(Duration::from_secs(retry)))
.await
.map_err(|_| OcrPollingError::PollTimeout)?;
}
status => {
return Err(OcrResponseError::OperationStatus(
status
.as_ref()
.map(ToString::to_string)
.unwrap_or_else(|| "None".into()),
)
.into());
}
}
}
}

229
litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs Normal file
View file

@ -0,0 +1,229 @@
use super::super::OcrAdapter;
use crate::Error;
use crate::auth::{InputSource, Sourced};
use crate::constants::AZURE_AI_OCR_PATH;
use crate::ocr::OcrClient;
use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection};
use crate::providers::azure_ai::auth::AzureAuthInputs;
use crate::url_utils::ApiUrl;
const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY";
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
#[derive(Clone, Debug)]
pub(crate) struct AzureMistralAdapter;
impl OcrAdapter for AzureMistralAdapter {
type ProviderResponse = MistralOcrResponse;
const PROVIDER: OcrProvider = OcrProvider::AzureAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<MistralOcrParams>(request)?;
let mut config = AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
config.azure_ad_token_provider = request.azure_ad_token_provider.clone();
let url = get_complete_url(request.connection.api_base.as_deref(), &credential_env)?;
let headers = validate_environment(&request.connection, &config, &credential_env).await?;
let retains_document = !request.document.source().starts_with("http://")
&& !request.document.source().starts_with("https://");
let document = inline_remote_document(
client.document_fetcher(),
request.document.clone(),
&request.connection,
)
.await?;
let body = mistral::transform_ocr_request(&request.model, document, &params)?;
transform_request_body(
client,
request,
&url,
&headers,
retains_document,
body,
|body| validate_inline_document(&body.document),
)
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
mistral::transform_ocr_response(&request.model, response)
}
}
fn get_complete_url(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, OcrError> {
let base = nonblank(api_base.map(str::to_string))
.or_else(|| nonblank(env_lookup(AZURE_AI_API_BASE_ENV)))
.ok_or_else(|| Error::Auth(
"Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter".into(),
))?;
let path: Vec<&str> = AZURE_AI_OCR_PATH.trim_matches('/').split('/').collect();
ApiUrl::parse(&base)
.and_then(|url| url.complete_path(&path))
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
pub(in crate::ocr::adapters) async fn validate_environment(
connection: &OcrConnection,
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
if config.azure_ad_token_provider.is_some() {
super::resolve_entra(config, env_lookup).await?;
}
super::validate_destination(connection, connection.extra_headers_source)?;
return Ok(connection.extra_headers.clone());
}
let key = nonblank(connection.api_key.clone())
.map(|value| Sourced::new(value, connection.api_key_source))
.or_else(|| {
nonblank(env_lookup(AZURE_AI_API_KEY_ENV))
.map(|value| Sourced::new(value, InputSource::Environment))
});
if let Some(key) = key {
super::validate_destination(connection, key.source())?;
return Ok(bearer_headers(connection, key.value()));
}
let key = super::resolve_entra(config, env_lookup)
.await?
.ok_or(Error::MissingAzureAiCredentials)?;
super::validate_destination(connection, key.source())?;
Ok(bearer_headers(connection, key.value()))
}
fn bearer_headers(connection: &OcrConnection, key: &str) -> Vec<(String, String)> {
std::iter::once(("Authorization".into(), format!("Bearer {key}")))
.chain(connection.extra_headers.clone())
.collect()
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completes_azure_path_and_preserves_query() {
assert_eq!(
get_complete_url(Some("https://example.com/?tenant=a"), &|_| None).unwrap(),
"https://example.com/providers/mistral/azure/ocr?tenant=a"
);
assert_eq!(
get_complete_url(
Some("https://example.com/providers/mistral/azure/ocr"),
&|_| None
)
.unwrap(),
"https://example.com/providers/mistral/azure/ocr"
);
}
#[tokio::test]
async fn supplied_authorization_precedes_keys() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
extra_headers: vec![("authorization".into(), "Bearer prepared".into())],
..Default::default()
};
assert_eq!(
validate_environment(&connection, &Default::default(), &|_| {
Some("environment-key".into())
})
.await
.unwrap(),
connection.extra_headers
);
}
#[tokio::test]
async fn request_key_precedes_environment_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
..Default::default()
};
assert_eq!(
validate_environment(&connection, &Default::default(), &|_| {
Some("environment-key".into())
})
.await
.unwrap()[0],
("Authorization".into(), "Bearer request-key".into())
);
}
#[tokio::test]
async fn request_endpoint_cannot_receive_environment_key() {
let connection = OcrConnection {
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let error = validate_environment(&connection, &Default::default(), &|name| {
(name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into())
})
.await
.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Azure endpoint")
);
}
#[tokio::test]
async fn request_endpoint_accepts_request_owned_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
api_key_source: InputSource::Request,
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let headers = validate_environment(&connection, &Default::default(), &|_| None)
.await
.unwrap();
assert_eq!(
headers[0],
("Authorization".into(), "Bearer request-key".into())
);
}
}

56
litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs Normal file
View file

@ -0,0 +1,56 @@
mod cohere;
mod document_intelligence;
mod mistral;
use std::sync::OnceLock;
use crate::Error;
use crate::auth::error::AuthConfigurationError;
use crate::auth::{InputSource, Sourced};
use crate::ocr::error::OcrError;
use crate::ocr::types::OcrConnection;
use crate::providers::azure_ai::auth::{AzureAuthInputs, AzureAuthService};
pub(crate) use cohere::AzureCohereAdapter;
pub(crate) use document_intelligence::AzureDocumentIntelligenceAdapter;
pub(crate) use mistral::AzureMistralAdapter;
pub(super) use mistral::validate_environment as validate_ai_environment;
async fn resolve_entra(
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<String>>, Error> {
static SERVICE: OnceLock<AzureAuthService> = OnceLock::new();
SERVICE
.get_or_init(AzureAuthService::default)
.get_azure_ad_token(config, env_lookup)
.await
.or_else(|error| match error {
crate::AuthError::EmptyAzureToken => Ok(None),
other => Err(other),
})
.map(|credential| {
credential.map(|credential| {
let source = credential.source();
let value = credential.value().secret().expose().to_string();
Sourced::new(value, source)
})
})
.map_err(Error::from)
}
fn validate_destination(
connection: &OcrConnection,
credential_source: InputSource,
) -> Result<(), OcrError> {
if connection.api_base.is_some()
&& connection.api_base_source == InputSource::Request
&& credential_source != InputSource::Request
{
return Err(Error::from(crate::AuthError::Configuration(
AuthConfigurationError::RequestAzureCredentialDestination,
))
.into());
}
Ok(())
}

123
litellm-rust/crates/core/src/ocr/adapters/cohere.rs Normal file
View file

@ -0,0 +1,123 @@
use super::OcrAdapter;
use crate::Error;
use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE};
use crate::ocr::OcrClient;
use crate::ocr::codecs::cohere::{
CohereParams, CohereResponse, transform_request, transform_response, validate_document,
};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{credential_env, transform_request_body};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection};
use crate::url_utils::ApiUrl;
pub(crate) struct CohereAdapter;
impl OcrAdapter for CohereAdapter {
type ProviderResponse = CohereResponse;
const PROVIDER: OcrProvider = OcrProvider::Cohere;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let params = super::super::wire::decode_request_value::<CohereParams>(
serde_json::Value::Object(request.optional_params.clone()),
"optional_params",
)?;
let headers = validate_environment(&request.connection, &credential_env)?;
let url = complete_url(
request
.connection
.api_base
.as_deref()
.unwrap_or(COHERE_PARSE_API_BASE),
)?;
let body = transform_request(&request.model, request.document.clone(), params)?;
transform_request_body(client, request, &url, &headers, true, body, |body| {
validate_document(&body.document)
})
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
transform_response(&request.model, response)
}
}
fn complete_url(base: &str) -> Result<String, OcrError> {
let parsed = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?;
if !matches!(parsed.scheme(), "http" | "https") {
return Err(invalid_api_base().into());
}
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&["v2", "parse"]))
.map(|url| url.into_string())
.map_err(|_| invalid_api_base().into())
}
fn invalid_api_base() -> OcrRequestError {
OcrRequestError::RequestField {
path: "api_base".into(),
}
}
fn validate_environment(
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| env_lookup(COHERE_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
.ok_or_else(|| {
Error::Auth("Missing COHERE_API_KEY - set it in the environment or pass api_key".into())
})?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completes_provider_urls_without_duplicate_paths_and_preserves_queries() {
for suffix in ["", "/v2", "/v2/parse"] {
assert_eq!(
complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(),
"https://example.com/v2/parse?tenant=a"
);
}
}
#[test]
fn rejects_invalid_urls_and_blank_keys() {
assert!(complete_url("relative/path").is_err());
assert!(complete_url("ftp://example.com").is_err());
assert!(matches!(
validate_environment(
&OcrConnection {
api_key: Some(" ".into()),
..Default::default()
},
&|_| None,
),
Err(OcrError::Public(Error::Auth(_)))
));
}
}

2
litellm-rust/crates/core/src/ocr/adapters/mistral.rs
View file

@ -33,7 +33,7 @@ impl OcrAdapter for MistralAdapter {
let url = get_complete_url(request.connection.api_base.as_deref())?;
let body =
mistral::transform_ocr_request(&request.model, request.document.clone(), &params)?;
transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await
transform_request_body(client, request, &url, &headers, true, body, |_| Ok(())).await
}
fn transform_ocr_response(

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

@ -5,12 +5,19 @@ use serde::de::DeserializeOwned;
use super::OcrClient;
use super::error::{OcrError, OcrResponseError};
use super::registry::OcrProvider;
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrResponseFormat};
use super::wire::DecodedOcrResponse;
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
mod azure;
mod cohere;
mod mistral;
mod reducto;
mod vertex;
pub(crate) use azure::{AzureCohereAdapter, AzureDocumentIntelligenceAdapter, AzureMistralAdapter};
pub(crate) use cohere::CohereAdapter;
pub(crate) use mistral::MistralAdapter;
pub(crate) use reducto::{ReductoLegacyAdapter, ReductoV3Adapter};
pub(crate) use vertex::{VertexDeepSeekAdapter, VertexMistralAdapter};
/// Converts a complete LiteLLM OCR call to provider HTTP and normalizes its response.
pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static {
@ -49,19 +56,34 @@ pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static {
_url: &str,
_headers: &[(String, String)],
request: &LiteLLMOcrRequest,
) -> impl Future<Output = Result<DecodedOcrResponse<Self::ProviderResponse>, OcrError>> + Send
{
let retain_native = request
.response_format()
.map(|format| format == OcrResponseFormat::Native);
async move { super::client::read_json_response(response, retain_native?).await }
) -> impl Future<
Output = Result<super::wire::DecodedOcrResponse<Self::ProviderResponse>, OcrError>,
> + Send {
async move {
let bytes =
super::client::read_response_bytes(response, request.connection.max_response_bytes)
.await?;
super::handler::post_call(&request.hooks, &bytes).await?;
Ok(super::wire::decode_response(
&bytes,
request.response_format()? == super::types::OcrResponseFormat::Native,
)?)
}
}
}
macro_rules! for_each_ocr_adapter {
($callback:ident) => {
$callback! {
Cohere, $crate::ocr::adapters::CohereAdapter, $crate::ocr::adapters::CohereAdapter, Cohere;
AzureCohere, $crate::ocr::adapters::AzureCohereAdapter, $crate::ocr::adapters::AzureCohereAdapter, AzureAi;
Mistral, $crate::ocr::adapters::MistralAdapter, $crate::ocr::adapters::MistralAdapter, Mistral;
AzureMistral, $crate::ocr::adapters::AzureMistralAdapter, $crate::ocr::adapters::AzureMistralAdapter, AzureAi;
AzureDocumentIntelligence, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, AzureAi;
ReductoLegacy, $crate::ocr::adapters::ReductoLegacyAdapter, $crate::ocr::adapters::ReductoLegacyAdapter, Reducto;
ReductoV3, $crate::ocr::adapters::ReductoV3Adapter, $crate::ocr::adapters::ReductoV3Adapter, Reducto;
VertexMistral, $crate::ocr::adapters::VertexMistralAdapter, $crate::ocr::adapters::VertexMistralAdapter, VertexAi;
VertexDeepSeek, $crate::ocr::adapters::VertexDeepSeekAdapter, $crate::ocr::adapters::VertexDeepSeekAdapter, VertexAi;
}
};
}

45
litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs Normal file
View file

@ -0,0 +1,45 @@
use super::super::OcrAdapter;
use crate::ocr::OcrClient;
use crate::ocr::codecs::reducto::{self, ReductoLegacyParams, ReductoResponse};
use crate::ocr::error::{OcrError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, build_http_request, credential_env,
guardrail_document, merge_extra_params,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
#[derive(Clone, Debug)]
pub(crate) struct ReductoLegacyAdapter;
impl OcrAdapter for ReductoLegacyAdapter {
type ProviderResponse = ReductoResponse;
const PROVIDER: OcrProvider = OcrProvider::Reducto;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let ParsedProviderParams {
known: params,
extra_params,
} = _prepare_ocr_request::<ReductoLegacyParams>(request)?;
let headers = super::validate_environment(&request.connection, &credential_env)?;
let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?;
let (document, headers) = guardrail_document(request, &url, &headers).await?;
let document =
super::prepare_document(client, document, &request.connection, &headers).await?;
let body = reducto::transform_legacy_ocr_request(&request.model, document, &params)?;
let body = merge_extra_params(&body, extra_params)?;
build_http_request(client, request, &url, &headers, &body)
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
reducto::transform_ocr_response(&request.model, response)
}
}

148
litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs Normal file
View file

@ -0,0 +1,148 @@
mod legacy;
mod v3;
use crate::Error;
use crate::constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX};
use crate::ocr::document::InlineDocument;
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::types::{OcrConnection, OcrDocument};
use crate::url_utils::ApiUrl;
pub(crate) use legacy::ReductoLegacyAdapter;
pub(crate) use v3::ReductoV3Adapter;
pub(super) fn get_complete_url(api_base: Option<&str>, path: &str) -> Result<String, OcrError> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(REDUCTO_API_BASE);
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&[path]))
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
pub(super) fn validate_environment(
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let api_key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
env_lookup(REDUCTO_API_KEY_ENV)
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
})
.ok_or(Error::MissingReductoApiKey)?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {api_key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
pub(super) async fn prepare_document(
client: &crate::ocr::OcrClient,
document: OcrDocument,
connection: &OcrConnection,
headers: &[(String, String)],
) -> Result<OcrDocument, OcrError> {
if document.source().starts_with(REDUCTO_ID_PREFIX) {
if document.source()[REDUCTO_ID_PREFIX.len()..]
.trim()
.is_empty()
{
return Err(OcrRequestError::RequestField {
path: "document file id".into(),
}
.into());
}
return Ok(document);
}
let inline = InlineDocument::parse(document.source())?.ok_or(OcrRequestError::ReductoSource)?;
let mime = inline.mime_type().to_string();
let bytes = inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
let part = reqwest::multipart::Part::bytes(bytes)
.file_name("document")
.mime_str(&mime)
.map_err(|_| OcrRequestError::InvalidDataUri)?;
let builder = client
.provider_http()
.post(get_complete_url(connection.api_base.as_deref(), "upload")?)
.multipart(reqwest::multipart::Form::new().part("file", part))
.timeout(connection.timeout);
let builder = crate::http_utils::with_headers(
builder,
headers,
crate::http_utils::HeaderPolicy::Except(&["content-type", "content-length"]),
);
let response = crate::http_utils::http_request(builder)
.await
.map_err(crate::error::TransportError::from)?;
let uploaded = crate::ocr::client::read_json_response::<
crate::ocr::codecs::reducto::ReductoUploadResponse,
>(response, false, connection.max_response_bytes)
.await?
.data;
let file_id = uploaded
.file_id
.as_deref()
.map(str::trim)
.filter(|id| !id.is_empty());
let Some(file_id) = file_id else {
return Err(OcrResponseError::ResponseField {
path: "file_id".into(),
}
.into());
};
Ok(document.with_source(file_id.to_string()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn explicit_key_precedes_environment_key() {
let connection = OcrConnection {
api_key: Some("passed-key".into()),
..Default::default()
};
let headers = validate_environment(&connection, &|_| Some("env-key".into())).unwrap();
assert_eq!(headers[0].1, "Bearer passed-key");
}
#[test]
fn blank_explicit_key_uses_environment_key() {
let connection = OcrConnection {
api_key: Some(" ".into()),
..Default::default()
};
let headers = validate_environment(&connection, &|_| Some(" env-key ".into())).unwrap();
assert_eq!(headers[0].1, "Bearer env-key");
}
#[test]
fn existing_authorization_skips_key_lookup() {
let connection = OcrConnection {
extra_headers: vec![("authorization".into(), "Bearer existing".into())],
..Default::default()
};
assert_eq!(
validate_environment(&connection, &|_| None).unwrap(),
connection.extra_headers
);
}
}

45
litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs Normal file
View file

@ -0,0 +1,45 @@
use super::super::OcrAdapter;
use crate::ocr::OcrClient;
use crate::ocr::codecs::reducto::{self, ReductoResponse, ReductoV3Params};
use crate::ocr::error::{OcrError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, build_http_request, credential_env,
guardrail_document, merge_extra_params,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
#[derive(Clone, Debug)]
pub(crate) struct ReductoV3Adapter;
impl OcrAdapter for ReductoV3Adapter {
type ProviderResponse = ReductoResponse;
const PROVIDER: OcrProvider = OcrProvider::Reducto;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let ParsedProviderParams {
known: params,
extra_params,
} = _prepare_ocr_request::<ReductoV3Params>(request)?;
let headers = super::validate_environment(&request.connection, &credential_env)?;
let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?;
let (document, headers) = guardrail_document(request, &url, &headers).await?;
let document =
super::prepare_document(client, document, &request.connection, &headers).await?;
let body = reducto::transform_v3_ocr_request(&request.model, document, &params)?;
let body = merge_extra_params(&body, extra_params)?;
build_http_request(client, request, &url, &headers, &body)
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
reducto::transform_ocr_response(&request.model, response)
}
}

140
litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs Normal file
View file

@ -0,0 +1,140 @@
use super::super::OcrAdapter;
use super::validate_destination;
use crate::Error;
use crate::auth::vertex::{self, VertexConfig};
use crate::ocr::OcrClient;
use crate::ocr::codecs::deepseek::{self, DeepSeekOcrParams, DeepSeekOcrResponse};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use crate::url_utils::ApiUrl;
const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com";
const MODEL_NAMESPACE: &str = "deepseek-ai";
const DEFAULT_LOCATION: &str = "us-central1";
#[derive(Clone, Debug)]
pub(crate) struct VertexDeepSeekAdapter;
impl OcrAdapter for VertexDeepSeekAdapter {
type ProviderResponse = DeepSeekOcrResponse;
const PROVIDER: OcrProvider = OcrProvider::VertexAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
validate_destination(&request.connection)?;
let ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<DeepSeekOcrParams>(request)?;
let config = VertexConfig::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
let authentication = client
.vertex_auth()
.validate_environment(
request.connection.extra_headers.clone(),
request.connection.api_key.as_deref(),
&config,
&credential_env,
)
.await
.map_err(Error::from)?;
let location = vertex::get_vertex_ai_location(&config, &credential_env)
.unwrap_or_else(|| DEFAULT_LOCATION.to_string());
let url = get_complete_url(
request.connection.api_base.as_deref(),
&authentication.project_id,
&location,
)?;
let document = request.document.clone();
let body =
deepseek::transform_ocr_request(&provider_model(&request.model), document, &params)?;
transform_request_body(
client,
request,
&url,
&authentication.headers,
false,
body,
|_| Ok(()),
)
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
deepseek::transform_ocr_response(&request.model, response)
}
}
fn provider_model(model: &str) -> String {
if model.starts_with(&format!("{MODEL_NAMESPACE}/")) {
model.to_string()
} else {
format!("{MODEL_NAMESPACE}/{model}")
}
}
fn get_complete_url(
api_base: Option<&str>,
project: &str,
location: &str,
) -> Result<String, OcrError> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(DEFAULT_API_BASE);
ApiUrl::parse(base)
.and_then(|url| {
url.complete_path(&[
"v1",
"projects",
project,
"locations",
location,
"endpoints",
"openapi",
"chat",
"completions",
])
})
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
#[cfg(test)]
mod tests {
use super::{get_complete_url, provider_model};
#[test]
fn adapter_owns_model_namespace_and_endpoint() {
assert_eq!(
provider_model("deepseek-ocr-maas"),
"deepseek-ai/deepseek-ocr-maas"
);
assert_eq!(
provider_model("deepseek-ai/deepseek-ocr-maas"),
"deepseek-ai/deepseek-ocr-maas"
);
assert_eq!(
get_complete_url(None, "proj-1", "europe-west4").unwrap(),
"https://aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/endpoints/openapi/chat/completions"
);
}
}

157
litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs Normal file
View file

@ -0,0 +1,157 @@
use super::super::OcrAdapter;
use super::validate_destination;
use crate::Error;
use crate::auth::vertex::{self, VertexConfig};
use crate::ocr::OcrClient;
use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use crate::url_utils::ApiUrl;
const DEFAULT_LOCATION: &str = "us-central1";
#[derive(Clone, Debug)]
pub(crate) struct VertexMistralAdapter;
impl OcrAdapter for VertexMistralAdapter {
type ProviderResponse = MistralOcrResponse;
const PROVIDER: OcrProvider = OcrProvider::VertexAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
validate_destination(&request.connection)?;
let ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<MistralOcrParams>(request)?;
let config = VertexConfig::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
let authentication = client
.vertex_auth()
.validate_environment(
request.connection.extra_headers.clone(),
request.connection.api_key.as_deref(),
&config,
&credential_env,
)
.await
.map_err(Error::from)?;
let location = vertex::get_vertex_ai_location(&config, &credential_env)
.unwrap_or_else(|| DEFAULT_LOCATION.to_string());
let url = get_complete_url(
request.connection.api_base.as_deref(),
&authentication.project_id,
&location,
&request.model,
)?;
let retains_document = !request.document.source().starts_with("http://")
&& !request.document.source().starts_with("https://");
let document = inline_remote_document(
client.document_fetcher(),
request.document.clone(),
&request.connection,
)
.await?;
let body = mistral::transform_ocr_request(&request.model, document, &params)?;
transform_request_body(
client,
request,
&url,
&authentication.headers,
retains_document,
body,
|body| validate_inline_document(&body.document),
)
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
mistral::transform_ocr_response(&request.model, response)
}
}
fn get_complete_url(
api_base: Option<&str>,
project: &str,
location: &str,
model: &str,
) -> Result<String, OcrError> {
validate_location(location)?;
let default_base = format!("https://{location}-aiplatform.googleapis.com");
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(&default_base);
let prediction = format!("{model}:rawPredict");
ApiUrl::parse(base)
.and_then(|url| {
url.complete_path(&[
"v1",
"projects",
project,
"locations",
location,
"publishers",
"mistralai",
"models",
&prediction,
])
})
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
fn validate_location(location: &str) -> Result<(), OcrError> {
let valid = !location.is_empty()
&& location
.bytes()
.all(|value| value.is_ascii_lowercase() || value.is_ascii_digit() || value == b'-')
&& location
.as_bytes()
.first()
.is_some_and(u8::is_ascii_alphanumeric)
&& location
.as_bytes()
.last()
.is_some_and(u8::is_ascii_alphanumeric);
if valid {
return Ok(());
}
Err(OcrRequestError::RequestField {
path: "vertex_location".into(),
}
.into())
}
#[cfg(test)]
mod tests {
use super::get_complete_url;
#[test]
fn endpoint_uses_location_project_and_model() {
assert_eq!(
get_complete_url(None, "proj-1", "europe-west4", "mistral-ocr-maas").unwrap(),
"https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
assert!(get_complete_url(None, "proj-1", "attacker.example/path", "model").is_err());
}
}

21
litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs Normal file
View file

@ -0,0 +1,21 @@
mod deepseek;
mod mistral;
use crate::Error;
use crate::auth::InputSource;
use crate::auth::error::AuthConfigurationError;
use crate::ocr::error::OcrError;
use crate::ocr::types::OcrConnection;
pub(crate) use deepseek::VertexDeepSeekAdapter;
pub(crate) use mistral::VertexMistralAdapter;
fn validate_destination(connection: &OcrConnection) -> Result<(), OcrError> {
if connection.api_base.is_some() && connection.api_base_source == InputSource::Request {
return Err(Error::from(crate::AuthError::Configuration(
AuthConfigurationError::RequestVertexCredentialDestination,
))
.into());
}
Ok(())
}

161
litellm-rust/crates/core/src/ocr/client.rs
View file

@ -1,24 +1,39 @@
use std::sync::OnceLock;
use std::time::Duration;
use bytes::{Bytes, BytesMut};
use serde::de::DeserializeOwned;
use super::error::OcrError;
use super::handler::perform_ocr_request;
use super::error::{OcrError, OcrResponseError};
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use super::wire::{DecodedOcrResponse, decode_response};
use crate::Error;
use crate::auth::vertex::VertexAuth;
use crate::constants::OCR_CONNECT_TIMEOUT_SECS;
use crate::error::TransportError;
use crate::media::MediaFetcher;
#[derive(Clone)]
pub struct OcrClient {
provider_http: reqwest::Client,
polling_http: reqwest::Client,
document_fetcher: MediaFetcher,
vertex_auth: VertexAuth,
}
impl OcrClient {
pub fn new(provider_http: reqwest::Client) -> Result<Self, TransportError> {
Ok(Self { provider_http })
let document_fetcher = MediaFetcher::new().map_err(TransportError::from)?;
Ok(Self {
provider_http,
polling_http: no_redirect_http()?,
document_fetcher,
vertex_auth: VertexAuth::default(),
})
}
pub fn shared() -> Result<Self, Error> {
shared_client()
}
#[tracing::instrument(
@ -28,20 +43,72 @@ impl OcrClient {
skip_all
)]
pub async fn perform(&self, request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
perform_ocr_request(self, request).await
use super::{
NativeOutcome, OcrAdmission, OcrCall, OcrCallStep, OcrHookHost, OcrHost,
OcrHostOperation, OcrHostResult,
};
let host = OcrHookHost::new(request.hooks.clone());
let mut request = Some(request);
let NativeOutcome::Completed(mut call) = OcrCall::admit(self.clone(), OcrAdmission::all())
else {
return Err(Error::InvalidRequest(
"native OCR host admission declined".into(),
));
};
let mut result = None;
loop {
match call.resume(result.take()).await? {
OcrCallStep::Host(OcrHostOperation::ProjectRequest) => {
result = Some(OcrHostResult::Request(Ok((
Box::new(request.take().ok_or_else(|| {
Error::InvalidRequest("OCR request was already projected".into())
})?),
false,
))))
}
OcrCallStep::Host(operation) => result = Some(host.invoke(operation).await),
OcrCallStep::Complete(response) => return Ok(response),
}
}
}
pub(crate) fn provider_http(&self) -> &reqwest::Client {
&self.provider_http
}
pub(crate) fn polling_http(&self) -> &reqwest::Client {
&self.polling_http
}
pub(crate) fn document_fetcher(&self) -> &MediaFetcher {
&self.document_fetcher
}
pub(crate) fn vertex_auth(&self) -> &VertexAuth {
&self.vertex_auth
}
#[cfg(test)]
pub(crate) fn for_test(provider_http: reqwest::Client) -> Self {
Self { provider_http }
pub(crate) fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self {
Self {
provider_http,
polling_http: no_redirect_http().expect("test polling client builds"),
document_fetcher: MediaFetcher::for_test(document_http),
vertex_auth: VertexAuth::default(),
}
}
}
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
fn no_redirect_http() -> Result<reqwest::Client, TransportError> {
reqwest::Client::builder()
.connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(TransportError::from)
}
pub(crate) fn shared_client() -> Result<OcrClient, Error> {
static CLIENT: OnceLock<Result<OcrClient, TransportError>> = OnceLock::new();
let client = CLIENT
.get_or_init(|| {
@ -52,18 +119,50 @@ pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error
.and_then(OcrClient::new)
})
.clone()?;
client.perform(request).await
Ok(client)
}
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
shared_client()?.perform(request).await
}
pub async fn read_json_response<T: DeserializeOwned>(
response: reqwest::Response,
native: bool,
max_response_bytes: usize,
) -> Result<DecodedOcrResponse<T>, OcrError> {
let bytes = read_response_bytes(response, max_response_bytes).await?;
Ok(decode_response(&bytes, native)?)
}
pub(crate) async fn read_response_bytes(
mut response: reqwest::Response,
max_response_bytes: usize,
) -> Result<Bytes, OcrError> {
let status = response.status();
let bytes = response
.bytes()
.await
.map_err(crate::error::TransportError::from)?;
let limit = if status.is_success() {
max_response_bytes
} else {
max_response_bytes.min(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1))
};
if status.is_success()
&& response
.content_length()
.is_some_and(|length| length > limit as u64)
{
return Err(OcrResponseError::TooLarge { limit }.into());
}
let mut bytes = BytesMut::new();
while let Some(chunk) = response.chunk().await.map_err(transport_error)? {
let remaining = limit.saturating_sub(bytes.len());
if status.is_success() && chunk.len() > remaining {
return Err(OcrResponseError::TooLarge { limit }.into());
}
bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
if !status.is_success() && bytes.len() == limit {
break;
}
}
if !status.is_success() {
return Err(crate::error::TransportError::Http {
status: status.as_u16(),
@ -71,5 +170,41 @@ pub async fn read_json_response<T: DeserializeOwned>(
}
.into());
}
Ok(decode_response(&bytes, native)?)
Ok(bytes.freeze())
}
pub(crate) fn transport_error(error: reqwest::Error) -> Error {
if error.is_timeout() {
return Error::Http {
status: 408,
body: "OCR request timed out".into(),
};
}
crate::error::TransportError::from(error).into()
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn request_timeout_has_an_http_408_status() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let _connection = listener.accept().await.unwrap();
tokio::time::sleep(Duration::from_secs(1)).await;
});
let error = reqwest::Client::new()
.get(format!("http://{address}"))
.timeout(Duration::from_millis(10))
.send()
.await
.unwrap_err();
assert!(matches!(
transport_error(error),
Error::Http { status: 408, .. }
));
server.abort();
}
}

254
litellm-rust/crates/core/src/ocr/codecs/cohere.rs Normal file
View file

@ -0,0 +1,254 @@
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value, json};
use crate::ocr::document::InlineDocument;
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
pub(crate) enum OutputFormat {
#[default]
Markdown,
Blocks,
}
#[derive(Deserialize)]
pub(crate) struct CohereParams {
#[serde(default)]
pub output_format: OutputFormat,
}
#[derive(Deserialize, Serialize)]
pub(crate) struct CohereRequest {
pub model: String,
pub document: OcrDocument,
pub output_format: OutputFormat,
}
pub(crate) fn validate_document(document: &OcrDocument) -> Result<(), OcrRequestError> {
let OcrDocument::ImageUrl { image_url, .. } = document else {
return Err(OcrRequestError::CohereImageOnly);
};
if image_url.is_empty() {
return Err(OcrRequestError::CohereImageOnly);
}
if let Some(inline) = InlineDocument::parse(image_url)? {
if !inline.mime_type().type_.eq_ignore_ascii_case("image") {
return Err(OcrRequestError::CohereImageOnly);
}
inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
}
Ok(())
}
#[derive(Deserialize)]
pub(crate) struct CohereResponse {
#[serde(default)]
pages: Vec<CoherePage>,
meta: Option<CohereMeta>,
}
#[derive(Deserialize)]
struct CoherePage {
index: Option<i64>,
markdown: Option<CohereMarkdown>,
blocks: Option<Vec<Map<String, Value>>>,
}
#[derive(Deserialize)]
struct CohereMarkdown {
#[serde(default)]
content: String,
images: Option<Vec<Map<String, Value>>>,
}
#[derive(Deserialize)]
struct CohereMeta {
billed_units: Option<CohereBilledUnits>,
}
#[derive(Deserialize)]
struct CohereBilledUnits {
pages: Option<i64>,
}
pub(crate) fn transform_response(
model: &str,
response: CohereResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
let pages_processed = response
.meta
.and_then(|meta| meta.billed_units)
.and_then(|units| units.pages)
.map(Ok)
.unwrap_or_else(|| {
i64::try_from(response.pages.len()).map_err(|_| OcrResponseError::NumericRange("pages"))
})?;
let pages = response
.pages
.into_iter()
.enumerate()
.map(|(position, page)| {
let index = page.index.map(Ok).unwrap_or_else(|| {
i64::try_from(position).map_err(|_| OcrResponseError::NumericRange("page index"))
})?;
let (content, images) = page
.markdown
.map(|markdown| {
let images =
markdown
.images
.filter(|images| !images.is_empty())
.map(|images| {
images
.into_iter()
.map(|mut image| {
if let Some(Value::Object(bbox)) =
image.get("bounding_box").cloned()
{
image.insert("bbox".into(), Value::Object(bbox));
}
Value::Object(image)
})
.collect::<Vec<_>>()
});
(markdown.content, images)
})
.unwrap_or_default();
let mut normalized = json!({"index": index, "markdown": content, "images": images});
if let Some(blocks) = page.blocks {
normalized["blocks"] = json!(blocks);
}
Ok(normalized)
})
.collect::<Result<Vec<_>, OcrResponseError>>()?;
Ok(LiteLLMOcrResponse {
pages,
model: model.into(),
document_annotation: None,
usage_info: Some(json!({"pages_processed": pages_processed})),
object: "ocr".into(),
extra_fields: Map::new(),
provider_native_response: None,
})
}
pub(crate) fn transform_request(
model: &str,
document: OcrDocument,
params: CohereParams,
) -> Result<CohereRequest, OcrRequestError> {
validate_document(&document)?;
Ok(CohereRequest {
model: model.into(),
document,
output_format: params.output_format,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn response_normalizes_markdown_images_blocks_and_billed_pages() {
let response = serde_json::from_value(json!({
"pages": [
{
"type":"markdown",
"index":4,
"markdown":{
"content":"receipt",
"images":[{
"id":"image",
"bounding_box":{"top_left_x":1,"bottom_right_x":48},
"bounding_box_normalized":{"top_left_x":0.04,"bottom_right_x":0.15},
"description":"scan",
"category":"logo"
}]
}
},
{"type":"blocks","blocks":[{"type":"text","text":{"content":"total"}}]}
],
"meta":{"api_version":{"version":"2"},"billed_units":{"pages":3}}
}))
.unwrap();
let normalized = transform_response("parse-v5.0", response).unwrap();
assert_eq!(normalized.pages[0]["index"], 4);
assert_eq!(normalized.pages[0]["markdown"], "receipt");
assert_eq!(normalized.pages[0]["images"][0]["bbox"]["top_left_x"], 1);
assert_eq!(
normalized.pages[0]["images"][0]["bounding_box_normalized"]["bottom_right_x"],
0.15
);
assert_eq!(normalized.pages[0]["images"][0]["description"], "scan");
assert_eq!(normalized.pages[0]["images"][0]["category"], "logo");
assert_eq!(normalized.pages[1]["index"], 1);
assert_eq!(normalized.pages[1]["markdown"], "");
assert_eq!(normalized.pages[1]["blocks"][0]["text"]["content"], "total");
assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 3);
}
#[test]
fn response_defaults_and_invalid_fields() {
for value in [
json!({}),
json!({"meta":null}),
json!({"pages":[],"meta":{"billed_units":null}}),
] {
let normalized =
transform_response("parse", serde_json::from_value(value).unwrap()).unwrap();
assert!(normalized.pages.is_empty());
assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 0);
}
for value in [
json!({"pages":null}),
json!({"pages":[{"markdown":"text"}]}),
json!({"pages":[{"index":"bad"}]}),
] {
assert!(serde_json::from_value::<CohereResponse>(value).is_err());
}
let normalized = transform_response(
"parse",
serde_json::from_value(json!({"pages":[{"markdown":null}]})).unwrap(),
)
.unwrap();
assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 1);
assert!(normalized.pages[0]["images"].is_null());
}
#[test]
fn request_requires_image_and_supported_output_format() {
for value in [
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
json!({"type":"image_url","image_url":""}),
json!({"type":"image_url","image_url":"data:application/pdf;base64,YQ=="}),
] {
assert_eq!(
validate_document(&serde_json::from_value(value).unwrap()),
Err(OcrRequestError::CohereImageOnly)
);
}
assert!(serde_json::from_value::<CohereParams>(json!({"output_format":"html"})).is_err());
for format in ["markdown", "blocks"] {
assert!(
serde_json::from_value::<CohereParams>(json!({"output_format":format})).is_ok()
);
}
let request = transform_request(
"parse-v5.0",
serde_json::from_value(json!({
"type":"image_url",
"image_url":"https://example.com/image.png"
}))
.unwrap(),
serde_json::from_value(json!({})).unwrap(),
)
.unwrap();
assert_eq!(
serde_json::to_value(request).unwrap()["output_format"],
"markdown"
);
}
}

5
litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs Normal file
View file

@ -0,0 +1,5 @@
mod transformation;
mod types;
pub(crate) use transformation::{transform_ocr_request, transform_ocr_response};
pub(crate) use types::{DeepSeekOcrParams, DeepSeekOcrResponse};

102
litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs Normal file
View file

@ -0,0 +1,102 @@
use serde::de::IntoDeserializer;
use serde_json::{Value, json};
use super::types::*;
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) fn transform_ocr_request(
provider_model: &str,
document: OcrDocument,
params: &DeepSeekOcrParams,
) -> Result<DeepSeekOcrRequest, OcrRequestError> {
if document.source().is_empty() {
return Err(OcrRequestError::MissingDocumentUrl);
}
let content = OcrDocument::ImageUrl {
image_url: document.source().to_string(),
extra_fields: serde_json::Map::new(),
};
Ok(DeepSeekOcrRequest {
model: provider_model.to_string(),
messages: vec![DeepSeekOcrMessage {
role: UserRole::User,
content: vec![content],
}],
params: params.clone(),
})
}
pub(crate) fn transform_ocr_response(
model: &str,
response: DeepSeekOcrResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
let content = response
.choices
.into_iter()
.next()
.and_then(|choice| choice.message.content)
.ok_or(OcrResponseError::EmptyContent)?;
let decoded = decode_content(content)?;
let pages = match decoded.result.pages {
Some(pages) if !pages.is_empty() => pages
.into_iter()
.map(|page| serde_json::to_value(page).expect("DeepSeek page serializes"))
.collect(),
_ => vec![json!({
"index":0,
"markdown":decoded.fallback_markdown,
"images":null
})],
};
Ok(LiteLLMOcrResponse {
pages,
model: decoded.result.model.unwrap_or_else(|| model.to_string()),
document_annotation: decoded.result.document_annotation,
usage_info: decoded.result.usage_info.or(response.usage),
object: "ocr".into(),
extra_fields: decoded.result.extra_fields,
provider_native_response: None,
})
}
struct DecodedContent {
result: DeepSeekOcrResult,
fallback_markdown: String,
}
fn decode_content(content: DeepSeekContent) -> Result<DecodedContent, OcrResponseError> {
let (result, fallback_markdown) = match content {
DeepSeekContent::Text(text) if text.is_empty() => {
return Err(OcrResponseError::EmptyContent);
}
DeepSeekContent::Text(text) => (decode_json_content(&text)?, text),
DeepSeekContent::Object(object) => {
let fallback =
serde_json::to_string(&object).map_err(|_| OcrResponseError::ResponseField {
path: "choices[0].message.content".into(),
})?;
(Some(object), fallback)
}
};
Ok(DecodedContent {
result: result.unwrap_or_default(),
fallback_markdown,
})
}
fn decode_json_content(text: &str) -> Result<Option<DeepSeekOcrResult>, OcrResponseError> {
if !text.trim_start().starts_with('{') {
return Ok(None);
}
let value = match serde_json::from_str::<Value>(text) {
Ok(value) => value,
Err(_) => return Ok(None),
};
serde_path_to_error::deserialize(value.into_deserializer())
.map(Some)
.map_err(|error| OcrResponseError::ResponseField {
path: format!("choices[0].message.content.{}", error.path()),
})
}

95
litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs Normal file
View file

@ -0,0 +1,95 @@
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub stream: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub n: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub stop: Option<StopSequences>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum StopSequences {
One(String),
Many(Vec<String>),
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrRequest {
pub model: String,
pub messages: Vec<DeepSeekOcrMessage>,
#[serde(flatten)]
pub params: DeepSeekOcrParams,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrMessage {
pub role: UserRole,
pub content: Vec<crate::ocr::types::OcrDocument>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub(crate) enum UserRole {
User,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DeepSeekOcrResponse {
#[serde(default)]
pub choices: Vec<DeepSeekChoice>,
pub usage: Option<Value>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DeepSeekChoice {
pub message: DeepSeekResponseMessage,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DeepSeekResponseMessage {
pub content: Option<DeepSeekContent>,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(untagged)]
pub(crate) enum DeepSeekContent {
Text(String),
Object(DeepSeekOcrResult),
}
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrResult {
#[serde(skip_serializing_if = "Option::is_none")]
pub pages: Option<Vec<DeepSeekPage>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub usage_info: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub document_annotation: Option<Value>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekPage {
#[serde(default)]
pub index: i64,
#[serde(default)]
pub markdown: String,
pub images: Option<Value>,
pub dimensions: Option<Value>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}

9
litellm-rust/crates/core/src/ocr/codecs/document_intelligence/mod.rs Normal file
View file

@ -0,0 +1,9 @@
mod params;
mod transformation;
mod types;
pub(crate) use params::{decode_input_params, map_ocr_params};
pub(crate) use transformation::{transform_ocr_request, transform_ocr_response};
pub(crate) use types::{
AzureDocumentIntelligenceOperation, DocumentIntelligenceParams, OperationStatus,
};

219
litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs Normal file
View file

@ -0,0 +1,219 @@
use std::collections::BTreeSet;
use serde_json::{Map, Value};
use super::types::{
DocumentIntelligenceInputParams, DocumentIntelligenceParams, FeaturesInput, PagesInput,
};
use crate::ocr::error::OcrRequestError;
use crate::ocr::prepare::ParsedProviderParams;
pub(crate) fn decode_input_params(
params: Map<String, Value>,
prefix: &str,
) -> Result<ParsedProviderParams<DocumentIntelligenceInputParams>, OcrRequestError> {
if let Some(Value::Array(pages)) = params.get("pages") {
if pages.iter().any(Value::is_boolean) {
return Err(OcrRequestError::Pages("boolean page index".into()));
}
if pages
.iter()
.any(|page| page.is_number() && page.as_i64().is_none())
{
return Err(OcrRequestError::Pages("page index is out of range".into()));
}
if !pages.iter().all(Value::is_i64) && !pages.iter().all(Value::is_string) {
return Err(OcrRequestError::Pages("mixed page element types".into()));
}
}
crate::ocr::wire::decode_request_value(Value::Object(params), prefix)
}
pub(crate) fn map_ocr_params(
params: DocumentIntelligenceInputParams,
) -> Result<DocumentIntelligenceParams, OcrRequestError> {
Ok(DocumentIntelligenceParams {
pages: params.pages.map(normalize_pages).transpose()?.flatten(),
features: params
.features
.map(normalize_features)
.transpose()?
.flatten(),
})
}
fn normalize_pages(pages: PagesInput) -> Result<Option<String>, OcrRequestError> {
let normalized = match pages {
PagesInput::ZeroBasedIndices(indices) => {
if indices.is_empty() {
return Ok(None);
}
indices
.into_iter()
.map(|page| {
if page < 0 {
return Err(OcrRequestError::Pages("negative page index".into()));
}
page.checked_add(1)
.ok_or_else(|| OcrRequestError::Pages("page index is out of range".into()))
})
.collect::<Result<BTreeSet<_>, _>>()?
.into_iter()
.map(|page| page.to_string())
.collect::<Vec<_>>()
.join(",")
}
PagesInput::NativeTokens(tokens) => {
if tokens.is_empty() {
return Ok(None);
}
tokens
.iter()
.map(|token| token.trim())
.collect::<Vec<_>>()
.join(",")
}
PagesInput::NativeRange(range) => range
.split(',')
.map(str::trim)
.collect::<Vec<_>>()
.join(","),
};
if !normalized.split(',').all(valid_page_token) {
return Err(OcrRequestError::Pages("invalid native page range".into()));
}
Ok(Some(normalized))
}
fn valid_page_token(token: &str) -> bool {
let mut parts = token.split('-');
let start = parts.next().unwrap_or_default();
if start.is_empty() || !start.chars().all(|character| character.is_ascii_digit()) {
return false;
}
match parts.next() {
None => true,
Some(end) => {
!end.is_empty()
&& end.chars().all(|character| character.is_ascii_digit())
&& parts.next().is_none()
}
}
}
fn normalize_features(features: FeaturesInput) -> Result<Option<String>, OcrRequestError> {
let tokens = match features {
FeaturesInput::Names(names) => names,
FeaturesInput::CommaSeparated(names) => names.split(',').map(str::to_string).collect(),
};
if tokens.is_empty() {
return Ok(None);
}
let normalized = tokens.iter().map(|token| token.trim()).collect::<Vec<_>>();
if !normalized.iter().all(|token| {
let Some((first, rest)) = token.as_bytes().split_first() else {
return false;
};
first.is_ascii_alphabetic() && rest.iter().all(u8::is_ascii_alphanumeric)
}) {
return Err(OcrRequestError::Features);
}
Ok(Some(normalized.join(",")))
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::{Value, json};
use super::*;
fn map(value: Value) -> Result<DocumentIntelligenceParams, OcrRequestError> {
let fields = value.as_object().unwrap().clone();
map_ocr_params(decode_input_params(fields, "optional_params")?.known)
}
#[test]
fn input_params_retain_unknown_fields() {
let parsed = decode_input_params(
json!({
"pages": [0],
"future_ocr_option": true,
"extra_body": {"provider_option": "value"}
})
.as_object()
.unwrap()
.clone(),
"optional_params",
)
.unwrap();
assert_eq!(
parsed.known.pages,
Some(PagesInput::ZeroBasedIndices(vec![0]))
);
assert_eq!(parsed.extra_params["future_ocr_option"], true);
assert_eq!(
parsed.extra_params["extra_body"],
json!({"provider_option": "value"})
);
assert_eq!(
serde_json::to_value(map_ocr_params(parsed.known).unwrap()).unwrap(),
json!({"pages": "1", "features": null})
);
}
#[rstest]
#[case(json!([0, 1, 2]), Some("1,2,3"))]
#[case(json!([2, 0, 0, 1]), Some("1,2,3"))]
#[case(json!([]), None)]
#[case(json!("3-9"), Some("3-9"))]
#[case(json!("1-3, 5"), Some("1-3,5"))]
#[case(json!(["1", "3-5"]), Some("1,3-5"))]
fn page_mapping_matches_python(#[case] input: Value, #[case] expected: Option<&str>) {
assert_eq!(
map(json!({"pages": input})).unwrap().pages.as_deref(),
expected
);
}
#[rstest]
#[case(json!("a,b"))]
#[case(json!([-1]))]
#[case(json!([true, false]))]
#[case(json!([1, "2"]))]
#[case(json!(5))]
fn invalid_page_mapping_matches_python(#[case] input: Value) {
assert!(map(json!({"pages": input})).is_err());
}
#[rstest]
#[case(json!(["keyValuePairs"]), "keyValuePairs")]
#[case(json!(["keyValuePairs", "languages"]), "keyValuePairs,languages")]
#[case(json!("keyValuePairs"), "keyValuePairs")]
#[case(json!("keyValuePairs,languages"), "keyValuePairs,languages")]
#[case(json!("keyValuePairs, languages"), "keyValuePairs,languages")]
fn feature_mapping_matches_python(#[case] input: Value, #[case] expected: &str) {
assert_eq!(
map(json!({"features": input})).unwrap().features.as_deref(),
Some(expected)
);
}
#[rstest]
#[case(json!("keyValuePairs&pages=9"))]
#[case(json!("key value pairs"))]
#[case(json!(""))]
#[case(json!([1, 2]))]
#[case(json!([["keyValuePairs"]]))]
#[case(json!({"feature":"keyValuePairs"}))]
#[case(json!(5))]
fn invalid_feature_mapping_matches_python(#[case] input: Value) {
assert!(map(json!({"features": input})).is_err());
}
#[test]
fn empty_feature_list_is_omitted() {
assert_eq!(map(json!({"features": []})).unwrap().features, None);
}
}

108
litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs Normal file
View file

@ -0,0 +1,108 @@
use base64::{Engine, engine::general_purpose::STANDARD};
use serde_json::{Map, Value, json};
use super::types::*;
use crate::constants::{AZURE_DI_DEFAULT_DPI, AZURE_DI_DEFAULT_HEIGHT, AZURE_DI_DEFAULT_WIDTH};
use crate::ocr::document::InlineDocument;
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) fn transform_ocr_request(
document: OcrDocument,
) -> Result<DocumentIntelligenceRequest, OcrRequestError> {
let source = document.source();
if source.is_empty() {
return Err(OcrRequestError::MissingDocumentUrl);
}
Ok(if let Some(document) = InlineDocument::parse(source)? {
DocumentIntelligenceRequest::Base64Source(
STANDARD.encode(document.decode(crate::constants::OCR_INLINE_MAX_BYTES)?),
)
} else {
DocumentIntelligenceRequest::UrlSource(source.to_string())
})
}
pub(crate) fn transform_ocr_response(
model: &str,
response: AzureDocumentIntelligenceOperation,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
if response.status != Some(OperationStatus::Succeeded) {
return Err(OcrResponseError::OperationStatus(
response
.status
.map(|status| status.to_string())
.unwrap_or_else(|| "None".into()),
));
}
let result = response.analyze_result.unwrap_or_default();
let pages = result
.pages
.into_iter()
.map(normalize_page)
.collect::<Result<Vec<_>, _>>()?;
let pages_processed = pages.len();
let mut extra_fields = Map::new();
extra_fields.insert("content".into(), option_value(result.content));
extra_fields.insert("tables".into(), option_value(result.tables));
extra_fields.insert("keyValuePairs".into(), option_value(result.key_value_pairs));
Ok(LiteLLMOcrResponse {
pages,
model: model.into(),
document_annotation: None,
usage_info: Some(json!({"pages_processed":pages_processed})),
object: "ocr".into(),
extra_fields,
provider_native_response: None,
})
}
fn normalize_page(page: AzureDocumentIntelligencePage) -> Result<Value, OcrResponseError> {
let index = page
.page_number
.unwrap_or(1)
.checked_sub(1)
.ok_or(OcrResponseError::NumericRange("page.pageNumber"))?;
let scale = if page.unit.as_deref().unwrap_or("inch") == "inch" {
AZURE_DI_DEFAULT_DPI as f64
} else {
1.0
};
let width = pixel_dimension(
page.width.unwrap_or(AZURE_DI_DEFAULT_WIDTH),
scale,
"page.width",
)?;
let height = pixel_dimension(
page.height.unwrap_or(AZURE_DI_DEFAULT_HEIGHT),
scale,
"page.height",
)?;
let markdown = page
.lines
.iter()
.map(|line| line.content.as_deref().unwrap_or_default())
.collect::<Vec<_>>()
.join("\n");
Ok(json!({
"index":index,
"markdown":markdown,
"images":null,
"dimensions":{"width":width,"height":height,"dpi":AZURE_DI_DEFAULT_DPI}
}))
}
fn pixel_dimension(value: f64, scale: f64, field: &'static str) -> Result<i64, OcrResponseError> {
let value = value * scale;
if !value.is_finite() || value < i64::MIN as f64 || value > i64::MAX as f64 {
return Err(OcrResponseError::NumericRange(field));
}
Ok(value.trunc() as i64)
}
fn option_value<T: serde::Serialize>(value: Option<T>) -> Value {
value
.and_then(|value| serde_json::to_value(value).ok())
.unwrap_or(Value::Null)
}

138
litellm-rust/crates/core/src/ocr/codecs/document_intelligence/types.rs Normal file
View file

@ -0,0 +1,138 @@
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value};
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum PagesInput {
ZeroBasedIndices(Vec<i64>),
NativeTokens(Vec<String>),
NativeRange(String),
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum FeaturesInput {
Names(Vec<String>),
CommaSeparated(String),
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub(crate) struct DocumentIntelligenceInputParams {
pub pages: Option<PagesInput>,
pub features: Option<FeaturesInput>,
}
#[derive(Clone, Debug, PartialEq, Serialize)]
pub(crate) struct DocumentIntelligenceParams {
pub pages: Option<String>,
pub features: Option<String>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) enum DocumentIntelligenceRequest {
#[serde(rename = "urlSource")]
UrlSource(String),
#[serde(rename = "base64Source")]
Base64Source(String),
}
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum OperationStatus {
Succeeded,
Running,
NotStarted,
Failed,
Unknown(String),
}
impl<'de> Deserialize<'de> for OperationStatus {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
Ok(match String::deserialize(deserializer)?.as_str() {
"succeeded" => Self::Succeeded,
"running" => Self::Running,
"notStarted" => Self::NotStarted,
"failed" => Self::Failed,
value => Self::Unknown(value.to_string()),
})
}
}
impl std::fmt::Display for OperationStatus {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(match self {
Self::Succeeded => "succeeded",
Self::Running => "running",
Self::NotStarted => "notStarted",
Self::Failed => "failed",
Self::Unknown(value) => value,
})
}
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct AzureDocumentIntelligenceOperation {
pub status: Option<OperationStatus>,
#[serde(rename = "analyzeResult")]
pub analyze_result: Option<AzureDocumentIntelligenceAnalyzeResult>,
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct AzureDocumentIntelligenceAnalyzeResult {
pub content: Option<String>,
#[serde(default)]
pub pages: Vec<AzureDocumentIntelligencePage>,
pub tables: Option<Vec<Map<String, Value>>>,
#[serde(rename = "keyValuePairs")]
pub key_value_pairs: Option<Vec<Map<String, Value>>>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct AzureDocumentIntelligencePage {
#[serde(rename = "pageNumber", default, deserialize_with = "optional_i64")]
pub page_number: Option<i64>,
#[serde(default, deserialize_with = "optional_f64")]
pub width: Option<f64>,
#[serde(default, deserialize_with = "optional_f64")]
pub height: Option<f64>,
pub unit: Option<String>,
#[serde(default)]
pub lines: Vec<AzureDocumentIntelligenceLine>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct AzureDocumentIntelligenceLine {
pub content: Option<String>,
}
fn optional_i64<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<i64>, D::Error> {
match Option::<Value>::deserialize(deserializer)? {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(number)) => number
.as_i64()
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected an integer")),
Some(Value::String(value)) => value
.parse::<i64>()
.map(Some)
.map_err(|_| serde::de::Error::custom("expected an integer")),
Some(_) => Err(serde::de::Error::custom("expected an integer")),
}
}
fn optional_f64<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<f64>, D::Error> {
match Option::<Value>::deserialize(deserializer)? {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(number)) => number
.as_f64()
.filter(|value| value.is_finite())
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected a finite number")),
Some(Value::String(value)) => value
.parse::<f64>()
.ok()
.filter(|value| value.is_finite())
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected a finite number")),
Some(_) => Err(serde::de::Error::custom("expected a number")),
}
}

199
litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs
View file

@ -36,8 +36,105 @@ mod tests {
use rstest::rstest;
use serde_json::{Value, json};
fn mapped_params(value: Value) -> Value {
serde_json::to_value(serde_json::from_value::<MistralOcrParams>(value).unwrap()).unwrap()
}
fn document() -> OcrDocument {
serde_json::from_value(
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
)
.unwrap()
}
#[rstest]
fn extract_header_is_a_supported_ocr_param() {
assert_eq!(
mapped_params(json!({"extract_header":true}))["extract_header"],
true
);
}
#[rstest]
fn extract_footer_is_a_supported_ocr_param() {
assert_eq!(
mapped_params(json!({"extract_footer":false}))["extract_footer"],
false
);
}
#[rstest]
fn existing_ocr_params_remain_supported() {
let mapped = mapped_params(json!({
"pages":[0,2],
"include_image_base64":true,
"image_limit":2,
"image_min_size":100,
"bbox_annotation_format":{"type":"json_schema"},
"document_annotation_format":{"type":"json_schema"}
}));
assert_eq!(mapped["pages"], json!([0, 2]));
assert_eq!(mapped["include_image_base64"], true);
assert_eq!(mapped["image_limit"], 2);
assert_eq!(mapped["image_min_size"], 100);
assert_eq!(mapped["bbox_annotation_format"]["type"], "json_schema");
assert_eq!(mapped["document_annotation_format"]["type"], "json_schema");
}
#[rstest]
fn map_ocr_params_forwards_extract_header() {
assert_eq!(
mapped_params(json!({"extract_header":true}))["extract_header"],
true
);
}
#[rstest]
fn map_ocr_params_forwards_extract_footer() {
assert_eq!(
mapped_params(json!({"extract_footer":true}))["extract_footer"],
true
);
}
#[rstest]
fn map_ocr_params_forwards_extract_header_and_footer() {
let mapped = mapped_params(json!({"extract_header":true,"extract_footer":false}));
assert_eq!(mapped["extract_header"], true);
assert_eq!(mapped["extract_footer"], false);
}
#[rstest]
fn map_ocr_params_drops_unknown_params() {
let mapped = mapped_params(json!({"extract_header":true,"unsupported_param":"value"}));
assert_eq!(mapped["extract_header"], true);
assert!(mapped.get("unsupported_param").is_none());
}
#[rstest]
#[case("table_format", json!("html"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("confidence_scores_granularity", json!("block"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("include_blocks", json!(true))]
#[case("id", json!("req-123"))]
fn new_ocr_params_are_supported(#[case] name: &str, #[case] value: Value) {
assert_eq!(mapped_params(json!({name:value.clone()}))[name], value);
}
#[rstest]
#[case("table_format", json!("html"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("include_blocks", json!(true))]
#[case("id", json!("req-123"))]
fn map_ocr_params_forwards_new_ocr_params(#[case] name: &str, #[case] value: Value) {
assert_eq!(mapped_params(json!({name:value.clone()}))[name], value);
}
#[rstest]
#[case("pages", json!([0, 2]))]
#[case("pages", json!("0,2-4"))]
#[case("include_image_base64", json!(true))]
#[case("image_limit", json!(2))]
#[case("image_min_size", json!(100))]
@ -53,50 +150,102 @@ mod tests {
fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) {
let params: MistralOcrParams =
serde_json::from_value(json!({name: value.clone()})).unwrap();
let document: OcrDocument = serde_json::from_value(
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
)
.unwrap();
let result =
serde_json::to_value(transform_ocr_request("model", document, &params).unwrap())
serde_json::to_value(transform_ocr_request("model", document(), &params).unwrap())
.unwrap();
assert_eq!(result["model"], "model");
assert_eq!(result[name], value);
}
#[test]
fn request_mapping_filters_unknown_fields() {
let params: MistralOcrParams = serde_json::from_value(json!({"unknown": true})).unwrap();
let document: OcrDocument = serde_json::from_value(
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
#[rstest]
#[case("table_format", json!("html"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("id", json!("req-123"))]
#[case("extract_header", json!(true))]
#[case("include_blocks", json!(true))]
#[case("pages", json!([0,1]))]
fn transform_ocr_request_includes_each_optional_param(
#[case] name: &str,
#[case] value: Value,
) {
let params: MistralOcrParams = serde_json::from_value(json!({name:value.clone()})).unwrap();
let result = serde_json::to_value(
transform_ocr_request("mistral-ocr-latest", document(), &params).unwrap(),
)
.unwrap();
let result =
serde_json::to_value(transform_ocr_request("model", document, &params).unwrap())
.unwrap();
assert!(result.get("unknown").is_none());
assert_eq!(result[name], value);
assert_eq!(result["model"], "mistral-ocr-latest");
}
#[test]
fn response_preserves_provider_fields() {
#[rstest]
fn transform_ocr_request_includes_multiple_new_params() {
let params: MistralOcrParams = serde_json::from_value(json!({
"table_format":"html",
"confidence_scores_granularity":"page",
"extract_header":true
}))
.unwrap();
let result = serde_json::to_value(
transform_ocr_request("mistral-ocr-latest", document(), &params).unwrap(),
)
.unwrap();
assert_eq!(result["table_format"], "html");
assert_eq!(result["confidence_scores_granularity"], "page");
assert_eq!(result["extract_header"], true);
}
#[rstest]
fn transform_ocr_response_preserves_blocks_and_confidence_scores() {
let response: MistralOcrResponse = serde_json::from_value(json!({
"pages":[{"index":0,"markdown":"hello","header":"head","confidence_scores":{"mean":0.99}}],
"pages":[{
"index":0,
"markdown":"hello",
"images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}],
"dimensions":{"width":612,"height":792,"dpi":72},
"blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}],
"confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97}
}],
"model":"returned-model",
"usage_info":{"pages_processed":1,"future_counter":5},
"future_response_field":"kept"
"document_annotation":"{\"language\":\"en\"}",
"usage_info":{"pages_processed":1}
}))
.unwrap();
let result = transform_ocr_response("model", response)
.unwrap()
.into_json();
assert_eq!(result["pages"][0]["header"], "head");
assert_eq!(result["usage_info"]["future_counter"], 5);
assert_eq!(result["future_response_field"], "kept");
assert_eq!(result["pages"][0]["blocks"][0]["type"], "title");
assert_eq!(result["pages"][0]["blocks"][0]["bbox"]["x"], 1);
assert_eq!(
result["pages"][0]["blocks"][0]["confidence_scores"]["mean"],
0.98
);
assert_eq!(
result["pages"][0]["confidence_scores"]["average_page_confidence_score"],
0.99
);
assert_eq!(result["pages"][0]["images"][0]["id"], "img-0");
assert_eq!(result["pages"][0]["dimensions"]["dpi"], 72);
assert_eq!(result["model"], "returned-model");
assert_eq!(result["document_annotation"], "{\"language\":\"en\"}");
assert_eq!(result["usage_info"]["pages_processed"], 1);
}
#[test]
fn response_rejects_null_pages() {
assert!(serde_json::from_value::<MistralOcrResponse>(json!({"pages":null})).is_err());
#[rstest]
fn transform_ocr_response_preserves_ocr4_page_fields() {
let page = json!({
"index":0,
"markdown":"table page",
"tables":[{"rows":2,"cols":3}],
"hyperlinks":["https://example.com"],
"header":"header",
"footer":"footer"
});
let response: MistralOcrResponse =
serde_json::from_value(json!({"pages":[page.clone()]})).unwrap();
let result = transform_ocr_response("model", response)
.unwrap()
.into_json();
assert_eq!(result["pages"][0], page);
}
}

9
litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs
View file

@ -3,10 +3,17 @@ use serde_json::{Map, Value};
use crate::ocr::types::OcrDocument;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum MistralOcrPages {
Range(String),
Indices(Vec<i64>),
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub(crate) struct MistralOcrParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub pages: Option<Vec<i64>>,
pub pages: Option<MistralOcrPages>,
#[serde(skip_serializing_if = "Option::is_none")]
pub include_image_base64: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]

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

@ -1 +1,5 @@
pub(crate) mod cohere;
pub(crate) mod deepseek;
pub(crate) mod document_intelligence;
pub(crate) mod mistral;
pub(crate) mod reducto;

9
litellm-rust/crates/core/src/ocr/codecs/reducto/mod.rs Normal file
View file

@ -0,0 +1,9 @@
mod transformation;
mod types;
pub(crate) use transformation::{
transform_legacy_ocr_request, transform_ocr_response, transform_v3_ocr_request,
};
pub(crate) use types::{
ReductoLegacyParams, ReductoResponse, ReductoUploadResponse, ReductoV3Params,
};

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