mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
Merge remote-tracking branch 'origin/main' into fix/databricks-gpt5-function-tools-reasoning
# Conflicts: # litellm/responses/bridge_check.py
This commit is contained in:
commit
fec5b3d139
2133 changed files with 245943 additions and 24514 deletions
|
|
@ -1293,6 +1293,19 @@ jobs:
|
|||
uv venv --python 3.12 .venv-mcp-peer
|
||||
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
|
||||
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
|
||||
- when:
|
||||
condition:
|
||||
equal: [sdk, << parameters.suite >>]
|
||||
steps:
|
||||
- run:
|
||||
name: Install redis-server 7.2.16
|
||||
command: |
|
||||
mkdir -p /tmp/redis-bin
|
||||
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
|
||||
docker cp "${cid}:/usr/local/bin/redis-server" /tmp/redis-bin/redis-server
|
||||
docker rm "$cid"
|
||||
/tmp/redis-bin/redis-server --version | grep -q 'v=7.2.16'
|
||||
echo 'export PATH="/tmp/redis-bin:$PATH"' >> "$BASH_ENV"
|
||||
- start_postgres:
|
||||
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
|
||||
server_args: "-c max_connections=300 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"
|
||||
|
|
|
|||
|
|
@ -1,90 +0,0 @@
|
|||
#!/usr/bin/env bash
|
||||
#
|
||||
# pre-push — enforce Conventional Branches
|
||||
# https://conventional-branch.github.io/
|
||||
#
|
||||
# Branch format: <type>/<description>
|
||||
# <type> must be one of: feature, bugfix, hotfix, release, chore
|
||||
#
|
||||
# Protected branches (always allowed):
|
||||
# - main
|
||||
# - dependabot/*
|
||||
# - gh-readonly-queue/*
|
||||
#
|
||||
# Tag pushes and branch deletions are skipped.
|
||||
# Bypass: git push --no-verify.
|
||||
|
||||
set -eu
|
||||
|
||||
ZERO_OID="0000000000000000000000000000000000000000"
|
||||
ZERO_OID_SHA256="0000000000000000000000000000000000000000000000000000000000000000"
|
||||
ALLOWED_TYPES="feature|bugfix|hotfix|release|chore"
|
||||
BRANCH_PATTERN="^(${ALLOWED_TYPES})/.+"
|
||||
|
||||
PROTECTED_NAMES="main"
|
||||
PROTECTED_PREFIXES="dependabot/ gh-readonly-queue/"
|
||||
|
||||
is_protected() {
|
||||
branch="$1"
|
||||
for name in $PROTECTED_NAMES; do
|
||||
if [ "$branch" = "$name" ]; then
|
||||
return 0
|
||||
fi
|
||||
done
|
||||
for prefix in $PROTECTED_PREFIXES; do
|
||||
case "$branch" in "$prefix"*) return 0 ;; esac
|
||||
done
|
||||
return 1
|
||||
}
|
||||
|
||||
invalid=""
|
||||
|
||||
while read -r local_ref local_oid remote_ref remote_oid; do
|
||||
# Branch deletion (no local commit being pushed).
|
||||
if [ "$local_oid" = "$ZERO_OID" ] || [ "$local_oid" = "$ZERO_OID_SHA256" ]; then
|
||||
continue
|
||||
fi
|
||||
|
||||
# Only validate branch pushes; ignore tags and other ref namespaces.
|
||||
case "$remote_ref" in
|
||||
refs/heads/*) ;;
|
||||
*) continue ;;
|
||||
esac
|
||||
|
||||
branch="${remote_ref#refs/heads/}"
|
||||
|
||||
if is_protected "$branch"; then
|
||||
continue
|
||||
fi
|
||||
|
||||
if ! printf '%s' "$branch" | grep -Eq "$BRANCH_PATTERN"; then
|
||||
invalid="$invalid $branch"
|
||||
fi
|
||||
done
|
||||
|
||||
if [ -n "$invalid" ]; then
|
||||
cat >&2 <<EOF
|
||||
✗ Branch name does not follow Conventional Branches.
|
||||
|
||||
Invalid:$invalid
|
||||
|
||||
Expected: <type>/<description>
|
||||
|
||||
Allowed types: feature, bugfix, hotfix, release, chore
|
||||
Examples:
|
||||
feature/weighted-round-robin
|
||||
bugfix/streaming-empty-chunks
|
||||
chore/bump-deps
|
||||
hotfix/auth-bypass
|
||||
|
||||
Protected (always allowed): main, dependabot/*, gh-readonly-queue/*.
|
||||
|
||||
See https://conventional-branch.github.io/
|
||||
|
||||
Rename with: git branch -m <new-name>
|
||||
To bypass (use sparingly): git push --no-verify
|
||||
EOF
|
||||
exit 1
|
||||
fi
|
||||
|
||||
exit 0
|
||||
2
.github/CODEOWNERS
vendored
2
.github/CODEOWNERS
vendored
|
|
@ -1,2 +1,4 @@
|
|||
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri
|
||||
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri
|
||||
/helm/ @yassin-berriai
|
||||
/terraform/ @yassin-berriai
|
||||
|
|
|
|||
2
.github/ci-coverage-allowlist.yml
vendored
2
.github/ci-coverage-allowlist.yml
vendored
|
|
@ -5,7 +5,7 @@ description: >-
|
|||
|
||||
test_paths:
|
||||
- reason: >-
|
||||
litellm.agent() end-to-end suite. It drives the real claude, codex and opencode CLIs and
|
||||
litellm.agent() end-to-end suite. It drives the real claude, codex, opencode and pi CLIs and
|
||||
deepagents against a live LiteLLM AI Gateway, so it needs those binaries on PATH plus
|
||||
LITELLM_PROXY_API_BASE / LITELLM_PROXY_API_KEY, and skips without them. Run manually
|
||||
before changing litellm/harness; the mocked coverage runs in tests/unit/harness and
|
||||
|
|
|
|||
69
.github/e2e-stack/assert_tests_ran.py
vendored
69
.github/e2e-stack/assert_tests_ran.py
vendored
|
|
@ -1,69 +0,0 @@
|
|||
import os
|
||||
import re
|
||||
import sys
|
||||
import xml.etree.ElementTree as ET
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e"))
|
||||
from coverage_registry.management_cases import MANAGEMENT_CASES
|
||||
|
||||
|
||||
def main() -> int:
|
||||
selected: Final = tuple(sys.argv[2:])
|
||||
try:
|
||||
report: Final = ET.parse(Path(sys.argv[1])).getroot()
|
||||
except (ET.ParseError, OSError):
|
||||
_ = sys.stdout.write("::error::could not read the test execution report\n")
|
||||
return 1
|
||||
cases: Final = tuple(report.iter("testcase"))
|
||||
expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT")
|
||||
passed: Final = frozenset(
|
||||
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
|
||||
)
|
||||
missing: Final = tuple(path for path in selected if path not in passed)
|
||||
required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected)
|
||||
passed_nodes: Final = frozenset(
|
||||
prop.get("value")
|
||||
for case in cases
|
||||
if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
|
||||
for prop in case.findall("./properties/property")
|
||||
if prop.get("name") == "management_node"
|
||||
)
|
||||
missing_nodes: Final = required_nodes - passed_nodes
|
||||
for node in sorted(missing_nodes):
|
||||
_ = sys.stdout.write(f"::error::required management case did not pass: {node}\n")
|
||||
for path in selected:
|
||||
collected: Final = sum(case.get("file") == path for case in cases)
|
||||
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
|
||||
_ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n")
|
||||
for case in cases:
|
||||
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error", "skipped")):
|
||||
continue
|
||||
outcome = "skipped" if case.find("skipped") is not None else "failed"
|
||||
_ = sys.stdout.write(f" {outcome}: {case.get('classname', '')}::{case.get('name', '')}\n")
|
||||
for prop in case.findall("./properties/property"):
|
||||
name = prop.get("name", "")
|
||||
value = prop.get("value", "")
|
||||
if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch(
|
||||
r"[A-Za-z0-9_.:<>-]{1,240}", value
|
||||
):
|
||||
_ = sys.stdout.write(f" {name}: {value}\n")
|
||||
if expected_count is not None and (
|
||||
len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases)
|
||||
):
|
||||
_ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n")
|
||||
return 1
|
||||
if (
|
||||
selected
|
||||
and not missing
|
||||
and not missing_nodes
|
||||
and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
|
||||
):
|
||||
return 0
|
||||
_ = sys.stdout.write("::error::every selected file must execute a passing test, with no failures or errors\n")
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
58
.github/scripts/assert_ci_coverage.py
vendored
58
.github/scripts/assert_ci_coverage.py
vendored
|
|
@ -32,6 +32,7 @@ IGNORE_ARG_RE: Final = re.compile(r"--ignore(?:-glob)?[= ](\S+)")
|
|||
DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*")
|
||||
COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE)
|
||||
GLOB_CHARS = frozenset("*?")
|
||||
DEDICATED_INTEGRATION_WORKFLOWS: Final = MappingProxyType({"image-scan.yml": "tests/integration/database/image/"})
|
||||
|
||||
# Trees whose jobs are sharded with no catch-all bucket, so every child that holds
|
||||
# tests has to be named by some shard or it runs nowhere. A child listed here is
|
||||
|
|
@ -560,6 +561,17 @@ def _integration_github_files(runner: pathlib.Path) -> frozenset[str]:
|
|||
return frozenset(ast.literal_eval(values))
|
||||
|
||||
|
||||
def _owns(workflow: str, relative_path: str) -> bool:
|
||||
owned_root: Final = DEDICATED_INTEGRATION_WORKFLOWS.get(workflow)
|
||||
return owned_root is not None and relative_path.startswith(owned_root)
|
||||
|
||||
|
||||
def _tokens_outside_owner(relative_path: str, workflow_tokens: Mapping[str, frozenset[str]]) -> frozenset[str]:
|
||||
return frozenset[str]().union(
|
||||
*(tokens for workflow, tokens in workflow_tokens.items() if not _owns(workflow, relative_path))
|
||||
)
|
||||
|
||||
|
||||
def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozenset[str], tuple[Finding, ...]]:
|
||||
runner: Final = repo_root / "tests/integration/run.py"
|
||||
if not runner.exists():
|
||||
|
|
@ -611,15 +623,20 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
|
|||
)
|
||||
!= 1
|
||||
)
|
||||
gha_tokens: Final = _invoked_test_tokens(
|
||||
scalar
|
||||
for path in (repo_root / ".github/workflows").glob("*.y*ml")
|
||||
for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
|
||||
workflow_tokens: Final = MappingProxyType(
|
||||
{
|
||||
path.name: _invoked_test_tokens(_scalars(yaml.safe_load(path.read_text()), path.name))
|
||||
for path in (repo_root / ".github/workflows").glob("*.y*ml")
|
||||
}
|
||||
)
|
||||
gha_tokens: Final = frozenset[str]().union(*workflow_tokens.values())
|
||||
circle_owned: Final = (
|
||||
frozenset(str(path.relative_to(repo_root)) for path in integration_root.rglob("test_*.py")) - github_files
|
||||
)
|
||||
findings: Final = (
|
||||
tuple(
|
||||
Finding(path, "integration contract is also selected by GitHub Actions")
|
||||
for path in paths
|
||||
Finding(path, "integration file outside GITHUB_FILES is selected by GitHub Actions")
|
||||
for path in sorted(circle_owned)
|
||||
if any(_token_covers(token, path) for token in gha_tokens)
|
||||
)
|
||||
+ tuple(
|
||||
|
|
@ -632,6 +649,11 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
|
|||
for path in sorted(github_files)
|
||||
if not (repo_root / path).is_file()
|
||||
)
|
||||
+ tuple(
|
||||
Finding(path, "GitHub-owned integration contract runs outside a dedicated integration workflow")
|
||||
for path in sorted(github_files)
|
||||
if any(_token_covers(token, path) for token in _tokens_outside_owner(path, workflow_tokens))
|
||||
)
|
||||
)
|
||||
browser_commands: Final = tuple(
|
||||
scalar.value
|
||||
|
|
@ -688,13 +710,9 @@ def main() -> int:
|
|||
scalars = _all_scalars()
|
||||
|
||||
integration_paths, ownership_findings = _integration_ownership()
|
||||
test_findings = (
|
||||
_uncovered_tests(
|
||||
allowlist,
|
||||
_invoked_selections(scalars)
|
||||
+ (Selection(included=integration_paths, ignored=frozenset()),),
|
||||
)
|
||||
+ ownership_findings
|
||||
test_findings = _uncovered_tests(
|
||||
allowlist,
|
||||
_invoked_selections(scalars) + (Selection(included=integration_paths, ignored=frozenset()),),
|
||||
)
|
||||
dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
|
||||
stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles())
|
||||
|
|
@ -711,13 +729,25 @@ def main() -> int:
|
|||
test_findings,
|
||||
"Add each to a job's test path, or list it in .github/ci-coverage-allowlist.yml with a reason.",
|
||||
)
|
||||
if ownership_findings:
|
||||
browser_remedy = (
|
||||
" CircleCI's browser group owns the integrationCritical Playwright contracts, so no GitHub workflow may "
|
||||
"select them and the default Playwright config must exclude them."
|
||||
if any("browser" in finding.detail or "Playwright" in finding.detail for finding in ownership_findings)
|
||||
else ""
|
||||
)
|
||||
_report(
|
||||
"integration contracts with the wrong CI owner",
|
||||
ownership_findings,
|
||||
"Follow tests/integration/AGENTS.md, which says which CI job owns each integration file." + browser_remedy,
|
||||
)
|
||||
if dockerfile_findings:
|
||||
_report(
|
||||
"Dockerfiles that no CI job builds",
|
||||
dockerfile_findings,
|
||||
"Build each in a workflow, or list it in .github/ci-coverage-allowlist.yml with a reason.",
|
||||
)
|
||||
if stale_findings or test_findings or dockerfile_findings:
|
||||
if stale_findings or test_findings or ownership_findings or dockerfile_findings:
|
||||
return 1
|
||||
|
||||
_write(
|
||||
|
|
|
|||
|
|
@ -24,12 +24,9 @@ jobs:
|
|||
- name: Update JSON Data
|
||||
run: |
|
||||
uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py"
|
||||
- name: Regenerate JSON Schema
|
||||
run: |
|
||||
uv run --frozen python ci_cd/generate_model_prices_schema.py
|
||||
- name: Create Pull Request
|
||||
run: |
|
||||
git add model_prices_and_context_window.json model_prices_and_context_window.schema.json
|
||||
git add model_prices_and_context_window.json
|
||||
git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')"
|
||||
gh pr create --title "Update model_prices_and_context_window.json file" \
|
||||
--body "Automated update for model_prices_and_context_window.json" \
|
||||
|
|
|
|||
|
|
@ -41,10 +41,6 @@ jobs:
|
|||
uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md"
|
||||
env:
|
||||
TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
|
||||
- name: Regenerate the JSON schema
|
||||
if: steps.existing.outputs.open_pr == ''
|
||||
run: |
|
||||
uv run --frozen python ci_cd/generate_model_prices_schema.py
|
||||
- name: Create a pull request when the registry changed
|
||||
if: steps.existing.outputs.open_pr == ''
|
||||
run: |
|
||||
|
|
@ -57,8 +53,7 @@ jobs:
|
|||
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
|
||||
git checkout -b "$branch"
|
||||
git add model_prices_and_context_window.json \
|
||||
litellm/model_prices_and_context_window_backup.json \
|
||||
model_prices_and_context_window.schema.json
|
||||
litellm/model_prices_and_context_window_backup.json
|
||||
git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')"
|
||||
gh auth setup-git
|
||||
git push origin "$branch"
|
||||
|
|
|
|||
6
.github/workflows/test-linting.yml
vendored
6
.github/workflows/test-linting.yml
vendored
|
|
@ -21,7 +21,7 @@ jobs:
|
|||
pull-requests: read
|
||||
actions: read
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
timeout-minutes: 20
|
||||
services:
|
||||
redis:
|
||||
image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
|
||||
|
|
@ -539,8 +539,8 @@ jobs:
|
|||
- name: test_unit_passed_gate
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_unit_passed_gate.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 tests/code_coverage_tests/test_e2e_idp_stack.py
|
||||
- name: test_e2e_idp_stack
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_idp_stack.py
|
||||
|
||||
- name: test_e2e_metadata
|
||||
env:
|
||||
|
|
|
|||
180
.github/workflows/test-mcp-oauth-e2e.yml
vendored
180
.github/workflows/test-mcp-oauth-e2e.yml
vendored
|
|
@ -1,180 +0,0 @@
|
|||
name: MCP OAuth happy path
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
paths:
|
||||
- '.github/workflows/test-mcp-oauth-e2e.yml'
|
||||
- '.github/e2e-stack/**'
|
||||
- 'tests/e2e/*.py'
|
||||
- 'tests/e2e/pytest.ini'
|
||||
- 'tests/e2e/idp_realm.json'
|
||||
- 'tests/e2e/mcp/**'
|
||||
- 'litellm/experimental_mcp_client/**'
|
||||
- 'litellm/proxy/_experimental/mcp_server/**'
|
||||
- 'litellm/proxy/auth/**'
|
||||
- 'litellm/proxy/management_endpoints/mcp_management_endpoints.py'
|
||||
- 'litellm/proxy/_types.py'
|
||||
- 'litellm/types/mcp_server/mcp_server_manager.py'
|
||||
- 'litellm/proxy/management_endpoints/*sso*.py'
|
||||
- 'litellm/proxy/management_endpoints/sso/**'
|
||||
- 'litellm/proxy/common_utils/encrypt_decrypt_utils.py'
|
||||
- 'litellm/proxy/proxy_server.py'
|
||||
- 'litellm/proxy/schema.prisma'
|
||||
- 'ui/litellm-dashboard/src/app/connect/**'
|
||||
- 'ui/litellm-dashboard/src/app/mcp/oauth/**'
|
||||
- 'pyproject.toml'
|
||||
- 'uv.lock'
|
||||
workflow_dispatch:
|
||||
|
||||
permissions: {}
|
||||
|
||||
concurrency:
|
||||
group: mcp-oauth-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
oauth:
|
||||
if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository
|
||||
runs-on: ubuntu-latest
|
||||
environment: e2e-changed
|
||||
timeout-minutes: 45
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16.6
|
||||
env:
|
||||
POSTGRES_USER: litellm
|
||||
POSTGRES_PASSWORD: dbpassword9090
|
||||
POSTGRES_DB: litellm
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U litellm"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 10
|
||||
env:
|
||||
DATABASE_HOST: 127.0.0.1
|
||||
DATABASE_PORT: '5432'
|
||||
DATABASE_USER: litellm
|
||||
DATABASE_PASSWORD: dbpassword9090
|
||||
DATABASE_NAME: litellm
|
||||
DATABASE_URL: postgresql://litellm:dbpassword9090@127.0.0.1:5432/litellm
|
||||
E2E_KEYCLOAK_URL: http://127.0.0.1:8081
|
||||
E2E_KEYCLOAK_ADMIN_USER: admin
|
||||
E2E_KEYCLOAK_ADMIN_PASSWORD: e2e-ephemeral-idp-not-a-secret
|
||||
E2E_FIXTURE_MODE: live
|
||||
E2E_PROVIDER_CACHE: '0'
|
||||
E2E_MCP_OAUTH_LIVE: '1'
|
||||
E2E_REQUIRED_TEST_COUNT: '4'
|
||||
steps:
|
||||
- name: Checkout the tested source
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || github.sha }}
|
||||
persist-credentials: false
|
||||
|
||||
- name: Require and materialize the upstream login
|
||||
env:
|
||||
STORAGE_STATE: ${{ secrets.E2E_LINEAR_STORAGE_STATE_B64 }}
|
||||
run: |
|
||||
umask 077
|
||||
python3 - <<'PY'
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import secrets
|
||||
from pathlib import Path
|
||||
encoded = os.environ.get("STORAGE_STATE", "")
|
||||
if not encoded:
|
||||
raise SystemExit("E2E_LINEAR_STORAGE_STATE_B64 is required; capture and provision a test-account login")
|
||||
state = json.loads(base64.b64decode(encoded, validate=True))
|
||||
if not isinstance(state, dict) or not state.get("cookies"):
|
||||
raise SystemExit("The captured login must contain browser cookies")
|
||||
directory = Path(os.environ["RUNNER_TEMP"]) / "mcp-oauth-private"
|
||||
directory.mkdir(mode=0o700)
|
||||
path = directory / "linear-state.json"
|
||||
path.write_text(json.dumps(state))
|
||||
with open(os.environ["GITHUB_ENV"], "a") as output:
|
||||
output.write(f"E2E_LINEAR_STORAGE_STATE={path}\n")
|
||||
for name in ("LITELLM_MASTER_KEY", "LITELLM_SALT_KEY"):
|
||||
value = "sk-e2e-" + secrets.token_hex(24)
|
||||
print(f"::add-mask::{value}")
|
||||
output.write(f"{name}={value}\n")
|
||||
PY
|
||||
|
||||
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: '3.13'
|
||||
- uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: '0.10.9'
|
||||
- uses: ./.github/actions/cache-cargo-build
|
||||
- name: Install the frozen E2E environment
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra proxy-runtime --extra extra_proxy --group ci --group proxy-dev --group e2e-dev
|
||||
uv run --no-sync python scripts/prisma_generate_if_needed.py
|
||||
uv run --no-sync playwright install --with-deps chromium
|
||||
|
||||
- name: Configure license access
|
||||
id: aws
|
||||
uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0
|
||||
with:
|
||||
role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
|
||||
aws-region: us-east-1
|
||||
role-session-name: mcp-oauth-${{ github.run_id }}
|
||||
role-duration-seconds: 900
|
||||
output-env-credentials: false
|
||||
output-credentials: true
|
||||
- name: Load the E2E license
|
||||
env:
|
||||
AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }}
|
||||
AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }}
|
||||
AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }}
|
||||
AWS_DEFAULT_REGION: us-east-1
|
||||
run: |
|
||||
license="$(aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license --query SecretString --output text)"
|
||||
test -n "${license}"
|
||||
echo "::add-mask::${license}"
|
||||
echo "LITELLM_LICENSE=${license}" >> "${GITHUB_ENV}"
|
||||
|
||||
- uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
|
||||
with:
|
||||
node-version-file: ui/litellm-dashboard/.nvmrc
|
||||
- name: Build the gateway consent UI at the tested commit
|
||||
run: |
|
||||
cd ui/litellm-dashboard
|
||||
../../scripts/with_dashboard_node.sh npm ci
|
||||
../../scripts/with_dashboard_node.sh npm run build
|
||||
mkdir -p ../../litellm/proxy/_experimental/out
|
||||
cp -r out/. ../../litellm/proxy/_experimental/out/
|
||||
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name index.html | while read -r page; do
|
||||
mkdir -p "${page%.html}"
|
||||
mv "${page}" "${page%.html}/index.html"
|
||||
done
|
||||
|
||||
- name: Prepare the isolated database and IdP
|
||||
run: |
|
||||
umask 077
|
||||
bash .github/e2e-stack/start-idp.sh
|
||||
uv run --no-sync python migrations/run.py > "${RUNNER_TEMP}/mcp-oauth-private/migrations.log" 2>&1
|
||||
|
||||
- name: Run every required OAuth variant without retries
|
||||
run: |
|
||||
umask 077
|
||||
uv run --no-sync pytest -c tests/e2e/pytest.ini tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
|
||||
--rootdir=. --reruns 0 --tb=short -o junit_family=xunit1 \
|
||||
--junitxml="${RUNNER_TEMP}/mcp-oauth-private/results.xml" \
|
||||
> "${RUNNER_TEMP}/mcp-oauth-private/pytest.log" 2>&1
|
||||
- name: Report JUnit results and reject skipped or missing cases
|
||||
if: always()
|
||||
run: |
|
||||
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \
|
||||
"${RUNNER_TEMP}/mcp-oauth-private/results.xml" tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
|
||||
- name: Remove private login and logs
|
||||
if: always()
|
||||
run: |
|
||||
docker rm -f e2e-keycloak >/dev/null 2>&1 || true
|
||||
rm -rf "${RUNNER_TEMP}/mcp-oauth-private"
|
||||
8
.github/workflows/test-rust.yml
vendored
8
.github/workflows/test-rust.yml
vendored
|
|
@ -16,6 +16,10 @@ on:
|
|||
- "litellm/proxy/guardrails/**"
|
||||
- "litellm/utils.py"
|
||||
- "litellm/ocr/**"
|
||||
- "litellm/chat_completions/dispatch.py"
|
||||
- "litellm/embeddings/dispatch.py"
|
||||
- "litellm/messages/dispatch.py"
|
||||
- "litellm/responses/dispatch.py"
|
||||
- "litellm/llms/base_llm/ocr/**"
|
||||
- "litellm/llms/custom_httpx/llm_http_handler.py"
|
||||
- "Makefile"
|
||||
|
|
@ -46,6 +50,10 @@ on:
|
|||
- "litellm/proxy/guardrails/**"
|
||||
- "litellm/utils.py"
|
||||
- "litellm/ocr/**"
|
||||
- "litellm/chat_completions/dispatch.py"
|
||||
- "litellm/embeddings/dispatch.py"
|
||||
- "litellm/messages/dispatch.py"
|
||||
- "litellm/responses/dispatch.py"
|
||||
- "litellm/llms/base_llm/ocr/**"
|
||||
- "litellm/llms/custom_httpx/llm_http_handler.py"
|
||||
- "Makefile"
|
||||
|
|
|
|||
14
.github/workflows/test-unit.yml
vendored
14
.github/workflows/test-unit.yml
vendored
|
|
@ -266,6 +266,7 @@ jobs:
|
|||
tests/unit/llms/featherless_ai
|
||||
tests/unit/llms/firecrawl
|
||||
tests/unit/llms/fireworks_ai
|
||||
tests/unit/llms/friendliai
|
||||
tests/unit/llms/gdc
|
||||
tests/unit/llms/gemini
|
||||
tests/unit/llms/gigachat
|
||||
|
|
@ -321,6 +322,7 @@ jobs:
|
|||
tests/unit/llms/pass_through
|
||||
tests/unit/llms/perplexity
|
||||
tests/unit/llms/pg_vector
|
||||
tests/unit/llms/pi
|
||||
tests/unit/llms/publicai
|
||||
tests/unit/llms/ragflow
|
||||
tests/unit/llms/recraft
|
||||
|
|
@ -410,6 +412,7 @@ jobs:
|
|||
tests/unit/embeddings
|
||||
tests/unit/endpoints
|
||||
tests/unit/files
|
||||
tests/unit/fine_tuning
|
||||
tests/unit/images
|
||||
tests/unit/interactions
|
||||
tests/unit/messages
|
||||
|
|
@ -456,11 +459,11 @@ jobs:
|
|||
tests/unit/telemetry tests/unit/proxy/telemetry
|
||||
tests/unit/tracing
|
||||
python-version: "3.12"
|
||||
workers: 2
|
||||
workers: 4
|
||||
reruns: 0
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 60
|
||||
dist: loadscope
|
||||
dist: load
|
||||
|
||||
- shard: proxy-auth
|
||||
test-path: |-
|
||||
|
|
@ -565,7 +568,6 @@ jobs:
|
|||
|
||||
- shard: mcp-elicitation
|
||||
test-path: |-
|
||||
--override-ini=pythonpath=tests
|
||||
tests/unit/experimental_mcp_client/test_mcp_client.py
|
||||
tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
|
||||
tests/unit/proxy/_experimental/mcp_server/test_interactions.py
|
||||
|
|
@ -573,9 +575,8 @@ jobs:
|
|||
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
|
||||
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
|
||||
tests/unit/proxy/_experimental/mcp_server/test_operations.py
|
||||
tests/integration/mcp/test_interactions.py
|
||||
python-version: "3.12"
|
||||
workers: 2
|
||||
workers: 4
|
||||
reruns: 0
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 60
|
||||
|
|
@ -660,6 +661,7 @@ jobs:
|
|||
- shard: proxy-extras
|
||||
test-path: >-
|
||||
tests/unit/litellm_proxy_extras
|
||||
tests/unit/db_scripts
|
||||
python-version: "3.12"
|
||||
workers: 2
|
||||
reruns: 0
|
||||
|
|
@ -698,7 +700,6 @@ jobs:
|
|||
tests/unit/responses/litellm_completion_transformation
|
||||
tests/unit/responses/test_additional_tools.py
|
||||
tests/unit/responses/test_custom_tool_call.py
|
||||
tests/unit/responses/test_dispatch.py
|
||||
tests/unit/responses/test_metadata_codex_callback.py
|
||||
tests/unit/responses/test_no_duplicate_spend_logs.py
|
||||
tests/unit/responses/test_null_test_fix.py
|
||||
|
|
@ -1453,7 +1454,6 @@ jobs:
|
|||
tests/unit/caching/test_redis_connection_pool.py \
|
||||
tests/unit/caching/test_redis_cluster_cache.py \
|
||||
tests/unit/caching/test_evicted_client_closer.py \
|
||||
tests/integration/sdk/test_redis_cluster_iam_auth.py \
|
||||
--tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 --durations=20 \
|
||||
--cov=./litellm --cov-report=xml:coverage.xml
|
||||
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ A test must only fail when litellm code changes. Never pin facts we don't own (a
|
|||
|
||||
`tests/unit/` mirrors `litellm/` in a parallel path (see `tests/unit/AGENTS.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
|
||||
|
||||
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md`
|
||||
End-to-end tests belong in `tests/e2e/`. A rule about one directory's files, like its harness or which CI runs it, goes in that directory's `AGENTS.md`, never only in a README or here. Read a directory's `AGENTS.md` before editing files in it
|
||||
|
||||
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
|
||||
|
||||
|
|
@ -72,7 +72,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests,
|
|||
|
||||
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
|
||||
|
||||
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
|
||||
Internal contributors branch a new PR off the repository's current default branch, with a branch name prefixed with litellm_
|
||||
|
||||
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
|
||||
|
||||
|
|
|
|||
|
|
@ -39,12 +39,12 @@ git clone https://github.com/YOUR_USERNAME/litellm.git
|
|||
cd litellm
|
||||
|
||||
# Create a new branch for your feature (see "Commit and Branch Conventions" below)
|
||||
git checkout -b feature/your-feature
|
||||
git checkout -b litellm_your_feature
|
||||
|
||||
# Install development dependencies
|
||||
make install-dev
|
||||
|
||||
# Install git hooks that enforce commit + branch conventions (one-time, opt-in)
|
||||
# Install git hooks that enforce commit conventions (one-time, opt-in)
|
||||
make install-hooks
|
||||
|
||||
# Verify your setup works
|
||||
|
|
@ -55,7 +55,7 @@ That's it! Your local development environment is ready.
|
|||
|
||||
## Commit and Branch Conventions
|
||||
|
||||
Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) and branches follow [Conventional Branches](https://conventional-branch.github.io/). Run `make install-hooks` once per clone to enable the local git hooks that enforce these — see the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for the full type list, examples, the protected-branch bypass list, and how to opt out.
|
||||
Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/). Internal contributors name branches with the `litellm_` prefix and no `/`, and fork branches can use any name. Run `make install-hooks` once per clone to enable the local hook that enforces the commit convention. See the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for details
|
||||
|
||||
### 2. Development Workflow
|
||||
|
||||
|
|
@ -78,8 +78,8 @@ uv run pytest tests/unit/<your_test_file>.py -v
|
|||
git add .
|
||||
git commit -m "feat(scope): your descriptive commit message"
|
||||
|
||||
# Push and create a PR (branch must follow Conventional Branches — see above)
|
||||
git push origin feature/your-feature
|
||||
# Push and create a PR
|
||||
git push origin litellm_your_feature
|
||||
```
|
||||
|
||||
## Adding Testing
|
||||
|
|
|
|||
4
Makefile
4
Makefile
|
|
@ -22,7 +22,7 @@ help:
|
|||
@echo " make install-proxy-dev-ci - Install proxy dev dependencies (CI-compatible)"
|
||||
@echo " make install-test-deps - Install the full local test environment"
|
||||
@echo " make install-helm-unittest - Install helm unittest plugin"
|
||||
@echo " make install-hooks - Install git hooks (Conventional Commits + Branches)"
|
||||
@echo " make install-hooks - Install git hooks (Conventional Commits)"
|
||||
@echo " make check - Run CI-equivalent lint on staged files, or on the diff vs the base branch when nothing is staged"
|
||||
@echo " make pre-commit - Legacy alias for make check"
|
||||
@echo " make format - Apply ruff format code formatting"
|
||||
|
|
@ -120,7 +120,7 @@ install-helm-unittest:
|
|||
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.8.2; \
|
||||
}
|
||||
|
||||
# Install git hooks that enforce Conventional Commits and Conventional Branches.
|
||||
# Install git hooks that enforce Conventional Commits.
|
||||
# Opt-in: not chained into install-dev.
|
||||
install-hooks:
|
||||
./scripts/install_git_hooks.sh
|
||||
|
|
|
|||
|
|
@ -287,7 +287,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
|
|||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode or Deep Agents on any model (Python SDK)</summary>
|
||||
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode, pi or Deep Agents on any model (Python SDK)</summary>
|
||||
|
||||
### Python SDK - Agents
|
||||
|
||||
|
|
@ -296,7 +296,7 @@ import litellm
|
|||
from litellm import Harness, sandbox
|
||||
|
||||
result = litellm.agent(
|
||||
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.DEEPAGENTS
|
||||
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.PI, Harness.DEEPAGENTS
|
||||
"Find why tests/test_router.py is flaky and fix it.",
|
||||
sandbox=sandbox.local("./repo"),
|
||||
model="litellm_proxy/claude-sonnet-4-5", # a model group on your AI Gateway
|
||||
|
|
@ -305,7 +305,7 @@ result = litellm.agent(
|
|||
print(result.text, result.cost, [f.path for f in result.files])
|
||||
```
|
||||
|
||||
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex` or `opencode`), or `deepagents langchain-litellm` for Deep Agents.
|
||||
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex`, `opencode` or `pi`), or `deepagents langchain-litellm` for Deep Agents.
|
||||
|
||||
[**Docs: Agent Harnesses**](https://docs.litellm.ai/docs/harness)
|
||||
|
||||
|
|
@ -384,6 +384,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
|
|||
| [ModelScope (`modelscope`)](https://docs.litellm.ai/docs/providers/modelscope) | ✅ | ✅ | ✅ | | ✅ | | | | | |
|
||||
| [Moonshot (`moonshot`)](https://docs.litellm.ai/docs/providers/moonshot) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Morph (`morph`)](https://docs.litellm.ai/docs/providers/morph) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Nace (`nace`)](https://console.nace.ai/docs/drex) | | | | | | | | | | |
|
||||
| [Nebius AI Studio (`nebius`)](https://docs.litellm.ai/docs/providers/nebius) | ✅ | ✅ | ✅ | ✅ | | | | | | |
|
||||
| [NLP Cloud (`nlp_cloud`)](https://docs.litellm.ai/docs/providers/nlp_cloud) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Novita AI (`novita`)](https://novita.ai/models/llm?utm_source=github_litellm&utm_medium=github_readme&utm_campaign=github_link) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
|
|
|
|||
|
|
@ -4,6 +4,33 @@ ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2
|
|||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
||||
# ---------- Rust extension ----------
|
||||
# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes
|
||||
FROM $LITELLM_BUILD_IMAGE AS rust-native
|
||||
|
||||
WORKDIR /app
|
||||
USER root
|
||||
|
||||
COPY --from=uvbin /uv /uvx /usr/local/bin/
|
||||
|
||||
RUN for i in 1 2 3; do \
|
||||
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \
|
||||
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
|
||||
sleep 5; \
|
||||
done
|
||||
|
||||
ENV UV_PYTHON_DOWNLOADS=0
|
||||
|
||||
COPY pyproject.toml README.md rust-toolchain.toml ./
|
||||
COPY .cargo ./.cargo
|
||||
COPY litellm-rust ./litellm-rust
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p litellm/rust_bridge /wheels && \
|
||||
touch litellm/__init__.py litellm/rust_bridge/__init__.py && \
|
||||
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
|
||||
python -c "import maturin; maturin.build_wheel('/wheels')"
|
||||
|
||||
# ---------- Builder ----------
|
||||
FROM $LITELLM_BUILD_IMAGE AS builder
|
||||
|
||||
|
|
@ -54,7 +81,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
COPY . .
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv sync --frozen --no-default-groups --no-editable \
|
||||
uv sync --frozen --no-install-project --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
--group admin-mcp \
|
||||
|
|
@ -63,6 +90,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra saml \
|
||||
--python python3.13
|
||||
|
||||
RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \
|
||||
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl
|
||||
|
||||
RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
|
|
|
|||
|
|
@ -165,6 +165,7 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
|
|||
BACKEND_MOUNT_PATHS: frozenset[str] = frozenset(
|
||||
{
|
||||
"/admin",
|
||||
"/metrics",
|
||||
"/swagger", # API documentation static assets belong to the backend
|
||||
"/mcp", # lazily-mounted MCP sub-app serves on the backend component
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
"""Guard the cost map on pull requests.
|
||||
|
||||
Every pull request whose diff against its merge base touches one of the three cost map files gets the file
|
||||
checks: the files parse, the backup copy matches the root file, and the JSON schema is in sync and validates the
|
||||
map. A pull request that leaves all three untouched skips them, since merging it keeps the base branch's copies
|
||||
and its head tree only carries whatever state the branch was cut from. Pull requests from the cost map sync bot
|
||||
(branches named litellm_cost_map_sync_*) always get the file checks and additionally may only touch those three
|
||||
files and may only add or update models.
|
||||
Every pull request whose diff against its merge base touches one of the two cost map files gets the file
|
||||
checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched
|
||||
skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the
|
||||
branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always
|
||||
get the file checks and additionally may only touch those two files or the legacy schema
|
||||
path and may only add or update models.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -19,13 +19,12 @@ from dataclasses import dataclass
|
|||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors
|
||||
|
||||
COST_MAP_PATH: Final = "model_prices_and_context_window.json"
|
||||
BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
|
||||
SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
|
||||
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH)
|
||||
LEGACY_SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
|
||||
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH)
|
||||
BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
|
||||
SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"})
|
||||
EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
CostMap = dict[str, object]
|
||||
|
|
@ -35,7 +34,6 @@ CostMap = dict[str, object]
|
|||
class Snapshot:
|
||||
cost_map: str
|
||||
backup: str
|
||||
schema: str
|
||||
|
||||
|
||||
def _parse_object(text: str, path: str) -> CostMap | str:
|
||||
|
|
@ -46,13 +44,6 @@ def _parse_object(text: str, path: str) -> CostMap | str:
|
|||
return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root"
|
||||
|
||||
|
||||
def _rendered_schema(cost_map: CostMap) -> str:
|
||||
try:
|
||||
return render(build_schema(cost_map))
|
||||
except SystemExit as error:
|
||||
return str(error)
|
||||
|
||||
|
||||
def _evaluation_output_price_failure(
|
||||
key: str, entry: object, evaluation_output_price_sources: Mapping[str, str]
|
||||
) -> str | None:
|
||||
|
|
@ -78,29 +69,13 @@ def _file_failures(
|
|||
*,
|
||||
evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES,
|
||||
) -> tuple[str, ...]:
|
||||
schema_text: Final = _rendered_schema(head_map)
|
||||
if not schema_text.startswith("{"):
|
||||
return (schema_text,)
|
||||
backup_failure: Final = (
|
||||
()
|
||||
if head.backup == head.cost_map
|
||||
else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",)
|
||||
)
|
||||
schema_failure: Final = (
|
||||
()
|
||||
if head.schema == schema_text
|
||||
else (
|
||||
f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; "
|
||||
"run `python ci_cd/generate_model_prices_schema.py` and commit the result",
|
||||
)
|
||||
)
|
||||
return (
|
||||
*backup_failure,
|
||||
*schema_failure,
|
||||
*(
|
||||
f"{COST_MAP_PATH} does not validate against its schema: {error}"
|
||||
for error in validation_errors(head_map, json.loads(schema_text))[:20]
|
||||
),
|
||||
*(
|
||||
failure
|
||||
for key, entry in head_map.items()
|
||||
|
|
@ -135,7 +110,7 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str
|
|||
*(
|
||||
f"bot PRs may only change the cost map files, not {path}"
|
||||
for path in changed_files
|
||||
if path not in GUARDED_PATHS
|
||||
if path not in GUARDED_PATHS and path != LEGACY_SCHEMA_PATH
|
||||
),
|
||||
*(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map),
|
||||
*(f"bot PRs may not remove fields: {ref}" for ref in removed_fields),
|
||||
|
|
|
|||
|
|
@ -1,409 +0,0 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Final, Optional
|
||||
|
||||
import jsonschema
|
||||
|
||||
REPO_ROOT = Path(__file__).parent.parent
|
||||
PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json"
|
||||
SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json"
|
||||
|
||||
SPECIAL_ROOT_KEYS = frozenset({"sample_spec", "fallback_generalizations"})
|
||||
|
||||
JsonSchema = dict
|
||||
|
||||
NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
|
||||
NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
|
||||
BOOLEAN: JsonSchema = {"type": "boolean"}
|
||||
STRING: JsonSchema = {"type": "string"}
|
||||
TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"}
|
||||
WEEKDAY_PATTERN: Final = (
|
||||
r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$"
|
||||
)
|
||||
|
||||
EXTRA_BOOLEAN_KEYS = frozenset(
|
||||
{
|
||||
"gemini_native_audio",
|
||||
"gemini_audio_only_live",
|
||||
"uses_embed_content",
|
||||
"use_openai_responses_path",
|
||||
"bedrock_converse_supports_strict_tools",
|
||||
"thinking_always_on",
|
||||
}
|
||||
)
|
||||
|
||||
HOURS_UTC: Final[JsonSchema] = {
|
||||
"description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.',
|
||||
"oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}],
|
||||
}
|
||||
|
||||
OFF_PEAK_WINDOW: Final[JsonSchema] = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"hours_utc": HOURS_UTC,
|
||||
"weekdays": {
|
||||
"type": "array",
|
||||
"description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.",
|
||||
"items": {
|
||||
"oneOf": [
|
||||
{"type": "integer", "minimum": 1, "maximum": 7},
|
||||
{"type": "string", "pattern": WEEKDAY_PATTERN},
|
||||
]
|
||||
},
|
||||
"minItems": 1,
|
||||
},
|
||||
},
|
||||
"required": ["hours_utc"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
OBJECT_KEYS: dict[str, JsonSchema] = {
|
||||
"off_peak_pricing": {
|
||||
"type": "object",
|
||||
"description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.",
|
||||
"properties": {
|
||||
"hours_utc": HOURS_UTC,
|
||||
"windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1},
|
||||
"weekday_timezone": {
|
||||
"type": "string",
|
||||
"description": "IANA zone the weekdays of each window are read on; defaults to UTC.",
|
||||
},
|
||||
"input_cost_per_token": NONNEG_NUMBER,
|
||||
"output_cost_per_token": NONNEG_NUMBER,
|
||||
"output_cost_per_reasoning_token": NONNEG_NUMBER,
|
||||
"cache_read_input_token_cost": NONNEG_NUMBER,
|
||||
"cache_creation_input_token_cost": NONNEG_NUMBER,
|
||||
},
|
||||
"anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"search_context_cost_per_query": {
|
||||
"type": "object",
|
||||
"description": "USD cost per web search query, keyed by search context size.",
|
||||
"properties": {
|
||||
"search_context_size_low": NONNEG_NUMBER,
|
||||
"search_context_size_medium": NONNEG_NUMBER,
|
||||
"search_context_size_high": NONNEG_NUMBER,
|
||||
},
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"guardrail_cost_per_unit": {
|
||||
"type": "object",
|
||||
"description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).",
|
||||
"additionalProperties": NONNEG_NUMBER,
|
||||
},
|
||||
"metadata": {
|
||||
"type": "object",
|
||||
"description": "Free-form notes about the entry (e.g. pricing derivation).",
|
||||
},
|
||||
"provider_specific_entry": {
|
||||
"type": "object",
|
||||
"description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).",
|
||||
},
|
||||
}
|
||||
|
||||
ARRAY_KEYS: dict[str, JsonSchema] = {
|
||||
"supported_audio_formats": {
|
||||
"type": "array",
|
||||
"description": "Audio container formats the model can return.",
|
||||
"items": {"type": "string", "enum": ["mp3", "wav"]},
|
||||
},
|
||||
"supported_endpoints": {
|
||||
"type": "array",
|
||||
"description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.",
|
||||
"items": STRING,
|
||||
},
|
||||
"supported_modalities": {
|
||||
"type": "array",
|
||||
"description": "Input modalities the model accepts.",
|
||||
"items": {"type": "string", "enum": ["text", "image", "audio", "video"]},
|
||||
},
|
||||
"supported_output_modalities": {
|
||||
"type": "array",
|
||||
"description": "Output modalities the model can produce.",
|
||||
"items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]},
|
||||
},
|
||||
"reasoning_effort_levels": {
|
||||
"type": "array",
|
||||
"description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.",
|
||||
"items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]},
|
||||
},
|
||||
"supported_regions": {
|
||||
"type": "array",
|
||||
"description": "Cloud regions the model is available in ('global' or region ids).",
|
||||
"items": STRING,
|
||||
},
|
||||
"tiered_pricing": {
|
||||
"type": "array",
|
||||
"description": "Context-length or result-count tiered rates; each tier's costs apply within its range.",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"range": {
|
||||
"type": "array",
|
||||
"description": "[min, max] prompt-token span this tier applies to.",
|
||||
"items": NONNEG_NUMBER,
|
||||
"minItems": 2,
|
||||
"maxItems": 2,
|
||||
},
|
||||
"max_results_range": {
|
||||
"type": "array",
|
||||
"description": "[min, max] result-count span this tier applies to (search models).",
|
||||
"items": NONNEG_NUMBER,
|
||||
"minItems": 2,
|
||||
"maxItems": 2,
|
||||
},
|
||||
"input_cost_per_token": NONNEG_NUMBER,
|
||||
"output_cost_per_token": NONNEG_NUMBER,
|
||||
"output_cost_per_reasoning_token": NONNEG_NUMBER,
|
||||
"cache_read_input_token_cost": NONNEG_NUMBER,
|
||||
"cache_creation_input_token_cost": NONNEG_NUMBER,
|
||||
"input_cost_per_query": NONNEG_NUMBER,
|
||||
},
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
INTEGER_KEYS: dict[str, JsonSchema] = {
|
||||
"max_tokens": {
|
||||
**NONNEG_INTEGER,
|
||||
"description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.",
|
||||
},
|
||||
"max_input_tokens": {
|
||||
**NONNEG_INTEGER,
|
||||
"description": "Maximum prompt/context tokens the model accepts.",
|
||||
},
|
||||
"max_output_tokens": {
|
||||
**NONNEG_INTEGER,
|
||||
"description": "Maximum tokens the model can generate in one response.",
|
||||
},
|
||||
"output_vector_size": {
|
||||
**NONNEG_INTEGER,
|
||||
"description": "Embedding dimension for embedding models.",
|
||||
},
|
||||
"prompt_cache_min_tokens": {
|
||||
**NONNEG_INTEGER,
|
||||
"description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
|
||||
},
|
||||
"tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
|
||||
"rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
|
||||
}
|
||||
|
||||
NUMBER_KEYS: dict[str, JsonSchema] = {
|
||||
"regional_processing_uplift_multiplier_eu": {
|
||||
"type": "number",
|
||||
"minimum": 1,
|
||||
"description": "Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%).",
|
||||
},
|
||||
"regional_processing_uplift_multiplier_us": {
|
||||
"type": "number",
|
||||
"minimum": 1,
|
||||
"description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).",
|
||||
},
|
||||
"regional_endpoint_uplift_multiplier": {
|
||||
"type": "number",
|
||||
"minimum": 1,
|
||||
"description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).",
|
||||
},
|
||||
}
|
||||
|
||||
COST_DESCRIPTIONS: dict[str, str] = {
|
||||
"input_cost_per_token": "USD per prompt token.",
|
||||
"output_cost_per_token": "USD per generated token.",
|
||||
"output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.",
|
||||
"google_maps_grounding_cost_per_query": (
|
||||
"USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit."
|
||||
),
|
||||
"cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.",
|
||||
"cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.",
|
||||
"input_cost_per_token_batches": "USD per prompt token via the provider's batch API.",
|
||||
"output_cost_per_token_batches": "USD per generated token via the provider's batch API.",
|
||||
}
|
||||
|
||||
|
||||
def cost_description(key: str) -> Optional[str]:
|
||||
if key in COST_DESCRIPTIONS:
|
||||
return COST_DESCRIPTIONS[key]
|
||||
if key.endswith("_flex"):
|
||||
return "Flex service-tier rate for the same-named base field."
|
||||
if key.endswith("_priority"):
|
||||
return "Priority service-tier rate for the same-named base field."
|
||||
if "_above_" in key:
|
||||
return "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
return None
|
||||
|
||||
|
||||
def cost_schema(key: str) -> JsonSchema:
|
||||
description = cost_description(key)
|
||||
return {**NONNEG_NUMBER, "description": description} if description else dict(NONNEG_NUMBER)
|
||||
|
||||
|
||||
def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
|
||||
return {
|
||||
"litellm_provider": {
|
||||
"type": "string",
|
||||
"description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.",
|
||||
},
|
||||
"mode": {
|
||||
"type": "string",
|
||||
"description": "Primary API surface / task type of the model.",
|
||||
"enum": list(modes),
|
||||
},
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "URL of the provider pricing/model page this entry was taken from.",
|
||||
},
|
||||
"deprecation_date": {
|
||||
"type": "string",
|
||||
"description": "Date the provider deprecates the model, YYYY-MM-DD.",
|
||||
"format": "date",
|
||||
"pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$",
|
||||
},
|
||||
"web_search_billing_unit": {
|
||||
"type": "string",
|
||||
"description": "Whether web search is billed per query or per prompt.",
|
||||
"enum": ["per_query", "per_prompt"],
|
||||
},
|
||||
"bedrock_output_config_effort_ceiling": {
|
||||
"type": "string",
|
||||
"description": "Highest reasoning effort the Bedrock output_config accepts for this model.",
|
||||
"enum": ["low", "medium", "high", "max", "xhigh"],
|
||||
},
|
||||
"default_reasoning_effort": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"Reasoning effort the provider applies when the request omits reasoning_effort. "
|
||||
"Gates whether a non-default temperature or the top_p/logprobs sampling params are "
|
||||
"accepted, which hold only when the effort resolves to 'none'."
|
||||
),
|
||||
"enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
|
||||
},
|
||||
"comment": STRING,
|
||||
"audio_transcription_config": STRING,
|
||||
"vertex_ai_audio_api": {
|
||||
"type": "string",
|
||||
"enum": ["lyria_predict", "lyria_interactions"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def classify(key: str, modes: tuple) -> Optional[JsonSchema]:
|
||||
curated = {**OBJECT_KEYS, **ARRAY_KEYS, **string_key_schemas(modes), **INTEGER_KEYS, **NUMBER_KEYS}
|
||||
if key in curated:
|
||||
return curated[key]
|
||||
if key.startswith("supports_") or key in EXTRA_BOOLEAN_KEYS:
|
||||
return BOOLEAN
|
||||
if "cost" in key:
|
||||
return cost_schema(key)
|
||||
return None
|
||||
|
||||
|
||||
def build_schema(prices: dict) -> JsonSchema:
|
||||
entries = {name: entry for name, entry in prices.items() if name not in SPECIAL_ROOT_KEYS}
|
||||
all_keys = tuple(sorted({key for entry in entries.values() for key in entry}))
|
||||
modes = tuple(sorted({entry["mode"] for entry in entries.values() if "mode" in entry}))
|
||||
unclassified = tuple(key for key in all_keys if classify(key, modes) is None)
|
||||
if unclassified:
|
||||
raise SystemExit(
|
||||
f"Unclassified keys in {PRICES_PATH.name}: {', '.join(unclassified)}. "
|
||||
f"Add them to the key tables in {Path(__file__).name} and rerun it."
|
||||
)
|
||||
entry_properties = {key: classify(key, modes) for key in all_keys}
|
||||
return {
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"title": "LiteLLM model_prices_and_context_window.json",
|
||||
"description": (
|
||||
"Schema for LiteLLM's model price and context window registry "
|
||||
"(https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). "
|
||||
"Every top-level key except 'sample_spec' and 'fallback_generalizations' is a model id, "
|
||||
"optionally prefixed with its provider (e.g. 'azure/gpt-5.4'), mapping to a model entry. "
|
||||
"All costs are USD per unit. New optional fields are added regularly, so consumers should "
|
||||
"ignore unknown fields rather than reject them."
|
||||
),
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"sample_spec": {
|
||||
"type": "object",
|
||||
"description": (
|
||||
"Documentation placeholder illustrating the entry shape; not a real model and not "
|
||||
"schema-conformant (several values are prose)."
|
||||
),
|
||||
},
|
||||
"fallback_generalizations": {
|
||||
"type": "object",
|
||||
"description": "Regex rules that generalize unknown model ids to known families; not a model entry.",
|
||||
"properties": {
|
||||
"rules": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": STRING,
|
||||
"pattern": STRING,
|
||||
"description": STRING,
|
||||
},
|
||||
"required": ["name", "pattern"],
|
||||
"additionalProperties": True,
|
||||
},
|
||||
}
|
||||
},
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
"additionalProperties": {"$ref": "#/$defs/modelEntry"},
|
||||
"$defs": {
|
||||
"modelEntry": {
|
||||
"type": "object",
|
||||
"description": (
|
||||
"Pricing, limits, and capability flags for one model. Fields other than litellm_provider "
|
||||
"are optional; boolean capability flags are simply omitted when unknown or false."
|
||||
),
|
||||
"required": ["litellm_provider"],
|
||||
"properties": entry_properties,
|
||||
"additionalProperties": True,
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def render(schema: JsonSchema) -> str:
|
||||
return json.dumps(schema, indent=2) + "\n"
|
||||
|
||||
|
||||
def validation_errors(prices: dict, schema: JsonSchema) -> tuple:
|
||||
validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER)
|
||||
return tuple(
|
||||
f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
|
||||
for error in validator.iter_errors(prices)
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
check = "--check" in sys.argv[1:]
|
||||
prices = json.loads(PRICES_PATH.read_text())
|
||||
rendered = render(build_schema(prices))
|
||||
errors = validation_errors(prices, json.loads(rendered))
|
||||
if errors:
|
||||
print(f"{PRICES_PATH.name} does not validate against the generated schema:")
|
||||
print("\n".join(errors[:20]))
|
||||
return 1
|
||||
if not check:
|
||||
SCHEMA_PATH.write_text(rendered)
|
||||
print(f"wrote {SCHEMA_PATH}")
|
||||
return 0
|
||||
if not SCHEMA_PATH.exists() or SCHEMA_PATH.read_text() != rendered:
|
||||
print(
|
||||
f"{SCHEMA_PATH.name} is out of sync with {PRICES_PATH.name}. "
|
||||
f"Run `python {Path(__file__).relative_to(REPO_ROOT)}` and commit the result."
|
||||
)
|
||||
return 1
|
||||
print(f"{SCHEMA_PATH.name} is in sync and {PRICES_PATH.name} validates against it")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
|
@ -6563,6 +6563,87 @@
|
|||
],
|
||||
"title": "litellm_project_model_rate_limit_used_metric",
|
||||
"type": "timeseries"
|
||||
},
|
||||
{
|
||||
"collapsed": false,
|
||||
"gridPos": {
|
||||
"h": 1,
|
||||
"w": 24,
|
||||
"x": 0,
|
||||
"y": 447
|
||||
},
|
||||
"id": 115,
|
||||
"panels": [],
|
||||
"title": "Agent tracing export (needs LITELLM_LENS_URL)",
|
||||
"type": "row"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"description": "Request records the Lens exporter delivered to the Lens service against the records it dropped before the service accepted them",
|
||||
"fieldConfig": {
|
||||
"defaults": {
|
||||
"color": {
|
||||
"mode": "palette-classic"
|
||||
},
|
||||
"custom": {
|
||||
"drawStyle": "line",
|
||||
"fillOpacity": 10,
|
||||
"lineWidth": 1,
|
||||
"showPoints": "never",
|
||||
"spanNulls": false
|
||||
},
|
||||
"unit": "short"
|
||||
},
|
||||
"overrides": []
|
||||
},
|
||||
"gridPos": {
|
||||
"h": 8,
|
||||
"w": 24,
|
||||
"x": 0,
|
||||
"y": 448
|
||||
},
|
||||
"id": 116,
|
||||
"options": {
|
||||
"legend": {
|
||||
"calcs": [],
|
||||
"displayMode": "list",
|
||||
"placement": "bottom",
|
||||
"showLegend": true
|
||||
},
|
||||
"tooltip": {
|
||||
"mode": "multi",
|
||||
"sort": "desc"
|
||||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"editorMode": "code",
|
||||
"expr": "sum(rate(litellm_lens_export_rows_written_total[$__rate_interval]))",
|
||||
"legendFormat": "written",
|
||||
"range": true,
|
||||
"refId": "A"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"editorMode": "code",
|
||||
"expr": "sum(rate(litellm_lens_export_rows_dropped_total[$__rate_interval]))",
|
||||
"legendFormat": "dropped",
|
||||
"range": true,
|
||||
"refId": "B"
|
||||
}
|
||||
],
|
||||
"title": "litellm_lens_export_rows_written and litellm_lens_export_rows_dropped rate",
|
||||
"type": "timeseries"
|
||||
}
|
||||
],
|
||||
"preload": false,
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
# LiteLLM All Prometheus Metrics dashboard
|
||||
|
||||
Every `litellm_*` metric family the proxy can expose on `/metrics` (141 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
|
||||
Every `litellm_*` metric family the proxy can expose on `/metrics` (143 families across 98 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes), project model rate limits, and the agent tracing export to Lens. Panel titles are the metric names so you can grep the JSON for the metric you care about
|
||||
|
||||
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/unit/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard
|
||||
|
||||
The first eleven rows need only `callbacks: ["prometheus"]`. The last three rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected
|
||||
The first eleven rows and the project model rate limits row need only `callbacks: ["prometheus"]`. The other rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, the agent tracing export row needs `LITELLM_LENS_URL`, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected
|
||||
|
||||
## Pre-requisites
|
||||
|
||||
|
|
|
|||
201
db_scripts/strip_vector_store_internal_metadata.py
Normal file
201
db_scripts/strip_vector_store_internal_metadata.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Iterator, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate, repeat, takewhile
|
||||
from typing import Final, Literal, TextIO
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError
|
||||
|
||||
from litellm.vector_stores.utils import is_litellm_written_metadata_key
|
||||
|
||||
DEFAULT_API_BASE: Final = "https://api.openai.com/v1"
|
||||
LIST_PAGE_SIZE: Final = "100"
|
||||
|
||||
|
||||
class VectorStore(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
id: str
|
||||
metadata: dict[str, JsonValue] | None = None
|
||||
|
||||
|
||||
class VectorStorePage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
data: tuple[VectorStore, ...]
|
||||
has_more: bool = False
|
||||
|
||||
|
||||
class OpenAIErrorBody(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
message: str = ""
|
||||
|
||||
|
||||
class OpenAIErrorResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
error: OpenAIErrorBody
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Options:
|
||||
api_key: str
|
||||
api_base: str
|
||||
vector_store_ids: tuple[str, ...]
|
||||
extra_keys: frozenset[str]
|
||||
apply: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StoreResult:
|
||||
store_id: str
|
||||
matched_keys: tuple[str, ...]
|
||||
status: Literal["dry_run", "updated", "failed"]
|
||||
error: str = ""
|
||||
|
||||
|
||||
def is_proxy_written_key(key: str, extra_keys: frozenset[str]) -> bool:
|
||||
return is_litellm_written_metadata_key(key) or key in extra_keys
|
||||
|
||||
|
||||
def matching_keys(store: VectorStore, extra_keys: frozenset[str]) -> tuple[str, ...]:
|
||||
return tuple(sorted(key for key in (store.metadata or {}) if is_proxy_written_key(key, extra_keys)))
|
||||
|
||||
|
||||
def fetch_page(client: httpx.Client, after: str | None) -> VectorStorePage:
|
||||
params: Final = {"limit": LIST_PAGE_SIZE} if after is None else {"limit": LIST_PAGE_SIZE, "after": after}
|
||||
return VectorStorePage.model_validate_json(client.get("/vector_stores", params=params).raise_for_status().content)
|
||||
|
||||
|
||||
def fetch_next_page(client: httpx.Client, page: VectorStorePage | None) -> VectorStorePage | None:
|
||||
if page is None or not page.has_more or not page.data:
|
||||
return None
|
||||
return fetch_page(client, after=page.data[-1].id)
|
||||
|
||||
|
||||
def list_all_stores(client: httpx.Client) -> Iterator[VectorStore]:
|
||||
pages: Final = accumulate(
|
||||
repeat(None), lambda page, _: fetch_next_page(client, page), initial=fetch_page(client, after=None)
|
||||
)
|
||||
for page in takewhile(lambda page: page is not None, pages):
|
||||
yield from page.data
|
||||
|
||||
|
||||
def retrieve_stores(client: httpx.Client, store_ids: tuple[str, ...]) -> Iterator[VectorStore]:
|
||||
for store_id in store_ids:
|
||||
yield VectorStore.model_validate_json(client.get(f"/vector_stores/{store_id}").raise_for_status().content)
|
||||
|
||||
|
||||
def describe_http_error(response: httpx.Response) -> str:
|
||||
try:
|
||||
message: Final = OpenAIErrorResponse.model_validate_json(response.content).error.message
|
||||
except ValidationError:
|
||||
return f"HTTP {response.status_code}"
|
||||
return f"HTTP {response.status_code}: {message}"
|
||||
|
||||
|
||||
def delete_keys(client: httpx.Client, store_id: str, keys: tuple[str, ...]) -> StoreResult:
|
||||
try:
|
||||
response: Final = client.post(f"/vector_stores/{store_id}", json={"metadata": dict.fromkeys(keys)})
|
||||
except httpx.TransportError as error:
|
||||
return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=type(error).__name__)
|
||||
if response.is_success:
|
||||
return StoreResult(store_id=store_id, matched_keys=keys, status="updated")
|
||||
return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=describe_http_error(response))
|
||||
|
||||
|
||||
def clean_store(client: httpx.Client, store: VectorStore, options: Options) -> StoreResult | None:
|
||||
keys: Final = matching_keys(store, options.extra_keys)
|
||||
if not keys:
|
||||
return None
|
||||
if not options.apply:
|
||||
return StoreResult(store_id=store.id, matched_keys=keys, status="dry_run")
|
||||
return delete_keys(client, store.id, keys)
|
||||
|
||||
|
||||
def format_result(result: StoreResult) -> str:
|
||||
key_list: Final = ", ".join(result.matched_keys)
|
||||
count: Final = len(result.matched_keys)
|
||||
match result.status:
|
||||
case "dry_run":
|
||||
return f"{result.store_id}: would delete {count} keys: {key_list}"
|
||||
case "updated":
|
||||
return f"{result.store_id}: deleted {count} keys: {key_list}"
|
||||
case "failed":
|
||||
return f"{result.store_id}: failed to delete {count} keys ({result.error}): {key_list}"
|
||||
|
||||
|
||||
def clean_and_report(
|
||||
client: httpx.Client, stores: Iterator[VectorStore], options: Options, out: TextIO
|
||||
) -> Iterator[StoreResult | None]:
|
||||
for store in stores:
|
||||
result: Final = clean_store(client, store, options)
|
||||
if result is not None:
|
||||
print(format_result(result), file=out)
|
||||
yield result
|
||||
|
||||
|
||||
def run(options: Options, client: httpx.Client, out: TextIO) -> int:
|
||||
stores: Final = (
|
||||
retrieve_stores(client, options.vector_store_ids) if options.vector_store_ids else list_all_stores(client)
|
||||
)
|
||||
scanned: Final = tuple(clean_and_report(client, stores, options, out))
|
||||
results: Final = tuple(result for result in scanned if result is not None)
|
||||
updated: Final = sum(1 for result in results if result.status == "updated")
|
||||
failed: Final = sum(1 for result in results if result.status == "failed")
|
||||
mode: Final = "applied" if options.apply else "dry run, pass --apply to delete"
|
||||
print(
|
||||
f"scanned={len(scanned)} matched={len(results)} updated={updated} failed={failed} ({mode})",
|
||||
file=out,
|
||||
)
|
||||
return 1 if failed else 0
|
||||
|
||||
|
||||
def parse_options(argv: Sequence[str] | None) -> Options:
|
||||
parser: Final = argparse.ArgumentParser(
|
||||
description="Remove LiteLLM proxy-written keys from OpenAI vector store metadata. Dry run by default."
|
||||
)
|
||||
parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY"), help="defaults to $OPENAI_API_KEY")
|
||||
parser.add_argument("--api-base", default=DEFAULT_API_BASE)
|
||||
parser.add_argument("--vector-store-id", action="append", default=[], help="repeatable, defaults to every store")
|
||||
parser.add_argument("--extra-key", action="append", default=[], help="repeatable, another key name to delete")
|
||||
parser.add_argument("--apply", action="store_true", help="delete the matching keys instead of listing them")
|
||||
args: Final = parser.parse_args(argv)
|
||||
if not args.api_key:
|
||||
parser.error("--api-key or OPENAI_API_KEY is required")
|
||||
return Options(
|
||||
api_key=args.api_key,
|
||||
api_base=args.api_base.rstrip("/"),
|
||||
vector_store_ids=tuple(args.vector_store_id),
|
||||
extra_keys=frozenset(args.extra_key),
|
||||
apply=args.apply,
|
||||
)
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None, transport: httpx.BaseTransport | None = None) -> int:
|
||||
options: Final = parse_options(argv)
|
||||
with httpx.Client(
|
||||
base_url=options.api_base,
|
||||
headers={"Authorization": f"Bearer {options.api_key}"},
|
||||
timeout=60.0,
|
||||
transport=transport,
|
||||
) as client:
|
||||
try:
|
||||
return run(options, client, sys.stdout)
|
||||
except httpx.HTTPStatusError as error:
|
||||
sys.stderr.write(f"could not read vector stores: {describe_http_error(error.response)}\n")
|
||||
return 2
|
||||
except httpx.TransportError as error:
|
||||
sys.stderr.write(f"could not reach {options.api_base}: {type(error).__name__}\n")
|
||||
return 2
|
||||
except ValidationError:
|
||||
sys.stderr.write(f"unexpected vector store response shape from {options.api_base}\n")
|
||||
return 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
|
@ -20,9 +20,15 @@ class EnterpriseCustomGuardrailHelper:
|
|||
the tag's value (the mode for that tag).
|
||||
- If no request tag matches, fall back to default mode(s).
|
||||
"""
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
StandardLoggingPayloadSetup,
|
||||
)
|
||||
from litellm.router_strategy.tag_based_routing import (
|
||||
without_baggage_mapped_tags,
|
||||
)
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
||||
|
|
@ -36,9 +42,13 @@ class EnterpriseCustomGuardrailHelper:
|
|||
|
||||
proxy_server_request = data.get("proxy_server_request", {})
|
||||
|
||||
request_tags = StandardLoggingPayloadSetup.get_request_tags(
|
||||
litellm_params=data,
|
||||
proxy_server_request=proxy_server_request,
|
||||
request_tags = without_baggage_mapped_tags(
|
||||
StandardLoggingPayloadSetup.get_request_tags(
|
||||
litellm_params=data,
|
||||
proxy_server_request=proxy_server_request,
|
||||
),
|
||||
data.get(get_metadata_variable_name_from_kwargs(data)),
|
||||
keep_serving_deployment_tags=True,
|
||||
)
|
||||
|
||||
# Check if any request tag matches a Mode tag key
|
||||
|
|
|
|||
|
|
@ -1,14 +1,16 @@
|
|||
from typing import Any, Optional
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
|
||||
async def enterprise_custom_auth(
|
||||
request: Request, api_key: str, user_custom_auth: Optional[Any]
|
||||
) -> Optional[UserAPIKeyAuth]:
|
||||
request: Request,
|
||||
api_key: str,
|
||||
user_custom_auth: Callable[[Request, str], Awaitable[object]] | None,
|
||||
) -> object:
|
||||
from litellm_enterprise.proxy.proxy_server import custom_auth_settings
|
||||
|
||||
if user_custom_auth is None:
|
||||
|
|
|
|||
|
|
@ -2,7 +2,9 @@
|
|||
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import replace as dataclasses_replace
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
|
|
@ -13,15 +15,22 @@ from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
|
|||
ManagedObjectRow,
|
||||
)
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
BATCH_LINE_ITEMS_MAX_PART_BYTES,
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
|
||||
MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
)
|
||||
from litellm.repositories.table_repositories import ManagedObjectRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.verification_token_repository import VerificationTokenRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.batches.batch_line_items import BatchFileProvider
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
|
|
@ -67,6 +76,19 @@ def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[Manage
|
|||
return ManagedObjectRepository(prisma_client).table
|
||||
|
||||
|
||||
def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
|
||||
return VerificationTokenRepository(prisma_client).table
|
||||
|
||||
|
||||
def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_TeamTable]":
|
||||
return TeamRepository(prisma_client).table
|
||||
|
||||
|
||||
def _row_metadata(row: object) -> dict[str, object]:
|
||||
metadata: Final = getattr(row, "metadata", None)
|
||||
return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {} # cast-ok: prisma Json object column
|
||||
|
||||
|
||||
class CheckBatchCost:
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -87,6 +109,10 @@ class CheckBatchCost:
|
|||
# the guaranteed-failing primary query on every subsequent cycle.
|
||||
self._has_batch_processed_column: bool = True
|
||||
self.batch_processed_support_confirmed: bool = False
|
||||
self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
|
||||
self._line_item_store_turn: Final = asyncio.Lock()
|
||||
self._poll_idle: Final = asyncio.Event()
|
||||
self._poll_idle.set()
|
||||
|
||||
@staticmethod
|
||||
def _is_missing_batch_processed_column_error(err: Exception) -> bool:
|
||||
|
|
@ -202,7 +228,7 @@ class CheckBatchCost:
|
|||
"so it will no longer be polled"
|
||||
)
|
||||
|
||||
async def _claim_job_for_costing(self, job: "ManagedObjectRow") -> bool:
|
||||
async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool:
|
||||
"""
|
||||
Atomically flip batch_processed from false to true, returning whether this pod won
|
||||
the row. Every pod and uvicorn worker schedules its own poller against the shared
|
||||
|
|
@ -216,10 +242,16 @@ class CheckBatchCost:
|
|||
"""
|
||||
if not self._has_batch_processed_column:
|
||||
return True
|
||||
from litellm_enterprise.proxy.hooks.managed_files import postgres_storable_json
|
||||
|
||||
try:
|
||||
claimed: Final = await _managed_object_table(self.prisma_client).update_many(
|
||||
where={"id": job.id, "batch_processed": False},
|
||||
data={"batch_processed": True},
|
||||
data={
|
||||
"batch_processed": True,
|
||||
"status": response.status if response.status != "completed" else "complete",
|
||||
"file_object": postgres_storable_json(response),
|
||||
},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
@ -228,7 +260,9 @@ class CheckBatchCost:
|
|||
return False
|
||||
return claimed > 0
|
||||
|
||||
async def _release_job_claim(self, job: "ManagedObjectRow") -> None:
|
||||
async def _release_job_claim(
|
||||
self, job: "ManagedObjectRow", original_status: str | None, original_file_object: object
|
||||
) -> None:
|
||||
"""Give a claimed row back once billing it failed, so a later poll cycle retries it.
|
||||
|
||||
Safe to match on batch_processed=True: while this poller is active the retrieve
|
||||
|
|
@ -240,7 +274,11 @@ class CheckBatchCost:
|
|||
try:
|
||||
await _managed_object_table(self.prisma_client).update_many(
|
||||
where={"id": job.id, "batch_processed": True},
|
||||
data={"batch_processed": False},
|
||||
data={
|
||||
"batch_processed": False,
|
||||
"status": original_status,
|
||||
"file_object": json.dumps(original_file_object),
|
||||
},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
@ -299,9 +337,9 @@ class CheckBatchCost:
|
|||
|
||||
async def _finalize_unbilled_terminal_job(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
|
||||
) -> None:
|
||||
) -> bool:
|
||||
"""Persist a terminal batch that has nothing billable, converting any raw
|
||||
provider file ids to managed ids, and take it out of the poll page."""
|
||||
provider file ids to managed ids, and take it out of the poll page. Returns whether it was persisted."""
|
||||
try:
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_is_base64_encoded_unified_file_id,
|
||||
|
|
@ -329,10 +367,12 @@ class CheckBatchCost:
|
|||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
|
||||
)
|
||||
return True
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}"
|
||||
)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _record_error(
|
||||
|
|
@ -537,8 +577,6 @@ class CheckBatchCost:
|
|||
|
||||
@staticmethod
|
||||
def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
|
||||
import json
|
||||
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
file_object: Final = job.file_object
|
||||
|
|
@ -581,7 +619,11 @@ class CheckBatchCost:
|
|||
from litellm.files.types import FileContentCallOptions
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
deployment_pricing_model_info,
|
||||
mask_api_base_credentials,
|
||||
resolve_vertex_location_for_cost,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_is_base64_encoded_unified_file_id,
|
||||
)
|
||||
|
|
@ -691,8 +733,7 @@ class CheckBatchCost:
|
|||
)
|
||||
except Exception as _e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckBatchCost: failed to create managed file ID for "
|
||||
f"{_file_attr}={_raw_file_id!r}: {_e}"
|
||||
f"CheckBatchCost: failed to create managed file ID for {_file_attr}={_raw_file_id!r}: {_e}"
|
||||
)
|
||||
|
||||
# Pass the deployment's router-registered pricing (litellm_params custom
|
||||
|
|
@ -711,6 +752,12 @@ class CheckBatchCost:
|
|||
custom_llm_provider=batch_file_provider,
|
||||
model_name=model_name,
|
||||
model_info=deployment_model_info,
|
||||
vertex_location=resolve_vertex_location_for_cost(
|
||||
custom_llm_provider=llm_provider,
|
||||
litellm_params=deployment_info.litellm_params.model_dump(exclude_none=True),
|
||||
optional_params=None,
|
||||
model=model_name,
|
||||
),
|
||||
)
|
||||
error_file_failed_requests: Final = await count_error_file_failed_requests(
|
||||
response,
|
||||
|
|
@ -728,8 +775,9 @@ class CheckBatchCost:
|
|||
failed_requests=output_file_result.failed_requests + error_file_failed_requests,
|
||||
)
|
||||
)
|
||||
logging_model: Final = batch_result.models[0] if batch_result.models else litellm_model_name or model_id
|
||||
logging_obj = LiteLLMLogging(
|
||||
model=batch_result.models[0],
|
||||
model=logging_model,
|
||||
messages=[{"role": "user", "content": "<retrieve_batch>"}],
|
||||
stream=False,
|
||||
call_type="aretrieve_batch",
|
||||
|
|
@ -760,7 +808,9 @@ class CheckBatchCost:
|
|||
custom_llm_provider=str(llm_provider) if llm_provider else None,
|
||||
)
|
||||
|
||||
if not await self._claim_job_for_costing(job):
|
||||
original_status: Final = job.status
|
||||
original_file_object: Final = job.file_object
|
||||
if not await self._claim_job_for_costing(job, response.model_copy(update={"usage": batch_result.usage})):
|
||||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod "
|
||||
"in this window, so its cost is already being tracked there"
|
||||
|
|
@ -779,7 +829,7 @@ class CheckBatchCost:
|
|||
batch_completion_cost=batch_result.completion_cost,
|
||||
)
|
||||
except Exception:
|
||||
await self._release_job_claim(job)
|
||||
await self._release_job_claim(job, original_status, original_file_object)
|
||||
raise
|
||||
|
||||
# Record batch duration (completed_at - created_at)
|
||||
|
|
@ -794,7 +844,200 @@ class CheckBatchCost:
|
|||
|
||||
return model_name, str(llm_provider) if llm_provider else None
|
||||
|
||||
async def check_batch_cost(self):
|
||||
async def _retire_unbilled_terminal_job(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str
|
||||
) -> None:
|
||||
"""Persist a terminal batch that has nothing billable, then store its lines, reading its input file by the id
|
||||
the provider reported rather than the managed id that was persisted."""
|
||||
provider_input_file_id: Final = response.input_file_id
|
||||
if await self._finalize_unbilled_terminal_job(job, response):
|
||||
response.input_file_id = provider_input_file_id
|
||||
self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=False)
|
||||
|
||||
def _store_batch_line_items_in_background(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
|
||||
) -> None:
|
||||
"""Store a retired batch's lines off the poll loop, one batch at a time and only between polls, so neither
|
||||
pricing the lines nor a slow or hung bucket holds up billing or retiring the other batches. While a poll
|
||||
cycle's worth of batches is still waiting, the lines of the next one are dropped with an error."""
|
||||
if not litellm.batch_line_item_callbacks:
|
||||
return
|
||||
if len(self._line_item_stores) >= MAX_OBJECTS_PER_POLL_CYCLE:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: {len(self._line_item_stores)} batches are still waiting to store their line items, "
|
||||
f"so the line items of batch {batch_id} (job {job.id}) were not stored"
|
||||
)
|
||||
return
|
||||
store: Final = asyncio.create_task(
|
||||
self._store_batch_line_items_in_turn(job, response, model_id, batch_id, billed=billed)
|
||||
)
|
||||
self._line_item_stores.add(store)
|
||||
store.add_done_callback(self._line_item_stores.discard)
|
||||
|
||||
async def _store_batch_line_items_in_turn(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
|
||||
) -> None:
|
||||
async with self._line_item_store_turn:
|
||||
await self._poll_idle.wait()
|
||||
await self._store_batch_line_items(job, response, model_id, batch_id, billed=billed)
|
||||
|
||||
async def _store_batch_line_items(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
|
||||
) -> None:
|
||||
"""Send every line of a finished batch to batch_line_item_callbacks, one part at a time. A billed batch's
|
||||
output file is read again here, so nothing waiting for its turn holds a batch's files. Failures, including an
|
||||
output file that can no longer be read, are logged and never raised, so the batch stays billed and retired
|
||||
whether or not its lines were stored."""
|
||||
if not litellm.batch_line_item_callbacks:
|
||||
return
|
||||
from litellm.batches.batch_line_items import (
|
||||
BatchFileProvider,
|
||||
BatchLineItemSource,
|
||||
batch_line_item_parts,
|
||||
dispatch_batch_line_items,
|
||||
line_item_metadata,
|
||||
)
|
||||
from litellm.batches.batch_utils import fetch_batch_managed_file_content
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
deployment_pricing_model_info,
|
||||
resolve_vertex_location_for_cost,
|
||||
)
|
||||
|
||||
try:
|
||||
deployment_info: Final = self.llm_router.get_deployment(model_id=model_id)
|
||||
if deployment_info is None:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: deployment {model_id} is no longer in the router, so the line items of batch "
|
||||
f"{batch_id} (job {job.id}) were not stored"
|
||||
)
|
||||
return
|
||||
litellm_model_name: Final = deployment_info.litellm_params.model
|
||||
model_name, llm_provider, _, _ = get_llm_provider(
|
||||
model=litellm_model_name,
|
||||
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
|
||||
)
|
||||
cost_provider: Final = cast(BatchFileProvider, llm_provider)
|
||||
credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
|
||||
file_access: Final = {
|
||||
**credentials,
|
||||
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
|
||||
}
|
||||
output_content: Final = (
|
||||
await fetch_batch_managed_file_content(
|
||||
response.output_file_id, custom_llm_provider=cost_provider, litellm_params=file_access
|
||||
)
|
||||
if billed and response.output_file_id
|
||||
else b""
|
||||
)
|
||||
error_content: Final = await self._batch_file_content(response.error_file_id, cost_provider, file_access)
|
||||
if not output_content and not error_content:
|
||||
return
|
||||
redact_bodies: Final = await self._creator_turned_off_message_logging(job)
|
||||
source: Final = BatchLineItemSource(
|
||||
batch_id=job.unified_object_id,
|
||||
finished_at=response.completed_at
|
||||
or response.failed_at
|
||||
or response.expired_at
|
||||
or response.cancelled_at
|
||||
or response.created_at,
|
||||
model=model_name,
|
||||
custom_llm_provider=llm_provider,
|
||||
cost_provider=cost_provider,
|
||||
model_info=deployment_pricing_model_info(model_id=model_id, deployment_model=litellm_model_name),
|
||||
metadata=line_item_metadata(await self._creator_attribution.batch_metadata(job, batch_id)),
|
||||
vertex_location=resolve_vertex_location_for_cost(
|
||||
custom_llm_provider=llm_provider,
|
||||
litellm_params=deployment_info.litellm_params.model_dump(exclude_none=True),
|
||||
optional_params=None,
|
||||
model=model_name,
|
||||
),
|
||||
)
|
||||
input_content: Final = (
|
||||
b""
|
||||
if redact_bodies
|
||||
else await self._batch_file_content(response.input_file_id, cost_provider, file_access)
|
||||
)
|
||||
parts: Final = batch_line_item_parts(
|
||||
source, input_content, (output_content, error_content), BATCH_LINE_ITEMS_MAX_PART_BYTES
|
||||
)
|
||||
handed = 0
|
||||
# pricing a part's lines is CPU-bound, so each part is built off the event loop
|
||||
while (payload := await asyncio.to_thread(next, parts, None)) is not None:
|
||||
await dispatch_batch_line_items(payload, redact_bodies=redact_bodies)
|
||||
handed += len(payload["lines"])
|
||||
if handed:
|
||||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: handed the {handed} line items of batch {batch_id} to batch_line_item_callbacks"
|
||||
)
|
||||
except Exception as error: # noqa: BLE001 # storing the lines is best effort, the batch is already billed
|
||||
verbose_proxy_logger.exception(
|
||||
f"CheckBatchCost: the line items of batch {batch_id} (job {job.id}) were not stored: {error}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _batch_file_content(
|
||||
file_id: str | None, custom_llm_provider: "BatchFileProvider", file_access: Mapping[str, object]
|
||||
) -> bytes:
|
||||
"""A batch file's raw JSONL, or nothing when the batch has no such file or the provider will not serve it."""
|
||||
from litellm.batches.batch_utils import fetch_batch_managed_file_content
|
||||
|
||||
if not file_id:
|
||||
return b""
|
||||
try:
|
||||
return await fetch_batch_managed_file_content(
|
||||
file_id, custom_llm_provider=custom_llm_provider, litellm_params=file_access
|
||||
)
|
||||
except Exception as error: # noqa: BLE001 # a missing input or error file leaves those lines out, not the rest
|
||||
verbose_proxy_logger.warning(f"CheckBatchCost: could not read batch file {file_id}: {error}")
|
||||
return b""
|
||||
|
||||
async def _creator_turned_off_message_logging(self, job: "ManagedObjectRow") -> bool:
|
||||
"""Whether the creating key's or team's logging settings turn off message logging, falling back to the
|
||||
proxy-wide setting. A batch whose saved credential is not a virtual key's hash, such as one created under JWT
|
||||
or OAuth2 auth or with the master key, has only its team's and the proxy's settings. A key or team that can no
|
||||
longer be read counts as turned off."""
|
||||
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata
|
||||
from litellm.proxy.proxy_server import proxy_config
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
api_key: Final = getattr(job, "api_key", None)
|
||||
team_id: Final = getattr(job, "team_id", None)
|
||||
has_key_row: Final = api_key is not None and is_valid_sha256_hash(api_key)
|
||||
try:
|
||||
key_row: Final = (
|
||||
await _token_table(self.prisma_client).find_unique(where={"token": api_key}) if has_key_row else None
|
||||
)
|
||||
team_row: Final = (
|
||||
await _team_table(self.prisma_client).find_unique(where={"team_id": team_id}) if team_id else None
|
||||
)
|
||||
except Exception as error:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: could not read the logging settings of batch job {job.id}'s creator, "
|
||||
f"so its line items are stored without request and response: {error}"
|
||||
)
|
||||
return True
|
||||
if (has_key_row and key_row is None) or (team_id and team_row is None):
|
||||
return True
|
||||
settings: Final = get_dynamic_logging_metadata(
|
||||
UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id),
|
||||
proxy_config,
|
||||
)
|
||||
callback_vars: Final = (settings.callback_vars if settings is not None else None) or {}
|
||||
creator_setting: Final = str_to_bool(callback_vars.get("turn_off_message_logging"))
|
||||
return creator_setting if creator_setting is not None else litellm.turn_off_message_logging is True
|
||||
|
||||
async def check_batch_cost(self) -> None:
|
||||
"""Run one poll cycle. Batch line items wait until it has billed and retired every batch it found."""
|
||||
self._poll_idle.clear()
|
||||
try:
|
||||
await self._poll_batches()
|
||||
finally:
|
||||
self._poll_idle.set()
|
||||
|
||||
async def _poll_batches(self):
|
||||
"""
|
||||
Check if the batch JOB has been tracked.
|
||||
- get all status="validating" and file_purpose="batch" jobs
|
||||
|
|
@ -888,10 +1131,7 @@ class CheckBatchCost:
|
|||
continue
|
||||
|
||||
## RETRIEVE THE BATCH JOB OUTPUT FILE
|
||||
if (
|
||||
response.status in PROVIDER_TERMINAL_BATCH_STATUSES
|
||||
and response.output_file_id is not None
|
||||
):
|
||||
if response.status in PROVIDER_TERMINAL_BATCH_STATUSES and response.output_file_id is not None:
|
||||
try:
|
||||
tracked = await self._track_completed_batch_cost(
|
||||
job=job,
|
||||
|
|
@ -908,7 +1148,7 @@ class CheckBatchCost:
|
|||
f"CheckBatchCost: output file {response.output_file_id} of batch {batch_id} "
|
||||
f"does not exist at the provider; retiring job {job.id} unbilled"
|
||||
)
|
||||
await self._finalize_unbilled_terminal_job(job, response)
|
||||
await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
|
||||
continue
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to track cost for batch {batch_id} "
|
||||
|
|
@ -922,22 +1162,21 @@ class CheckBatchCost:
|
|||
# Track this job for the final metrics summary
|
||||
processed_models.append(tracked)
|
||||
|
||||
# mark the job as complete
|
||||
try:
|
||||
update_data: dict = {
|
||||
"status": response.status if response.status != "completed" else "complete",
|
||||
"file_object": response.model_dump_json(),
|
||||
}
|
||||
if self._has_batch_processed_column:
|
||||
update_data["batch_processed"] = True
|
||||
await _managed_object_table(self.prisma_client).update(
|
||||
where={"id": job.id},
|
||||
data=update_data,
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
|
||||
)
|
||||
if not self._has_batch_processed_column:
|
||||
try:
|
||||
await _managed_object_table(self.prisma_client).update(
|
||||
where={"id": job.id},
|
||||
data={
|
||||
"status": response.status if response.status != "completed" else "complete",
|
||||
"file_object": response.model_dump_json(),
|
||||
},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
|
||||
)
|
||||
continue
|
||||
self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True)
|
||||
|
||||
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
|
|
@ -950,7 +1189,7 @@ class CheckBatchCost:
|
|||
f"has not appeared yet; leaving job {job.id} for the next poll cycle"
|
||||
)
|
||||
continue
|
||||
await self._finalize_unbilled_terminal_job(job, response)
|
||||
await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
|
||||
|
||||
# Record polling run metrics (always, even if nothing was processed)
|
||||
if prom_logger:
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import base64
|
|||
import json
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
|
|
@ -416,7 +417,7 @@ def _storable_json_value(value: JsonValue) -> JsonValue:
|
|||
return value
|
||||
|
||||
|
||||
def _postgres_storable_json(file_object: BaseModel) -> str:
|
||||
def postgres_storable_json(file_object: BaseModel) -> str:
|
||||
dumped: Final = file_object.model_dump_json()
|
||||
if "\\u0000" not in dumped:
|
||||
return dumped
|
||||
|
|
@ -479,6 +480,13 @@ def _with_managed_background_marker(request: Mapping[str, object]) -> dict[str,
|
|||
|
||||
|
||||
_MANAGED_FILES_TARGET: Final = "managed_files"
|
||||
_PROVIDER_FILE_ID_KWARGS: Final = MappingProxyType(
|
||||
{
|
||||
CallTypes.acreate_batch: ("input_file_id",),
|
||||
CallTypes.acreate_fine_tuning_job: ("training_file", "validation_file"),
|
||||
}
|
||||
)
|
||||
_PROVIDER_FILE_ID_KEYS: Final = tuple(dict.fromkeys(chain.from_iterable(_PROVIDER_FILE_ID_KWARGS.values())))
|
||||
_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0)
|
||||
|
||||
|
||||
|
|
@ -650,7 +658,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
# FIX: Update status and file_object on every operation to keep state in sync
|
||||
update_columns: Final = {
|
||||
"file_object": _postgres_storable_json(stored_object),
|
||||
"file_object": postgres_storable_json(stored_object),
|
||||
"status": file_object.status,
|
||||
"updated_by": user_api_key_dict.user_id,
|
||||
}
|
||||
|
|
@ -666,7 +674,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
data={
|
||||
"create": {
|
||||
"unified_object_id": unified_object_id,
|
||||
"file_object": _postgres_storable_json(stored_object),
|
||||
"file_object": postgres_storable_json(stored_object),
|
||||
"model_object_id": model_object_id,
|
||||
"file_purpose": file_purpose,
|
||||
"created_by": resolve_resource_owner_id(user_api_key_dict),
|
||||
|
|
@ -702,7 +710,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
"unified_object_id": response_id,
|
||||
"model_object_id": decrypted_response_id(response_id),
|
||||
"file_purpose": "response",
|
||||
"file_object": _postgres_storable_json(_without_model_output(response)),
|
||||
"file_object": postgres_storable_json(_without_model_output(response)),
|
||||
"status": PENDING_RESPONSE_STATUSES[0],
|
||||
"created_by": resolve_resource_owner_id(user_api_key_dict),
|
||||
"team_id": user_api_key_dict.team_id,
|
||||
|
|
@ -1562,6 +1570,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if isinstance(file_id, str) and file_id
|
||||
)
|
||||
await self.check_file_ids_access(fine_tuning_file_ids, user_api_key_dict)
|
||||
if fine_tuning_file_ids:
|
||||
data["model_file_id_mapping"] = await self.get_model_file_id_mapping(
|
||||
fine_tuning_file_ids, user_api_key_dict.parent_otel_span
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
|
|
@ -1576,24 +1588,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if request_kwargs is None:
|
||||
return healthy_deployments
|
||||
|
||||
input_file_id = cast(Optional[str], request_kwargs.get("input_file_id"))
|
||||
model_file_id_mapping = cast(
|
||||
Optional[Dict[str, Dict[str, str]]],
|
||||
request_kwargs.get("model_file_id_mapping"),
|
||||
request: Final = cast(Mapping[str, object], request_kwargs)
|
||||
model_file_id_mapping: Final = cast(Mapping[str, Mapping[str, str]], request.get("model_file_id_mapping") or {})
|
||||
requested_file_ids: Final = (request.get(key) for key in _PROVIDER_FILE_ID_KEYS)
|
||||
model_ids_per_file: Final = tuple(
|
||||
frozenset(model_file_id_mapping[file_id])
|
||||
for file_id in requested_file_ids
|
||||
if isinstance(file_id, str) and file_id in model_file_id_mapping
|
||||
)
|
||||
allowed_model_ids: Final[frozenset[str]] = (
|
||||
model_ids_per_file[0].intersection(*model_ids_per_file[1:]) if model_ids_per_file else frozenset()
|
||||
)
|
||||
allowed_model_ids = []
|
||||
if input_file_id and model_file_id_mapping:
|
||||
model_id_dict = model_file_id_mapping.get(input_file_id, {})
|
||||
allowed_model_ids = list(model_id_dict.keys())
|
||||
|
||||
if len(allowed_model_ids) == 0:
|
||||
return healthy_deployments
|
||||
|
||||
return [
|
||||
holders: Final[list[dict[str, object]]] = [
|
||||
deployment
|
||||
for deployment in healthy_deployments
|
||||
if deployment.get("model_info", {}).get("id") in allowed_model_ids
|
||||
]
|
||||
return holders if holders or "input_file_id" in request else healthy_deployments
|
||||
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
|
||||
|
|
@ -1601,25 +1616,21 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
"""
|
||||
Allow modifying the request just before it's sent to the deployment.
|
||||
"""
|
||||
accessor_key: Optional[str] = None
|
||||
if call_type and call_type == CallTypes.acreate_batch:
|
||||
accessor_key = "input_file_id"
|
||||
elif call_type and call_type == CallTypes.acreate_fine_tuning_job:
|
||||
accessor_key = "training_file"
|
||||
else:
|
||||
accessor_keys: Final = _PROVIDER_FILE_ID_KWARGS.get(call_type) if call_type else None
|
||||
if accessor_keys is None:
|
||||
return kwargs
|
||||
|
||||
if accessor_key:
|
||||
model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping"))
|
||||
# model_info may be at top-level or nested under litellm_metadata
|
||||
# (batch/file operations use litellm_metadata)
|
||||
model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None))
|
||||
if model_id is None:
|
||||
model_id = cast(
|
||||
Optional[str],
|
||||
kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
|
||||
)
|
||||
for accessor_key in accessor_keys:
|
||||
input_file_id = cast(Optional[str], kwargs.get(accessor_key))
|
||||
model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping"))
|
||||
# model_info may be at top-level or nested under litellm_metadata
|
||||
# (batch/file operations use litellm_metadata)
|
||||
model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None))
|
||||
if model_id is None:
|
||||
model_id = cast(
|
||||
Optional[str],
|
||||
kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
|
||||
)
|
||||
mapped_file_id: Optional[str] = None
|
||||
if input_file_id and model_file_id_mapping and model_id:
|
||||
mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None)
|
||||
|
|
@ -1813,7 +1824,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
),
|
||||
)
|
||||
|
||||
async def get_model_file_id_mapping(self, file_ids: List[str], litellm_parent_otel_span: Span) -> dict:
|
||||
async def get_model_file_id_mapping(self, file_ids: Sequence[str], litellm_parent_otel_span: Span) -> dict:
|
||||
"""
|
||||
Get model-specific file IDs for a list of proxy file IDs.
|
||||
Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.77"
|
||||
version = "0.1.78"
|
||||
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.77"
|
||||
version = "0.1.78"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -20,6 +20,33 @@ RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${
|
|||
make -j"$(nproc)" pgbouncer && \
|
||||
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
|
||||
|
||||
# ---------- Rust extension ----------
|
||||
# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes
|
||||
FROM $LITELLM_BUILD_IMAGE AS rust-native
|
||||
|
||||
WORKDIR /app
|
||||
USER root
|
||||
|
||||
COPY --from=uvbin /uv /uvx /usr/local/bin/
|
||||
|
||||
RUN for i in 1 2 3; do \
|
||||
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \
|
||||
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
|
||||
sleep 5; \
|
||||
done
|
||||
|
||||
ENV UV_PYTHON_DOWNLOADS=0
|
||||
|
||||
COPY pyproject.toml README.md rust-toolchain.toml ./
|
||||
COPY .cargo ./.cargo
|
||||
COPY litellm-rust ./litellm-rust
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p litellm/rust_bridge /wheels && \
|
||||
touch litellm/__init__.py litellm/rust_bridge/__init__.py && \
|
||||
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
|
||||
python -c "import maturin; maturin.build_wheel('/wheels')"
|
||||
|
||||
# ---------- Builder ----------
|
||||
FROM $LITELLM_BUILD_IMAGE AS builder
|
||||
|
||||
|
|
@ -66,11 +93,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
# Stage 2 — copy source and install the project + workspace members.
|
||||
# Stage 2 — copy source and install the workspace members.
|
||||
COPY . .
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv sync --frozen --no-default-groups --no-editable \
|
||||
uv sync --frozen --no-install-project --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
--group admin-mcp \
|
||||
|
|
@ -79,6 +106,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \
|
||||
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl
|
||||
|
||||
# PYTHONPATH=/app makes the source tree shadow the installed package, so the
|
||||
# compiled Rust extension must live next to the source or it is never imported.
|
||||
RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
|
||||
|
|
|
|||
|
|
@ -70,11 +70,18 @@ proxyConfigMap:
|
|||
proxy_config:
|
||||
general_settings:
|
||||
master_key: os.environ/PROXY_MASTER_KEY
|
||||
model_list: []
|
||||
```
|
||||
|
||||
The chart ships no models, so add them under `proxy_config.model_list` or from the Admin UI with `store_model_in_db: true`
|
||||
|
||||
```
|
||||
proxy_config:
|
||||
model_list:
|
||||
- model_name: gpt-3.5-turbo
|
||||
- model_name: gpt-5.6
|
||||
litellm_params:
|
||||
model: gpt-3.5-turbo
|
||||
api_key: eXaMpLeOnLy
|
||||
model: openai/gpt-5.6
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
```
|
||||
|
||||
#### Example using existing `proxyConfigMap` instead of creating it:
|
||||
|
|
@ -126,6 +133,11 @@ Set `billingMetrics.caSecretName` only when the collector is a private or test o
|
|||
| `db.secret.name` | If `db.useExisting` is `true`, the name of the Kubernetes Secret that contains credentials. | `postgres` |
|
||||
| `db.secret.usernameKey` | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the username for authenticating with the Postgres instance. | `username` |
|
||||
| `db.secret.passwordKey` | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the password associates with the above user. | `password` |
|
||||
| `db.username` | If `db.useExisting` is `true`, the database user as a plain value. When set, `db.secret.usernameKey` is not read. | `""` |
|
||||
| `db.port` | If `db.useExisting` is `true` with token auth, the Postgres port used to build the connection URL. Password auth takes the port from `db.url`. | `5432` |
|
||||
| `db.schema` | With token auth, the Postgres schema holding the LiteLLM tables, rendered as `DATABASE_SCHEMA` when the tables are not in `public`. Ignored with password auth, where the schema stays in `db.url` as `?schema=`. | `""` |
|
||||
| `db.useIAMAuth` | If `db.useExisting` is `true` and `db.deployStandalone` is `false`, authenticate with an AWS RDS IAM token instead of a password. See [Token auth](#token-auth-to-an-existing-postgres-rds-iam-or-microsoft-entra-id). | `false` |
|
||||
| `db.useAzureEntraAuth` | If `db.useExisting` is `true` and `db.deployStandalone` is `false`, authenticate with a Microsoft Entra ID token instead of a password. Mutually exclusive with `db.useIAMAuth`. | `false` |
|
||||
| `db.useStackgresOperator` | Not yet implemented. | `false` |
|
||||
| `db.deployStandalone` | Deploy a standalone, single instance deployment of Postgres, using the Bitnami postgresql chart. This is useful for getting started but doesn't provide HA or (by default) data backups. | `true` |
|
||||
| `postgresql.*` | If `db.deployStandalone` is `true`, configuration passed to the Bitnami postgresql chart. See the [Bitnami Documentation](https://github.com/bitnami/charts/tree/main/bitnami/postgresql) for full configuration details. See [values.yaml](./values.yaml) for the default configuration. | See [values.yaml](./values.yaml) |
|
||||
|
|
@ -156,6 +168,22 @@ data:
|
|||
type: Opaque
|
||||
```
|
||||
|
||||
#### Token auth to an existing Postgres (RDS IAM or Microsoft Entra ID)
|
||||
|
||||
When the database accepts short-lived tokens instead of a password (AWS RDS or Aurora with IAM database authentication, or Azure Database for PostgreSQL with Microsoft Entra ID), there is no password to put in a Secret. Turn on `db.useIAMAuth` or `db.useAzureEntraAuth` and the chart stops referencing `db.secret.passwordKey` and `db.url`. It sets `IAM_TOKEN_DB_AUTH` (or `AZURE_POSTGRESQL_AUTH`) and emits `DATABASE_HOST`, `DATABASE_PORT`, `DATABASE_USERNAME` and `DATABASE_NAME` for the proxy, the migrations job and the collector sidecar. Each of them mints its own token at startup and builds `DATABASE_URL` from those parts, and the in-container pgbouncer (`db.connectionPool`) refreshes the token before it expires.
|
||||
|
||||
```yaml
|
||||
db:
|
||||
useExisting: true
|
||||
deployStandalone: false
|
||||
useIAMAuth: true
|
||||
endpoint: litellm.cluster-abc123.us-west-1.rds.amazonaws.com
|
||||
database: litellm
|
||||
username: litellm_iam # the role granted rds_iam; no db.secret needed
|
||||
```
|
||||
|
||||
The pod needs cloud credentials for the token: on EKS use IRSA or Pod Identity on the chart's `serviceAccount` (and `migrationJob.serviceAccountName` when the Helm hook runs before the ServiceAccount exists), on AKS use workload identity. Set `AWS_RDS_REGION` under `envVars` when the RDS hostname is not a standard regional endpoint (the token must be signed for the region the database lives in). TLS options that would otherwise go in `db.url` are set as `DATABASE_SSLMODE` and `DATABASE_SSLROOTCERT` under `envVars`. Leaving `db.username` empty reads the user from `db.secret.usernameKey` as before.
|
||||
|
||||
#### Examples for `environmentSecrets` and `environemntConfigMaps`
|
||||
|
||||
```yaml
|
||||
|
|
|
|||
|
|
@ -162,6 +162,77 @@ taken before the change, which by that point no longer exists.
|
|||
{{- end -}}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Whether the existing Postgres is reached with a short-lived token (AWS RDS IAM
|
||||
or Microsoft Entra ID) instead of a password.
|
||||
*/}}
|
||||
{{- define "litellm.db.tokenAuth" -}}
|
||||
{{- if or .Values.db.useIAMAuth .Values.db.useAzureEntraAuth -}}true{{- end -}}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.validateDbTokenAuth" -}}
|
||||
{{- if and .Values.db.useIAMAuth .Values.db.useAzureEntraAuth -}}
|
||||
{{- fail "db.useIAMAuth and db.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" -}}
|
||||
{{- end -}}
|
||||
{{- if and (include "litellm.db.tokenAuth" .) (or (not .Values.db.useExisting) .Values.db.deployStandalone) -}}
|
||||
{{- fail "db.useIAMAuth and db.useAzureEntraAuth need db.useExisting: true and db.deployStandalone: false (token auth targets an existing RDS or Azure Database for PostgreSQL; the bundled postgresql subchart only supports password auth)" -}}
|
||||
{{- end -}}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Database env for an existing Postgres, shared by the proxy container, the
|
||||
collector sidecar and the migrations Job. Under token auth every entrypoint
|
||||
mints its own token and assembles DATABASE_URL itself, so db.url is left out.
|
||||
*/}}
|
||||
{{- define "litellm.existingDbEnv" -}}
|
||||
{{- $tokenAuth := include "litellm.db.tokenAuth" . -}}
|
||||
- name: DATABASE_USERNAME
|
||||
{{- if .Values.db.username }}
|
||||
value: {{ .Values.db.username | quote }}
|
||||
{{- else }}
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.db.secret.name }}
|
||||
key: {{ .Values.db.secret.usernameKey }}
|
||||
{{- end }}
|
||||
{{- if .Values.db.useIAMAuth }}
|
||||
- name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
{{- else if .Values.db.useAzureEntraAuth }}
|
||||
- name: AZURE_POSTGRESQL_AUTH
|
||||
value: "true"
|
||||
{{- else }}
|
||||
- name: DATABASE_PASSWORD
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.db.secret.name }}
|
||||
key: {{ .Values.db.secret.passwordKey }}
|
||||
{{- end }}
|
||||
- name: DATABASE_HOST
|
||||
{{- if .Values.db.secret.endpointKey }}
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.db.secret.name }}
|
||||
key: {{ .Values.db.secret.endpointKey }}
|
||||
{{- else }}
|
||||
value: {{ .Values.db.endpoint }}
|
||||
{{- end }}
|
||||
{{- if $tokenAuth }}
|
||||
- name: DATABASE_PORT
|
||||
value: {{ .Values.db.port | quote }}
|
||||
{{- end }}
|
||||
- name: DATABASE_NAME
|
||||
value: {{ .Values.db.database }}
|
||||
{{- if and $tokenAuth .Values.db.schema }}
|
||||
- name: DATABASE_SCHEMA
|
||||
value: {{ .Values.db.schema | quote }}
|
||||
{{- end }}
|
||||
{{- if not $tokenAuth }}
|
||||
- name: DATABASE_URL
|
||||
value: {{ .Values.db.url | quote }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Environment shared by the proxy container and the opt-in collector sidecar:
|
||||
database, pgbouncer, master key, redis, user envVars. Both containers must see
|
||||
|
|
@ -169,6 +240,7 @@ the same DATABASE_URL and REDIS_* so the sidecar reaches the pod's pgbouncer
|
|||
and the same spend transaction buffer.
|
||||
*/}}
|
||||
{{- define "litellm.proxyEnv" -}}
|
||||
{{- include "litellm.validateDbTokenAuth" . -}}
|
||||
- name: HOST
|
||||
value: "{{ .Values.listen | default "0.0.0.0" }}"
|
||||
- name: PORT
|
||||
|
|
@ -189,29 +261,7 @@ and the same spend transaction buffer.
|
|||
- name: DATABASE_NAME
|
||||
value: litellm
|
||||
{{- else if .Values.db.useExisting }}
|
||||
- name: DATABASE_USERNAME
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.db.secret.name }}
|
||||
key: {{ .Values.db.secret.usernameKey }}
|
||||
- name: DATABASE_PASSWORD
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.db.secret.name }}
|
||||
key: {{ .Values.db.secret.passwordKey }}
|
||||
- name: DATABASE_HOST
|
||||
{{- if .Values.db.secret.endpointKey }}
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.db.secret.name }}
|
||||
key: {{ .Values.db.secret.endpointKey }}
|
||||
{{- else }}
|
||||
value: {{ .Values.db.endpoint }}
|
||||
{{- end }}
|
||||
- name: DATABASE_NAME
|
||||
value: {{ .Values.db.database }}
|
||||
- name: DATABASE_URL
|
||||
value: {{ .Values.db.url | quote }}
|
||||
{{ include "litellm.existingDbEnv" . }}
|
||||
{{- end }}
|
||||
{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
|
||||
- name: DATABASE_READER_HOST
|
||||
|
|
|
|||
|
|
@ -68,30 +68,9 @@ spec:
|
|||
{{- end }}
|
||||
workingDir: "/app"
|
||||
env:
|
||||
{{- include "litellm.validateDbTokenAuth" . -}}
|
||||
{{- if .Values.db.useExisting }}
|
||||
- name: DATABASE_USERNAME
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.db.secret.name }}
|
||||
key: {{ .Values.db.secret.usernameKey }}
|
||||
- name: DATABASE_PASSWORD
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.db.secret.name }}
|
||||
key: {{ .Values.db.secret.passwordKey }}
|
||||
- name: DATABASE_HOST
|
||||
{{- if .Values.db.secret.endpointKey }}
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.db.secret.name }}
|
||||
key: {{ .Values.db.secret.endpointKey }}
|
||||
{{- else }}
|
||||
value: {{ .Values.db.endpoint }}
|
||||
{{- end }}
|
||||
- name: DATABASE_NAME
|
||||
value: {{ .Values.db.database }}
|
||||
- name: DATABASE_URL
|
||||
value: {{ .Values.db.url | quote }}
|
||||
{{- include "litellm.existingDbEnv" . | nindent 12 }}
|
||||
{{- else if .Values.db.deployStandalone }}
|
||||
- name: DATABASE_URL
|
||||
value: postgresql://{{ .Values.postgresql.auth.username }}:{{ .Values.postgresql.auth.password }}@{{ .Release.Name }}-postgresql/{{ .Values.postgresql.auth.database }}
|
||||
|
|
|
|||
288
helm/litellm-helm/tests/database_auth_tests.yaml
Normal file
288
helm/litellm-helm/tests/database_auth_tests.yaml
Normal file
|
|
@ -0,0 +1,288 @@
|
|||
suite: test existing database token auth
|
||||
templates:
|
||||
- deployment.yaml
|
||||
- migrations-job.yaml
|
||||
- configmap-litellm.yaml
|
||||
tests:
|
||||
- it: should keep password auth from the secret when no token auth is set
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db.useExisting: true
|
||||
db.deployStandalone: false
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_USERNAME
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: postgres
|
||||
key: username
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: postgres
|
||||
key: password
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL
|
||||
value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: IAM_TOKEN_DB_AUTH
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: AZURE_POSTGRESQL_AUTH
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PORT
|
||||
any: true
|
||||
|
||||
- it: should render RDS IAM auth without a password secret on the proxy and the collector sidecar
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db.useExisting: true
|
||||
db.deployStandalone: false
|
||||
db.useIAMAuth: true
|
||||
db.username: litellm_iam
|
||||
db.endpoint: writer.cluster-abc.us-west-1.rds.amazonaws.com
|
||||
db.port: 5433
|
||||
collector.enabled: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_USERNAME
|
||||
value: litellm_iam
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_HOST
|
||||
value: writer.cluster-abc.us-west-1.rds.amazonaws.com
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PORT
|
||||
value: "5433"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_NAME
|
||||
value: litellm
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: AZURE_POSTGRESQL_AUTH
|
||||
any: true
|
||||
- contains:
|
||||
path: spec.template.spec.containers[1].env
|
||||
content:
|
||||
name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[1].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
|
||||
- it: should give the migrations job the same token auth env as the proxy
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
migrationJob.enabled: true
|
||||
db.useExisting: true
|
||||
db.deployStandalone: false
|
||||
db.useIAMAuth: true
|
||||
db.username: litellm_iam
|
||||
db.endpoint: writer.cluster-abc.us-west-1.rds.amazonaws.com
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_USERNAME
|
||||
value: litellm_iam
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_HOST
|
||||
value: writer.cluster-abc.us-west-1.rds.amazonaws.com
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PORT
|
||||
value: "5432"
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL
|
||||
any: true
|
||||
|
||||
- it: should carry db.schema as DATABASE_SCHEMA under token auth and omit it when empty
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db.useExisting: true
|
||||
db.deployStandalone: false
|
||||
db.useIAMAuth: true
|
||||
db.username: litellm_iam
|
||||
db.schema: tenant1
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_SCHEMA
|
||||
value: tenant1
|
||||
|
||||
- it: should leave the schema to db.url under password auth even when db.schema is set
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db.useExisting: true
|
||||
db.deployStandalone: false
|
||||
db.schema: tenant1
|
||||
db.url: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)?schema=tenant1
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_SCHEMA
|
||||
any: true
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL
|
||||
value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)?schema=tenant1
|
||||
|
||||
- it: should not emit DATABASE_SCHEMA when db.schema is empty
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db.useExisting: true
|
||||
db.deployStandalone: false
|
||||
db.useIAMAuth: true
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_SCHEMA
|
||||
any: true
|
||||
|
||||
- it: should read the user from the secret when db.username is empty under token auth
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db.useExisting: true
|
||||
db.deployStandalone: false
|
||||
db.useIAMAuth: true
|
||||
db.secret.name: rds-identity
|
||||
db.secret.usernameKey: user
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_USERNAME
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: rds-identity
|
||||
key: user
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
|
||||
- it: should render Entra auth with AZURE_POSTGRESQL_AUTH and no password
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db.useExisting: true
|
||||
db.deployStandalone: false
|
||||
db.useAzureEntraAuth: true
|
||||
db.username: litellm@example.onmicrosoft.com
|
||||
db.endpoint: litellm.postgres.database.azure.com
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: AZURE_POSTGRESQL_AUTH
|
||||
value: "true"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_USERNAME
|
||||
value: litellm@example.onmicrosoft.com
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: IAM_TOKEN_DB_AUTH
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL
|
||||
any: true
|
||||
|
||||
- it: should fail when both token sources are enabled
|
||||
template: deployment.yaml
|
||||
set:
|
||||
migrationJob.enabled: false
|
||||
db.useExisting: true
|
||||
db.deployStandalone: false
|
||||
db.useIAMAuth: true
|
||||
db.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "db.useIAMAuth and db.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source"
|
||||
|
||||
- it: should fail when token auth is set while the bundled postgres is still deployed
|
||||
template: deployment.yaml
|
||||
set:
|
||||
migrationJob.enabled: false
|
||||
db.useExisting: true
|
||||
db.useIAMAuth: true
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorPattern: "need db.useExisting: true and db.deployStandalone: false"
|
||||
|
||||
- it: should fail when token auth is set without useExisting
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
migrationJob.enabled: true
|
||||
db.deployStandalone: false
|
||||
db.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorPattern: "need db.useExisting: true and db.deployStandalone: false"
|
||||
|
|
@ -270,6 +270,32 @@ tests:
|
|||
- matchRegex:
|
||||
path: data["config.yaml"]
|
||||
pattern: '(?m)^\s*health_check_details:\s*false$'
|
||||
- it: should not render any placeholder models when the user sets proxy_config without model_list
|
||||
template: configmap-litellm.yaml
|
||||
set:
|
||||
proxy_config.general_settings.store_model_in_db: true
|
||||
asserts:
|
||||
- notMatchRegex:
|
||||
path: data["config.yaml"]
|
||||
pattern: 'model_name:'
|
||||
- matchRegex:
|
||||
path: data["config.yaml"]
|
||||
pattern: 'store_model_in_db: true'
|
||||
- matchRegex:
|
||||
path: data["config.yaml"]
|
||||
pattern: 'master_key: os.environ/PROXY_MASTER_KEY'
|
||||
- it: should render a user-supplied model_list as the only models
|
||||
template: configmap-litellm.yaml
|
||||
set:
|
||||
proxy_config.model_list:
|
||||
- model_name: my-model
|
||||
litellm_params:
|
||||
model: openai/my-model
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
asserts:
|
||||
- matchRegex:
|
||||
path: data["config.yaml"]
|
||||
pattern: 'model_list:\n- litellm_params:\n api_key: os.environ/OPENAI_API_KEY\n model: openai/my-model\n model_name: my-model\n$'
|
||||
- it: should allow overriding liveness, readiness, and startup probes
|
||||
template: deployment.yaml
|
||||
set:
|
||||
|
|
|
|||
|
|
@ -180,17 +180,7 @@ proxyConfigMap:
|
|||
# Examples: https://github.com/BerriAI/litellm/tree/main/litellm/proxy/example_config_yaml
|
||||
# Reference: https://docs.litellm.ai/docs/proxy/configs
|
||||
proxy_config:
|
||||
model_list:
|
||||
# At least one model must exist for the proxy to start.
|
||||
- model_name: gpt-3.5-turbo
|
||||
litellm_params:
|
||||
model: gpt-3.5-turbo
|
||||
api_key: eXaMpLeOnLy
|
||||
- model_name: fake-openai-endpoint
|
||||
litellm_params:
|
||||
model: openai/fake
|
||||
api_key: fake-key
|
||||
api_base: https://exampleopenaiendpoint-production.up.railway.app/
|
||||
model_list: []
|
||||
general_settings:
|
||||
master_key: os.environ/PROXY_MASTER_KEY
|
||||
|
||||
|
|
@ -376,6 +366,30 @@ db:
|
|||
endpoint: localhost
|
||||
database: litellm
|
||||
url: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)
|
||||
# Authenticate with a short-lived token instead of a password. useIAMAuth is
|
||||
# AWS RDS / Aurora IAM database authentication (the pod needs AWS credentials,
|
||||
# e.g. IRSA or EKS Pod Identity, and the database user needs GRANT rds_iam);
|
||||
# useAzureEntraAuth is Azure Database for PostgreSQL with a Microsoft Entra ID
|
||||
# identity (workload identity). Mutually exclusive, and both need
|
||||
# useExisting: true with deployStandalone: false. The chart then sets
|
||||
# IAM_TOKEN_DB_AUTH or AZURE_POSTGRESQL_AUTH, emits DATABASE_HOST, DATABASE_PORT,
|
||||
# DATABASE_USERNAME and DATABASE_NAME and skips DATABASE_PASSWORD and db.url:
|
||||
# the proxy, the migrations job and the collector sidecar each mint their own
|
||||
# token and build DATABASE_URL from those parts. Set DATABASE_SSLMODE /
|
||||
# DATABASE_SSLROOTCERT under envVars for TLS options you would otherwise put
|
||||
# in db.url.
|
||||
useIAMAuth: false
|
||||
useAzureEntraAuth: false
|
||||
# Postgres port, used to build the token-auth URL (password auth takes the
|
||||
# port from db.url)
|
||||
port: 5432
|
||||
# Postgres schema for the LiteLLM tables, rendered as DATABASE_SCHEMA under
|
||||
# token auth only (password auth carries it as ?schema= in db.url)
|
||||
schema: ""
|
||||
# Database user as a plain value. Empty means read it from
|
||||
# db.secret.usernameKey. With token auth the user name is the only credential
|
||||
# the chart needs, so setting it here removes the need for db.secret entirely.
|
||||
username: ""
|
||||
secret:
|
||||
name: postgres
|
||||
usernameKey: username
|
||||
|
|
|
|||
|
|
@ -475,6 +475,27 @@ harmless no-op for the Job and authoritative for the app pods.
|
|||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- $sslCertReqs := $root.Values.redis.sslCertReqs }}
|
||||
{{- if and $sslCertReqs (not (has $sslCertReqs (list "none" "optional" "required"))) }}
|
||||
{{- fail (printf "redis.sslCertReqs must be one of none, optional, required (got %q)" $sslCertReqs) }}
|
||||
{{- end }}
|
||||
{{- if and $urlKey (or $root.Values.redis.ssl $sslCertReqs $root.Values.redis.sslCaCerts) }}
|
||||
{{- fail "redis.ssl, redis.sslCertReqs and redis.sslCaCerts cannot be combined with redis.passwordSecret.urlKey, use a rediss:// URL with ssl_cert_reqs / ssl_ca_certs query parameters instead" }}
|
||||
{{- end }}
|
||||
{{- if or $root.Values.redis.host $hostKey }}
|
||||
{{- if $root.Values.redis.ssl }}
|
||||
- name: REDIS_SSL
|
||||
value: "true"
|
||||
{{- end }}
|
||||
{{- if $sslCertReqs }}
|
||||
- name: REDIS_SSL_CERT_REQS
|
||||
value: {{ $sslCertReqs | quote }}
|
||||
{{- end }}
|
||||
{{- if $root.Values.redis.sslCaCerts }}
|
||||
- name: REDIS_SSL_CA_CERTS
|
||||
value: {{ $root.Values.redis.sslCaCerts | quote }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- with $component.extraEnv }}
|
||||
{{ toYaml . }}
|
||||
{{- end }}
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ templates:
|
|||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- backend/deployment.yaml
|
||||
- migrations-job.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
|
|
@ -320,3 +321,168 @@ tests:
|
|||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"
|
||||
|
||||
- it: redis omits TLS env vars by default
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.host: redis.example.com
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL_CERT_REQS
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL_CA_CERTS
|
||||
any: true
|
||||
|
||||
- it: redis ssl false omits REDIS_SSL
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.host: redis.example.com
|
||||
redis.ssl: false
|
||||
redis.sslCertReqs: required
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL
|
||||
any: true
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL_CERT_REQS
|
||||
value: "required"
|
||||
|
||||
- it: gateway renders redis TLS env vars
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.host: redis.example.com
|
||||
redis.ssl: true
|
||||
redis.sslCertReqs: none
|
||||
redis.sslCaCerts: /etc/redis-tls/ca.crt
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL
|
||||
value: "true"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL_CERT_REQS
|
||||
value: "none"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL_CA_CERTS
|
||||
value: "/etc/redis-tls/ca.crt"
|
||||
|
||||
- it: backend renders redis TLS env vars
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
redis.host: redis.example.com
|
||||
redis.ssl: true
|
||||
redis.sslCertReqs: required
|
||||
redis.sslCaCerts: /etc/redis-tls/ca.crt
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL
|
||||
value: "true"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL_CERT_REQS
|
||||
value: "required"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL_CA_CERTS
|
||||
value: "/etc/redis-tls/ca.crt"
|
||||
|
||||
- it: migration job renders redis TLS env vars
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
db.deployStandalone: false
|
||||
db.useExisting: true
|
||||
db.endpoint: postgres.example.com
|
||||
db.secret.name: pg-secret
|
||||
redis.host: redis.example.com
|
||||
redis.ssl: true
|
||||
redis.sslCertReqs: optional
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL
|
||||
value: "true"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL_CERT_REQS
|
||||
value: "optional"
|
||||
|
||||
- it: redis TLS values reject a secret-backed URL
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.passwordSecret.name: redis-secret
|
||||
redis.passwordSecret.urlKey: connection-url
|
||||
redis.sslCaCerts: /etc/redis-tls/ca.crt
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "redis.ssl, redis.sslCertReqs and redis.sslCaCerts cannot be combined with redis.passwordSecret.urlKey, use a rediss:// URL with ssl_cert_reqs / ssl_ca_certs query parameters instead"
|
||||
|
||||
- it: hostKey mode renders redis TLS env vars
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.passwordSecret.name: redis-secret
|
||||
redis.passwordSecret.hostKey: hostname
|
||||
redis.ssl: true
|
||||
redis.sslCaCerts: /etc/redis-tls/ca.crt
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL
|
||||
value: "true"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL_CA_CERTS
|
||||
value: "/etc/redis-tls/ca.crt"
|
||||
|
||||
- it: redis TLS env vars are omitted when redis is not configured
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.ssl: true
|
||||
redis.sslCertReqs: none
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: REDIS_SSL_CERT_REQS
|
||||
any: true
|
||||
|
||||
- it: redis sslCertReqs rejects unknown values
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
redis.host: redis.example.com
|
||||
redis.ssl: true
|
||||
redis.sslCertReqs: strict
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: 'redis.sslCertReqs must be one of none, optional, required (got "strict")'
|
||||
|
|
|
|||
|
|
@ -275,10 +275,21 @@ database:
|
|||
# self-hosted Redis Cluster). The chart emits REDIS_CLUSTER_NODES from
|
||||
# `host` / `port` as the single seed; the cluster client discovers the
|
||||
# remaining nodes from CLUSTER SLOTS at startup.
|
||||
#
|
||||
# For a TLS Redis endpoint set `ssl: true`, which the chart emits as REDIS_SSL.
|
||||
# `sslCertReqs` (none / optional / required) becomes REDIS_SSL_CERT_REQS and
|
||||
# `sslCaCerts` (path to a CA bundle mounted in the pod, e.g. via
|
||||
# `extraVolumeMounts`) becomes REDIS_SSL_CA_CERTS. All three are omitted when
|
||||
# unset. They apply to `host` / `hostKey` mode only. With `passwordSecret.urlKey`
|
||||
# the proxy reads TLS settings from the URL itself, so use `rediss://` and the
|
||||
# `ssl_cert_reqs` / `ssl_ca_certs` query parameters there instead.
|
||||
redis:
|
||||
cluster: false
|
||||
host: ""
|
||||
port: 6379
|
||||
ssl: false
|
||||
sslCertReqs: ""
|
||||
sslCaCerts: ""
|
||||
passwordSecret:
|
||||
name: "" # Leave empty for auth-less Redis
|
||||
passwordKey: password
|
||||
|
|
|
|||
50
litellm-proxy-extras/litellm_proxy_extras/_redaction.py
Normal file
50
litellm-proxy-extras/litellm_proxy_extras/_redaction.py
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
import functools
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
from urllib.parse import unquote, urlsplit
|
||||
|
||||
_REDACTED: Final = "REDACTED"
|
||||
_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _secret_shape_redactor() -> Callable[[str], str]:
|
||||
try:
|
||||
from litellm._logging import redact_secrets
|
||||
except ImportError:
|
||||
return lambda text: text
|
||||
return redact_secrets
|
||||
|
||||
|
||||
def _url_passwords(url: str) -> frozenset[str]:
|
||||
try:
|
||||
parts: Final = urlsplit(url)
|
||||
except ValueError:
|
||||
return frozenset()
|
||||
query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
|
||||
raw_query_passwords: Final = tuple(
|
||||
value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
|
||||
)
|
||||
raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
|
||||
return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
|
||||
|
||||
|
||||
def _configured_database_passwords() -> frozenset[str]:
|
||||
database_url: Final = os.getenv("DATABASE_URL")
|
||||
direct_url: Final = os.getenv("DIRECT_URL")
|
||||
database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
|
||||
direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
|
||||
return database_passwords | direct_passwords
|
||||
|
||||
|
||||
def _redact_credentials(text: str) -> str:
|
||||
"""Mask configured database passwords before passing the text to LiteLLM redaction."""
|
||||
passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
|
||||
alternation: Final = "|".join(re.escape(password) for password in passwords)
|
||||
password_pattern: Final = (
|
||||
re.compile(rf"(?<![A-Za-z0-9_])(?:{alternation})(?![A-Za-z0-9_])", re.IGNORECASE) if passwords else None
|
||||
)
|
||||
result: Final = password_pattern.sub(_REDACTED, text) if password_pattern is not None else text
|
||||
return _secret_shape_redactor()(result)
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
import random
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
|
|
@ -7,6 +8,7 @@ from typing import TYPE_CHECKING, Final
|
|||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras._redaction import _redact_credentials
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
MIGRATION_LOCK_TIMEOUT_ENV_VAR,
|
||||
migration_lock_timeout,
|
||||
|
|
@ -14,6 +16,26 @@ from litellm_proxy_extras.prisma_toolchain import (
|
|||
|
||||
MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
|
||||
PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
|
||||
_SERVER_ANSWERED_MARKERS: Final = ("FATAL:", "SSL", "certificate")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AdvisoryLock:
|
||||
key: int
|
||||
name: str
|
||||
holder: str
|
||||
|
||||
|
||||
COORDINATOR_LOCK: Final = AdvisoryLock(
|
||||
MIGRATION_LOCK_KEY,
|
||||
"v2 migration coordinator lock",
|
||||
"another v2 migration resolver",
|
||||
)
|
||||
PRISMA_LOCK: Final = AdvisoryLock(
|
||||
PRISMA_ADVISORY_LOCK_KEY,
|
||||
"Prisma migration lock",
|
||||
"another Prisma migration",
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import psycopg
|
||||
|
|
@ -49,6 +71,28 @@ def _try_lock(connection: "psycopg.Connection[tuple[object, ...]]", key: int = M
|
|||
return row is not None and row.acquired
|
||||
|
||||
|
||||
def _coordination_failure(database_url: str, exc: "psycopg.Error", *, connecting: bool) -> str:
|
||||
import psycopg
|
||||
|
||||
prefix: Final = "Lost or could not establish v2 migration coordination with the database: "
|
||||
raw: Final = str(exc)
|
||||
unquoted: Final = re.sub(r"\"[^\"]*\"|'[^']*'", "", raw)
|
||||
detail: Final = _redact_credentials(raw).strip()
|
||||
if (
|
||||
not connecting
|
||||
or not isinstance(exc, psycopg.OperationalError)
|
||||
or any(marker in unquoted for marker in _SERVER_ANSWERED_MARKERS)
|
||||
):
|
||||
return prefix + detail
|
||||
parsed_url: Final = urlsplit(database_url)
|
||||
try:
|
||||
port: Final = parsed_url.port
|
||||
except ValueError:
|
||||
return prefix + f"P1001: Can't reach database server. {detail}"
|
||||
host: Final = parsed_url.hostname or "localhost"
|
||||
return prefix + f"P1001: Can't reach database server at `{host}:{port or 5432}`. {detail}"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MigrationCoordinator:
|
||||
connection: "psycopg.Connection[tuple[object, ...]]"
|
||||
|
|
@ -75,39 +119,45 @@ class MigrationCoordinator:
|
|||
|
||||
|
||||
@contextmanager
|
||||
def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, None]:
|
||||
def migration_lock(
|
||||
database_url: str,
|
||||
lock: AdvisoryLock = COORDINATOR_LOCK,
|
||||
) -> Generator[MigrationCoordinator, None, None]:
|
||||
import psycopg
|
||||
|
||||
wait_seconds: Final = migration_lock_timeout()
|
||||
deadline: Final = time.monotonic() + wait_seconds
|
||||
try:
|
||||
with psycopg.connect(database_url, connect_timeout=10, autocommit=True) as connection:
|
||||
connection: Final = psycopg.connect(database_url, connect_timeout=10, autocommit=True)
|
||||
except psycopg.Error as exc:
|
||||
raise RuntimeError(_coordination_failure(database_url, exc, connecting=True)) from None
|
||||
try:
|
||||
with connection:
|
||||
coordinator: Final = MigrationCoordinator(connection)
|
||||
logger.info("Waiting for the v2 migration coordinator lock (up to %ss)", wait_seconds)
|
||||
logger.info("Waiting for the %s (up to %ss)", lock.name, wait_seconds)
|
||||
while time.monotonic() < deadline:
|
||||
with connection.transaction():
|
||||
if _try_lock(connection):
|
||||
logger.info("Acquired the v2 migration coordinator lock")
|
||||
if _try_lock(connection, lock.key):
|
||||
logger.info("Acquired the %s", lock.name)
|
||||
|
||||
yield coordinator
|
||||
coordinator.check_connection()
|
||||
return
|
||||
time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
|
||||
except psycopg.Error as exc:
|
||||
raise RuntimeError(f"Lost or could not establish v2 migration coordination with the database: {exc}") from exc
|
||||
raise RuntimeError(_coordination_failure(database_url, exc, connecting=False)) from None
|
||||
raise RuntimeError(
|
||||
f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. "
|
||||
f"Timed out waiting for {lock.holder} after {wait_seconds}s. "
|
||||
f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
|
||||
"""Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one
|
||||
pinned connection while a Prisma command runs with its session advisory lock disabled, so a
|
||||
transaction pooler can never strand that lock on an idle backend after the command exits."""
|
||||
with migration_lock(database_url) as coordinator:
|
||||
coordinator.acquire_prisma_lock()
|
||||
"""Hold only Prisma's key as an xact lock on a pinned connection with Prisma's session lock disabled.
|
||||
|
||||
This keeps deploy independent of the coordinator key held by request-log index builds."""
|
||||
with migration_lock(database_url, PRISMA_LOCK) as coordinator:
|
||||
coordinator.release_statement_snapshot()
|
||||
yield
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,23 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
CREATE SEQUENCE IF NOT EXISTS "litellm_environment_variables_revision_seq";
|
||||
GRANT USAGE ON SEQUENCE "litellm_environment_variables_revision_seq" TO PUBLIC;
|
||||
|
||||
CREATE OR REPLACE FUNCTION litellm_bump_environment_variables_revision() RETURNS TRIGGER AS $$
|
||||
DECLARE
|
||||
next_revision BIGINT := nextval(format('%I.litellm_environment_variables_revision_seq', TG_TABLE_SCHEMA)::regclass);
|
||||
BEGIN
|
||||
IF TG_OP = 'UPDATE' THEN
|
||||
NEW."reload_revision" := GREATEST(next_revision, OLD."reload_revision" + 1);
|
||||
ELSE
|
||||
NEW."reload_revision" := next_revision;
|
||||
END IF;
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
DROP TRIGGER IF EXISTS "litellm_environment_variables_revision" ON "LiteLLM_Config";
|
||||
CREATE TRIGGER "litellm_environment_variables_revision"
|
||||
BEFORE INSERT OR UPDATE ON "LiteLLM_Config"
|
||||
FOR EACH ROW WHEN (NEW."param_name" = 'environment_variables')
|
||||
EXECUTE PROCEDURE litellm_bump_environment_variables_revision();
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
CREATE INDEX CONCURRENTLY IF NOT EXISTS "idx_background_interaction_settlement_settled_at"
|
||||
ON "LiteLLM_BackgroundInteractionSettlement" ("settled_at");
|
||||
|
|
@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement {
|
|||
outcome String?
|
||||
|
||||
@@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at")
|
||||
@@index([settled_at], map: "idx_background_interaction_settlement_settled_at")
|
||||
}
|
||||
|
||||
model LiteLLM_Lens {
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
import functools
|
||||
import glob
|
||||
import os
|
||||
import random
|
||||
|
|
@ -13,10 +12,10 @@ from contextlib import AbstractContextManager, nullcontext
|
|||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
|
||||
from urllib.parse import unquote, urlsplit
|
||||
|
||||
from litellm_proxy_extras import prisma_toolchain
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras._redaction import _redact_credentials
|
||||
from litellm_proxy_extras.migration_lock import held_migration_lock
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
|
|
@ -212,51 +211,6 @@ def _max_migration_timestamp(names) -> int:
|
|||
return max(_migration_timestamp(n) for n in names)
|
||||
|
||||
|
||||
_REDACTED: Final = "REDACTED"
|
||||
_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _secret_shape_redactor() -> Callable[[str], str]:
|
||||
try:
|
||||
from litellm._logging import redact_secrets
|
||||
except ImportError:
|
||||
return lambda text: text
|
||||
return redact_secrets
|
||||
|
||||
|
||||
def _url_passwords(url: str) -> frozenset[str]:
|
||||
try:
|
||||
parts: Final = urlsplit(url)
|
||||
except ValueError:
|
||||
return frozenset()
|
||||
query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
|
||||
raw_query_passwords: Final = tuple(
|
||||
value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
|
||||
)
|
||||
raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
|
||||
return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
|
||||
|
||||
|
||||
def _configured_database_passwords() -> frozenset[str]:
|
||||
database_url: Final = os.getenv("DATABASE_URL")
|
||||
direct_url: Final = os.getenv("DIRECT_URL")
|
||||
database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
|
||||
direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
|
||||
return database_passwords | direct_passwords
|
||||
|
||||
|
||||
def _redact_credentials(text: str) -> str:
|
||||
"""Mask configured database passwords before passing the text to LiteLLM redaction."""
|
||||
passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
|
||||
alternation: Final = "|".join(re.escape(password) for password in passwords)
|
||||
password_pattern: Final = (
|
||||
re.compile(rf"(?<![A-Za-z0-9_])(?:{alternation})(?![A-Za-z0-9_])", re.IGNORECASE) if passwords else None
|
||||
)
|
||||
result: Final = password_pattern.sub(_REDACTED, text) if password_pattern is not None else text
|
||||
return _secret_shape_redactor()(result)
|
||||
|
||||
|
||||
def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
|
||||
if isinstance(command, tuple):
|
||||
return tuple(_redact_credentials(str(argument)) for argument in command)
|
||||
|
|
@ -789,16 +743,23 @@ class ProxyExtrasDBManager:
|
|||
(connection_limit, pool_timeout, schema, pgbouncer, sslaccept, ...) and
|
||||
translate Prisma's TLS params back, since libpq reads ``sslcert`` as a
|
||||
client certificate where Prisma reads it as the CA."""
|
||||
from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
|
||||
from urllib.parse import parse_qsl, quote, unquote, urlencode, urlparse, urlunparse
|
||||
|
||||
parsed: Final = urlparse(url)
|
||||
netloc: Final = parsed.netloc
|
||||
userinfo, separator, hostport = netloc.rpartition("@")
|
||||
user, password_separator, password = userinfo.partition(":")
|
||||
canonical_userinfo: Final = quote(unquote(user), safe="") + (
|
||||
password_separator + quote(unquote(password), safe="") if password_separator else ""
|
||||
)
|
||||
canonical_netloc: Final = f"{canonical_userinfo}{separator}{hostport}" if separator else netloc
|
||||
if not parsed.query:
|
||||
return url
|
||||
return url if canonical_netloc == netloc else urlunparse(parsed._replace(netloc=canonical_netloc))
|
||||
pairs: Final = tuple(parse_qsl(parsed.query))
|
||||
kept: Final = tuple((k, v) for k, v in pairs if k in LIBPQ_URL_PARAMS)
|
||||
sslaccept: Final = next((v for k, v in pairs if k == "sslaccept"), None)
|
||||
libpq_pairs: Final = ProxyExtrasDBManager._libpq_tls_params(kept, sslaccept)
|
||||
return urlunparse(parsed._replace(query=urlencode(libpq_pairs, quote_via=quote)))
|
||||
return urlunparse(parsed._replace(netloc=canonical_netloc, query=urlencode(libpq_pairs, quote_via=quote)))
|
||||
|
||||
@staticmethod
|
||||
def _libpq_tls_params(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.109"
|
||||
version = "0.4.110"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.109"
|
||||
version = "0.4.110"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
1
litellm-rust/Cargo.lock
generated
1
litellm-rust/Cargo.lock
generated
|
|
@ -4083,6 +4083,7 @@ dependencies = [
|
|||
name = "litellm-inference-responses"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"litellm-auth",
|
||||
|
|
|
|||
|
|
@ -94,6 +94,7 @@ const OPENAI_COMPATIBLE_PROVIDERS: &[&str] = &[
|
|||
"docker_model_runner",
|
||||
"ragflow",
|
||||
"pinstripes",
|
||||
"openinfer",
|
||||
"darkbloom",
|
||||
"meta",
|
||||
"cognition",
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ const OPENROUTER_API_HEADERS: [&str; 2] = [
|
|||
"x-openrouter-experimental-metadata",
|
||||
];
|
||||
const ANTHROPIC_API_HEADER_PROVIDERS: &str =
|
||||
"anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter";
|
||||
"anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter,deepseek";
|
||||
|
||||
pub async fn create(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ use wiremock::{
|
|||
#[rstest]
|
||||
#[case::anthropic("anthropic/test-model", "/v1/messages", true)]
|
||||
#[case::azure("azure_ai/test-model", "/anthropic/v1/messages", false)]
|
||||
#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", false)]
|
||||
#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", true)]
|
||||
#[tokio::test]
|
||||
async fn messages_reaches_the_provider_and_preserves_json_or_sse(
|
||||
#[case] model: &str,
|
||||
|
|
|
|||
|
|
@ -6,9 +6,11 @@ use litellm_llms::{
|
|||
bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
|
||||
openai_like::chat::transformation::OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG,
|
||||
};
|
||||
use litellm_llms_types::formats::chat_completions::ChatMessage;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::Error;
|
||||
use crate::types::{ChatCompletionsCall, ResolvedChatCompletionsRequest};
|
||||
|
||||
const HEADER_CONTEXT: &str = "chat completions";
|
||||
|
||||
|
|
@ -42,3 +44,58 @@ pub(super) fn string_headers(
|
|||
) -> Result<Vec<(String, String)>, Error> {
|
||||
shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from)
|
||||
}
|
||||
|
||||
pub(super) fn resolve_request(
|
||||
call: ChatCompletionsCall,
|
||||
config: &'static dyn BaseConfig,
|
||||
) -> Result<ResolvedChatCompletionsRequest, Error> {
|
||||
let messages = serde_json::from_value::<Vec<ChatMessage>>(call.messages).map_err(|err| {
|
||||
Error::InvalidRequest(litellm_llms::ErrorDetail::invalid(
|
||||
"chat completions messages",
|
||||
err,
|
||||
))
|
||||
})?;
|
||||
if messages.is_empty() {
|
||||
return Err(Error::InvalidRequest(
|
||||
"chat completions requires at least one message".into(),
|
||||
));
|
||||
}
|
||||
if let Some(reason) = config.unsupported_reason(&messages, &call.optional_params) {
|
||||
return Err(Error::Unsupported(reason.0));
|
||||
}
|
||||
Ok(ResolvedChatCompletionsRequest {
|
||||
messages,
|
||||
optional_params: call.optional_params,
|
||||
api_key: call.api_key,
|
||||
api_base: call.api_base,
|
||||
extra_headers: call.extra_headers,
|
||||
timeout: call.timeout,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::chat_completions_provider;
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic(LlmProviders::Anthropic, ("anthropic-version", "2023-06-01"))]
|
||||
#[case::bedrock(LlmProviders::Bedrock, ("Content-Type", "application/json"))]
|
||||
fn resolves_supported_provider_configs(
|
||||
#[case] provider: LlmProviders,
|
||||
#[case] expected_header: (&str, &str),
|
||||
) {
|
||||
let config = chat_completions_provider(provider)
|
||||
.expect("provider has a chat completions config")
|
||||
.config();
|
||||
assert!(config.default_headers().contains(&expected_header));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::openai(LlmProviders::Openai)]
|
||||
fn rejects_providers_without_a_chat_config(#[case] provider: LlmProviders) {
|
||||
assert!(chat_completions_provider(provider).is_none());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,44 +3,66 @@ use std::time::Duration;
|
|||
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest};
|
||||
use litellm_http::{Client, outbound::OutboundRequest, request::truncate_error_body};
|
||||
use litellm_http::{
|
||||
Client,
|
||||
outbound::OutboundRequest,
|
||||
request::{truncate_error_body, with_default_headers},
|
||||
};
|
||||
use litellm_llms::base_llm::{
|
||||
auth::{Authenticated, resolve_auth},
|
||||
auth::{Authenticated, ValidatedEnvironment, resolve_auth},
|
||||
chat::transformation::BaseConfig,
|
||||
chat::transformation::ProviderChatResponseData,
|
||||
};
|
||||
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
|
||||
use litellm_secrets::source::Secrets;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::Error;
|
||||
use crate::{constants::CHAT_COMPLETIONS_TIMEOUT_SECS, types::ProviderChatCompletionsRequest};
|
||||
use crate::{
|
||||
common_utils::string_headers, constants::CHAT_COMPLETIONS_TIMEOUT_SECS,
|
||||
types::ResolvedChatCompletionsRequest,
|
||||
};
|
||||
|
||||
#[expect(clippy::too_many_arguments)] // mirrors BaseLLMHTTPHandler.completion, which takes each stage input separately
|
||||
pub(super) async fn execute(
|
||||
http: &Client,
|
||||
auth: &AuthServices,
|
||||
request: ProviderChatCompletionsRequest,
|
||||
config: &'static dyn BaseConfig,
|
||||
provider: litellm_inference::provider::ResolvedProvider<'_>,
|
||||
call: ResolvedChatCompletionsRequest,
|
||||
secrets: Secrets,
|
||||
cache: Option<litellm_cache_response::ScopedCache>,
|
||||
cache_options: Option<litellm_cache_response::CachePolicy>,
|
||||
interceptors: &impl Interceptors<Error>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
let ProviderChatCompletionsRequest {
|
||||
model,
|
||||
custom_llm_provider,
|
||||
config,
|
||||
url,
|
||||
body,
|
||||
optional_params,
|
||||
environment,
|
||||
secrets,
|
||||
timeout,
|
||||
api_key,
|
||||
} = request;
|
||||
let validated = config.validate_environment(
|
||||
string_headers(call.extra_headers)?,
|
||||
call.api_key.as_deref(),
|
||||
provider.model,
|
||||
&call.optional_params,
|
||||
&|key| secrets.get(key),
|
||||
)?;
|
||||
let url = config.get_complete_url(
|
||||
call.api_base.as_deref(),
|
||||
provider.model,
|
||||
&call.optional_params,
|
||||
&|key| secrets.get(key),
|
||||
)?;
|
||||
let body = config
|
||||
.transform_request(provider.model, call.messages, call.optional_params.clone())?
|
||||
.body;
|
||||
let environment = ValidatedEnvironment {
|
||||
headers: with_default_headers(validated.headers, config.default_headers()),
|
||||
auth: validated.auth,
|
||||
};
|
||||
let provider_name = <&'static str>::from(provider.provider);
|
||||
let context = RequestContext {
|
||||
model: model.clone(),
|
||||
custom_llm_provider,
|
||||
optional_params: Value::Object(optional_params),
|
||||
model: provider.model.to_owned(),
|
||||
custom_llm_provider: provider_name.to_owned(),
|
||||
optional_params: Value::Object(call.optional_params),
|
||||
secret_fields: Vec::new(),
|
||||
api_key,
|
||||
api_key: call.api_key.map(litellm_auth::SecretValue::new),
|
||||
};
|
||||
let authenticated = resolve_auth(auth, environment, &|key| secrets.get(key)).await?;
|
||||
let identity = litellm_host::interceptors::ProviderIdentity {
|
||||
|
|
@ -76,7 +98,7 @@ pub(super) async fn execute(
|
|||
},
|
||||
wire.url,
|
||||
&wire.body,
|
||||
timeout,
|
||||
call.timeout,
|
||||
)?;
|
||||
|
||||
let response = litellm_inference::outbound::send(outbound, http)
|
||||
|
|
@ -121,7 +143,7 @@ pub(super) async fn execute(
|
|||
))
|
||||
})?;
|
||||
config
|
||||
.transform_response(&model, ProviderChatResponseData { body })
|
||||
.transform_response(provider.model, ProviderChatResponseData { body })
|
||||
.map_err(Error::from)
|
||||
.map_err(as_response_error)
|
||||
},
|
||||
|
|
@ -134,7 +156,7 @@ pub(super) async fn execute(
|
|||
///
|
||||
/// A config reports the same variants on either side of the call: a missing
|
||||
/// field or an unsupported block can mean "this request cannot be translated"
|
||||
/// during prepare and "this response cannot be normalized" here. Only the
|
||||
/// before the request is sent and "this response cannot be normalized" here. Only the
|
||||
/// second kind has already been billed, and a host that keeps a reference
|
||||
/// implementation must not retry those, so collapse them to one variant that
|
||||
/// can only mean the provider was already called.
|
||||
|
|
@ -178,9 +200,10 @@ mod tests {
|
|||
|
||||
use super::*;
|
||||
use crate::{
|
||||
prepare::{prepare_provider_request, resolve_request},
|
||||
types::ChatCompletionsRequest,
|
||||
common_utils::{chat_completions_provider, resolve_request},
|
||||
types::{ChatCompletionsCall, ResolvedChatCompletionsRequest},
|
||||
};
|
||||
use litellm_inference::provider::resolve_llm_provider;
|
||||
|
||||
const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
|
||||
|
||||
|
|
@ -217,22 +240,33 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
fn prepared(api_base: &str) -> ProviderChatCompletionsRequest {
|
||||
prepare_provider_request(
|
||||
resolve_request(ChatCompletionsRequest {
|
||||
model: "anthropic/claude-sonnet-4-5",
|
||||
fn prepared(
|
||||
api_base: &str,
|
||||
) -> (
|
||||
litellm_inference::provider::ResolvedProvider<'static>,
|
||||
&'static dyn BaseConfig,
|
||||
ResolvedChatCompletionsRequest,
|
||||
) {
|
||||
let provider =
|
||||
resolve_llm_provider("anthropic/claude-sonnet-4-5", None, "chat completions").unwrap();
|
||||
let config = chat_completions_provider(provider.provider)
|
||||
.expect("provider has a chat completions config")
|
||||
.config();
|
||||
let request = resolve_request(
|
||||
ChatCompletionsCall {
|
||||
model: "anthropic/claude-sonnet-4-5".into(),
|
||||
messages: json!([{"role": "user", "content": "hi"}]),
|
||||
optional_params: json!({"max_tokens": 16}).as_object().unwrap().clone(),
|
||||
api_key: Some("sk-test"),
|
||||
api_base: Some(api_base),
|
||||
api_key: Some("sk-test".into()),
|
||||
api_base: Some(api_base.into()),
|
||||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
timeout: None,
|
||||
})
|
||||
.unwrap(),
|
||||
std::sync::Arc::new(|_: &str| None),
|
||||
},
|
||||
config,
|
||||
)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
(provider, config, request)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -247,10 +281,14 @@ mod tests {
|
|||
.await;
|
||||
let interceptors = RecordingHooks::default();
|
||||
|
||||
let (provider, config, request) = prepared(&upstream.uri());
|
||||
execute(
|
||||
&Client::plain_for_test(),
|
||||
&AuthServices::default(),
|
||||
prepared(&upstream.uri()),
|
||||
config,
|
||||
provider,
|
||||
request,
|
||||
std::sync::Arc::new(|_: &str| None),
|
||||
None,
|
||||
None,
|
||||
&interceptors,
|
||||
|
|
@ -288,10 +326,14 @@ mod tests {
|
|||
.await;
|
||||
let interceptors = RecordingHooks::default();
|
||||
|
||||
let (provider, config, request) = prepared(&upstream.uri());
|
||||
let error = execute(
|
||||
&Client::plain_for_test(),
|
||||
&AuthServices::default(),
|
||||
prepared(&upstream.uri()),
|
||||
config,
|
||||
provider,
|
||||
request,
|
||||
std::sync::Arc::new(|_: &str| None),
|
||||
None,
|
||||
None,
|
||||
&interceptors,
|
||||
|
|
|
|||
|
|
@ -5,14 +5,13 @@ pub use litellm_inference::RouteError as Error;
|
|||
mod common_utils;
|
||||
pub mod constants;
|
||||
pub(crate) mod handler;
|
||||
mod prepare;
|
||||
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
|
||||
use prepare::{prepare_provider_request, resolve_request};
|
||||
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_inference::provider::resolve_llm_provider;
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use std::sync::Arc;
|
||||
use types::ChatCompletionsRequest;
|
||||
use types::{ChatCompletionsCall, ChatCompletionsRequest};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ChatCompletionsRoute {
|
||||
|
|
@ -55,7 +54,7 @@ impl ChatCompletionsRoute {
|
|||
} = options.into();
|
||||
litellm_host::lifecycle::observe_unary(
|
||||
observers.clone(),
|
||||
self.run_call(
|
||||
self.run(
|
||||
request.into(),
|
||||
cache_options,
|
||||
interceptors,
|
||||
|
|
@ -65,30 +64,52 @@ impl ChatCompletionsRoute {
|
|||
.await
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "chat_completions",
|
||||
model = %call.model,
|
||||
provider,
|
||||
resolved_model,
|
||||
stream = false,
|
||||
outcome
|
||||
))]
|
||||
async fn run(
|
||||
&self,
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
call: ChatCompletionsCall,
|
||||
cache_options: Option<litellm_cache_response::CachePolicy>,
|
||||
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
let resolved = resolve_request(request)?;
|
||||
let snapshot = self
|
||||
.secrets
|
||||
.resolve(&resolved.config.secret_names())
|
||||
.await?;
|
||||
let prepared = prepare_provider_request(resolved, snapshot)?;
|
||||
litellm_inference::diagnostic::provider(&prepared.model, &prepared.custom_llm_provider);
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<ChatCompletionsResponse, Error>> =
|
||||
Box::pin(handler::execute(
|
||||
litellm_inference::diagnostic::unary(async {
|
||||
let model = call.model.clone();
|
||||
let custom_llm_provider = call.custom_llm_provider.clone();
|
||||
let provider =
|
||||
resolve_llm_provider(&model, custom_llm_provider.as_deref(), "chat completions")?;
|
||||
let config = common_utils::chat_completions_provider(provider.provider)
|
||||
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider.provider).to_string()))?
|
||||
.config();
|
||||
let call = common_utils::resolve_request(call, config)?;
|
||||
litellm_inference::diagnostic::provider(
|
||||
provider.model,
|
||||
<&'static str>::from(provider.provider),
|
||||
);
|
||||
let secrets = self.secrets.resolve(&config.secret_names()).await?;
|
||||
let execute: futures_util::future::BoxFuture<
|
||||
'_,
|
||||
Result<ChatCompletionsResponse, Error>,
|
||||
> = Box::pin(handler::execute(
|
||||
&self.http,
|
||||
&self.auth,
|
||||
prepared,
|
||||
config,
|
||||
provider,
|
||||
call,
|
||||
secrets,
|
||||
self.cache.clone(),
|
||||
cache_options,
|
||||
interceptors,
|
||||
observers,
|
||||
));
|
||||
execute.await
|
||||
execute.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,638 +0,0 @@
|
|||
use litellm_auth::SecretValue;
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_http::request::with_default_headers;
|
||||
use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig};
|
||||
use litellm_llms_types::formats::chat_completions::ChatMessage;
|
||||
use litellm_secrets::source::Secrets;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
Error,
|
||||
common_utils::{chat_completions_provider, string_headers},
|
||||
};
|
||||
use crate::types::{
|
||||
ChatCompletionsRequest, ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest,
|
||||
};
|
||||
use litellm_inference::provider::resolve_llm_provider;
|
||||
|
||||
pub(super) struct ResolvedProvider {
|
||||
pub(super) model: String,
|
||||
pub(super) custom_llm_provider: String,
|
||||
pub(super) config: &'static dyn BaseConfig,
|
||||
}
|
||||
|
||||
pub(super) fn resolve_provider_config<'a>(
|
||||
model: &'a str,
|
||||
custom_llm_provider: Option<&'a str>,
|
||||
) -> Result<ResolvedProvider, Error> {
|
||||
let provider_info = resolve_llm_provider(model, custom_llm_provider, "chat completions")?;
|
||||
let config = chat_completions_provider(provider_info.provider)
|
||||
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider_info.provider).to_string()))?
|
||||
.config();
|
||||
Ok(ResolvedProvider {
|
||||
model: provider_info.model.to_string(),
|
||||
custom_llm_provider: <&str>::from(provider_info.provider).to_string(),
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn parse_messages(messages: Value) -> Result<Vec<ChatMessage>, Error> {
|
||||
serde_json::from_value(messages).map_err(|err| {
|
||||
Error::InvalidRequest(litellm_llms::ErrorDetail::invalid(
|
||||
"chat completions messages",
|
||||
err,
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn resolve_request(
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
) -> Result<ResolvedChatCompletionsRequest<'_>, Error> {
|
||||
let ResolvedProvider {
|
||||
model,
|
||||
custom_llm_provider,
|
||||
config,
|
||||
} = resolve_provider_config(request.model, request.custom_llm_provider)?;
|
||||
let messages = parse_messages(request.messages)?;
|
||||
if messages.is_empty() {
|
||||
return Err(Error::InvalidRequest(
|
||||
"chat completions requires at least one message".into(),
|
||||
));
|
||||
}
|
||||
if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) {
|
||||
return Err(Error::Unsupported(reason.0));
|
||||
}
|
||||
Ok(ResolvedChatCompletionsRequest {
|
||||
model,
|
||||
custom_llm_provider,
|
||||
config,
|
||||
messages,
|
||||
optional_params: request.optional_params,
|
||||
api_key: request.api_key,
|
||||
api_base: request.api_base,
|
||||
extra_headers: request.extra_headers,
|
||||
timeout: request.timeout,
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_environment(
|
||||
request: &ResolvedChatCompletionsRequest<'_>,
|
||||
model: &str,
|
||||
config: &dyn BaseConfig,
|
||||
secrets: &dyn Lookup,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
let env_lookup = |key: &str| secrets.get(key);
|
||||
let forwarded = string_headers(request.extra_headers.clone())?;
|
||||
let validated = config.validate_environment(
|
||||
forwarded,
|
||||
request.api_key,
|
||||
model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
Ok(ValidatedEnvironment {
|
||||
headers: with_default_headers(validated.headers, config.default_headers()),
|
||||
auth: validated.auth,
|
||||
})
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
|
||||
pub(super) fn prepare_provider_request(
|
||||
request: ResolvedChatCompletionsRequest<'_>,
|
||||
secrets: Secrets,
|
||||
) -> Result<ProviderChatCompletionsRequest, Error> {
|
||||
let environment =
|
||||
validate_environment(&request, &request.model, request.config, secrets.as_ref())?;
|
||||
let model = request.model;
|
||||
let config = request.config;
|
||||
let env_lookup = |key: &str| secrets.get(key);
|
||||
let url = config.get_complete_url(
|
||||
request.api_base,
|
||||
&model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let transformed =
|
||||
config.transform_request(&model, request.messages, request.optional_params.clone())?;
|
||||
|
||||
Ok(ProviderChatCompletionsRequest {
|
||||
model,
|
||||
custom_llm_provider: request.custom_llm_provider,
|
||||
config,
|
||||
url,
|
||||
body: transformed.body,
|
||||
optional_params: request.optional_params,
|
||||
environment,
|
||||
secrets,
|
||||
timeout: request.timeout,
|
||||
api_key: request.api_key.map(|key| SecretValue::new(key.to_string())),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_auth::CredentialPlacement;
|
||||
use litellm_llms::base_llm::auth::{AuthScheme, resolve_auth};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use super::{prepare_provider_request, resolve_request};
|
||||
use crate::{
|
||||
Error,
|
||||
types::{ChatCompletionsRequest, ProviderChatCompletionsRequest},
|
||||
};
|
||||
|
||||
fn prepare_chat_completions_call(
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
) -> Result<ProviderChatCompletionsRequest, Error> {
|
||||
prepare_provider_request(
|
||||
resolve_request(request)?,
|
||||
std::sync::Arc::new(|_: &str| None),
|
||||
)
|
||||
}
|
||||
|
||||
/// The headers as they go on the wire, credential applied.
|
||||
fn wire_headers(prepared: &ProviderChatCompletionsRequest) -> Vec<(String, String)> {
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.build()
|
||||
.unwrap()
|
||||
.block_on(resolve_auth(
|
||||
&litellm_auth::AuthServices::default(),
|
||||
prepared.environment.clone(),
|
||||
&|_| None,
|
||||
))
|
||||
.unwrap()
|
||||
.headers
|
||||
}
|
||||
|
||||
fn request<'a>(
|
||||
model: &'a str,
|
||||
provider: Option<&'a str>,
|
||||
messages: Value,
|
||||
optional_params: Value,
|
||||
) -> ChatCompletionsRequest<'a> {
|
||||
ChatCompletionsRequest {
|
||||
model,
|
||||
messages,
|
||||
optional_params: match optional_params {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("params must be an object, got {other}"),
|
||||
},
|
||||
api_key: Some("sk-test"),
|
||||
api_base: None,
|
||||
custom_llm_provider: provider,
|
||||
extra_headers: None,
|
||||
timeout: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn preparation_error(request: ChatCompletionsRequest<'_>) -> Error {
|
||||
prepare_chat_completions_call(request)
|
||||
.err()
|
||||
.expect("request preparation should fail")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_the_provider_from_the_model_prefix() {
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert_eq!(prepared.model, "claude-sonnet-4-5");
|
||||
assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages");
|
||||
assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_an_explicit_provider_prefix_from_the_model() {
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert_eq!(prepared.model, "claude-sonnet-4-5");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn adds_the_auth_and_default_headers() {
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert!(
|
||||
wire_headers(&prepared).contains(&("x-api-key".to_string(), "sk-test".to_string()))
|
||||
);
|
||||
assert!(
|
||||
wire_headers(&prepared)
|
||||
.contains(&("anthropic-version".to_string(), "2023-06-01".to_string()))
|
||||
);
|
||||
assert!(matches!(
|
||||
prepared.environment.auth,
|
||||
AuthScheme::Credential {
|
||||
placement: CredentialPlacement::Header("x-api-key"),
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_deployment_credential_replaces_a_caller_supplied_auth_header() {
|
||||
// Python builds `{**headers, **anthropic_headers}`, so the deployment's key
|
||||
// overwrites a forwarded one. Honouring the caller's would let whoever sends
|
||||
// the request choose the Anthropic principal it bills to.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"X-Api-Key".to_string(),
|
||||
json!("sk-caller"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let headers = wire_headers(&prepared);
|
||||
let keys: Vec<_> = headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
|
||||
.collect();
|
||||
assert_eq!(keys.len(), 1, "got {:?}", headers);
|
||||
assert_eq!(keys[0].1, "sk-test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() {
|
||||
// Anthropic's `validate_environment` pops `x-api-key` and sets `authorization`
|
||||
// for an OAuth token, so re-adding the key here would put the credential into
|
||||
// a header the host removed on purpose.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([
|
||||
(
|
||||
"Authorization".to_string(),
|
||||
json!("Bearer sk-ant-oat01-token"),
|
||||
),
|
||||
("X-Api-Key".to_string(), json!("sk-caller")),
|
||||
]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
assert!(
|
||||
!wire_headers(&prepared)
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"),
|
||||
"the resolved key must not be applied over an OAuth bearer, got {:?}",
|
||||
wire_headers(&prepared)
|
||||
);
|
||||
assert!(
|
||||
wire_headers(&prepared)
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer sk-ant-oat01-token")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() {
|
||||
// Only an OAuth bearer replaces the credential. Python sends the deployment's
|
||||
// `x-api-key` alongside any other forwarded `authorization`, so deferring on
|
||||
// the mere presence of that header would drop the deployment's auth.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([
|
||||
("Authorization".to_string(), json!("Bearer unrelated")),
|
||||
("X-Api-Key".to_string(), json!("sk-caller")),
|
||||
]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let headers = wire_headers(&prepared);
|
||||
let keys: Vec<_> = headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
|
||||
.collect();
|
||||
assert_eq!(keys.len(), 1, "got {:?}", headers);
|
||||
assert_eq!(keys[0].1, "sk-test");
|
||||
assert!(
|
||||
wire_headers(&prepared)
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer unrelated"),
|
||||
"the unrelated authorization must survive, got {:?}",
|
||||
wire_headers(&prepared)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
fn rejects_empty_messages_before_resolving_credentials() {
|
||||
let call = ChatCompletionsRequest {
|
||||
api_key: None,
|
||||
..request("claude-sonnet-4-5", Some("anthropic"), json!([]), json!({}))
|
||||
};
|
||||
assert!(matches!(preparation_error(call), Error::InvalidRequest(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_an_unknown_provider() {
|
||||
assert_eq!(
|
||||
preparation_error(request(
|
||||
"openai/gpt-4o",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)),
|
||||
Error::InvalidProvider("openai".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_a_model_with_no_resolvable_provider() {
|
||||
assert!(matches!(
|
||||
preparation_error(request(
|
||||
"claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)),
|
||||
Error::InvalidProvider(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_an_empty_or_malformed_message_list() {
|
||||
assert_eq!(
|
||||
preparation_error(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([]),
|
||||
json!({}),
|
||||
)),
|
||||
Error::InvalidRequest(
|
||||
"chat completions requires at least one message"
|
||||
.to_string()
|
||||
.into()
|
||||
)
|
||||
);
|
||||
assert!(matches!(
|
||||
preparation_error(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!("not a list"),
|
||||
json!({}),
|
||||
)),
|
||||
Error::InvalidRequest(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_string_extra_headers() {
|
||||
let mut call = request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
|
||||
assert_eq!(
|
||||
preparation_error(call),
|
||||
Error::Headers(litellm_http::request::HeaderError {
|
||||
context: "chat completions",
|
||||
name: "x-trace".to_string(),
|
||||
actual: "number",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepares_a_bedrock_call_without_resolving_credentials() {
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"maxTokens": 16}),
|
||||
);
|
||||
call.api_key = None;
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
assert_eq!(
|
||||
prepared.url,
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse"
|
||||
);
|
||||
assert!(matches!(
|
||||
&prepared.environment.auth,
|
||||
AuthScheme::AwsSigV4 { region, service: "bedrock", .. } if region == "us-east-1"
|
||||
));
|
||||
// SigV4 signs the serialized body, so prepare must not have added an
|
||||
// Authorization header; the signer does it over the bytes sent.
|
||||
assert!(
|
||||
!prepared
|
||||
.environment
|
||||
.headers
|
||||
.iter()
|
||||
.any(|(name, _)| name.eq_ignore_ascii_case("authorization"))
|
||||
);
|
||||
assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
|
||||
// Python signs only the AWS header set and reattaches the rest, so a header
|
||||
// the caller forwarded rides along without joining the canonical request.
|
||||
// Signing it makes Converse 403 on a deployment that works on Python.
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"maxTokens": 16,
|
||||
"aws_access_key_id": "AKIDEXAMPLE",
|
||||
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
|
||||
}),
|
||||
);
|
||||
// A key would resolve to a bearer token and never reach the signer.
|
||||
call.api_key = None;
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"x-request-id".to_string(),
|
||||
json!("abc-123"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let authenticated = resolve_auth(
|
||||
&litellm_auth::AuthServices::default(),
|
||||
prepared.environment,
|
||||
&|_| None,
|
||||
)
|
||||
.await
|
||||
.expect("resolves");
|
||||
let signed = crate::handler::outbound_request(
|
||||
authenticated,
|
||||
prepared.url,
|
||||
&prepared.body,
|
||||
prepared.timeout,
|
||||
)
|
||||
.expect("signs");
|
||||
|
||||
let authorization = signed
|
||||
.header("authorization")
|
||||
.expect("carries an authorization header")
|
||||
.to_string();
|
||||
assert!(
|
||||
authorization.starts_with("AWS4-HMAC-SHA256"),
|
||||
"expected a SigV4 signature, got {authorization}"
|
||||
);
|
||||
assert!(
|
||||
!authorization.contains("x-request-id"),
|
||||
"forwarded header reached SignedHeaders: {authorization}"
|
||||
);
|
||||
// It still goes on the wire, it is just not part of the signature.
|
||||
assert!(
|
||||
signed
|
||||
.headers()
|
||||
.iter()
|
||||
.any(|(name, value)| name == "x-request-id" && value == "abc-123"),
|
||||
"forwarded header was dropped instead of reattached"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::authorization("Authorization")]
|
||||
#[case::amz_date("x-amz-date")]
|
||||
#[case::security_token("x-amz-security-token")]
|
||||
#[case::date("Date")]
|
||||
#[tokio::test]
|
||||
async fn rejects_a_forwarded_header_the_signer_computes(#[case] forwarded: &str) {
|
||||
let call = ChatCompletionsRequest {
|
||||
api_key: None,
|
||||
extra_headers: Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])),
|
||||
..request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"maxTokens": 16,
|
||||
"aws_access_key_id": "AKIDEXAMPLE",
|
||||
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
|
||||
}),
|
||||
)
|
||||
};
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let authenticated = resolve_auth(
|
||||
&litellm_auth::AuthServices::default(),
|
||||
prepared.environment,
|
||||
&|_| None,
|
||||
)
|
||||
.await
|
||||
.expect("resolves");
|
||||
let error = crate::handler::outbound_request(
|
||||
authenticated,
|
||||
prepared.url,
|
||||
&prepared.body,
|
||||
prepared.timeout,
|
||||
)
|
||||
.expect_err("conflicting signing headers must fail");
|
||||
assert!(
|
||||
matches!(error, Error::Unsupported(_)),
|
||||
"{forwarded} returned {error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
|
||||
// `get_request_headers` assigns `headers["Authorization"]` unconditionally
|
||||
// once a bearer token resolves, so the deployment's identity wins on
|
||||
// Python. Keeping the caller's would authorize and bill the call as a
|
||||
// different principal, and only when the deployment carries `rust: true`.
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"maxTokens": 16}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"Authorization".to_string(),
|
||||
json!("Bearer caller-supplied"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let headers = wire_headers(&prepared);
|
||||
let authorizations: Vec<_> = headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("authorization"))
|
||||
.map(|(_, value)| value.as_str())
|
||||
.collect();
|
||||
assert_eq!(
|
||||
authorizations,
|
||||
vec!["Bearer sk-test"],
|
||||
"the deployment token must be the only authorization on the wire"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
|
||||
// The opposite precedence, and deliberate: Anthropic's own transform
|
||||
// honours a forwarded OAuth bearer, so the Bedrock fix above must not be
|
||||
// generalized into a rule that the configured key always wins.
|
||||
//
|
||||
// An OAuth bearer is the whole of that exception. This forwarded a plain
|
||||
// `x-api-key` until round 17, which read as the same claim and was not:
|
||||
// Python overwrites a forwarded `x-api-key` with the deployment's.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"authorization".to_string(),
|
||||
json!("Bearer sk-ant-oat01-forwarded"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let headers = wire_headers(&prepared);
|
||||
let keys: Vec<_> = headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
|
||||
.map(|(_, value)| value.as_str())
|
||||
.collect();
|
||||
assert!(keys.is_empty(), "got {:?}", headers);
|
||||
assert!(
|
||||
wire_headers(&prepared)
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer sk-ant-oat01-forwarded")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
|
||||
// The configured bearer identity has its own account and quota boundary,
|
||||
// so a request carrying one must not be signed as whatever principal the
|
||||
// host's AWS credentials resolve to.
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"maxTokens": 16}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert!(matches!(
|
||||
&prepared.environment.auth,
|
||||
AuthScheme::Credential { placement: CredentialPlacement::Bearer, secret }
|
||||
if secret.expose() == "sk-test"
|
||||
));
|
||||
assert!(
|
||||
wire_headers(&prepared)
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer sk-test"),
|
||||
"prepare did not carry the bearer token"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,4 +1,3 @@
|
|||
use litellm_host::observation::ObservationSender;
|
||||
use std::convert::Infallible;
|
||||
|
||||
use litellm_host::{
|
||||
|
|
@ -7,10 +6,7 @@ use litellm_host::{
|
|||
};
|
||||
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
|
||||
|
||||
use super::{
|
||||
ChatCompletionsRoute, Error,
|
||||
types::{ChatCompletionsCall, ChatCompletionsRequest},
|
||||
};
|
||||
use super::{ChatCompletionsRoute, Error, types::ChatCompletionsCall};
|
||||
|
||||
pub struct ChatCompletions;
|
||||
|
||||
|
|
@ -37,44 +33,12 @@ impl ChatCompletionsRoute {
|
|||
call,
|
||||
observers,
|
||||
move |call, _, interceptors, observers| async move {
|
||||
self.run_call(call, cache_options, &interceptors, observers.as_ref())
|
||||
self.run(call, cache_options, &interceptors, observers.as_ref())
|
||||
.await
|
||||
.map(CallOutput::Complete)
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "chat_completions",
|
||||
model = %call.model,
|
||||
provider,
|
||||
resolved_model,
|
||||
stream = false,
|
||||
outcome
|
||||
))]
|
||||
pub(super) async fn run_call(
|
||||
&self,
|
||||
call: ChatCompletionsCall,
|
||||
cache_options: Option<litellm_cache_response::CachePolicy>,
|
||||
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
litellm_inference::diagnostic::unary(async {
|
||||
let request = ChatCompletionsRequest {
|
||||
model: &call.model,
|
||||
messages: call.messages,
|
||||
optional_params: call.optional_params,
|
||||
api_key: call.api_key.as_deref(),
|
||||
api_base: call.api_base.as_deref(),
|
||||
custom_llm_provider: call.custom_llm_provider.as_deref(),
|
||||
extra_headers: call.extra_headers,
|
||||
timeout: call.timeout,
|
||||
};
|
||||
self.run(request, cache_options, interceptors, observers)
|
||||
.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl litellm_inference::caching::Cachable for ChatCompletions {
|
||||
|
|
|
|||
|
|
@ -1,8 +1,5 @@
|
|||
use litellm_secrets::source::Secrets;
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_auth::SecretValue;
|
||||
use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig};
|
||||
use litellm_llms_types::formats::chat_completions::ChatMessage;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
|
@ -49,31 +46,11 @@ impl From<ChatCompletionsRequest<'_>> for ChatCompletionsCall {
|
|||
}
|
||||
}
|
||||
|
||||
pub struct ResolvedChatCompletionsRequest<'a> {
|
||||
pub model: String,
|
||||
pub custom_llm_provider: String,
|
||||
pub config: &'static dyn BaseConfig,
|
||||
pub struct ResolvedChatCompletionsRequest {
|
||||
pub messages: Vec<ChatMessage>,
|
||||
pub optional_params: Map<String, Value>,
|
||||
pub api_key: Option<&'a str>,
|
||||
pub api_base: Option<&'a str>,
|
||||
pub api_key: Option<String>,
|
||||
pub api_base: Option<String>,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
pub struct ProviderChatCompletionsRequest {
|
||||
pub model: String,
|
||||
pub custom_llm_provider: String,
|
||||
pub config: &'static dyn BaseConfig,
|
||||
pub url: String,
|
||||
pub body: Value,
|
||||
/// The route's parameters before the provider transformation, reported to the host
|
||||
/// beside the wire request.
|
||||
pub optional_params: Map<String, Value>,
|
||||
/// The forwarded and default headers plus how the call authenticates; the credential
|
||||
/// itself is applied when the request is sent.
|
||||
pub environment: ValidatedEnvironment,
|
||||
pub secrets: Secrets,
|
||||
pub timeout: Option<Duration>,
|
||||
pub api_key: Option<SecretValue>,
|
||||
}
|
||||
|
|
|
|||
439
litellm-rust/crates/inference-chat/tests/request.rs
Normal file
439
litellm-rust/crates/inference-chat/tests/request.rs
Normal file
|
|
@ -0,0 +1,439 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_inference_chat::{Error, types::ChatCompletionsRequest};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Map, Value, json};
|
||||
use wiremock::ResponseTemplate;
|
||||
|
||||
mod support;
|
||||
use support::*;
|
||||
|
||||
const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
|
||||
|
||||
fn request<'a>(
|
||||
model: &'a str,
|
||||
provider: Option<&'a str>,
|
||||
messages: Value,
|
||||
optional_params: Value,
|
||||
) -> ChatCompletionsRequest<'a> {
|
||||
let Value::Object(optional_params) = optional_params else {
|
||||
panic!("params must be an object");
|
||||
};
|
||||
ChatCompletionsRequest {
|
||||
model,
|
||||
messages,
|
||||
optional_params,
|
||||
api_key: Some("sk-test"),
|
||||
api_base: None,
|
||||
custom_llm_provider: provider,
|
||||
extra_headers: None,
|
||||
timeout: Some(Duration::from_secs(10)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn complete(
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
) -> Result<litellm_llms_types::formats::chat_completions::ChatCompletionsResponse, Error> {
|
||||
chat_completions_route().execute(request, &(), None).await
|
||||
}
|
||||
|
||||
fn object(value: Value) -> Map<String, Value> {
|
||||
let Value::Object(map) = value else {
|
||||
panic!("expected an object");
|
||||
};
|
||||
map
|
||||
}
|
||||
|
||||
fn anthropic_response() -> ResponseTemplate {
|
||||
ResponseTemplate::new(200).set_body_raw(ANTHROPIC_MESSAGE, "application/json")
|
||||
}
|
||||
|
||||
fn bedrock_response() -> ResponseTemplate {
|
||||
json_response(json!({
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}
|
||||
}))
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn resolves_the_provider_from_the_model_prefix() {
|
||||
let upstream = upstream([anthropic_response()]).await;
|
||||
let base = upstream.uri();
|
||||
complete(ChatCompletionsRequest {
|
||||
api_base: Some(&base),
|
||||
..request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
)
|
||||
})
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
let sent = only_request(&upstream).await;
|
||||
assert_eq!(sent.url.path(), "/v1/messages");
|
||||
assert_eq!(sent.json()["model"], "claude-sonnet-4-5");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn strips_an_explicit_provider_prefix_from_the_model() {
|
||||
let upstream = upstream([anthropic_response()]).await;
|
||||
let base = upstream.uri();
|
||||
complete(ChatCompletionsRequest {
|
||||
api_base: Some(&base),
|
||||
..request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)
|
||||
})
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
assert_eq!(
|
||||
only_request(&upstream).await.json()["model"],
|
||||
"claude-sonnet-4-5"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn adds_the_auth_and_default_headers() {
|
||||
let upstream = upstream([anthropic_response()]).await;
|
||||
let base = upstream.uri();
|
||||
complete(ChatCompletionsRequest {
|
||||
api_base: Some(&base),
|
||||
..request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)
|
||||
})
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
let sent = only_request(&upstream).await;
|
||||
assert_eq!(sent.header_values("x-api-key"), ["sk-test"]);
|
||||
assert_eq!(sent.header("anthropic-version"), Some("2023-06-01"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_deployment_credential_replaces_a_caller_supplied_auth_header() {
|
||||
let upstream = upstream([anthropic_response()]).await;
|
||||
let base = upstream.uri();
|
||||
complete(ChatCompletionsRequest {
|
||||
api_base: Some(&base),
|
||||
extra_headers: Some(object(json!({"X-Api-Key": "sk-caller"}))),
|
||||
..request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)
|
||||
})
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
assert_eq!(
|
||||
only_request(&upstream).await.header_values("x-api-key"),
|
||||
["sk-test"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() {
|
||||
let upstream = upstream([anthropic_response()]).await;
|
||||
let base = upstream.uri();
|
||||
complete(ChatCompletionsRequest {
|
||||
api_base: Some(&base),
|
||||
extra_headers: Some(object(json!({
|
||||
"Authorization": "Bearer sk-ant-oat01-token",
|
||||
"X-Api-Key": "sk-caller"
|
||||
}))),
|
||||
..request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)
|
||||
})
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
let sent = only_request(&upstream).await;
|
||||
assert_eq!(sent.header_values("x-api-key"), ["sk-caller"]);
|
||||
assert_eq!(
|
||||
sent.header_values("authorization"),
|
||||
["Bearer sk-ant-oat01-token"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() {
|
||||
let upstream = upstream([anthropic_response()]).await;
|
||||
let base = upstream.uri();
|
||||
complete(ChatCompletionsRequest {
|
||||
api_base: Some(&base),
|
||||
extra_headers: Some(object(json!({
|
||||
"Authorization": "Bearer unrelated",
|
||||
"X-Api-Key": "sk-caller"
|
||||
}))),
|
||||
..request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)
|
||||
})
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
let sent = only_request(&upstream).await;
|
||||
assert_eq!(sent.header_values("x-api-key"), ["sk-test"]);
|
||||
assert_eq!(sent.header_values("authorization"), ["Bearer unrelated"]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rejects_empty_messages_before_resolving_credentials() {
|
||||
let error = complete(ChatCompletionsRequest {
|
||||
api_key: None,
|
||||
..request("claude-sonnet-4-5", Some("anthropic"), json!([]), json!({}))
|
||||
})
|
||||
.await
|
||||
.expect_err("empty messages are rejected");
|
||||
assert!(matches!(error, Error::InvalidRequest(_)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rejects_an_unknown_provider() {
|
||||
assert_eq!(
|
||||
complete(request(
|
||||
"openai/gpt-4o",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
))
|
||||
.await
|
||||
.unwrap_err(),
|
||||
Error::InvalidProvider("openai".into())
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rejects_a_model_with_no_resolvable_provider() {
|
||||
assert!(matches!(
|
||||
complete(request(
|
||||
"claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
))
|
||||
.await,
|
||||
Err(Error::InvalidProvider(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty(json!([]), Some("chat completions requires at least one message"))]
|
||||
#[case::malformed(json!("not a list"), None)]
|
||||
#[tokio::test]
|
||||
async fn rejects_an_empty_or_malformed_message_list(
|
||||
#[case] messages: Value,
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
let error = complete(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
messages,
|
||||
json!({}),
|
||||
))
|
||||
.await
|
||||
.unwrap_err();
|
||||
if let Some(expected) = expected {
|
||||
assert_eq!(error, Error::InvalidRequest(expected.to_string().into()));
|
||||
} else {
|
||||
assert!(matches!(error, Error::InvalidRequest(_)));
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rejects_non_string_extra_headers() {
|
||||
assert_eq!(
|
||||
complete(ChatCompletionsRequest {
|
||||
extra_headers: Some(Map::from_iter([("x-trace".to_string(), json!(7))])),
|
||||
..request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)
|
||||
})
|
||||
.await
|
||||
.unwrap_err(),
|
||||
Error::Headers(litellm_http::request::HeaderError {
|
||||
context: "chat completions",
|
||||
name: "x-trace".to_string(),
|
||||
actual: "number",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
|
||||
let upstream = upstream([bedrock_response()]).await;
|
||||
let base = upstream.uri();
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"maxTokens": 16,
|
||||
"aws_access_key_id": "AKIDEXAMPLE",
|
||||
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
|
||||
}),
|
||||
);
|
||||
call.api_key = None;
|
||||
call.api_base = Some(&base);
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"x-request-id".to_string(),
|
||||
json!("abc-123"),
|
||||
)]));
|
||||
complete(call).await.expect("call succeeds");
|
||||
|
||||
let sent = only_request(&upstream).await;
|
||||
let authorization = sent.header("authorization").expect("signed request");
|
||||
assert!(authorization.starts_with("AWS4-HMAC-SHA256"));
|
||||
assert!(!authorization.contains("x-request-id"));
|
||||
assert_eq!(sent.header("x-request-id"), Some("abc-123"));
|
||||
assert_eq!(sent.url.path(), "/model/anthropic.claude-v2/converse");
|
||||
let body: Value = serde_json::from_slice(&sent.body).unwrap();
|
||||
assert_eq!(body["inferenceConfig"], json!({"maxTokens": 16}));
|
||||
assert_eq!(
|
||||
body["messages"],
|
||||
json!([{"role": "user", "content": [{"text": "hi"}]}])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::authorization("Authorization")]
|
||||
#[case::amz_date("x-amz-date")]
|
||||
#[case::security_token("x-amz-security-token")]
|
||||
#[case::date("Date")]
|
||||
#[tokio::test]
|
||||
async fn rejects_a_forwarded_header_the_signer_computes(#[case] forwarded: &str) {
|
||||
let upstream = wiremock::MockServer::start().await;
|
||||
let base = upstream.uri();
|
||||
let call = ChatCompletionsRequest {
|
||||
api_key: None,
|
||||
api_base: Some(&base),
|
||||
extra_headers: Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])),
|
||||
..request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"maxTokens": 16,
|
||||
"aws_access_key_id": "AKIDEXAMPLE",
|
||||
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
|
||||
}),
|
||||
)
|
||||
};
|
||||
let error = complete(call)
|
||||
.await
|
||||
.expect_err("conflicting signing headers fail");
|
||||
assert!(matches!(error, Error::Unsupported(_)), "{error:?}");
|
||||
assert!(received(&upstream).await.is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
|
||||
let upstream = upstream([bedrock_response()]).await;
|
||||
let base = upstream.uri();
|
||||
complete(ChatCompletionsRequest {
|
||||
api_base: Some(&base),
|
||||
extra_headers: Some(Map::from_iter([(
|
||||
"Authorization".to_string(),
|
||||
json!("Bearer caller-supplied"),
|
||||
)])),
|
||||
..request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"maxTokens": 16}),
|
||||
)
|
||||
})
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
assert_eq!(
|
||||
only_request(&upstream).await.header_values("authorization"),
|
||||
["Bearer sk-test"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
|
||||
let upstream = upstream([anthropic_response()]).await;
|
||||
let base = upstream.uri();
|
||||
complete(ChatCompletionsRequest {
|
||||
api_base: Some(&base),
|
||||
extra_headers: Some(Map::from_iter([(
|
||||
"authorization".to_string(),
|
||||
json!("Bearer sk-ant-oat01-forwarded"),
|
||||
)])),
|
||||
..request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)
|
||||
})
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
let sent = only_request(&upstream).await;
|
||||
assert!(sent.header_values("x-api-key").is_empty());
|
||||
assert_eq!(
|
||||
sent.header_values("authorization"),
|
||||
["Bearer sk-ant-oat01-forwarded"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
|
||||
let upstream = upstream([bedrock_response()]).await;
|
||||
let base = upstream.uri();
|
||||
complete(ChatCompletionsRequest {
|
||||
api_base: Some(&base),
|
||||
..request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"maxTokens": 16}),
|
||||
)
|
||||
})
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
assert_eq!(
|
||||
only_request(&upstream).await.header_values("authorization"),
|
||||
["Bearer sk-test"]
|
||||
);
|
||||
}
|
||||
|
|
@ -7,8 +7,8 @@
|
|||
- delegate authentication policy, beta selection, payload rewriting, and response interpretation to those adapters
|
||||
- keep provider policy out of request preparation and transport handlers
|
||||
- calling a concrete provider helper for every provider is still a policy dependency
|
||||
- Route types such as `MessagesCall`, prepared requests, and response wrappers containing live streams describe execution. Reuse the shared Messages payload types inside them instead of defining another request or response schema here
|
||||
- Route types such as `MessagesCall`, provider calls, and response wrappers containing live streams describe execution. Reuse the shared Messages payload types inside them instead of defining another request or response schema here
|
||||
- The route returns `litellm_host::call::CallOutput`: a completed response, or a stream head and chunks
|
||||
- Per-call dependencies are grouped in `litellm_inference::context::CallContext`. `src/lib.rs` explicitly sequences cache lookup, provider execution, result acceptance (`CallContext::result_ready`), and cache storage
|
||||
- Per-call dependencies are grouped in `litellm_inference::context::CallContext`. `handler::execute` explicitly sequences the provider call build in Python's handler order, cache lookup, provider execution, result acceptance (`CallContext::result_ready`), and cache storage
|
||||
- Preserve the order of validation, normalization, caller-requested parameter removal, and provider transformation when that order affects observable behavior
|
||||
- Test provider dispatch, auth precedence, header handling, transformations, and responses through behavior, not source structure
|
||||
|
|
|
|||
|
|
@ -82,17 +82,17 @@ mod tests {
|
|||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
|
||||
#[rstest]
|
||||
#[case::tencent("tencent")]
|
||||
#[case::anthropic("anthropic")]
|
||||
#[case::azure_ai("azure_ai")]
|
||||
#[case::bedrock("bedrock")]
|
||||
#[case::deepseek("deepseek")]
|
||||
#[case::edenai("edenai")]
|
||||
#[case::minimax("minimax")]
|
||||
#[case::vertex_ai("vertex_ai")]
|
||||
fn provider_keeps_its_python_name(#[case] name: &str) {
|
||||
let provider =
|
||||
messages_provider(name.parse::<LlmProviders>().unwrap(), "claude-sonnet-4-5").unwrap();
|
||||
#[case::tencent(LlmProviders::Tencent, "tencent")]
|
||||
#[case::anthropic(LlmProviders::Anthropic, "anthropic")]
|
||||
#[case::azure_ai(LlmProviders::AzureAi, "azure_ai")]
|
||||
#[case::bedrock(LlmProviders::Bedrock, "bedrock")]
|
||||
#[case::deepseek(LlmProviders::Deepseek, "deepseek")]
|
||||
#[case::edenai(LlmProviders::Edenai, "edenai")]
|
||||
#[case::minimax(LlmProviders::Minimax, "minimax")]
|
||||
#[case::vertex_ai(LlmProviders::VertexAi, "vertex_ai")]
|
||||
#[case::openrouter(LlmProviders::Openrouter, "openrouter")]
|
||||
fn provider_keeps_its_python_name(#[case] provider: LlmProviders, #[case] name: &str) {
|
||||
let provider = messages_provider(provider, "claude-sonnet-4-5").unwrap();
|
||||
assert_eq!(provider.as_str(), name);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,148 +1,271 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures_util::{StreamExt, TryStreamExt, stream::BoxStream};
|
||||
use litellm_host::interceptors::{Interceptors, ProviderIdentity, RequestContext, WireRequest};
|
||||
use litellm_http::transport::Error as TransportError;
|
||||
use futures_util::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream};
|
||||
use litellm_auth::{AuthServices, SecretValue};
|
||||
use litellm_core_utils::{
|
||||
dot_notation_indexing::delete_nested_value,
|
||||
get_provider_specific_headers::get_provider_specific_headers, settings::Lookup,
|
||||
};
|
||||
use litellm_host::interceptors::{
|
||||
Cost, ExecutionFacts, Interceptors, ProviderIdentity, RequestContext, ResultSource, WireRequest,
|
||||
};
|
||||
use litellm_http::{request::with_default_headers, transport::Error as TransportError};
|
||||
use litellm_llms::base_llm::{
|
||||
auth::{Authenticated, resolve_auth},
|
||||
auth::{Authenticated, ValidatedEnvironment, resolve_auth},
|
||||
messages::{
|
||||
context::MessagesTransformContext,
|
||||
streaming::{self, ByteStream, StreamDecoder},
|
||||
transformation::BaseMessagesConfig,
|
||||
},
|
||||
};
|
||||
use litellm_llms_types::formats::messages::MessagesResponse;
|
||||
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
|
||||
use litellm_secrets::source::Secrets;
|
||||
use litellm_tracing::ByteChunk;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
Error, MessagesCallResponse, MessagesRoute, common_utils::truncate_error_body,
|
||||
prepare::ProviderMessagesRequest,
|
||||
Error, MessagesCall, MessagesCallResponse, MessagesRoute,
|
||||
common_utils::{MessagesProvider, string_headers, truncate_error_body},
|
||||
constants::MESSAGES_TIMEOUT_SECS,
|
||||
route::Messages,
|
||||
types::invalid_request,
|
||||
};
|
||||
use crate::constants::MESSAGES_TIMEOUT_SECS;
|
||||
use litellm_inference::{context::CallContext, outbound::outbound_request};
|
||||
use litellm_inference::{caching::CallCache, context::CallContext, outbound::outbound_request};
|
||||
|
||||
pub(super) struct ProviderCall {
|
||||
pub identity: ProviderIdentity,
|
||||
pub wire: WireRequest,
|
||||
pub(super) provider: super::common_utils::MessagesProvider,
|
||||
struct ProviderCall {
|
||||
identity: ProviderIdentity,
|
||||
wire: WireRequest,
|
||||
provider: MessagesProvider,
|
||||
signer: Option<litellm_auth_aws::SigV4Signer>,
|
||||
timeout: Option<Duration>,
|
||||
stream: bool,
|
||||
}
|
||||
|
||||
impl ProviderCall {
|
||||
pub fn cacheable(&self) -> bool {
|
||||
fn cacheable(&self) -> bool {
|
||||
self.signer.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
impl MessagesRoute {
|
||||
pub(super) async fn prepare_outbound(
|
||||
&self,
|
||||
request: ProviderMessagesRequest,
|
||||
context: &CallContext<'_, impl Interceptors<Error>>,
|
||||
) -> Result<ProviderCall, Error> {
|
||||
let ProviderMessagesRequest {
|
||||
provider,
|
||||
url,
|
||||
body,
|
||||
environment,
|
||||
timeout,
|
||||
api_key,
|
||||
} = request;
|
||||
let request_context = RequestContext {
|
||||
model: body.model.clone(),
|
||||
custom_llm_provider: provider.as_str().to_string(),
|
||||
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
|
||||
secret_fields: Vec::new(),
|
||||
api_key,
|
||||
};
|
||||
let authenticated =
|
||||
resolve_auth(&self.auth, environment, &|key| std::env::var(key).ok()).await?;
|
||||
let identity = ProviderIdentity {
|
||||
model: request_context.model.clone(),
|
||||
provider: request_context.custom_llm_provider.clone(),
|
||||
};
|
||||
let wire = context
|
||||
.interceptors
|
||||
.before_provider_request(
|
||||
WireRequest {
|
||||
url,
|
||||
headers: authenticated.headers,
|
||||
body: provider
|
||||
.config()
|
||||
.wire_body(serde_json::to_value(&body).map_err(serialize_failure)?),
|
||||
},
|
||||
request_context,
|
||||
)
|
||||
.await?;
|
||||
let stream = match wire.body.get("stream") {
|
||||
None | Some(Value::Null) => false,
|
||||
Some(Value::Bool(stream)) => *stream,
|
||||
Some(value) => {
|
||||
return Err(Error::InvalidRequest(
|
||||
litellm_llms::ErrorDetail::InvalidValue {
|
||||
field: "stream",
|
||||
expected: "a boolean",
|
||||
actual: value.clone(),
|
||||
},
|
||||
));
|
||||
}
|
||||
};
|
||||
Ok(ProviderCall {
|
||||
identity,
|
||||
wire,
|
||||
provider,
|
||||
signer: authenticated.signer,
|
||||
timeout,
|
||||
stream,
|
||||
pub(super) async fn execute(
|
||||
route: &MessagesRoute,
|
||||
provider: MessagesProvider,
|
||||
model: String,
|
||||
call: MessagesCall,
|
||||
secrets: Secrets,
|
||||
context: &CallContext<'_, impl Interceptors<Error>>,
|
||||
) -> Result<MessagesCallResponse, Error> {
|
||||
let request = provider_call(
|
||||
&route.auth,
|
||||
provider,
|
||||
model,
|
||||
call,
|
||||
secrets.as_ref(),
|
||||
context,
|
||||
)
|
||||
.boxed()
|
||||
.await?;
|
||||
let cache = CallCache::<Messages>::from_wire(
|
||||
route.cache.as_ref().filter(|_| request.cacheable()),
|
||||
context.cache,
|
||||
&request.identity,
|
||||
&request.wire,
|
||||
);
|
||||
let identity = request.identity.clone();
|
||||
let (output, source) = match cache.lookup().await {
|
||||
Some(hit) => hit,
|
||||
None => (
|
||||
call_provider(&route.http, request, context).await?,
|
||||
ResultSource::Provider,
|
||||
),
|
||||
};
|
||||
let reported = match &output {
|
||||
MessagesCallResponse::Complete(response) => provider.config().reported_cost(response),
|
||||
MessagesCallResponse::Stream { .. } => None,
|
||||
};
|
||||
context
|
||||
.result_ready(ExecutionFacts {
|
||||
provider: identity,
|
||||
cost: Cost::settle(&source, reported),
|
||||
source: source.clone(),
|
||||
})
|
||||
}
|
||||
.await?;
|
||||
Ok(cache.finish(output, &source).await)
|
||||
}
|
||||
|
||||
pub(super) async fn call_provider(
|
||||
&self,
|
||||
request: ProviderCall,
|
||||
context: &CallContext<'_, impl Interceptors<Error>>,
|
||||
) -> Result<MessagesCallResponse, Error> {
|
||||
let ProviderCall {
|
||||
identity,
|
||||
wire,
|
||||
provider,
|
||||
signer,
|
||||
timeout,
|
||||
stream,
|
||||
} = request;
|
||||
let provider_name = provider.as_str();
|
||||
log_request_body(provider_name, stream, &wire.body);
|
||||
let response = send(
|
||||
&self.http,
|
||||
Authenticated {
|
||||
headers: wire.headers,
|
||||
signer,
|
||||
async fn provider_call(
|
||||
auth: &AuthServices,
|
||||
provider: MessagesProvider,
|
||||
model: String,
|
||||
call: MessagesCall,
|
||||
secrets: &(dyn Lookup + Send + Sync),
|
||||
context: &CallContext<'_, impl Interceptors<Error>>,
|
||||
) -> Result<ProviderCall, Error> {
|
||||
let MessagesCall {
|
||||
body,
|
||||
api_key,
|
||||
api_base,
|
||||
litellm_params,
|
||||
extra_headers,
|
||||
provider_specific_header,
|
||||
timeout,
|
||||
shaping,
|
||||
..
|
||||
} = call;
|
||||
let connection_params = litellm_router_types::LitellmParams {
|
||||
api_base: api_base.clone().or(litellm_params.api_base.clone()),
|
||||
..litellm_params
|
||||
};
|
||||
let config = provider.config();
|
||||
let env_lookup = |key: &str| secrets.get(key);
|
||||
|
||||
let sanitized = config.shape_request(
|
||||
MessagesRequest { model, ..body },
|
||||
shaping.settings.reasoning_auto_summary,
|
||||
)?;
|
||||
let trimmed =
|
||||
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
|
||||
|
||||
let scoped =
|
||||
get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str());
|
||||
let forwarded = string_headers(Some(
|
||||
extra_headers.into_iter().flatten().chain(scoped).collect(),
|
||||
))?;
|
||||
let validated = config.validate_environment(
|
||||
forwarded,
|
||||
api_key.as_deref(),
|
||||
&trimmed.model,
|
||||
&trimmed.params,
|
||||
&connection_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
|
||||
let body = config.transform_anthropic_messages_request(
|
||||
trimmed,
|
||||
&MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params),
|
||||
)?;
|
||||
let environment = ValidatedEnvironment {
|
||||
headers: config.request_headers(
|
||||
with_default_headers(validated.headers, config.default_headers()),
|
||||
&body,
|
||||
),
|
||||
auth: validated.auth,
|
||||
};
|
||||
|
||||
let url = config.get_complete_url(
|
||||
api_base.as_deref(),
|
||||
&body.model,
|
||||
&connection_params,
|
||||
body.params.stream == Some(true),
|
||||
&env_lookup,
|
||||
)?;
|
||||
|
||||
let request_context = RequestContext {
|
||||
model: body.model.clone(),
|
||||
custom_llm_provider: provider.as_str().to_string(),
|
||||
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
|
||||
secret_fields: Vec::new(),
|
||||
api_key: api_key.map(SecretValue::new),
|
||||
};
|
||||
let authenticated = resolve_auth(auth, environment, &|key| std::env::var(key).ok()).await?;
|
||||
let identity = ProviderIdentity {
|
||||
model: request_context.model.clone(),
|
||||
provider: request_context.custom_llm_provider.clone(),
|
||||
};
|
||||
let wire = context
|
||||
.interceptors
|
||||
.before_provider_request(
|
||||
WireRequest {
|
||||
url,
|
||||
headers: authenticated.headers,
|
||||
body: config.wire_body(serde_json::to_value(&body).map_err(serialize_failure)?),
|
||||
},
|
||||
&wire.url,
|
||||
&wire.body,
|
||||
timeout,
|
||||
request_context,
|
||||
)
|
||||
.await?;
|
||||
if !response.status().is_success() {
|
||||
return Err(provider_error(response).await);
|
||||
}
|
||||
let config = provider.config();
|
||||
if stream {
|
||||
return Ok(streaming_response(
|
||||
response,
|
||||
config.stream_decoder(),
|
||||
provider_name,
|
||||
let stream = match wire.body.get("stream") {
|
||||
None | Some(Value::Null) => false,
|
||||
Some(Value::Bool(stream)) => *stream,
|
||||
Some(value) => {
|
||||
return Err(Error::InvalidRequest(
|
||||
litellm_llms::ErrorDetail::InvalidValue {
|
||||
field: "stream",
|
||||
expected: "a boolean",
|
||||
actual: value.clone(),
|
||||
},
|
||||
));
|
||||
}
|
||||
let text = response.text().await.map_err(network)?;
|
||||
log_response_body(&text);
|
||||
context.response_received(&text).await?;
|
||||
decode_response(config, &identity.model, &text)
|
||||
.map(|message| MessagesCallResponse::Complete(Box::new(message)))
|
||||
};
|
||||
Ok(ProviderCall {
|
||||
identity,
|
||||
wire,
|
||||
provider,
|
||||
signer: authenticated.signer,
|
||||
timeout: timeout.or_else(|| shaping.settings.request_timeout()),
|
||||
stream,
|
||||
})
|
||||
}
|
||||
|
||||
fn without_additional_drop_params(
|
||||
request: MessagesRequest,
|
||||
paths: &[String],
|
||||
) -> Result<MessagesRequest, Error> {
|
||||
if paths.is_empty() {
|
||||
return Ok(request);
|
||||
}
|
||||
let params = serde_json::to_value(request.params).map_err(invalid_request)?;
|
||||
let trimmed = paths
|
||||
.iter()
|
||||
.fold(params, |params, path| delete_nested_value(params, path));
|
||||
Ok(MessagesRequest {
|
||||
params: serde_json::from_value(trimmed).map_err(invalid_request)?,
|
||||
..request
|
||||
})
|
||||
}
|
||||
|
||||
async fn call_provider(
|
||||
http: &litellm_http::Client,
|
||||
request: ProviderCall,
|
||||
context: &CallContext<'_, impl Interceptors<Error>>,
|
||||
) -> Result<MessagesCallResponse, Error> {
|
||||
let ProviderCall {
|
||||
identity,
|
||||
wire,
|
||||
provider,
|
||||
signer,
|
||||
timeout,
|
||||
stream,
|
||||
} = request;
|
||||
let provider_name = provider.as_str();
|
||||
log_request_body(provider_name, stream, &wire.body);
|
||||
let response = send(
|
||||
http,
|
||||
Authenticated {
|
||||
headers: wire.headers,
|
||||
signer,
|
||||
},
|
||||
&wire.url,
|
||||
&wire.body,
|
||||
timeout,
|
||||
)
|
||||
.await?;
|
||||
if !response.status().is_success() {
|
||||
return Err(provider_error(response).await);
|
||||
}
|
||||
let config = provider.config();
|
||||
if stream {
|
||||
return Ok(streaming_response(
|
||||
response,
|
||||
config.stream_decoder(),
|
||||
provider_name,
|
||||
));
|
||||
}
|
||||
let text = response.text().await.map_err(network)?;
|
||||
log_response_body(&text);
|
||||
context.response_received(&text).await?;
|
||||
decode_response(config, &identity.model, &text)
|
||||
.map(|message| MessagesCallResponse::Complete(Box::new(message)))
|
||||
}
|
||||
|
||||
fn serialize_failure(err: serde_json::Error) -> Error {
|
||||
|
|
@ -271,10 +394,482 @@ fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) {
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use litellm_llms_types::headers::ProviderSpecificHeaders;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Map, Value, json};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
|
||||
|
||||
use super::*;
|
||||
use crate::{MessagesSettings, MessagesShaping, common_utils::messages_provider};
|
||||
use litellm_inference::provider::resolve_llm_provider;
|
||||
|
||||
#[fixture]
|
||||
fn shaping() -> MessagesShaping {
|
||||
MessagesShaping::default()
|
||||
}
|
||||
|
||||
fn body(value: Value) -> MessagesRequest {
|
||||
serde_json::from_value(value).unwrap()
|
||||
}
|
||||
|
||||
async fn provider_call_for(call: MessagesCall) -> Result<ProviderCall, Error> {
|
||||
provider_call_with_secrets(call, &|_: &str| None).await
|
||||
}
|
||||
|
||||
async fn provider_call_with_secrets(
|
||||
call: MessagesCall,
|
||||
secrets: &(dyn Lookup + Send + Sync),
|
||||
) -> Result<ProviderCall, Error> {
|
||||
let provider = resolve_llm_provider(
|
||||
&call.body.model,
|
||||
call.custom_llm_provider.as_deref(),
|
||||
"messages",
|
||||
)?;
|
||||
let config = messages_provider(provider.provider, provider.model)
|
||||
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider.provider).to_string()))?;
|
||||
let model = provider.model.to_owned();
|
||||
let interceptors = ();
|
||||
let context = CallContext::new(&interceptors, litellm_inference::CallOptions::default());
|
||||
provider_call(
|
||||
&AuthServices::default(),
|
||||
config,
|
||||
model,
|
||||
call,
|
||||
secrets,
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn wire_body(fields: Value, shaping: MessagesShaping) -> Result<Value, Error> {
|
||||
provider_call_for(MessagesCall {
|
||||
body: body(fields),
|
||||
api_key: Some("sk-test".into()),
|
||||
api_base: Some("https://anthropic.test".into()),
|
||||
custom_llm_provider: Some("anthropic".into()),
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: None,
|
||||
shaping,
|
||||
})
|
||||
.await
|
||||
.map(|call| call.wire.body)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::api_key(
|
||||
&[("ANTHROPIC_API_KEY", "sk-secret")],
|
||||
&[("x-api-key", "sk-secret")],
|
||||
"https://api.anthropic.com/v1/messages"
|
||||
)]
|
||||
#[case::auth_token(
|
||||
&[("ANTHROPIC_AUTH_TOKEN", "token")],
|
||||
&[("authorization", "Bearer token")],
|
||||
"https://api.anthropic.com/v1/messages"
|
||||
)]
|
||||
#[case::api_base(
|
||||
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_API_BASE", "https://gateway.test")],
|
||||
&[("x-api-key", "sk-secret")],
|
||||
"https://gateway.test/v1/messages"
|
||||
)]
|
||||
#[case::sdk_base_url(
|
||||
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_BASE_URL", "https://sdk.test")],
|
||||
&[("x-api-key", "sk-secret")],
|
||||
"https://sdk.test/v1/messages"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn credentials_and_base_come_from_the_resolved_secrets(
|
||||
shaping: MessagesShaping,
|
||||
#[case] secrets: &[(&str, &str)],
|
||||
#[case] expected_auth: &[(&str, &str)],
|
||||
#[case] expected_url: &str,
|
||||
) {
|
||||
let lookup = |name: &str| {
|
||||
secrets
|
||||
.iter()
|
||||
.find(|(key, _)| *key == name)
|
||||
.map(|(_, value)| value.to_string())
|
||||
};
|
||||
let call = provider_call_with_secrets(
|
||||
MessagesCall {
|
||||
body: body(
|
||||
json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
|
||||
),
|
||||
api_key: None,
|
||||
api_base: None,
|
||||
custom_llm_provider: Some("anthropic".into()),
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: None,
|
||||
shaping,
|
||||
},
|
||||
&lookup,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let auth: Vec<(&str, &str)> = call
|
||||
.wire
|
||||
.headers
|
||||
.iter()
|
||||
.filter(|(name, _)| matches!(name.as_str(), "x-api-key" | "authorization"))
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()))
|
||||
.collect();
|
||||
assert_eq!(
|
||||
(auth.as_slice(), call.wire.url.as_str()),
|
||||
(expected_auth, expected_url)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn missing_credentials_are_reported_before_the_request_is_transformed(
|
||||
shaping: MessagesShaping,
|
||||
) {
|
||||
let result = provider_call_for(MessagesCall {
|
||||
body: body(
|
||||
json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}]}),
|
||||
),
|
||||
api_key: None,
|
||||
api_base: None,
|
||||
custom_llm_provider: Some("anthropic".into()),
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: None,
|
||||
shaping,
|
||||
})
|
||||
.await;
|
||||
assert!(
|
||||
matches!(
|
||||
result,
|
||||
Err(Error::Auth(litellm_auth::Error::MissingApiKey { .. }))
|
||||
),
|
||||
"{:?}",
|
||||
result.err()
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::nothing_configured(None, None, None)]
|
||||
#[case::configured_setting(None, Some(1800.0), Some(Duration::from_secs(1800)))]
|
||||
#[case::call_timeout_wins(
|
||||
Some(Duration::from_secs(30)),
|
||||
Some(1800.0),
|
||||
Some(Duration::from_secs(30))
|
||||
)]
|
||||
#[case::call_timeout_only(Some(Duration::from_secs(30)), None, Some(Duration::from_secs(30)))]
|
||||
#[case::unusable_setting_ignored(None, Some(-1.0), None)]
|
||||
#[case::zero_setting_ignored(None, Some(0.0), None)]
|
||||
#[tokio::test]
|
||||
async fn timeout_falls_back_to_the_configured_request_timeout(
|
||||
shaping: MessagesShaping,
|
||||
#[case] call_timeout: Option<Duration>,
|
||||
#[case] request_timeout: Option<f64>,
|
||||
#[case] expected: Option<Duration>,
|
||||
) {
|
||||
let call = provider_call_for(MessagesCall {
|
||||
body: body(json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})),
|
||||
api_key: Some("sk-test".into()),
|
||||
api_base: Some("https://anthropic.test".into()),
|
||||
custom_llm_provider: Some("anthropic".into()),
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: call_timeout,
|
||||
shaping: MessagesShaping {
|
||||
settings: MessagesSettings { request_timeout, ..shaping.settings },
|
||||
..shaping
|
||||
},
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(call.timeout, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::top_level_and_nested_paths(
|
||||
json!({
|
||||
"max_tokens": 1024,
|
||||
"thinking": {"type": "enabled", "budget_tokens": 2048},
|
||||
"context_management": {"edits": [{"type": "clear_thinking_20251015"}]},
|
||||
"metadata": {"user_id": "u1"},
|
||||
"tools": [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}]
|
||||
}),
|
||||
&["thinking", "context_management", "tools[*].input_examples"],
|
||||
json!({
|
||||
"max_tokens": 1024,
|
||||
"metadata": {"user_id": "u1"},
|
||||
"tools": [{"name": "lookup", "input_schema": {"type": "object"}}]
|
||||
}),
|
||||
)]
|
||||
#[case::no_paths(
|
||||
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
|
||||
&[],
|
||||
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
|
||||
)]
|
||||
#[case::model_and_messages_are_never_dropped(
|
||||
json!({"max_tokens": 16}),
|
||||
&["model", "messages", "messages[0].content"],
|
||||
json!({"max_tokens": 16}),
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn wire_body_drops_configured_paths(
|
||||
shaping: MessagesShaping,
|
||||
#[case] fields: Value,
|
||||
#[case] additional_drop_params: &[&str],
|
||||
#[case] expected_fields: Value,
|
||||
) {
|
||||
let with_messages = |fields: Value| -> Value {
|
||||
let Value::Object(fields) = fields else {
|
||||
unreachable!()
|
||||
};
|
||||
Value::Object(
|
||||
[
|
||||
("model".to_string(), json!("claude-test")),
|
||||
(
|
||||
"messages".to_string(),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
),
|
||||
]
|
||||
.into_iter()
|
||||
.chain(fields)
|
||||
.collect(),
|
||||
)
|
||||
};
|
||||
let shaping = MessagesShaping {
|
||||
settings: MessagesSettings {
|
||||
additional_drop_params: additional_drop_params
|
||||
.iter()
|
||||
.map(ToString::to_string)
|
||||
.collect(),
|
||||
..shaping.settings
|
||||
},
|
||||
..shaping
|
||||
};
|
||||
assert_eq!(
|
||||
wire_body(with_messages(fields), shaping).await,
|
||||
Ok(with_messages(expected_fields))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::model_prefix_picks_the_provider(
|
||||
"azure_ai/claude-test",
|
||||
None,
|
||||
&[("x-priority", "extra"), ("x-scoped", "azure_ai")]
|
||||
)]
|
||||
#[case::explicit_provider(
|
||||
"claude-test",
|
||||
Some("anthropic"),
|
||||
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
|
||||
)]
|
||||
#[case::provider_prefix_on_an_anthropic_model(
|
||||
"anthropic/claude-test",
|
||||
None,
|
||||
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn provider_specific_headers_follow_the_resolved_provider(
|
||||
shaping: MessagesShaping,
|
||||
#[case] model: &str,
|
||||
#[case] custom_llm_provider: Option<&str>,
|
||||
#[case] expected: &[(&str, &str)],
|
||||
) {
|
||||
let configured: ProviderSpecificHeaders = serde_json::from_value(json!([
|
||||
{"custom_llm_provider": "azure_ai", "extra_headers": {"x-scoped": "azure_ai"}},
|
||||
{"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "anthropic", "x-priority": "scoped"}}
|
||||
]))
|
||||
.unwrap();
|
||||
let call = provider_call_for(MessagesCall {
|
||||
body: body(
|
||||
json!({"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
|
||||
),
|
||||
api_key: Some("sk-test".into()),
|
||||
api_base: Some("https://resource.services.ai.azure.com".into()),
|
||||
custom_llm_provider: custom_llm_provider.map(Into::into),
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: Some(Map::from_iter([("x-priority".into(), json!("extra"))])),
|
||||
provider_specific_header: Some(configured),
|
||||
timeout: None,
|
||||
shaping,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let caller_headers: Vec<(&str, &str)> = call
|
||||
.wire
|
||||
.headers
|
||||
.iter()
|
||||
.filter(|(name, _)| matches!(name.as_str(), "x-priority" | "x-scoped"))
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()))
|
||||
.collect();
|
||||
assert_eq!(caller_headers, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn wire_body_carries_the_provider_stripped_model(shaping: MessagesShaping) {
|
||||
assert_eq!(
|
||||
wire_body(
|
||||
json!({
|
||||
"model": "anthropic/claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 16
|
||||
}),
|
||||
shaping,
|
||||
)
|
||||
.await,
|
||||
Ok(json!({
|
||||
"model": "claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 16
|
||||
}))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn dropped_thinking_display_is_not_restored_by_auto_summary(shaping: MessagesShaping) {
|
||||
let shaping = MessagesShaping {
|
||||
settings: MessagesSettings {
|
||||
reasoning_auto_summary: true,
|
||||
additional_drop_params: vec!["thinking.display".to_string()],
|
||||
..shaping.settings
|
||||
},
|
||||
..shaping
|
||||
};
|
||||
assert_eq!(
|
||||
wire_body(
|
||||
json!({
|
||||
"model": "claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 4096,
|
||||
"thinking": {"type": "enabled", "budget_tokens": 2048}
|
||||
}),
|
||||
shaping,
|
||||
)
|
||||
.await,
|
||||
Ok(json!({
|
||||
"model": "claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 4096,
|
||||
"thinking": {"type": "enabled", "budget_tokens": 2048}
|
||||
}))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn dropping_an_invalid_metadata_user_id_does_not_skip_its_validation(
|
||||
shaping: MessagesShaping,
|
||||
) {
|
||||
let shaping = MessagesShaping {
|
||||
settings: MessagesSettings {
|
||||
additional_drop_params: vec!["metadata.user_id".to_string()],
|
||||
..shaping.settings
|
||||
},
|
||||
..shaping
|
||||
};
|
||||
assert!(matches!(
|
||||
wire_body(
|
||||
json!({
|
||||
"model": "claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 16,
|
||||
"metadata": {"user_id": 123}
|
||||
}),
|
||||
shaping,
|
||||
)
|
||||
.await,
|
||||
Err(Error::InvalidRequest(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn wire_body_rejects_invalid_metadata_before_the_call(shaping: MessagesShaping) {
|
||||
assert_eq!(
|
||||
wire_body(
|
||||
json!({
|
||||
"model": "claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 16,
|
||||
"metadata": {"user_id": 123}
|
||||
}),
|
||||
shaping,
|
||||
)
|
||||
.await,
|
||||
Err(Error::InvalidRequest(
|
||||
litellm_llms::ErrorDetail::InvalidValue {
|
||||
field: "metadata.user_id",
|
||||
expected: "a string",
|
||||
actual: json!(123),
|
||||
}
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::mantle(
|
||||
"bedrock_mantle/claude-test",
|
||||
"https://bedrock-mantle.us-west-2.api.aws",
|
||||
"/anthropic/v1/messages",
|
||||
"bedrock"
|
||||
)]
|
||||
#[case::platform(
|
||||
"bedrock/claude_platform/claude-test",
|
||||
"https://aws-external-anthropic.us-west-2.api.aws",
|
||||
"/v1/messages",
|
||||
"aws-external-anthropic"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn explicit_endpoint_region_is_used_by_the_signer(
|
||||
shaping: MessagesShaping,
|
||||
#[case] model: &str,
|
||||
#[case] base: &str,
|
||||
#[case] path: &str,
|
||||
#[case] service: &str,
|
||||
) {
|
||||
let call = provider_call_with_secrets(MessagesCall {
|
||||
body: body(json!({"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})),
|
||||
api_key: None,
|
||||
api_base: Some(base.into()),
|
||||
custom_llm_provider: None,
|
||||
litellm_params: crate::litellm_params(Map::from_iter([
|
||||
("model".into(), json!(model)),
|
||||
("aws_access_key_id".into(), json!("AKIDEXAMPLE")),
|
||||
("aws_secret_access_key".into(), json!("secret")),
|
||||
]))
|
||||
.unwrap(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: None,
|
||||
shaping,
|
||||
}, &|name: &str| match name {
|
||||
"AWS_REGION_NAME" => Some("us-east-1".into()),
|
||||
"ANTHROPIC_AWS_WORKSPACE_ID" => Some("workspace".into()),
|
||||
_ => None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(call.wire.url, format!("{base}{path}"));
|
||||
let signed = outbound_request(
|
||||
Authenticated {
|
||||
headers: call.wire.headers,
|
||||
signer: call.signer,
|
||||
},
|
||||
call.wire.url,
|
||||
&call.wire.body,
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let authorization = signed.header("authorization").unwrap_or_default();
|
||||
assert!(
|
||||
authorization.contains(&format!("/us-west-2/{service}/aws4_request")),
|
||||
"{authorization}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::event(
|
||||
|
|
|
|||
|
|
@ -1,18 +1,19 @@
|
|||
mod common_utils;
|
||||
mod constants;
|
||||
mod handler;
|
||||
mod prepare;
|
||||
pub mod route;
|
||||
mod types;
|
||||
|
||||
use futures_util::FutureExt;
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_host::interceptors::{Cost, ExecutionFacts, Interceptors, ResultSource};
|
||||
|
||||
use litellm_inference::{caching::CallCache, context::CallContext};
|
||||
use litellm_host::interceptors::Interceptors;
|
||||
use litellm_inference::context::CallContext;
|
||||
use litellm_inference::provider::resolve_llm_provider;
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::common_utils::messages_provider;
|
||||
|
||||
pub use litellm_inference::RouteError as Error;
|
||||
pub use types::{
|
||||
MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params,
|
||||
|
|
@ -74,39 +75,23 @@ impl MessagesRoute {
|
|||
context: CallContext<'_, impl Interceptors<Error>>,
|
||||
) -> Result<MessagesCallResponse, Error> {
|
||||
litellm_inference::diagnostic::call(async {
|
||||
let prepared = prepare::prepare(call, self.secrets.as_ref()).await?;
|
||||
let provider = resolve_llm_provider(
|
||||
&call.body.model,
|
||||
call.custom_llm_provider.as_deref(),
|
||||
"messages",
|
||||
)?;
|
||||
let config = messages_provider(provider.provider, provider.model).ok_or_else(|| {
|
||||
Error::InvalidProvider(<&str>::from(provider.provider).to_string())
|
||||
})?;
|
||||
litellm_inference::diagnostic::provider(
|
||||
&prepared.body.model,
|
||||
prepared.provider.as_str(),
|
||||
provider.model,
|
||||
<&'static str>::from(provider.provider),
|
||||
);
|
||||
let request = self.prepare_outbound(prepared, &context).boxed().await?;
|
||||
let cache = CallCache::<route::Messages>::from_wire(
|
||||
self.cache.as_ref().filter(|_| request.cacheable()),
|
||||
context.cache,
|
||||
&request.identity,
|
||||
&request.wire,
|
||||
);
|
||||
let identity = request.identity.clone();
|
||||
let config = request.provider.config();
|
||||
let (output, source) = match cache.lookup().await {
|
||||
Some(hit) => hit,
|
||||
None => (
|
||||
self.call_provider(request, &context).await?,
|
||||
ResultSource::Provider,
|
||||
),
|
||||
};
|
||||
let reported = match &output {
|
||||
MessagesCallResponse::Complete(response) => config.reported_cost(response),
|
||||
MessagesCallResponse::Stream { .. } => None,
|
||||
};
|
||||
context
|
||||
.result_ready(ExecutionFacts {
|
||||
provider: identity,
|
||||
cost: Cost::settle(&source, reported),
|
||||
source: source.clone(),
|
||||
})
|
||||
.await?;
|
||||
Ok(cache.finish(output, &source).await)
|
||||
let model = provider.model.to_owned();
|
||||
let secrets = self.secrets.resolve(config.config().secret_names()).await?;
|
||||
handler::execute(self, config, model, call, secrets, &context)
|
||||
.boxed()
|
||||
.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,519 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_auth::SecretValue;
|
||||
use litellm_core_utils::{
|
||||
dot_notation_indexing::delete_nested_value,
|
||||
get_provider_specific_headers::get_provider_specific_headers, settings::Lookup,
|
||||
};
|
||||
use litellm_http::request::with_default_headers;
|
||||
use litellm_llms::base_llm::{
|
||||
auth::ValidatedEnvironment, messages::context::MessagesTransformContext,
|
||||
};
|
||||
use litellm_llms_types::formats::messages::MessagesRequest;
|
||||
use litellm_secrets::source::SecretSource;
|
||||
|
||||
use super::{
|
||||
Error, MessagesCall,
|
||||
common_utils::{MessagesProvider, messages_provider, string_headers},
|
||||
types::invalid_request,
|
||||
};
|
||||
use litellm_inference::provider::resolve_llm_provider;
|
||||
|
||||
struct ResolvedProvider {
|
||||
model: String,
|
||||
provider: MessagesProvider,
|
||||
}
|
||||
|
||||
pub(super) struct ProviderMessagesRequest {
|
||||
pub(super) provider: MessagesProvider,
|
||||
pub(super) url: String,
|
||||
pub(super) body: MessagesRequest,
|
||||
pub(super) environment: ValidatedEnvironment,
|
||||
pub(super) timeout: Option<Duration>,
|
||||
/// The caller's own credential, reported to the host beside the wire request.
|
||||
pub(super) api_key: Option<SecretValue>,
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
|
||||
pub(super) async fn prepare(
|
||||
call: MessagesCall,
|
||||
secrets: &dyn SecretSource,
|
||||
) -> Result<ProviderMessagesRequest, Error> {
|
||||
let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?;
|
||||
let secrets = secrets
|
||||
.resolve(resolved.provider.config().secret_names())
|
||||
.await?;
|
||||
prepare_provider_request(call, resolved, secrets.as_ref())
|
||||
}
|
||||
|
||||
fn resolve_provider(
|
||||
model: &str,
|
||||
custom_llm_provider: Option<&str>,
|
||||
) -> Result<ResolvedProvider, Error> {
|
||||
let resolved = resolve_llm_provider(model, custom_llm_provider, "messages")?;
|
||||
let provider = messages_provider(resolved.provider, resolved.model)
|
||||
.ok_or_else(|| Error::InvalidProvider(<&str>::from(resolved.provider).to_string()))?;
|
||||
Ok(ResolvedProvider {
|
||||
model: resolved.model.to_string(),
|
||||
provider,
|
||||
})
|
||||
}
|
||||
|
||||
fn prepare_provider_request(
|
||||
call: MessagesCall,
|
||||
resolved: ResolvedProvider,
|
||||
secrets: &dyn Lookup,
|
||||
) -> Result<ProviderMessagesRequest, Error> {
|
||||
let ResolvedProvider { model, provider } = resolved;
|
||||
let MessagesCall {
|
||||
body,
|
||||
api_key,
|
||||
api_base,
|
||||
litellm_params,
|
||||
extra_headers,
|
||||
provider_specific_header,
|
||||
timeout,
|
||||
shaping,
|
||||
..
|
||||
} = call;
|
||||
let connection_params = litellm_router_types::LitellmParams {
|
||||
api_base: api_base.clone().or(litellm_params.api_base.clone()),
|
||||
..litellm_params
|
||||
};
|
||||
let config = provider.config();
|
||||
let env_lookup = |key: &str| secrets.get(key);
|
||||
|
||||
let sanitized = config.shape_request(
|
||||
MessagesRequest { model, ..body },
|
||||
shaping.settings.reasoning_auto_summary,
|
||||
)?;
|
||||
let trimmed =
|
||||
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
|
||||
let transformed = config.transform_anthropic_messages_request(
|
||||
trimmed,
|
||||
&MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params),
|
||||
)?;
|
||||
|
||||
let scoped =
|
||||
get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str());
|
||||
let forwarded = string_headers(Some(
|
||||
extra_headers.into_iter().flatten().chain(scoped).collect(),
|
||||
))?;
|
||||
let forwarded = config.request_headers(
|
||||
with_default_headers(forwarded, config.default_headers()),
|
||||
&transformed,
|
||||
);
|
||||
let environment = config.validate_environment(
|
||||
forwarded,
|
||||
api_key.as_deref(),
|
||||
&transformed.model,
|
||||
&connection_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
|
||||
let url = config.get_complete_url(
|
||||
api_base.as_deref(),
|
||||
&transformed.model,
|
||||
&connection_params,
|
||||
transformed.params.stream == Some(true),
|
||||
&env_lookup,
|
||||
)?;
|
||||
|
||||
Ok(ProviderMessagesRequest {
|
||||
provider,
|
||||
url,
|
||||
body: transformed,
|
||||
environment,
|
||||
timeout,
|
||||
api_key: api_key.map(SecretValue::new),
|
||||
})
|
||||
}
|
||||
|
||||
fn without_additional_drop_params(
|
||||
request: MessagesRequest,
|
||||
paths: &[String],
|
||||
) -> Result<MessagesRequest, Error> {
|
||||
if paths.is_empty() {
|
||||
return Ok(request);
|
||||
}
|
||||
let params = serde_json::to_value(request.params).map_err(invalid_request)?;
|
||||
let trimmed = paths
|
||||
.iter()
|
||||
.fold(params, |params, path| delete_nested_value(params, path));
|
||||
Ok(MessagesRequest {
|
||||
params: serde_json::from_value(trimmed).map_err(invalid_request)?,
|
||||
..request
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_llms::base_llm::auth::resolve_auth;
|
||||
use litellm_llms_types::headers::ProviderSpecificHeaders;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use super::*;
|
||||
use crate::{MessagesSettings, MessagesShaping};
|
||||
|
||||
#[fixture]
|
||||
fn shaping() -> MessagesShaping {
|
||||
MessagesShaping::default()
|
||||
}
|
||||
|
||||
fn body(value: Value) -> MessagesRequest {
|
||||
serde_json::from_value(value).unwrap()
|
||||
}
|
||||
|
||||
fn prepare(call: MessagesCall) -> Result<ProviderMessagesRequest, Error> {
|
||||
prepare_with_secrets(call, &|_: &str| None)
|
||||
}
|
||||
|
||||
fn prepare_with_secrets(
|
||||
call: MessagesCall,
|
||||
secrets: &dyn Lookup,
|
||||
) -> Result<ProviderMessagesRequest, Error> {
|
||||
let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?;
|
||||
prepare_provider_request(call, resolved, secrets)
|
||||
}
|
||||
|
||||
/// The headers as they go on the wire, credential applied.
|
||||
fn wire_headers(prepared: &ProviderMessagesRequest) -> Vec<(String, String)> {
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.build()
|
||||
.unwrap()
|
||||
.block_on(resolve_auth(
|
||||
&litellm_auth::AuthServices::default(),
|
||||
prepared.environment.clone(),
|
||||
&|_| None,
|
||||
))
|
||||
.unwrap()
|
||||
.headers
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::api_key(
|
||||
&[("ANTHROPIC_API_KEY", "sk-secret")],
|
||||
&[("x-api-key", "sk-secret")],
|
||||
"https://api.anthropic.com/v1/messages"
|
||||
)]
|
||||
#[case::auth_token(
|
||||
&[("ANTHROPIC_AUTH_TOKEN", "token")],
|
||||
&[("authorization", "Bearer token")],
|
||||
"https://api.anthropic.com/v1/messages"
|
||||
)]
|
||||
#[case::api_base(
|
||||
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_API_BASE", "https://gateway.test")],
|
||||
&[("x-api-key", "sk-secret")],
|
||||
"https://gateway.test/v1/messages"
|
||||
)]
|
||||
#[case::sdk_base_url(
|
||||
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_BASE_URL", "https://sdk.test")],
|
||||
&[("x-api-key", "sk-secret")],
|
||||
"https://sdk.test/v1/messages"
|
||||
)]
|
||||
fn credentials_and_base_come_from_the_resolved_secrets(
|
||||
shaping: MessagesShaping,
|
||||
#[case] secrets: &[(&str, &str)],
|
||||
#[case] expected_auth: &[(&str, &str)],
|
||||
#[case] expected_url: &str,
|
||||
) {
|
||||
let lookup = |name: &str| {
|
||||
secrets
|
||||
.iter()
|
||||
.find(|(key, _)| *key == name)
|
||||
.map(|(_, value)| value.to_string())
|
||||
};
|
||||
let prepared = prepare_with_secrets(
|
||||
MessagesCall {
|
||||
body: body(
|
||||
json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
|
||||
),
|
||||
api_key: None,
|
||||
api_base: None,
|
||||
custom_llm_provider: Some("anthropic".into()),
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: None,
|
||||
shaping,
|
||||
},
|
||||
&lookup,
|
||||
)
|
||||
.unwrap();
|
||||
let headers = wire_headers(&prepared);
|
||||
let auth: Vec<(&str, &str)> = headers
|
||||
.iter()
|
||||
.filter(|(name, _)| matches!(name.as_str(), "x-api-key" | "authorization"))
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()))
|
||||
.collect();
|
||||
assert_eq!(
|
||||
(auth.as_slice(), prepared.url.as_str()),
|
||||
(expected_auth, expected_url)
|
||||
);
|
||||
}
|
||||
|
||||
fn prepared_body(fields: Value, shaping: MessagesShaping) -> Result<Value, Error> {
|
||||
prepare(MessagesCall {
|
||||
body: body(fields),
|
||||
api_key: Some("sk-test".into()),
|
||||
api_base: Some("https://anthropic.test".into()),
|
||||
custom_llm_provider: Some("anthropic".into()),
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: None,
|
||||
shaping,
|
||||
})
|
||||
.map(|prepared| serde_json::to_value(prepared.body).unwrap())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::top_level_and_nested_paths(
|
||||
json!({
|
||||
"max_tokens": 1024,
|
||||
"thinking": {"type": "enabled", "budget_tokens": 2048},
|
||||
"context_management": {"edits": [{"type": "clear_thinking_20251015"}]},
|
||||
"metadata": {"user_id": "u1"},
|
||||
"tools": [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}]
|
||||
}),
|
||||
&["thinking", "context_management", "tools[*].input_examples"],
|
||||
json!({
|
||||
"max_tokens": 1024,
|
||||
"metadata": {"user_id": "u1"},
|
||||
"tools": [{"name": "lookup", "input_schema": {"type": "object"}}]
|
||||
}),
|
||||
)]
|
||||
#[case::no_paths(
|
||||
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
|
||||
&[],
|
||||
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
|
||||
)]
|
||||
#[case::model_and_messages_are_never_dropped(
|
||||
json!({"max_tokens": 16}),
|
||||
&["model", "messages", "messages[0].content"],
|
||||
json!({"max_tokens": 16}),
|
||||
)]
|
||||
fn prepared_body_drops_configured_paths(
|
||||
shaping: MessagesShaping,
|
||||
#[case] fields: Value,
|
||||
#[case] additional_drop_params: &[&str],
|
||||
#[case] expected_fields: Value,
|
||||
) {
|
||||
let with_messages = |fields: Value| -> Value {
|
||||
let Value::Object(fields) = fields else {
|
||||
unreachable!()
|
||||
};
|
||||
Value::Object(
|
||||
[
|
||||
("model".to_string(), json!("claude-test")),
|
||||
(
|
||||
"messages".to_string(),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
),
|
||||
]
|
||||
.into_iter()
|
||||
.chain(fields)
|
||||
.collect(),
|
||||
)
|
||||
};
|
||||
let shaping = MessagesShaping {
|
||||
settings: MessagesSettings {
|
||||
additional_drop_params: additional_drop_params
|
||||
.iter()
|
||||
.map(ToString::to_string)
|
||||
.collect(),
|
||||
..shaping.settings
|
||||
},
|
||||
..shaping
|
||||
};
|
||||
assert_eq!(
|
||||
prepared_body(with_messages(fields), shaping),
|
||||
Ok(with_messages(expected_fields))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::model_prefix_picks_the_provider(
|
||||
"azure_ai/claude-test",
|
||||
None,
|
||||
&[("x-priority", "extra"), ("x-scoped", "azure_ai")]
|
||||
)]
|
||||
#[case::explicit_provider(
|
||||
"claude-test",
|
||||
Some("anthropic"),
|
||||
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
|
||||
)]
|
||||
#[case::provider_prefix_on_an_anthropic_model(
|
||||
"anthropic/claude-test",
|
||||
None,
|
||||
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
|
||||
)]
|
||||
fn provider_specific_headers_follow_the_resolved_provider(
|
||||
shaping: MessagesShaping,
|
||||
#[case] model: &str,
|
||||
#[case] custom_llm_provider: Option<&str>,
|
||||
#[case] expected: &[(&str, &str)],
|
||||
) {
|
||||
let configured: ProviderSpecificHeaders = serde_json::from_value(json!([
|
||||
{"custom_llm_provider": "azure_ai", "extra_headers": {"x-scoped": "azure_ai"}},
|
||||
{"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "anthropic", "x-priority": "scoped"}}
|
||||
]))
|
||||
.unwrap();
|
||||
let prepared = prepare(MessagesCall {
|
||||
body: body(
|
||||
json!({"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
|
||||
),
|
||||
api_key: Some("sk-test".into()),
|
||||
api_base: Some("https://resource.services.ai.azure.com".into()),
|
||||
custom_llm_provider: custom_llm_provider.map(Into::into),
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: Some(Map::from_iter([("x-priority".into(), json!("extra"))])),
|
||||
provider_specific_header: Some(configured),
|
||||
timeout: None,
|
||||
shaping,
|
||||
})
|
||||
.unwrap();
|
||||
let caller_headers: Vec<(&str, &str)> = prepared
|
||||
.environment
|
||||
.headers
|
||||
.iter()
|
||||
.filter(|(name, _)| matches!(name.as_str(), "x-priority" | "x-scoped"))
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()))
|
||||
.collect();
|
||||
assert_eq!(caller_headers, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn prepared_body_carries_the_provider_stripped_model(shaping: MessagesShaping) {
|
||||
assert_eq!(
|
||||
prepared_body(
|
||||
json!({
|
||||
"model": "anthropic/claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 16
|
||||
}),
|
||||
shaping,
|
||||
),
|
||||
Ok(json!({
|
||||
"model": "claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 16
|
||||
}))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn dropped_thinking_display_is_not_restored_by_auto_summary(shaping: MessagesShaping) {
|
||||
let shaping = MessagesShaping {
|
||||
settings: MessagesSettings {
|
||||
reasoning_auto_summary: true,
|
||||
additional_drop_params: vec!["thinking.display".to_string()],
|
||||
..shaping.settings
|
||||
},
|
||||
..shaping
|
||||
};
|
||||
assert_eq!(
|
||||
prepared_body(
|
||||
json!({
|
||||
"model": "claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 4096,
|
||||
"thinking": {"type": "enabled", "budget_tokens": 2048}
|
||||
}),
|
||||
shaping,
|
||||
),
|
||||
Ok(json!({
|
||||
"model": "claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 4096,
|
||||
"thinking": {"type": "enabled", "budget_tokens": 2048}
|
||||
}))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn dropping_an_invalid_metadata_user_id_does_not_skip_its_validation(shaping: MessagesShaping) {
|
||||
let shaping = MessagesShaping {
|
||||
settings: MessagesSettings {
|
||||
additional_drop_params: vec!["metadata.user_id".to_string()],
|
||||
..shaping.settings
|
||||
},
|
||||
..shaping
|
||||
};
|
||||
assert!(matches!(
|
||||
prepared_body(
|
||||
json!({
|
||||
"model": "claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 16,
|
||||
"metadata": {"user_id": 123}
|
||||
}),
|
||||
shaping,
|
||||
),
|
||||
Err(Error::InvalidRequest(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn prepared_body_rejects_invalid_metadata_before_the_call(shaping: MessagesShaping) {
|
||||
assert_eq!(
|
||||
prepared_body(
|
||||
json!({
|
||||
"model": "claude-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 16,
|
||||
"metadata": {"user_id": 123}
|
||||
}),
|
||||
shaping,
|
||||
),
|
||||
Err(Error::InvalidRequest(
|
||||
litellm_llms::ErrorDetail::InvalidValue {
|
||||
field: "metadata.user_id",
|
||||
expected: "a string",
|
||||
actual: json!(123),
|
||||
}
|
||||
))
|
||||
);
|
||||
}
|
||||
#[rstest]
|
||||
#[case::mantle(
|
||||
"bedrock_mantle/claude-test",
|
||||
"https://bedrock-mantle.us-west-2.api.aws",
|
||||
"/anthropic/v1/messages",
|
||||
"bedrock"
|
||||
)]
|
||||
#[case::platform(
|
||||
"bedrock/claude_platform/claude-test",
|
||||
"https://aws-external-anthropic.us-west-2.api.aws",
|
||||
"/v1/messages",
|
||||
"aws-external-anthropic"
|
||||
)]
|
||||
fn explicit_endpoint_region_is_used_by_the_signer(
|
||||
shaping: MessagesShaping,
|
||||
#[case] model: &str,
|
||||
#[case] base: &str,
|
||||
#[case] path: &str,
|
||||
#[case] service: &str,
|
||||
) {
|
||||
let prepared = prepare_with_secrets(MessagesCall {
|
||||
body: body(json!({"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})),
|
||||
api_key: None,
|
||||
api_base: Some(base.into()),
|
||||
custom_llm_provider: None,
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: None,
|
||||
shaping,
|
||||
}, &|name: &str| match name {
|
||||
"AWS_REGION_NAME" => Some("us-east-1".into()),
|
||||
"ANTHROPIC_AWS_WORKSPACE_ID" => Some("workspace".into()),
|
||||
_ => None,
|
||||
}).unwrap();
|
||||
assert_eq!(prepared.url, format!("{base}{path}"));
|
||||
assert!(
|
||||
matches!(prepared.environment.auth, litellm_llms::base_llm::auth::AuthScheme::AwsSigV4 { region, service: actual, .. } if region == "us-west-2" && actual == service)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -60,6 +60,16 @@ pub struct MessagesSettings {
|
|||
pub reasoning_auto_summary: bool,
|
||||
#[serde(default)]
|
||||
pub additional_drop_params: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub request_timeout: Option<f64>,
|
||||
}
|
||||
|
||||
impl MessagesSettings {
|
||||
pub fn request_timeout(&self) -> Option<Duration> {
|
||||
self.request_timeout
|
||||
.and_then(|secs| Duration::try_from_secs_f64(secs).ok())
|
||||
.filter(|timeout| !timeout.is_zero())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
@ -120,13 +130,15 @@ mod tests {
|
|||
},
|
||||
"drop_params": true,
|
||||
"reasoning_auto_summary": true,
|
||||
"additional_drop_params": ["metadata.user_id", "thinking"]
|
||||
"additional_drop_params": ["metadata.user_id", "thinking"],
|
||||
"request_timeout": 1800.0
|
||||
}),
|
||||
MessagesShaping {
|
||||
settings: MessagesSettings {
|
||||
drop_params: true,
|
||||
reasoning_auto_summary: true,
|
||||
additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()],
|
||||
request_timeout: Some(1800.0),
|
||||
},
|
||||
capabilities: MessagesModelCapabilities {
|
||||
supports_reasoning: true,
|
||||
|
|
|
|||
|
|
@ -343,6 +343,77 @@ async fn feature_betas_join_the_callers_betas_in_one_sorted_header(
|
|||
assert_eq!(sent, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic("anthropic", &[AnthropicBeta::DangerousToolUse20260903])]
|
||||
#[case::deepseek("deepseek", &[])]
|
||||
#[case::minimax("minimax", &[])]
|
||||
#[tokio::test]
|
||||
async fn safeguards_add_the_dangerous_tool_use_beta_only_where_the_beta_map_has_a_column(
|
||||
call: MessagesCall,
|
||||
#[case] provider: &str,
|
||||
#[case] derived: &[AnthropicBeta],
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]);
|
||||
|
||||
run_message(with_fields(
|
||||
MessagesCall {
|
||||
custom_llm_provider: Some(provider.into()),
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]),
|
||||
..call
|
||||
},
|
||||
json!({"safeguards": safeguards}),
|
||||
))
|
||||
.await;
|
||||
|
||||
let request = only_request(&upstream).await;
|
||||
let expected: BetaSet = derived
|
||||
.iter()
|
||||
.cloned()
|
||||
.chain([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())])
|
||||
.collect();
|
||||
assert_eq!(sent_betas(&request), expected);
|
||||
assert_eq!(
|
||||
request.body_json::<Value>().unwrap()["safeguards"],
|
||||
safeguards
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn null_safeguards_are_dropped_instead_of_sent_without_their_beta(
|
||||
call: MessagesCall,
|
||||
#[values("anthropic", "azure_ai", "bedrock_mantle", "deepseek", "minimax")] provider: &str,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
run_message(with_fields(
|
||||
MessagesCall {
|
||||
custom_llm_provider: Some(provider.into()),
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]),
|
||||
..call
|
||||
},
|
||||
json!({"safeguards": null}),
|
||||
))
|
||||
.await;
|
||||
|
||||
let request = only_request(&upstream).await;
|
||||
assert_eq!(
|
||||
sent_betas(&request),
|
||||
BetaSet::from_iter([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())])
|
||||
);
|
||||
assert!(
|
||||
!request
|
||||
.body_json::<Map<String, Value>>()
|
||||
.unwrap()
|
||||
.contains_key("safeguards")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: MessagesCall) {
|
||||
|
|
|
|||
|
|
@ -44,8 +44,8 @@ pub fn consumed_optional_param_names(
|
|||
model: &str,
|
||||
custom_llm_provider: Option<&str>,
|
||||
) -> Result<Vec<&'static str>, Error> {
|
||||
let (model, config) = resolve_provider_config(model, custom_llm_provider)?;
|
||||
let provider_fields = config.get_supported_ocr_params(&model);
|
||||
let (provider, config) = resolve_provider_config(model, custom_llm_provider)?;
|
||||
let provider_fields = config.get_supported_ocr_params(provider.model);
|
||||
let auth_fields: &[&str] = match config {
|
||||
OcrConfigKind::AwsTextract | OcrConfigKind::AwsTextractAnalyze => AWS_AUTH_OPTION_FIELDS,
|
||||
OcrConfigKind::AzureAi
|
||||
|
|
|
|||
|
|
@ -1,81 +0,0 @@
|
|||
use litellm_host::observation::ObservationSender;
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_host::interceptors::Interceptors;
|
||||
use litellm_llms::base_llm::ocr::{error::Error, handler::OcrClient};
|
||||
use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
|
||||
|
||||
use super::{
|
||||
handler::perform_ocr_request,
|
||||
types::{LiteLLMOcrRequest, OcrDocumentInput, ResolvedOcrRequest},
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct OcrRoute {
|
||||
client: OcrClient,
|
||||
}
|
||||
|
||||
impl OcrRoute {
|
||||
pub fn new(client: OcrClient) -> Self {
|
||||
Self { client }
|
||||
}
|
||||
|
||||
pub async fn execute(
|
||||
&self,
|
||||
request: LiteLLMOcrRequest,
|
||||
interceptors: &impl Interceptors<Error>,
|
||||
observers: Option<ObservationSender>,
|
||||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
litellm_host::lifecycle::observe_unary(
|
||||
observers.clone(),
|
||||
self.run(request, interceptors, observers.as_ref()),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "ocr",
|
||||
model = %request.model,
|
||||
resolved_model = %request.model,
|
||||
provider = <&str>::from(request.config.provider()),
|
||||
stream = false,
|
||||
outcome
|
||||
))]
|
||||
pub(super) async fn run(
|
||||
&self,
|
||||
request: LiteLLMOcrRequest,
|
||||
interceptors: &impl Interceptors<Error>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
litellm_inference::diagnostic::unary(async {
|
||||
let caller_document = matches!(&request.document, OcrDocumentInput::Document(_));
|
||||
let prepared = prepare_request_document(request).await?;
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<LiteLLMOcrResponse, Error>> =
|
||||
Box::pin(perform_ocr_request(
|
||||
&self.client,
|
||||
prepared,
|
||||
interceptors,
|
||||
caller_document,
|
||||
observers,
|
||||
));
|
||||
execute.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
|
||||
async fn prepare_request_document(
|
||||
request: LiteLLMOcrRequest<OcrDocumentInput>,
|
||||
) -> Result<ResolvedOcrRequest, Error> {
|
||||
if let OcrDocumentInput::Document(_) = &request.document {
|
||||
return request.map_document(super::document::prepare_document);
|
||||
}
|
||||
let logger = litellm_tracing::Logger::current();
|
||||
let span = tracing::Span::current();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
logger.scope(|| span.in_scope(|| request.map_document(super::document::prepare_document)))
|
||||
})
|
||||
.await
|
||||
.map_err(|error| Error::DocumentTask(Arc::new(error)))?
|
||||
}
|
||||
|
|
@ -1,10 +1,10 @@
|
|||
use std::{collections::BTreeMap as Map, io::Read, path::Path};
|
||||
use std::{collections::BTreeMap as Map, io::Read, path::Path, sync::Arc};
|
||||
|
||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||
use litellm_llms::base_llm::ocr::{error::Error, transformation::OCR_INLINE_MAX_BYTES};
|
||||
use litellm_llms_types::formats::ocr::OcrDocument;
|
||||
|
||||
use crate::types::OcrDocumentInput;
|
||||
use crate::types::{LiteLLMOcrRequest, OcrDocumentInput, ResolvedOcrRequest};
|
||||
|
||||
pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, Error> {
|
||||
match input {
|
||||
|
|
@ -24,6 +24,21 @@ pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, Error> {
|
|||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_document(
|
||||
request: LiteLLMOcrRequest<OcrDocumentInput>,
|
||||
) -> Result<ResolvedOcrRequest, Error> {
|
||||
if let OcrDocumentInput::Document(_) = &request.document {
|
||||
return request.map_document(prepare_document);
|
||||
}
|
||||
let logger = litellm_tracing::Logger::current();
|
||||
let span = tracing::Span::current();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
logger.scope(|| span.in_scope(|| request.map_document(prepare_document)))
|
||||
})
|
||||
.await
|
||||
.map_err(|error| Error::DocumentTask(Arc::new(error)))?
|
||||
}
|
||||
|
||||
pub fn read_path_document(path: &Path, mime_type: Option<&str>) -> Result<OcrDocument, Error> {
|
||||
let mut bytes = Vec::new();
|
||||
std::fs::File::open(path)
|
||||
|
|
|
|||
|
|
@ -4,29 +4,45 @@ use litellm_host::{lifecycle::ExecutionEvent, observation::ObservationSender};
|
|||
use litellm_llms::base_llm::ocr::{
|
||||
error::Error,
|
||||
handler::{CallHooks, OcrClient},
|
||||
transformation::PreparedOcrRequest,
|
||||
transformation::{OcrConnection, PreparedOcrRequest},
|
||||
};
|
||||
use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
|
||||
use litellm_secrets::source::Secrets;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{arguments::is_secret_param, prepare::prepare_request, provider_config::OcrConfigKind};
|
||||
use crate::types::ResolvedOcrRequest;
|
||||
use super::{arguments::is_secret_param, provider_config::OcrConfigKind};
|
||||
use crate::types::{LiteLLMOcrRequest, ResolvedOcrRequest};
|
||||
|
||||
pub(crate) async fn perform_ocr_request(
|
||||
pub(crate) async fn execute(
|
||||
client: &OcrClient,
|
||||
request: ResolvedOcrRequest,
|
||||
config: OcrConfigKind,
|
||||
call: ResolvedOcrRequest,
|
||||
secrets: Secrets,
|
||||
host: &impl Interceptors<Error>,
|
||||
caller_document: bool,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
request.response_format()?;
|
||||
let config = request.config;
|
||||
let secrets = client
|
||||
.secret_source()
|
||||
.resolve(&config.secret_names())
|
||||
.await
|
||||
.map_err(|error| Error::Secret(std::sync::Arc::new(error)))?;
|
||||
let request = prepare_request(request, caller_document, client, secrets);
|
||||
call.response_format()?;
|
||||
let LiteLLMOcrRequest {
|
||||
model,
|
||||
document,
|
||||
credentials,
|
||||
transport,
|
||||
optional_params,
|
||||
input_sources,
|
||||
azure_ad_token_provider,
|
||||
..
|
||||
} = call;
|
||||
let resolved = config.resolve_credentials(credentials, &secrets);
|
||||
let request = PreparedOcrRequest {
|
||||
model,
|
||||
document,
|
||||
connection: OcrConnection::new(resolved, transport, client.settings().clone(), secrets),
|
||||
caller_document,
|
||||
optional_params,
|
||||
input_sources,
|
||||
azure_ad_token_provider,
|
||||
};
|
||||
let interceptors = OcrCallHooks::new(host, &request, config, observers);
|
||||
config.ocr(client, &request, &interceptors).await
|
||||
}
|
||||
|
|
@ -86,3 +102,257 @@ impl<H: Interceptors<Error>> CallHooks<Error> for OcrCallHooks<'_, H> {
|
|||
Box::pin(self.interceptors.after_provider_response(raw))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_core_utils::call_arguments::{CallArguments, compose_body, parse_options};
|
||||
use litellm_llms_types::formats::ocr::OcrResponseFormat;
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
OcrRoute,
|
||||
types::LiteLLMOcrRequest,
|
||||
wire::{OcrWireRequest, decode_request},
|
||||
};
|
||||
|
||||
fn client() -> OcrClient {
|
||||
OcrClient::for_test(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::no_redirect_for_test(),
|
||||
)
|
||||
.with_secrets(litellm_inference_testing::no_secrets())
|
||||
}
|
||||
|
||||
fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest {
|
||||
decode_request(OcrWireRequest {
|
||||
model: model.into(),
|
||||
document,
|
||||
api_key: Some(litellm_auth::SecretValue::new("test-key")),
|
||||
api_base: Some(base.into()),
|
||||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
optional_params: options.as_object().unwrap().clone(),
|
||||
input_sources: Default::default(),
|
||||
timeout_seconds: Some(2.0),
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn run_request(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
|
||||
OcrRoute::new(client()).execute(request, &(), None).await
|
||||
}
|
||||
|
||||
async fn upstream(request_count: u64, response: Value) -> MockServer {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(any())
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(response))
|
||||
.up_to_n_times(request_count)
|
||||
.mount(&server)
|
||||
.await;
|
||||
server
|
||||
}
|
||||
|
||||
fn image(url: &str) -> Value {
|
||||
json!({"type": "image_url", "image_url": url})
|
||||
}
|
||||
|
||||
fn body(request: &wiremock::Request) -> Value {
|
||||
serde_json::from_slice(&request.body).unwrap()
|
||||
}
|
||||
|
||||
fn assert_mistral_request(request: &wiremock::Request) {
|
||||
assert_eq!(
|
||||
request
|
||||
.headers
|
||||
.get("authorization")
|
||||
.unwrap()
|
||||
.to_str()
|
||||
.unwrap(),
|
||||
"Bearer test-key"
|
||||
);
|
||||
assert_eq!(
|
||||
request
|
||||
.headers
|
||||
.get("content-type")
|
||||
.unwrap()
|
||||
.to_str()
|
||||
.unwrap(),
|
||||
"application/json"
|
||||
);
|
||||
assert_eq!(
|
||||
body(request),
|
||||
json!({
|
||||
"model": "mistral-ocr-maas",
|
||||
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
|
||||
"pages": [0, 2],
|
||||
"include_image_base64": true,
|
||||
"unknown": "preserved"
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn cohere_body_keeps_native_document_fields_and_untyped_overrides() {
|
||||
let upstream = upstream(
|
||||
1,
|
||||
json!({"pages": [{"index": 0, "markdown": {"content": "hello"}}]}),
|
||||
)
|
||||
.await;
|
||||
let request = request(
|
||||
"cohere/parse",
|
||||
&upstream.uri(),
|
||||
image("https://example.com/original.png"),
|
||||
json!({
|
||||
"output_format": "markdown", "timeout": 30,
|
||||
"extra_body": {
|
||||
"output_format": {"future": true},
|
||||
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
|
||||
"provider_options": {"nested": [false, 0, null]}}
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
run_request(request).await.unwrap();
|
||||
|
||||
let [request] = upstream
|
||||
.received_requests()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_into()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
body(&request),
|
||||
json!({
|
||||
"model": "parse", "output_format": {"future": true},
|
||||
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
|
||||
"provider_options": {"nested": [false, 0, null]}}
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn explicit_null_options_use_defaults_before_http() {
|
||||
let upstream = upstream(
|
||||
1,
|
||||
json!({"pages": [{"index": 0, "markdown": {"content": "hello"}}]}),
|
||||
)
|
||||
.await;
|
||||
let request = request(
|
||||
"cohere/parse",
|
||||
&upstream.uri(),
|
||||
image("https://example.com/a.png"),
|
||||
json!({"output_format": null, "req_format": null}),
|
||||
);
|
||||
assert_eq!(
|
||||
request.response_format().unwrap(),
|
||||
OcrResponseFormat::Litellm
|
||||
);
|
||||
|
||||
run_request(request).await.unwrap();
|
||||
|
||||
let [request] = upstream
|
||||
.received_requests()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_into()
|
||||
.unwrap();
|
||||
let body = body(&request);
|
||||
assert_eq!(body["output_format"], "markdown");
|
||||
assert!(body.get("req_format").is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn direct_and_vertex_mistral_build_the_same_request_and_share_normalization() {
|
||||
let upstream = upstream(
|
||||
2,
|
||||
json!({
|
||||
"pages": [{"index": 0, "markdown": "hello"}],
|
||||
"extra": "preserved"
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let options = json!({
|
||||
"pages": [0, 2],
|
||||
"include_image_base64": true,
|
||||
"vertex_project": "project-1",
|
||||
"vertex_location": "us-central1",
|
||||
"unknown": "preserved"
|
||||
});
|
||||
let document =
|
||||
json!({"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"});
|
||||
let direct_call = request(
|
||||
"mistral/mistral-ocr-maas",
|
||||
&upstream.uri(),
|
||||
document.clone(),
|
||||
options.clone(),
|
||||
);
|
||||
let vertex_call = request(
|
||||
"vertex_ai/mistral-ocr-maas",
|
||||
&upstream.uri(),
|
||||
document,
|
||||
options,
|
||||
);
|
||||
assert_eq!(direct_call.transport.timeout, Some(Duration::from_secs(2)));
|
||||
assert_eq!(vertex_call.transport.timeout, Some(Duration::from_secs(2)));
|
||||
let direct_response = run_request(direct_call).await.unwrap().into_json();
|
||||
let vertex_response = run_request(vertex_call).await.unwrap().into_json();
|
||||
|
||||
let [direct_http, vertex_http] = upstream
|
||||
.received_requests()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_into()
|
||||
.unwrap();
|
||||
assert_eq!(direct_http.url.path(), "/v1/ocr");
|
||||
assert_eq!(
|
||||
vertex_http.url.path(),
|
||||
"/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
|
||||
);
|
||||
assert_mistral_request(&direct_http);
|
||||
assert_mistral_request(&vertex_http);
|
||||
assert_eq!(direct_response, vertex_response);
|
||||
assert_eq!(direct_response["model"], "mistral-ocr-maas");
|
||||
assert_eq!(direct_response["object"], "ocr");
|
||||
assert_eq!(direct_response["extra"], "preserved");
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct KnownParams {
|
||||
pages: Option<Vec<i64>>,
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn parsed_provider_params_separates_known_and_extra_params() {
|
||||
let arguments: CallArguments = serde_json::from_value(json!({
|
||||
"pages": [0, 2],
|
||||
"future_ocr_option": true,
|
||||
"extra_body": {"provider_option": "value"}
|
||||
}))
|
||||
.unwrap();
|
||||
let known: KnownParams = parse_options(&arguments).unwrap();
|
||||
assert_eq!(known.pages, Some(vec![0, 2]));
|
||||
assert_eq!(arguments["future_ocr_option"], true);
|
||||
assert_eq!(arguments["extra_body"], json!({"provider_option": "value"}));
|
||||
assert_eq!(
|
||||
arguments
|
||||
.iter()
|
||||
.filter(|(name, _)| name.as_str() != "pages")
|
||||
.count(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
compose_body(&arguments, &json!({"pages": known.pages}), &["pages"]).unwrap(),
|
||||
json!({
|
||||
"pages": [0, 2], "future_ocr_option": true, "provider_option": "value"
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,10 +1,82 @@
|
|||
pub mod arguments;
|
||||
mod client;
|
||||
pub use client::OcrRoute;
|
||||
pub mod document;
|
||||
pub(crate) mod handler;
|
||||
pub(crate) mod prepare;
|
||||
pub mod provider_config;
|
||||
pub mod route;
|
||||
pub mod types;
|
||||
pub mod wire;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_host::interceptors::Interceptors;
|
||||
use litellm_host::observation::ObservationSender;
|
||||
use litellm_llms::base_llm::ocr::{error::Error, handler::OcrClient};
|
||||
use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
|
||||
|
||||
use crate::{
|
||||
handler::execute,
|
||||
types::{LiteLLMOcrRequest, OcrDocumentInput},
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct OcrRoute {
|
||||
client: OcrClient,
|
||||
}
|
||||
|
||||
impl OcrRoute {
|
||||
pub fn new(client: OcrClient) -> Self {
|
||||
Self { client }
|
||||
}
|
||||
|
||||
pub async fn execute(
|
||||
&self,
|
||||
request: LiteLLMOcrRequest,
|
||||
interceptors: &impl Interceptors<Error>,
|
||||
observers: Option<ObservationSender>,
|
||||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
litellm_host::lifecycle::observe_unary(
|
||||
observers.clone(),
|
||||
self.run(request, interceptors, observers.as_ref()),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "ocr",
|
||||
model = %call.model,
|
||||
resolved_model = %call.model,
|
||||
provider = <&str>::from(call.config.provider()),
|
||||
stream = false,
|
||||
outcome
|
||||
))]
|
||||
async fn run(
|
||||
&self,
|
||||
call: LiteLLMOcrRequest,
|
||||
interceptors: &impl Interceptors<Error>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
litellm_inference::diagnostic::unary(async {
|
||||
let caller_document = matches!(&call.document, OcrDocumentInput::Document(_));
|
||||
let call = document::resolve_document(call).await?;
|
||||
let config = call.config;
|
||||
let secrets = self
|
||||
.client
|
||||
.secret_source()
|
||||
.resolve(&config.secret_names())
|
||||
.await
|
||||
.map_err(|error| Error::Secret(Arc::new(error)))?;
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<LiteLLMOcrResponse, Error>> =
|
||||
Box::pin(execute(
|
||||
&self.client,
|
||||
config,
|
||||
call,
|
||||
secrets,
|
||||
interceptors,
|
||||
caller_document,
|
||||
observers,
|
||||
));
|
||||
execute.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,309 +0,0 @@
|
|||
use litellm_auth::{InputSource, SecretValue, Sourced};
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use litellm_llms::base_llm::ocr::{
|
||||
handler::OcrClient,
|
||||
transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest},
|
||||
};
|
||||
use litellm_secrets::source::Secrets;
|
||||
|
||||
use crate::types::{LiteLLMOcrRequest, ResolvedOcrRequest};
|
||||
|
||||
pub(crate) fn prepare_request(
|
||||
request: ResolvedOcrRequest,
|
||||
caller_document: bool,
|
||||
client: &OcrClient,
|
||||
secrets: Secrets,
|
||||
) -> PreparedOcrRequest {
|
||||
let credentials = request.credentials.clone();
|
||||
let (preferred_api_key_env, api_base_env) = match request.config.provider() {
|
||||
LlmProviders::Mistral => (
|
||||
Some("MISTRAL_AZURE_API_KEY"),
|
||||
Some("MISTRAL_AZURE_API_BASE"),
|
||||
),
|
||||
LlmProviders::AzureAi => (None, Some("AZURE_AI_API_BASE")),
|
||||
_ => (None, None),
|
||||
};
|
||||
let secret = |name: &str| secrets.truthy(name);
|
||||
let dynamic_api_key = credentials.dynamic_api_key.or_else(|| {
|
||||
credentials.api_key.clone().or_else(|| {
|
||||
preferred_api_key_env
|
||||
.into_iter()
|
||||
.chain(request.config.get_api_key_env_var())
|
||||
.find_map(secret)
|
||||
.map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment))
|
||||
})
|
||||
});
|
||||
let dynamic_api_base = credentials.dynamic_api_base.or_else(|| {
|
||||
credentials.api_base.clone().or_else(|| {
|
||||
api_base_env
|
||||
.and_then(secret)
|
||||
.map(|value| Sourced::new(value, InputSource::Environment))
|
||||
})
|
||||
});
|
||||
let resolved = request
|
||||
.config
|
||||
.resolve_connection_params(OcrCredentialInputs {
|
||||
dynamic_api_key,
|
||||
dynamic_api_base,
|
||||
..credentials
|
||||
});
|
||||
let LiteLLMOcrRequest {
|
||||
model,
|
||||
document,
|
||||
transport,
|
||||
optional_params,
|
||||
input_sources,
|
||||
azure_ad_token_provider,
|
||||
..
|
||||
} = request;
|
||||
PreparedOcrRequest {
|
||||
model,
|
||||
document,
|
||||
connection: OcrConnection::new(resolved, transport, client.settings().clone(), secrets),
|
||||
caller_document,
|
||||
optional_params,
|
||||
input_sources,
|
||||
azure_ad_token_provider,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_core_utils::call_arguments::{CallArguments, compose_body, parse_options};
|
||||
use litellm_host::interceptors::WireRequest;
|
||||
use litellm_llms::{
|
||||
base_llm::ocr::{
|
||||
error::Error,
|
||||
handler::{CallHooks, OcrClient},
|
||||
transformation::BaseOcrConfig,
|
||||
},
|
||||
cohere::ocr::transformation::CohereParseConfig,
|
||||
mistral::ocr::transformation::MistralOcrConfig,
|
||||
vertex_ai::ocr::transformation::VertexAiOcrConfig,
|
||||
};
|
||||
use litellm_llms_types::formats::ocr::OcrResponseFormat;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
document::prepare_document,
|
||||
types::LiteLLMOcrRequest,
|
||||
wire::{OcrWireRequest, decode_request},
|
||||
};
|
||||
|
||||
/// Stands in for a host with no interceptors registered.
|
||||
struct NoHooks;
|
||||
|
||||
impl CallHooks<Error> for NoHooks {
|
||||
fn before_provider_request(
|
||||
&self,
|
||||
wire: WireRequest,
|
||||
) -> BoxFuture<'_, Result<WireRequest, Error>> {
|
||||
Box::pin(async move { Ok(wire) })
|
||||
}
|
||||
|
||||
fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
}
|
||||
|
||||
fn client() -> OcrClient {
|
||||
OcrClient::for_test(
|
||||
litellm_http::Client::plain_for_test(),
|
||||
litellm_http::Client::no_redirect_for_test(),
|
||||
)
|
||||
}
|
||||
|
||||
fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest {
|
||||
decode_request(OcrWireRequest {
|
||||
model: model.into(),
|
||||
document,
|
||||
api_key: Some(litellm_auth::SecretValue::new("test-key")),
|
||||
api_base: Some(base.into()),
|
||||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
optional_params: options.as_object().unwrap().clone(),
|
||||
input_sources: Default::default(),
|
||||
timeout_seconds: Some(2.0),
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn prepared(request: LiteLLMOcrRequest) -> PreparedOcrRequest {
|
||||
prepare_request(
|
||||
request.map_document(prepare_document).unwrap(),
|
||||
true,
|
||||
&client(),
|
||||
std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment),
|
||||
)
|
||||
}
|
||||
|
||||
fn image(url: &str) -> Value {
|
||||
json!({"type": "image_url", "image_url": url})
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn cohere_body_keeps_native_document_fields_and_untyped_overrides() {
|
||||
let request = request(
|
||||
"cohere/parse",
|
||||
"https://example.com",
|
||||
image("https://example.com/original.png"),
|
||||
json!({
|
||||
"output_format": "markdown", "timeout": 30,
|
||||
"extra_body": {
|
||||
"output_format": {"future": true},
|
||||
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
|
||||
"provider_options": {"nested": [false, 0, null]}}
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let http = CohereParseConfig
|
||||
.prepare_request(&prepared(request), &client(), &NoHooks)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let body: Value = serde_json::from_slice(http.body()).unwrap();
|
||||
assert_eq!(
|
||||
body,
|
||||
json!({
|
||||
"model": "parse", "output_format": {"future": true},
|
||||
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
|
||||
"provider_options": {"nested": [false, 0, null]}}
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn explicit_null_options_use_defaults_before_http() {
|
||||
let request = request(
|
||||
"cohere/parse",
|
||||
"https://example.com",
|
||||
image("https://example.com/a.png"),
|
||||
json!({"output_format": null, "req_format": null}),
|
||||
);
|
||||
assert_eq!(
|
||||
request.response_format().unwrap(),
|
||||
OcrResponseFormat::Litellm
|
||||
);
|
||||
|
||||
let http = CohereParseConfig
|
||||
.prepare_request(&prepared(request), &client(), &NoHooks)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let body: Value = serde_json::from_slice(http.body()).unwrap();
|
||||
assert_eq!(body["output_format"], "markdown");
|
||||
assert!(body.get("req_format").is_none());
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn direct_and_vertex_mistral_build_the_same_request_and_share_normalization() {
|
||||
let options = json!({
|
||||
"pages": [0, 2],
|
||||
"include_image_base64": true,
|
||||
"vertex_project": "project-1",
|
||||
"vertex_location": "us-central1",
|
||||
"unknown": "preserved"
|
||||
});
|
||||
let document =
|
||||
json!({"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"});
|
||||
let direct = prepared(request(
|
||||
"mistral/mistral-ocr-maas",
|
||||
"https://mistral.test",
|
||||
document.clone(),
|
||||
options.clone(),
|
||||
));
|
||||
let vertex = prepared(request(
|
||||
"vertex_ai/mistral-ocr-maas",
|
||||
"https://vertex.test",
|
||||
document,
|
||||
options,
|
||||
));
|
||||
|
||||
let direct_http = MistralOcrConfig
|
||||
.prepare_request(&direct, &client(), &NoHooks)
|
||||
.await
|
||||
.unwrap();
|
||||
let vertex_http = VertexAiOcrConfig
|
||||
.prepare_request(&vertex, &client(), &NoHooks)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr");
|
||||
assert_eq!(
|
||||
vertex_http.url(),
|
||||
"https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
|
||||
);
|
||||
for http in [&direct_http, &vertex_http] {
|
||||
assert_eq!(http.header("authorization").unwrap(), "Bearer test-key");
|
||||
assert_eq!(http.header("content-type").unwrap(), "application/json");
|
||||
assert_eq!(http.timeout(), Some(Duration::from_secs(2)));
|
||||
let body: Value = serde_json::from_slice(http.body()).unwrap();
|
||||
assert_eq!(
|
||||
body,
|
||||
json!({
|
||||
"model": "mistral-ocr-maas",
|
||||
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
|
||||
"pages": [0, 2],
|
||||
"include_image_base64": true,
|
||||
"unknown": "preserved"
|
||||
})
|
||||
);
|
||||
}
|
||||
let payload = serde_json::to_vec(
|
||||
&json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}),
|
||||
)
|
||||
.unwrap();
|
||||
let direct_response = MistralOcrConfig
|
||||
.transform_ocr_response(&direct.model, &payload, OcrResponseFormat::Litellm)
|
||||
.unwrap()
|
||||
.into_json();
|
||||
let vertex_response = VertexAiOcrConfig
|
||||
.transform_ocr_response(&vertex.model, &payload, OcrResponseFormat::Litellm)
|
||||
.unwrap()
|
||||
.into_json();
|
||||
assert_eq!(direct_response, vertex_response);
|
||||
assert_eq!(direct_response["model"], "mistral-ocr-maas");
|
||||
assert_eq!(direct_response["object"], "ocr");
|
||||
assert_eq!(direct_response["extra"], "preserved");
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct KnownParams {
|
||||
pages: Option<Vec<i64>>,
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
fn parsed_provider_params_separates_known_and_extra_params() {
|
||||
let arguments: CallArguments = serde_json::from_value(json!({
|
||||
"pages": [0, 2],
|
||||
"future_ocr_option": true,
|
||||
"extra_body": {"provider_option": "value"}
|
||||
}))
|
||||
.unwrap();
|
||||
let known: KnownParams = parse_options(&arguments).unwrap();
|
||||
assert_eq!(known.pages, Some(vec![0, 2]));
|
||||
assert_eq!(arguments["future_ocr_option"], true);
|
||||
assert_eq!(arguments["extra_body"], json!({"provider_option": "value"}));
|
||||
assert_eq!(
|
||||
arguments
|
||||
.iter()
|
||||
.filter(|(name, _)| name.as_str() != "pages")
|
||||
.count(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
compose_body(&arguments, &json!({"pages": known.pages}), &["pages"]).unwrap(),
|
||||
json!({
|
||||
"pages": [0, 2], "future_ocr_option": true, "provider_option": "value"
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
|
||||
use litellm_auth::{InputSource, SecretValue, Sourced};
|
||||
use litellm_core_utils::get_llm_provider_logic::{LlmProviders, get_custom_llm_provider};
|
||||
use litellm_inference::provider::{ResolvedProvider, resolve_llm_provider};
|
||||
use litellm_llms::{
|
||||
aws_textract::ocr::{
|
||||
analyze_transformation::TextractAnalyzeDocumentConfig, common_utils::TextractOperation,
|
||||
|
|
@ -25,6 +26,7 @@ use litellm_llms::{
|
|||
},
|
||||
};
|
||||
use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat};
|
||||
use litellm_secrets::source::Secrets;
|
||||
|
||||
macro_rules! with_config {
|
||||
($kind:expr, $config:ident => $body:expr) => {
|
||||
|
|
@ -115,7 +117,18 @@ impl OcrConfigKind {
|
|||
}
|
||||
|
||||
pub(crate) fn secret_names(self) -> Vec<&'static str> {
|
||||
with_config!(self, config => config.secret_names())
|
||||
let names = with_config!(self, config => config.secret_names());
|
||||
let fallback_names: &[&str] = match self.provider() {
|
||||
LlmProviders::Mistral => &["MISTRAL_AZURE_API_KEY", "MISTRAL_AZURE_API_BASE"],
|
||||
LlmProviders::AzureAi => &["AZURE_AI_API_BASE"],
|
||||
_ => &[],
|
||||
};
|
||||
let additional_names = fallback_names
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|name| !names.contains(name))
|
||||
.collect::<Vec<_>>();
|
||||
names.into_iter().chain(additional_names).collect()
|
||||
}
|
||||
|
||||
pub(crate) fn get_health_check_document(self) -> OcrDocument {
|
||||
|
|
@ -129,6 +142,43 @@ impl OcrConfigKind {
|
|||
with_config!(self, config => config.resolve_connection_params(inputs))
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_credentials(
|
||||
self,
|
||||
credentials: OcrCredentialInputs,
|
||||
secrets: &Secrets,
|
||||
) -> ResolvedOcrCredentials {
|
||||
let (preferred_api_key_env, api_base_env) = match self.provider() {
|
||||
LlmProviders::Mistral => (
|
||||
Some("MISTRAL_AZURE_API_KEY"),
|
||||
Some("MISTRAL_AZURE_API_BASE"),
|
||||
),
|
||||
LlmProviders::AzureAi => (None, Some("AZURE_AI_API_BASE")),
|
||||
_ => (None, None),
|
||||
};
|
||||
let secret = |name: &str| secrets.truthy(name);
|
||||
let dynamic_api_key = credentials.dynamic_api_key.or_else(|| {
|
||||
credentials.api_key.clone().or_else(|| {
|
||||
preferred_api_key_env
|
||||
.into_iter()
|
||||
.chain(self.get_api_key_env_var())
|
||||
.find_map(secret)
|
||||
.map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment))
|
||||
})
|
||||
});
|
||||
let dynamic_api_base = credentials.dynamic_api_base.or_else(|| {
|
||||
credentials.api_base.clone().or_else(|| {
|
||||
api_base_env
|
||||
.and_then(secret)
|
||||
.map(|value| Sourced::new(value, InputSource::Environment))
|
||||
})
|
||||
});
|
||||
self.resolve_connection_params(OcrCredentialInputs {
|
||||
dynamic_api_key,
|
||||
dynamic_api_base,
|
||||
..credentials
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn ocr(
|
||||
self,
|
||||
client: &OcrClient,
|
||||
|
|
@ -164,7 +214,7 @@ pub fn passthrough_response(
|
|||
endpoint: &str,
|
||||
body: &[u8],
|
||||
) -> Result<Option<LiteLLMOcrResponse>, Error> {
|
||||
let (model, config) = resolve_provider_config(model, Some("azure_ai"))?;
|
||||
let (provider, config) = resolve_provider_config(model, Some("azure_ai"))?;
|
||||
let segments: Vec<&str> = endpoint
|
||||
.split('/')
|
||||
.filter(|segment| !segment.is_empty())
|
||||
|
|
@ -173,7 +223,7 @@ pub fn passthrough_response(
|
|||
OcrConfigKind::AzureAi => segments == AZURE_AI_OCR_PATH,
|
||||
OcrConfigKind::AzureCohere => segments == AZURE_COHERE_PARSE_PATH,
|
||||
OcrConfigKind::AzureDocumentIntelligence => {
|
||||
segments == AzureDocumentIntelligenceOcrConfig::analyze_path(&model)?
|
||||
segments == AzureDocumentIntelligenceOcrConfig::analyze_path(provider.model)?
|
||||
}
|
||||
other => {
|
||||
let provider: &'static str = other.provider().into();
|
||||
|
|
@ -183,24 +233,28 @@ pub fn passthrough_response(
|
|||
if !is_ocr_endpoint {
|
||||
return Ok(None);
|
||||
}
|
||||
with_config!(config, config => config.transform_ocr_response(&model, body, OcrResponseFormat::Litellm))
|
||||
with_config!(config, config => config.transform_ocr_response(provider.model, body, OcrResponseFormat::Litellm))
|
||||
.map(Some)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_provider_config(
|
||||
model: &str,
|
||||
custom_llm_provider: Option<&str>,
|
||||
) -> Result<(String, OcrConfigKind), Error> {
|
||||
let provider =
|
||||
get_custom_llm_provider(model, custom_llm_provider).unwrap_or(CustomLlmProvider {
|
||||
model,
|
||||
custom_llm_provider: LlmProviders::Mistral.into(),
|
||||
});
|
||||
let llm_provider = provider
|
||||
.custom_llm_provider
|
||||
.parse::<LlmProviders>()
|
||||
.map_err(|_| Error::InvalidProvider(provider.custom_llm_provider.to_string()))?;
|
||||
let config = match llm_provider {
|
||||
pub(crate) fn resolve_provider_config<'a>(
|
||||
model: &'a str,
|
||||
custom_llm_provider: Option<&'a str>,
|
||||
) -> Result<(ResolvedProvider<'a>, OcrConfigKind), Error> {
|
||||
let provider = match resolve_llm_provider(model, custom_llm_provider, "ocr") {
|
||||
Ok(provider) => provider,
|
||||
Err(_) if get_custom_llm_provider(model, custom_llm_provider).is_none() => {
|
||||
ResolvedProvider {
|
||||
model,
|
||||
provider: LlmProviders::Mistral,
|
||||
}
|
||||
}
|
||||
Err(litellm_inference::RouteError::InvalidProvider(provider)) => {
|
||||
return Err(Error::InvalidProvider(provider));
|
||||
}
|
||||
Err(error) => return Err(Error::InvalidProvider(error.to_string())),
|
||||
};
|
||||
let config = match provider.provider {
|
||||
LlmProviders::AwsTextract => match TextractOperation::from_model(provider.model)? {
|
||||
TextractOperation::DetectDocumentText => OcrConfigKind::AwsTextract,
|
||||
TextractOperation::AnalyzeDocument => OcrConfigKind::AwsTextractAnalyze,
|
||||
|
|
@ -227,11 +281,11 @@ pub(crate) fn resolve_provider_config(
|
|||
LlmProviders::VertexAi => OcrConfigKind::VertexAi,
|
||||
_ => {
|
||||
return Err(Error::InvalidProvider(
|
||||
provider.custom_llm_provider.to_string(),
|
||||
<&'static str>::from(provider.provider).to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
Ok((provider.model.to_string(), config))
|
||||
Ok((provider, config))
|
||||
}
|
||||
|
||||
fn is_document_intelligence_model(model: &str) -> bool {
|
||||
|
|
@ -243,7 +297,7 @@ fn is_document_intelligence_model(model: &str) -> bool {
|
|||
mod tests {
|
||||
use std::collections::HashSet;
|
||||
|
||||
use litellm_auth::{InputSource, Sourced};
|
||||
use litellm_auth::Sourced;
|
||||
use litellm_llms::{
|
||||
base_llm::ocr::document::InlineDocument, cohere::ocr::transformation::validate_document,
|
||||
};
|
||||
|
|
@ -484,8 +538,8 @@ mod tests {
|
|||
#[case] expected_config: OcrConfigKind,
|
||||
) {
|
||||
let expected_model = qualified_model.split_once('/').unwrap().1;
|
||||
let (model, config) = resolve_provider_config(qualified_model, None).unwrap();
|
||||
assert_eq!(model, expected_model);
|
||||
let (provider, config) = resolve_provider_config(qualified_model, None).unwrap();
|
||||
assert_eq!(provider.model, expected_model);
|
||||
assert_eq!(config, expected_config);
|
||||
}
|
||||
|
||||
|
|
@ -528,7 +582,7 @@ mod tests {
|
|||
expected_config
|
||||
);
|
||||
assert_eq!(
|
||||
resolve_provider_config(model, None).unwrap().0,
|
||||
resolve_provider_config(model, None).unwrap().0.model,
|
||||
model.split_once('/').unwrap().1
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -75,6 +75,8 @@ impl OcrConnectionInputs {
|
|||
|
||||
pub struct LiteLLMOcrRequest<D = OcrDocumentInput> {
|
||||
pub model: String,
|
||||
pub(crate) requested_model: String,
|
||||
pub(crate) custom_llm_provider: Option<String>,
|
||||
pub document: D,
|
||||
pub credentials: OcrCredentialInputs,
|
||||
pub transport: OcrTransportConfig,
|
||||
|
|
@ -91,7 +93,8 @@ impl LiteLLMOcrRequest {
|
|||
custom_llm_provider: Option<&str>,
|
||||
optional_params: CallArguments,
|
||||
) -> Result<Self, Error> {
|
||||
let (model, config) = resolve_provider_config(&model, custom_llm_provider)?;
|
||||
let (provider, config) = resolve_provider_config(&model, custom_llm_provider)?;
|
||||
let resolved_model = provider.model.to_owned();
|
||||
let default_transport = OcrTransportConfig::default();
|
||||
let max_response_bytes = optional_params
|
||||
.get("max_response_bytes")
|
||||
|
|
@ -116,7 +119,9 @@ impl LiteLLMOcrRequest {
|
|||
.collect();
|
||||
|
||||
Ok(Self {
|
||||
model,
|
||||
model: resolved_model,
|
||||
requested_model: model,
|
||||
custom_llm_provider: custom_llm_provider.map(str::to_owned),
|
||||
document: document.into(),
|
||||
credentials: OcrCredentialInputs::default(),
|
||||
transport,
|
||||
|
|
@ -135,6 +140,8 @@ impl<D> LiteLLMOcrRequest<D> {
|
|||
) -> Result<LiteLLMOcrRequest<T>, E> {
|
||||
Ok(LiteLLMOcrRequest {
|
||||
model: self.model,
|
||||
requested_model: self.requested_model,
|
||||
custom_llm_provider: self.custom_llm_provider,
|
||||
document: map(self.document)?,
|
||||
credentials: self.credentials,
|
||||
transport: self.transport,
|
||||
|
|
@ -148,6 +155,8 @@ impl<D> LiteLLMOcrRequest<D> {
|
|||
pub fn with_document<T>(self, document: T) -> LiteLLMOcrRequest<T> {
|
||||
LiteLLMOcrRequest {
|
||||
model: self.model,
|
||||
requested_model: self.requested_model,
|
||||
custom_llm_provider: self.custom_llm_provider,
|
||||
document,
|
||||
credentials: self.credentials,
|
||||
transport: self.transport,
|
||||
|
|
|
|||
|
|
@ -66,6 +66,19 @@ async fn interception_runs_once_and_observers_receive_ordered_success_events() {
|
|||
assert_eq!(received(&upstream).await.len(), 1);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::mutated_model("model-b")]
|
||||
#[tokio::test]
|
||||
async fn provider_request_uses_the_request_model_after_it_is_mutated(#[case] model: &str) {
|
||||
let upstream = upstream([pages_response()]).await;
|
||||
let mut request = ocr_request("mistral/model-a", &upstream.uri(), json!({}));
|
||||
request.model = model.into();
|
||||
|
||||
perform(request).await.unwrap();
|
||||
|
||||
assert_eq!(only_request(&upstream).await.json()["model"], json!(model));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() {
|
||||
|
|
|
|||
|
|
@ -2,3 +2,4 @@
|
|||
- This crate owns Responses API call orchestration. Payloads belong in `litellm-llms-types`, transformations in `llms/src/base_llm/responses` and `llms/src/<provider>/responses`
|
||||
- The HTTP route returns `litellm_host::call::CallOutput`: a completed response, or a stream head and chunks, cached through `litellm_inference::caching::execute_streaming`
|
||||
- WebSocket sessions (`src/websocket.rs`) stay separate from the HTTP call driver because a connection can accept multiple requests while receiving events
|
||||
- `src/ids.rs` owns the LiteLLM-managed ids handed back to callers (`resp_`, `encitem_`, `litellm_enc:`, `cntr_`) and must stay fully compatible with `litellm/responses/utils.py`
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
base64.workspace = true
|
||||
bytes.workspace = true
|
||||
futures-util.workspace = true
|
||||
litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] }
|
||||
|
|
|
|||
|
|
@ -3,39 +3,61 @@ use litellm_host::observation::ObservationSender;
|
|||
use std::time::Duration;
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use litellm_host::interceptors::{Interceptors, RawResponse, WireRequest};
|
||||
use litellm_llms::base_llm::auth::{Authenticated, resolve_auth};
|
||||
use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest};
|
||||
use litellm_llms::base_llm::auth::{AuthScheme, Authenticated, resolve_auth};
|
||||
use litellm_secrets::source::Secrets;
|
||||
|
||||
use super::{
|
||||
Error,
|
||||
types::{ProviderResponsesRequest, ResponsesOutput, ResponsesStreamHead},
|
||||
Error, ResponsesRoute,
|
||||
types::{ResponsesCall, ResponsesOutput, ResponsesProvider, ResponsesStreamHead},
|
||||
};
|
||||
|
||||
pub(super) async fn execute(
|
||||
http: &litellm_http::Client,
|
||||
auth: &litellm_auth::AuthServices,
|
||||
request: ProviderResponsesRequest,
|
||||
cache: Option<litellm_cache_response::ScopedCache>,
|
||||
route: &ResponsesRoute,
|
||||
ResponsesProvider { identity, config }: ResponsesProvider,
|
||||
call: ResponsesCall,
|
||||
secrets: Secrets,
|
||||
cache_options: Option<litellm_cache_response::CachePolicy>,
|
||||
interceptors: &impl Interceptors<Error>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<ResponsesOutput, Error> {
|
||||
let authenticated = resolve_auth(auth, request.environment, &|_| None).await?;
|
||||
let identity = litellm_host::interceptors::ProviderIdentity {
|
||||
model: request.context.model.clone(),
|
||||
provider: request.context.custom_llm_provider.clone(),
|
||||
let lookup = |name: &str| secrets.get(name);
|
||||
let environment = config.validate_environment(
|
||||
litellm_http::request::string_headers("responses", call.extra_headers)?,
|
||||
call.api_key.as_deref(),
|
||||
&lookup,
|
||||
)?;
|
||||
let url = config.get_complete_url(call.api_base.as_deref(), &lookup);
|
||||
let context = RequestContext {
|
||||
model: identity.model.clone(),
|
||||
custom_llm_provider: identity.provider.clone(),
|
||||
optional_params: serde_json::Value::Object(call.optional_params.clone()),
|
||||
secret_fields: Vec::new(),
|
||||
api_key: match &environment.auth {
|
||||
AuthScheme::Credential { secret, .. } => Some(secret.clone()),
|
||||
_ => None,
|
||||
},
|
||||
};
|
||||
let body = config.transform_responses_api_request(
|
||||
&identity.model,
|
||||
call.input,
|
||||
call.optional_params,
|
||||
)?;
|
||||
let authenticated = resolve_auth(&route.auth, environment, &|_| None).await?;
|
||||
let wire = interceptors
|
||||
.before_provider_request(
|
||||
WireRequest {
|
||||
url: request.url,
|
||||
url,
|
||||
headers: authenticated.headers,
|
||||
body: request.body,
|
||||
body,
|
||||
},
|
||||
request.context,
|
||||
context,
|
||||
)
|
||||
.await?;
|
||||
let cache = cache.filter(|_| authenticated.signer.is_none());
|
||||
let cache = route
|
||||
.cache
|
||||
.clone()
|
||||
.filter(|_| authenticated.signer.is_none());
|
||||
let cache_request = litellm_inference::caching::CacheRequest::from_wire(
|
||||
identity,
|
||||
cache.as_ref().map(|_| &wire),
|
||||
|
|
@ -59,9 +81,9 @@ pub(super) async fn execute(
|
|||
},
|
||||
wire.url,
|
||||
&wire.body,
|
||||
Some(request.timeout.unwrap_or(Duration::from_secs(600))),
|
||||
Some(call.timeout.unwrap_or(Duration::from_secs(600))),
|
||||
)?;
|
||||
let response = litellm_inference::outbound::send(outbound, http)
|
||||
let response = litellm_inference::outbound::send(outbound, &route.http)
|
||||
.await
|
||||
.map_err(network)?;
|
||||
let status = response.status().as_u16();
|
||||
|
|
@ -103,8 +125,7 @@ pub(super) async fn execute(
|
|||
.map_err(Error::post_call)?;
|
||||
let value = serde_json::from_str(&body)
|
||||
.map_err(|error| Error::InvalidResponse(error.to_string().into()))?;
|
||||
request
|
||||
.config
|
||||
config
|
||||
.transform_response_api_response(value)
|
||||
.map(ResponsesOutput::Complete)
|
||||
.map_err(Error::from)
|
||||
|
|
|
|||
499
litellm-rust/crates/inference-responses/src/ids.rs
Normal file
499
litellm-rust/crates/inference-responses/src/ids.rs
Normal file
|
|
@ -0,0 +1,499 @@
|
|||
//! LiteLLM-managed identifiers on Responses payloads, fully compatible with `litellm/responses/utils.py`
|
||||
//!
|
||||
//! LiteLLM rewrites some ids in a Responses API result before handing it back, so that a later
|
||||
//! request carrying those ids can be routed to the provider and deployment that produced them,
|
||||
//! and then restored to what the provider originally issued. The native and Python paths serve
|
||||
//! the same traffic, so an id minted by either one must decode on the other.
|
||||
//!
|
||||
//! Each managed id is a fixed prefix followed by standard, padded base64 of a `;`-separated
|
||||
//! plaintext:
|
||||
//!
|
||||
//! | Id | Applied to | Decoded plaintext |
|
||||
//! |---|---|---|
|
||||
//! | `resp_<b64>` | the response `id` | `litellm:custom_llm_provider:{provider};model_id:{model_id};response_id:{id}` |
|
||||
//! | `encitem_<b64>` | the `id` of an output item with `encrypted_content` | `litellm:model_id:{model_id};item_id:{id}` |
|
||||
//! | `litellm_enc:<b64>;{content}` | that item's `encrypted_content` | `<b64>` is `model_id:{model_id}`, and the original content follows the `;` unencoded |
|
||||
//! | `cntr_<b64>` | `container_id` on output items, `code_interpreter_call`, and message annotations | `litellm:custom_llm_provider:{provider};model_id:{model_id};container_id:{id}` |
|
||||
//!
|
||||
//! For example, OpenAI's `resp_0a1b2c3d4e5f` served without a deployment id becomes
|
||||
//! `resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDpOb25lO3Jlc3BvbnNlX2lkOnJlc3BfMGExYjJjM2Q0ZTVm`,
|
||||
//! which decodes to `litellm:custom_llm_provider:openai;model_id:None;response_id:resp_0a1b2c3d4e5f`.
|
||||
//!
|
||||
//! How they fit together:
|
||||
//! - `resp_` lets a follow-up's `previous_response_id` find the provider and deployment again.
|
||||
//! [`encode_response`] stamps it on every response, unless the id is already managed.
|
||||
//! - `encitem_` and `litellm_enc:` carry encrypted-content affinity: encrypted reasoning only
|
||||
//! decrypts under the API key that produced it, so the router pins replays to that deployment.
|
||||
//! They are stamped only when the router sets `encrypted_content_affinity_enabled` and a
|
||||
//! deployment `model_id` is known (see [`IdScope`]). Items without an `id`, such as Codex's,
|
||||
//! carry the deployment in the wrapped content alone.
|
||||
//! - `cntr_` lets container calls find the provider and deployment that own a container. A
|
||||
//! container id that already names a provider is left as it is.
|
||||
//! - On the way in, [`restore_input`] turns `encitem_` ids and `litellm_enc:` content in the
|
||||
//! request input back into the provider's own bytes before the request is sent. The routing
|
||||
//! decision itself stays in Python's `EncryptedContentAffinityCheck`.
|
||||
//!
|
||||
//! Python's quirks are kept on purpose:
|
||||
//! - A missing provider or `model_id` is written as `None` in `resp_` ids (an f-string of
|
||||
//! `None`), but as an empty string in `cntr_` ids. [`decode_container_id`] reads both as absent.
|
||||
//! - Decoding follows Python's non-strict `base64.b64decode`: characters outside the alphabet
|
||||
//! are skipped, trailing bits are allowed, and padding stripped in transit is added back.
|
||||
//! - A `resp_` id counts as managed when it decodes to at least three `;`-separated parts.
|
||||
|
||||
use base64::{
|
||||
Engine, alphabet,
|
||||
engine::{GeneralPurpose, general_purpose::PAD},
|
||||
};
|
||||
use litellm_llms_types::formats::responses::ResponsesApiResponse;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
const RESPONSE_PREFIX: &str = "resp_";
|
||||
const ITEM_PREFIX: &str = "encitem_";
|
||||
const CONTAINER_PREFIX: &str = "cntr_";
|
||||
const WRAPPED_CONTENT_PREFIX: &str = "litellm_enc:";
|
||||
|
||||
/// Python's non-strict `base64.b64decode`: characters outside the alphabet are skipped and
|
||||
/// non-zero trailing bits are accepted
|
||||
const LENIENT: GeneralPurpose = GeneralPurpose::new(
|
||||
&alphabet::STANDARD,
|
||||
PAD.with_decode_allow_trailing_bits(true),
|
||||
);
|
||||
|
||||
/// Where a response came from, as encoded into the ids LiteLLM hands back
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct IdScope {
|
||||
pub provider: Option<String>,
|
||||
pub model_id: Option<String>,
|
||||
pub encrypted_content_affinity: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct DecodedContainerId {
|
||||
pub provider: Option<String>,
|
||||
pub model_id: Option<String>,
|
||||
pub container_id: String,
|
||||
}
|
||||
|
||||
/// Rewrites the response id, encrypted-content items and container ids the way the Python
|
||||
/// Responses path does before returning to the caller. An already managed id is left untouched
|
||||
pub fn encode_response(response: &mut ResponsesApiResponse, scope: &IdScope) {
|
||||
if is_managed_response_id(&response.id) {
|
||||
return;
|
||||
}
|
||||
response.id = response_id(scope, &response.id);
|
||||
if let Some(model_id) = scope
|
||||
.model_id
|
||||
.as_deref()
|
||||
.filter(|_| scope.encrypted_content_affinity)
|
||||
.filter(|model_id| !model_id.is_empty())
|
||||
{
|
||||
response
|
||||
.output
|
||||
.iter_mut()
|
||||
.filter_map(Value::as_object_mut)
|
||||
.for_each(|item| encode_encrypted_item(item, model_id));
|
||||
}
|
||||
response
|
||||
.output
|
||||
.iter_mut()
|
||||
.for_each(|item| encode_containers(item, scope));
|
||||
}
|
||||
|
||||
/// Restores encrypted-content item ids and contents that LiteLLM encoded on an earlier response,
|
||||
/// so the provider receives what it originally issued
|
||||
pub fn restore_input(input: &mut Value) {
|
||||
let Some(items) = input.as_array_mut() else {
|
||||
return;
|
||||
};
|
||||
for item in items.iter_mut().filter_map(Value::as_object_mut) {
|
||||
if let Some(Value::String(id)) = item.get_mut("id")
|
||||
&& let Some((_, original)) = decode_encrypted_item_id(id)
|
||||
{
|
||||
*id = original;
|
||||
}
|
||||
if let Some(Value::String(content)) = item.get_mut("encrypted_content")
|
||||
&& let Some(original) = unwrap_encrypted_content(content)
|
||||
{
|
||||
*content = original;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn response_id(scope: &IdScope, id: &str) -> String {
|
||||
let assembled = format!(
|
||||
"litellm:custom_llm_provider:{};model_id:{};response_id:{id}",
|
||||
python_str(scope.provider.as_deref()),
|
||||
python_str(scope.model_id.as_deref()),
|
||||
);
|
||||
format!("{RESPONSE_PREFIX}{}", LENIENT.encode(assembled))
|
||||
}
|
||||
|
||||
pub fn is_managed_response_id(id: &str) -> bool {
|
||||
decode_text(&id.replace(RESPONSE_PREFIX, ""))
|
||||
.is_some_and(|decoded| decoded.split(';').count() >= 3)
|
||||
}
|
||||
|
||||
pub fn encrypted_item_id(model_id: &str, id: &str) -> String {
|
||||
format!(
|
||||
"{ITEM_PREFIX}{}",
|
||||
LENIENT.encode(format!("litellm:model_id:{model_id};item_id:{id}"))
|
||||
)
|
||||
}
|
||||
|
||||
pub fn decode_encrypted_item_id(id: &str) -> Option<(String, String)> {
|
||||
let decoded = decode_text(&padded(id.strip_prefix(ITEM_PREFIX)?))?;
|
||||
let (model_id, item_id) = decoded.split_once(';')?;
|
||||
Some((
|
||||
model_id.replace("litellm:model_id:", ""),
|
||||
item_id.replace("item_id:", ""),
|
||||
))
|
||||
}
|
||||
|
||||
pub fn wrap_encrypted_content(content: &str, model_id: &str) -> String {
|
||||
format!(
|
||||
"{WRAPPED_CONTENT_PREFIX}{};{content}",
|
||||
LENIENT.encode(format!("model_id:{model_id}"))
|
||||
)
|
||||
}
|
||||
|
||||
/// The original content of a wrapped `encrypted_content`, or `None` when it is not wrapped
|
||||
pub fn unwrap_encrypted_content(content: &str) -> Option<String> {
|
||||
if !content.starts_with(WRAPPED_CONTENT_PREFIX) {
|
||||
return None;
|
||||
}
|
||||
let (metadata, original) = content.split_once(';')?;
|
||||
decode_text(&padded(&metadata.replace(WRAPPED_CONTENT_PREFIX, "")))?;
|
||||
Some(original.to_owned())
|
||||
}
|
||||
|
||||
pub fn container_id(scope: &IdScope, id: &str) -> String {
|
||||
let assembled = format!(
|
||||
"litellm:custom_llm_provider:{};model_id:{};container_id:{id}",
|
||||
scope.provider.as_deref().unwrap_or_default(),
|
||||
scope.model_id.as_deref().unwrap_or_default(),
|
||||
);
|
||||
format!("{CONTAINER_PREFIX}{}", LENIENT.encode(assembled))
|
||||
}
|
||||
|
||||
pub fn decode_container_id(id: &str) -> DecodedContainerId {
|
||||
let unmanaged = || DecodedContainerId {
|
||||
provider: None,
|
||||
model_id: None,
|
||||
container_id: id.to_owned(),
|
||||
};
|
||||
if !id.starts_with(CONTAINER_PREFIX) {
|
||||
return unmanaged();
|
||||
}
|
||||
let Some(decoded) = decode_text(&id.replace(CONTAINER_PREFIX, "")) else {
|
||||
return unmanaged();
|
||||
};
|
||||
let parsed = decoded
|
||||
.strip_prefix("litellm:custom_llm_provider:")
|
||||
.and_then(|rest| rest.split_once(';'))
|
||||
.and_then(|(provider, rest)| {
|
||||
let (model_id, rest) = rest.strip_prefix("model_id:")?.split_once(';')?;
|
||||
let container_id = rest.strip_prefix("container_id:")?;
|
||||
(!container_id.is_empty() && !container_id.contains('\n')).then_some((
|
||||
provider,
|
||||
model_id,
|
||||
container_id,
|
||||
))
|
||||
});
|
||||
let Some((provider, model_id, container_id)) = parsed else {
|
||||
return unmanaged();
|
||||
};
|
||||
let present = |value: &str| (!matches!(value, "" | "None")).then(|| value.to_owned());
|
||||
DecodedContainerId {
|
||||
provider: present(provider),
|
||||
model_id: present(model_id),
|
||||
container_id: container_id.to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
fn encode_encrypted_item(item: &mut Map<String, Value>, model_id: &str) {
|
||||
let Some(Value::String(content)) = item.get_mut("encrypted_content") else {
|
||||
return;
|
||||
};
|
||||
if content.is_empty() {
|
||||
return;
|
||||
}
|
||||
*content = wrap_encrypted_content(content, model_id);
|
||||
if let Some(Value::String(id)) = item.get_mut("id")
|
||||
&& !id.is_empty()
|
||||
{
|
||||
*id = encrypted_item_id(model_id, id);
|
||||
}
|
||||
}
|
||||
|
||||
fn encode_containers(item: &mut Value, scope: &IdScope) {
|
||||
let Some(item) = item.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
encode_container_field(item, scope);
|
||||
if let Some(nested) = item
|
||||
.get_mut("code_interpreter_call")
|
||||
.and_then(Value::as_object_mut)
|
||||
{
|
||||
encode_container_field(nested, scope);
|
||||
}
|
||||
if item.get("type").and_then(Value::as_str) != Some("message") {
|
||||
return;
|
||||
}
|
||||
let Some(content) = item.get_mut("content").and_then(Value::as_array_mut) else {
|
||||
return;
|
||||
};
|
||||
content
|
||||
.iter_mut()
|
||||
.filter_map(|part| part.get_mut("annotations").and_then(Value::as_array_mut))
|
||||
.flatten()
|
||||
.for_each(|annotation| encode_containers(annotation, scope));
|
||||
}
|
||||
|
||||
fn encode_container_field(item: &mut Map<String, Value>, scope: &IdScope) {
|
||||
if let Some(Value::String(id)) = item.get_mut("container_id")
|
||||
&& decode_container_id(id).provider.is_none()
|
||||
{
|
||||
*id = container_id(scope, id);
|
||||
}
|
||||
}
|
||||
|
||||
fn python_str(value: Option<&str>) -> &str {
|
||||
value.unwrap_or("None")
|
||||
}
|
||||
|
||||
fn padded(value: &str) -> String {
|
||||
let missing = value.len() % 4;
|
||||
let padding = if missing == 0 { 0 } else { 4 - missing };
|
||||
format!("{value}{}", "=".repeat(padding))
|
||||
}
|
||||
|
||||
fn decode_text(value: &str) -> Option<String> {
|
||||
let alphabet: String = value
|
||||
.chars()
|
||||
.filter(|character| {
|
||||
character.is_ascii_alphanumeric() || matches!(character, '+' | '/' | '=')
|
||||
})
|
||||
.collect();
|
||||
String::from_utf8(LENIENT.decode(alphabet).ok()?).ok()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn scope(provider: Option<&str>, model_id: Option<&str>, affinity: bool) -> IdScope {
|
||||
IdScope {
|
||||
provider: provider.map(str::to_owned),
|
||||
model_id: model_id.map(str::to_owned),
|
||||
encrypted_content_affinity: affinity,
|
||||
}
|
||||
}
|
||||
|
||||
fn response(id: &str, output: Value) -> ResponsesApiResponse {
|
||||
serde_json::from_value(json!({"id": id, "model": "gpt-6.1-sol", "output": output})).unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::openai_without_model_id(
|
||||
Some("openai"),
|
||||
None,
|
||||
"resp_0a1b2c3d4e5f",
|
||||
"resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDpOb25lO3Jlc3BvbnNlX2lkOnJlc3BfMGExYjJjM2Q0ZTVm"
|
||||
)]
|
||||
#[case::openrouter_with_model_id(
|
||||
Some("openrouter"),
|
||||
Some("deploy-1"),
|
||||
"gen-123",
|
||||
"resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5yb3V0ZXI7bW9kZWxfaWQ6ZGVwbG95LTE7cmVzcG9uc2VfaWQ6Z2VuLTEyMw=="
|
||||
)]
|
||||
fn response_ids_match_python(
|
||||
#[case] provider: Option<&str>,
|
||||
#[case] model_id: Option<&str>,
|
||||
#[case] id: &str,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
assert_eq!(response_id(&scope(provider, model_id, false), id), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::raw_short_id("resp_0a1b2c3d4e5f", false)]
|
||||
#[case::raw_long_id("resp_68a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8a9b0c1d2e3", false)]
|
||||
#[case::two_segments("resp_YTti", false)]
|
||||
#[case::three_segments("resp_YTtiO2M=", true)]
|
||||
#[case::managed(
|
||||
"resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDpOb25lO3Jlc3BvbnNlX2lkOng=",
|
||||
true
|
||||
)]
|
||||
fn managed_response_ids_match_python(#[case] id: &str, #[case] expected: bool) {
|
||||
assert_eq!(is_managed_response_id(id), expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encrypted_items_match_python() {
|
||||
assert_eq!(
|
||||
encrypted_item_id("deploy-1", "rs_0a1b"),
|
||||
"encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOnJzXzBhMWI="
|
||||
);
|
||||
assert_eq!(
|
||||
wrap_encrypted_content("gAAAA", "deploy-1"),
|
||||
"litellm_enc:bW9kZWxfaWQ6ZGVwbG95LTE=;gAAAA"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::padded("encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOnJzXzBhMWI=", Some(("deploy-1", "rs_0a1b")))]
|
||||
#[case::padding_stripped_in_transit("encitem_bGl0ZWxsbTptb2RlbF9pZDptO2l0ZW1faWQ6cnNfMQ", Some(("m", "rs_1")))]
|
||||
#[case::not_base64("encitem_notbase64!", None)]
|
||||
#[case::unmanaged("rs_0a1b", None)]
|
||||
fn encrypted_item_ids_decode_like_python(
|
||||
#[case] id: &str,
|
||||
#[case] expected: Option<(&str, &str)>,
|
||||
) {
|
||||
assert_eq!(
|
||||
decode_encrypted_item_id(id),
|
||||
expected.map(|(model_id, item_id)| (model_id.to_owned(), item_id.to_owned()))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unpadded_metadata("litellm_enc:bW9kZWxfaWQ6bQ;gAAAA", Some("gAAAA"))]
|
||||
#[case::no_separator("litellm_enc:nosemicolon", None)]
|
||||
#[case::unwrapped("gAAAA", None)]
|
||||
fn encrypted_content_unwraps_like_python(
|
||||
#[case] content: &str,
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
assert_eq!(unwrap_encrypted_content(content).as_deref(), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::with_model_id(
|
||||
Some("deploy-1"),
|
||||
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDpkZXBsb3ktMTtjb250YWluZXJfaWQ6Y250cl9hYmM="
|
||||
)]
|
||||
#[case::without_model_id(
|
||||
None,
|
||||
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDo7Y29udGFpbmVyX2lkOmNudHJfYWJj"
|
||||
)]
|
||||
fn container_ids_match_python(#[case] model_id: Option<&str>, #[case] expected: &str) {
|
||||
assert_eq!(
|
||||
container_id(&scope(Some("openai"), model_id, false), "cntr_abc"),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::raw("cntr_abc", None, None, "cntr_abc")]
|
||||
#[case::managed(
|
||||
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDptO2NvbnRhaW5lcl9pZDpjbnRyXzE=",
|
||||
Some("openai"),
|
||||
Some("m"),
|
||||
"cntr_1"
|
||||
)]
|
||||
#[case::empty_segments(
|
||||
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOjttb2RlbF9pZDo7Y29udGFpbmVyX2lkOmNudHJfMQ==",
|
||||
None,
|
||||
None,
|
||||
"cntr_1"
|
||||
)]
|
||||
#[case::python_none_provider(
|
||||
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOk5vbmU7bW9kZWxfaWQ6O2NvbnRhaW5lcl9pZDp4",
|
||||
None,
|
||||
None,
|
||||
"x"
|
||||
)]
|
||||
fn container_ids_decode_like_python(
|
||||
#[case] id: &str,
|
||||
#[case] provider: Option<&str>,
|
||||
#[case] model_id: Option<&str>,
|
||||
#[case] container_id: &str,
|
||||
) {
|
||||
assert_eq!(
|
||||
decode_container_id(id),
|
||||
DecodedContainerId {
|
||||
provider: provider.map(str::to_owned),
|
||||
model_id: model_id.map(str::to_owned),
|
||||
container_id: container_id.to_owned(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encodes_response_items_under_affinity() {
|
||||
let mut encoded = response(
|
||||
"resp_1",
|
||||
json!([
|
||||
{"type": "reasoning", "id": "rs_0a1b", "encrypted_content": "gAAAA", "summary": []},
|
||||
{"type": "reasoning", "encrypted_content": "gBBBB", "summary": []},
|
||||
{"type": "code_interpreter_call", "id": "ci_1", "container_id": "cntr_abc"},
|
||||
{"type": "message", "id": "msg_1", "content": [{"type": "output_text", "text": "", "annotations": [
|
||||
{"type": "container_file_citation", "container_id": "cntr_abc"}
|
||||
]}]}
|
||||
]),
|
||||
);
|
||||
let scope = scope(Some("openai"), Some("deploy-1"), true);
|
||||
encode_response(&mut encoded, &scope);
|
||||
let container = container_id(&scope, "cntr_abc");
|
||||
assert_eq!(encoded.id, response_id(&scope, "resp_1"));
|
||||
assert_eq!(
|
||||
Value::Array(encoded.output),
|
||||
json!([
|
||||
{
|
||||
"type": "reasoning",
|
||||
"id": "encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOnJzXzBhMWI=",
|
||||
"encrypted_content": "litellm_enc:bW9kZWxfaWQ6ZGVwbG95LTE=;gAAAA",
|
||||
"summary": []
|
||||
},
|
||||
{"type": "reasoning", "encrypted_content": "litellm_enc:bW9kZWxfaWQ6ZGVwbG95LTE=;gBBBB", "summary": []},
|
||||
{"type": "code_interpreter_call", "id": "ci_1", "container_id": container},
|
||||
{"type": "message", "id": "msg_1", "content": [{"type": "output_text", "text": "", "annotations": [
|
||||
{"type": "container_file_citation", "container_id": container}
|
||||
]}]}
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::affinity_disabled(scope(Some("openai"), Some("deploy-1"), false))]
|
||||
#[case::no_model_id(scope(Some("openai"), None, true))]
|
||||
fn leaves_encrypted_items_without_affinity(#[case] scope: IdScope) {
|
||||
let output = json!([{"type": "reasoning", "id": "rs_1", "encrypted_content": "gAAAA", "summary": []}]);
|
||||
let mut encoded = response("resp_1", output.clone());
|
||||
encode_response(&mut encoded, &scope);
|
||||
assert_eq!(Value::Array(encoded.output), output);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn leaves_managed_responses_untouched() {
|
||||
let scope = scope(Some("openai"), Some("deploy-1"), true);
|
||||
let id = response_id(&scope, "resp_1");
|
||||
let output = json!([{"type": "code_interpreter_call", "container_id": "cntr_abc"}]);
|
||||
let mut encoded = response(&id, output.clone());
|
||||
encode_response(&mut encoded, &scope);
|
||||
assert_eq!((encoded.id, Value::Array(encoded.output)), (id, output));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restores_encrypted_items_in_input() {
|
||||
let mut input = json!([
|
||||
{
|
||||
"type": "reasoning",
|
||||
"id": "encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOnJzXzBhMWI=",
|
||||
"encrypted_content": "litellm_enc:bW9kZWxfaWQ6ZGVwbG95LTE=;gAAAA",
|
||||
"summary": []
|
||||
},
|
||||
{"role": "user", "content": "hi"}
|
||||
]);
|
||||
restore_input(&mut input);
|
||||
assert_eq!(
|
||||
input,
|
||||
json!([
|
||||
{"type": "reasoning", "id": "rs_0a1b", "encrypted_content": "gAAAA", "summary": []},
|
||||
{"role": "user", "content": "hi"}
|
||||
])
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,19 +1,36 @@
|
|||
use litellm_host::observation::ObservationSender;
|
||||
pub use litellm_inference::RouteError as Error;
|
||||
|
||||
pub mod ids;
|
||||
pub mod route;
|
||||
pub mod types;
|
||||
pub mod websocket;
|
||||
|
||||
mod handler;
|
||||
mod prepare;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_core_utils::get_llm_provider_logic::{
|
||||
CustomLlmProvider, LlmProviders, get_custom_llm_provider,
|
||||
};
|
||||
use litellm_host::interceptors::Interceptors;
|
||||
use litellm_host::interceptors::ProviderIdentity;
|
||||
use litellm_llms::{
|
||||
base_llm::responses::transformation::BaseResponsesApiConfig,
|
||||
openai::responses::transformation::OpenAiResponsesApiConfig,
|
||||
openrouter::responses::transformation::OPENROUTER_RESPONSES_CONFIG,
|
||||
};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use types::{ResponsesCall, ResponsesOutput};
|
||||
use types::{ResponsesCall, ResponsesOutput, ResponsesProvider};
|
||||
|
||||
fn responses_provider(provider: LlmProviders) -> Option<&'static dyn BaseResponsesApiConfig> {
|
||||
match provider {
|
||||
LlmProviders::Openai => Some(&OpenAiResponsesApiConfig),
|
||||
LlmProviders::Openrouter => Some(&OPENROUTER_RESPONSES_CONFIG),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ResponsesRoute {
|
||||
|
|
@ -77,34 +94,68 @@ impl ResponsesRoute {
|
|||
observers: Option<&ObservationSender>,
|
||||
) -> Result<ResponsesOutput, Error> {
|
||||
litellm_inference::diagnostic::call(async {
|
||||
self.run_provider(call, cache_options, interceptors, observers)
|
||||
.await
|
||||
let explicit = call
|
||||
.custom_llm_provider
|
||||
.as_deref()
|
||||
.filter(|provider| !provider.is_empty());
|
||||
let (provider_name, resolved_model) =
|
||||
match get_custom_llm_provider(&call.model, explicit) {
|
||||
Some(CustomLlmProvider {
|
||||
model,
|
||||
custom_llm_provider,
|
||||
}) => (custom_llm_provider, model),
|
||||
None => ("openai", call.model.as_str()),
|
||||
};
|
||||
if resolved_model.is_empty() {
|
||||
return Err(Error::InvalidProvider(call.model));
|
||||
}
|
||||
let provider = provider_name
|
||||
.parse::<LlmProviders>()
|
||||
.map_err(|_| Error::Unsupported("native HTTP responses provider"))?;
|
||||
let config = responses_provider(provider)
|
||||
.ok_or(Error::Unsupported("native HTTP responses provider"))?;
|
||||
let wire_model = match provider {
|
||||
LlmProviders::Openai if resolved_model.contains('/') => {
|
||||
return Err(Error::InvalidProvider(call.model));
|
||||
}
|
||||
LlmProviders::Openai => resolved_model.into(),
|
||||
LlmProviders::Openrouter if explicit.is_some() && !resolved_model.contains('/') => {
|
||||
format!("openrouter/{resolved_model}")
|
||||
}
|
||||
LlmProviders::Openrouter => resolved_model.into(),
|
||||
_ => return Err(Error::Unsupported("native HTTP responses provider")),
|
||||
};
|
||||
let provider = ResponsesProvider {
|
||||
identity: ProviderIdentity {
|
||||
model: wire_model,
|
||||
provider: provider_name.into(),
|
||||
},
|
||||
config,
|
||||
};
|
||||
litellm_inference::diagnostic::provider(
|
||||
&provider.identity.model,
|
||||
&provider.identity.provider,
|
||||
);
|
||||
let secrets = self
|
||||
.secrets
|
||||
.resolve(
|
||||
provider
|
||||
.config
|
||||
.secret_names(call.api_key.as_deref(), call.api_base.as_deref()),
|
||||
)
|
||||
.await?;
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<ResponsesOutput, Error>> =
|
||||
Box::pin(handler::execute(
|
||||
self,
|
||||
provider,
|
||||
call,
|
||||
secrets,
|
||||
cache_options,
|
||||
interceptors,
|
||||
observers,
|
||||
));
|
||||
execute.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_provider(
|
||||
&self,
|
||||
call: ResponsesCall,
|
||||
cache_options: Option<litellm_cache_response::CachePolicy>,
|
||||
interceptors: &impl Interceptors<Error>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<ResponsesOutput, Error> {
|
||||
let request = prepare::prepare(call, self.secrets.as_ref()).await?;
|
||||
litellm_inference::diagnostic::provider(
|
||||
&request.context.model,
|
||||
&request.context.custom_llm_provider,
|
||||
);
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<ResponsesOutput, Error>> =
|
||||
Box::pin(handler::execute(
|
||||
&self.http,
|
||||
&self.auth,
|
||||
request,
|
||||
self.cache.clone(),
|
||||
cache_options,
|
||||
interceptors,
|
||||
observers,
|
||||
));
|
||||
execute.await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,162 +0,0 @@
|
|||
use litellm_core_utils::get_llm_provider_logic::{
|
||||
CustomLlmProvider, LlmProviders, get_custom_llm_provider,
|
||||
};
|
||||
use litellm_host::interceptors::{ProviderIdentity, RequestContext};
|
||||
use litellm_llms::{
|
||||
base_llm::responses::transformation::BaseResponsesApiConfig,
|
||||
openai::responses::transformation::OpenAiResponsesApiConfig,
|
||||
openrouter::responses::transformation::OPENROUTER_RESPONSES_CONFIG,
|
||||
};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
Error,
|
||||
types::{ProviderResponsesRequest, ResponsesCall},
|
||||
};
|
||||
|
||||
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
|
||||
pub(super) async fn prepare(
|
||||
call: ResponsesCall,
|
||||
secrets: &dyn SecretSource,
|
||||
) -> Result<ProviderResponsesRequest, Error> {
|
||||
let ResponsesProvider { identity, config } =
|
||||
resolve_provider(&call.model, call.custom_llm_provider.as_deref())?;
|
||||
let provider = identity.provider.as_str();
|
||||
let model = identity.model.as_str();
|
||||
let snapshot = secrets
|
||||
.resolve(config.secret_names(call.api_key.as_deref(), call.api_base.as_deref()))
|
||||
.await?;
|
||||
let lookup = |name: &str| snapshot.get(name);
|
||||
let environment = config.validate_environment(
|
||||
litellm_http::request::string_headers("responses", call.extra_headers)?,
|
||||
call.api_key.as_deref(),
|
||||
&lookup,
|
||||
)?;
|
||||
let context = RequestContext {
|
||||
model: model.into(),
|
||||
custom_llm_provider: provider.into(),
|
||||
optional_params: Value::Object(call.optional_params.clone()),
|
||||
secret_fields: Vec::new(),
|
||||
api_key: match &environment.auth {
|
||||
litellm_llms::base_llm::auth::AuthScheme::Credential { secret, .. } => {
|
||||
Some(secret.clone())
|
||||
}
|
||||
_ => None,
|
||||
},
|
||||
};
|
||||
let body = config.transform_responses_api_request(model, call.input, call.optional_params)?;
|
||||
Ok(ProviderResponsesRequest {
|
||||
url: config.get_complete_url(call.api_base.as_deref(), &lookup),
|
||||
config,
|
||||
environment,
|
||||
body,
|
||||
context,
|
||||
timeout: call.timeout,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) struct ResponsesProvider {
|
||||
pub identity: ProviderIdentity,
|
||||
pub config: &'static dyn BaseResponsesApiConfig,
|
||||
}
|
||||
|
||||
/// Python's `get_llm_provider`: a bare model is OpenAI's, and with an explicit OpenRouter
|
||||
/// provider a model without a vendor segment keeps the `openrouter/` prefix as its wire id.
|
||||
pub(super) fn resolve_provider(
|
||||
model: &str,
|
||||
custom_llm_provider: Option<&str>,
|
||||
) -> Result<ResponsesProvider, Error> {
|
||||
let explicit = custom_llm_provider.filter(|provider| !provider.is_empty());
|
||||
let (provider, resolved) = match get_custom_llm_provider(model, explicit) {
|
||||
Some(CustomLlmProvider {
|
||||
model,
|
||||
custom_llm_provider,
|
||||
}) => (custom_llm_provider, model),
|
||||
None => ("openai", model),
|
||||
};
|
||||
if resolved.is_empty() {
|
||||
return Err(Error::InvalidProvider(model.into()));
|
||||
}
|
||||
let (config, wire_model): (&'static dyn BaseResponsesApiConfig, String) =
|
||||
match provider.parse::<LlmProviders>() {
|
||||
Ok(LlmProviders::Openai) if resolved.contains('/') => {
|
||||
return Err(Error::InvalidProvider(model.into()));
|
||||
}
|
||||
Ok(LlmProviders::Openai) => (&OpenAiResponsesApiConfig, resolved.into()),
|
||||
Ok(LlmProviders::Openrouter) if explicit.is_some() && !resolved.contains('/') => (
|
||||
&OPENROUTER_RESPONSES_CONFIG,
|
||||
format!("openrouter/{resolved}"),
|
||||
),
|
||||
Ok(LlmProviders::Openrouter) => (&OPENROUTER_RESPONSES_CONFIG, resolved.into()),
|
||||
_ => return Err(Error::Unsupported("native HTTP responses provider")),
|
||||
};
|
||||
Ok(ResponsesProvider {
|
||||
identity: ProviderIdentity {
|
||||
model: wire_model,
|
||||
provider: provider.into(),
|
||||
},
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
/// Pinned against Python's `get_llm_provider` on 2026-10-11.
|
||||
#[rstest]
|
||||
#[case::bare("gpt-5", None, "openai", "gpt-5")]
|
||||
#[case::openai_prefix("openai/gpt-5", None, "openai", "gpt-5")]
|
||||
#[case::openai_explicit("gpt-5", Some("openai"), "openai", "gpt-5")]
|
||||
#[case::openrouter_prefix(
|
||||
"openrouter/anthropic/claude",
|
||||
None,
|
||||
"openrouter",
|
||||
"anthropic/claude"
|
||||
)]
|
||||
#[case::openrouter_explicit(
|
||||
"openrouter/anthropic/claude",
|
||||
Some("openrouter"),
|
||||
"openrouter",
|
||||
"anthropic/claude"
|
||||
)]
|
||||
#[case::openrouter_explicit_unprefixed(
|
||||
"anthropic/claude",
|
||||
Some("openrouter"),
|
||||
"openrouter",
|
||||
"anthropic/claude"
|
||||
)]
|
||||
#[case::openrouter_native_id(
|
||||
"openrouter/auto",
|
||||
Some("openrouter"),
|
||||
"openrouter",
|
||||
"openrouter/auto"
|
||||
)]
|
||||
#[case::openrouter_native_id_implicit("openrouter/auto", None, "openrouter", "auto")]
|
||||
fn providers_and_wire_models_match_python(
|
||||
#[case] model: &str,
|
||||
#[case] custom_llm_provider: Option<&str>,
|
||||
#[case] provider: &str,
|
||||
#[case] wire_model: &str,
|
||||
) {
|
||||
let resolved = resolve_provider(model, custom_llm_provider).unwrap();
|
||||
assert_eq!(resolved.identity.provider, provider);
|
||||
assert_eq!(resolved.identity.model, wire_model);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::other_prefix("other/model", None)]
|
||||
#[case::other_explicit("model", Some("other"))]
|
||||
#[case::openai_with_vendor_segment("openai/other/model", None)]
|
||||
#[case::empty_model("openai/", None)]
|
||||
#[case::conflicting_prefix("openrouter/anthropic/claude", Some("openai"))]
|
||||
fn everything_else_is_rejected(#[case] model: &str, #[case] custom_llm_provider: Option<&str>) {
|
||||
assert!(matches!(
|
||||
resolve_provider(model, custom_llm_provider),
|
||||
Err(Error::InvalidProvider(_) | Error::Unsupported(_))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
@ -2,9 +2,8 @@ use std::time::Duration;
|
|||
|
||||
use bytes::Bytes;
|
||||
use litellm_host::call::CallOutput;
|
||||
use litellm_llms::base_llm::{
|
||||
auth::ValidatedEnvironment, responses::transformation::BaseResponsesApiConfig,
|
||||
};
|
||||
use litellm_host::interceptors::ProviderIdentity;
|
||||
use litellm_llms::base_llm::responses::transformation::BaseResponsesApiConfig;
|
||||
use litellm_llms_types::formats::responses::ResponsesApiResponse;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
|
@ -21,17 +20,13 @@ pub struct ResponsesCall {
|
|||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
pub(crate) struct ResponsesProvider {
|
||||
pub identity: ProviderIdentity,
|
||||
pub config: &'static dyn BaseResponsesApiConfig,
|
||||
}
|
||||
|
||||
pub struct ResponsesStreamHead {
|
||||
pub headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
pub type ResponsesOutput = CallOutput<ResponsesApiResponse, ResponsesStreamHead, Bytes, Error>;
|
||||
|
||||
pub(super) struct ProviderResponsesRequest {
|
||||
pub config: &'static dyn BaseResponsesApiConfig,
|
||||
pub environment: ValidatedEnvironment,
|
||||
pub url: String,
|
||||
pub body: Value,
|
||||
pub context: litellm_host::interceptors::RequestContext,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ use std::sync::Arc;
|
|||
use futures_util::TryStreamExt;
|
||||
use litellm_host::{call::HostedCompletion, lifecycle::CallEvent};
|
||||
use litellm_inference_responses::{
|
||||
Error,
|
||||
route::Responses,
|
||||
types::{ResponsesCall, ResponsesOutput},
|
||||
};
|
||||
|
|
@ -32,6 +33,110 @@ fn call() -> ResponsesCall {
|
|||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::bare_openai("gpt-5", None, "gpt-5")]
|
||||
#[case::openai_prefix("openai/gpt-5", None, "gpt-5")]
|
||||
#[case::empty_explicit_provider("openai/gpt-5", Some(""), "gpt-5")]
|
||||
#[case::explicit_openai("gpt-5", Some("openai"), "gpt-5")]
|
||||
#[case::openrouter_prefix("openrouter/anthropic/claude", None, "anthropic/claude")]
|
||||
#[case::explicit_openrouter_prefix(
|
||||
"openrouter/anthropic/claude",
|
||||
Some("openrouter"),
|
||||
"anthropic/claude"
|
||||
)]
|
||||
#[case::explicit_openrouter_unprefixed("anthropic/claude", Some("openrouter"), "anthropic/claude")]
|
||||
#[case::explicit_openrouter_native_id("auto", Some("openrouter"), "openrouter/auto")]
|
||||
#[case::implicit_openrouter_native_id("openrouter/auto", None, "auto")]
|
||||
#[tokio::test]
|
||||
async fn provider_resolution_preserves_the_wire_model(
|
||||
#[case] model: &str,
|
||||
#[case] custom_llm_provider: Option<&str>,
|
||||
#[case] expected_wire_model: &str,
|
||||
) {
|
||||
let response = json!({
|
||||
"id": "response-1",
|
||||
"model": expected_wire_model,
|
||||
"output": [{"type": "message", "content": []}],
|
||||
"usage": {"total_tokens": 7}
|
||||
});
|
||||
let upstream = upstream([json_response(response)]).await;
|
||||
responses_route(no_secrets())
|
||||
.execute(
|
||||
ResponsesCall {
|
||||
model: model.into(),
|
||||
api_key: Some("test-key".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
custom_llm_provider: custom_llm_provider.map(str::to_owned),
|
||||
..call()
|
||||
},
|
||||
&(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
only_request(&upstream).await.json()["model"],
|
||||
expected_wire_model
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::explicit_unknown_provider("test-model", Some("unknown-provider"))]
|
||||
#[case::unsupported_responses_provider("test-model", Some("anthropic"))]
|
||||
#[tokio::test]
|
||||
async fn unsupported_providers_fail_before_transport(
|
||||
#[case] model: &str,
|
||||
#[case] custom_llm_provider: Option<&str>,
|
||||
) {
|
||||
let upstream = upstream([]).await;
|
||||
let result = responses_route(no_secrets())
|
||||
.execute(
|
||||
ResponsesCall {
|
||||
model: model.into(),
|
||||
api_base: Some(upstream.uri()),
|
||||
custom_llm_provider: custom_llm_provider.map(str::to_owned),
|
||||
..call()
|
||||
},
|
||||
&(),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::Unsupported("native HTTP responses provider"))
|
||||
));
|
||||
assert!(received(&upstream).await.is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::openai_with_vendor_segment("openai/other/model", None)]
|
||||
#[case::empty_model("openai/", None)]
|
||||
#[case::conflicting_prefix("openrouter/anthropic/claude", Some("openai"))]
|
||||
#[tokio::test]
|
||||
async fn invalid_provider_models_fail_before_transport(
|
||||
#[case] model: &str,
|
||||
#[case] custom_llm_provider: Option<&str>,
|
||||
) {
|
||||
let upstream = upstream([]).await;
|
||||
let result = responses_route(no_secrets())
|
||||
.execute(
|
||||
ResponsesCall {
|
||||
model: model.into(),
|
||||
api_base: Some(upstream.uri()),
|
||||
custom_llm_provider: custom_llm_provider.map(str::to_owned),
|
||||
..call()
|
||||
},
|
||||
&(),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(matches!(result, Err(Error::InvalidProvider(_))));
|
||||
assert!(received(&upstream).await.is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::direct(false)]
|
||||
#[case::hosted(true)]
|
||||
|
|
|
|||
|
|
@ -1,28 +1,57 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_http::{Client, request::truncate_error_body};
|
||||
use litellm_llms::base_llm::auth::resolve_auth;
|
||||
use litellm_http::{
|
||||
Client,
|
||||
request::{string_headers, truncate_error_body, with_default_headers},
|
||||
};
|
||||
use litellm_llms::base_llm::{
|
||||
audio_transcription::transformation::BaseAudioTranscriptionConfig,
|
||||
auth::{ValidatedEnvironment, resolve_auth},
|
||||
};
|
||||
use litellm_secrets::source::Secrets;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::Error;
|
||||
use crate::{
|
||||
constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS, types::ProviderAudioTranscriptionRequest,
|
||||
};
|
||||
use crate::{constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS, types::AudioTranscriptionRequest};
|
||||
use litellm_inference::provider::ResolvedProvider;
|
||||
|
||||
pub async fn execute_audio_transcription_provider_call(
|
||||
pub(super) async fn execute(
|
||||
http: &Client,
|
||||
auth: &litellm_auth::AuthServices,
|
||||
request: ProviderAudioTranscriptionRequest,
|
||||
config: &'static dyn BaseAudioTranscriptionConfig,
|
||||
provider: ResolvedProvider<'_>,
|
||||
call: AudioTranscriptionRequest<'_>,
|
||||
secrets: Secrets,
|
||||
) -> Result<Value, Error> {
|
||||
let env_lookup = |key: &str| request.secrets.get(key);
|
||||
let authenticated = resolve_auth(auth, request.environment.clone(), &env_lookup).await?;
|
||||
let env_lookup = |key: &str| secrets.get(key);
|
||||
let forwarded = string_headers("audio transcription", call.extra_headers)?;
|
||||
let validated = config.validate_environment(
|
||||
forwarded,
|
||||
provider.model,
|
||||
&call.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let environment = ValidatedEnvironment {
|
||||
headers: with_default_headers(validated.headers, config.default_headers()),
|
||||
auth: validated.auth,
|
||||
};
|
||||
let url = config.get_complete_url(
|
||||
call.api_base,
|
||||
provider.model,
|
||||
&call.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let filtered_params = config.map_transcription_params(&call.optional_params);
|
||||
let body = config
|
||||
.transform_audio_transcription_request(provider.model, call.audio, filtered_params)?
|
||||
.body;
|
||||
let authenticated = resolve_auth(auth, environment, &env_lookup).await?;
|
||||
let outbound = litellm_inference::outbound::outbound_request(
|
||||
authenticated,
|
||||
request.url.clone(),
|
||||
&request.body,
|
||||
url,
|
||||
&body,
|
||||
Some(
|
||||
request
|
||||
.timeout
|
||||
call.timeout
|
||||
.unwrap_or(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)),
|
||||
),
|
||||
)?;
|
||||
|
|
@ -47,8 +76,7 @@ pub async fn execute_audio_transcription_provider_call(
|
|||
error,
|
||||
))
|
||||
})?;
|
||||
Ok(request
|
||||
.config
|
||||
.transform_audio_transcription_response(&request.model, response_json)?
|
||||
Ok(config
|
||||
.transform_audio_transcription_response(provider.model, response_json)?
|
||||
.into_json())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,16 +2,29 @@ pub mod types;
|
|||
pub use litellm_inference::RouteError as Error;
|
||||
mod constants;
|
||||
mod handler;
|
||||
mod prepare;
|
||||
pub use handler::execute_audio_transcription_provider_call;
|
||||
|
||||
use litellm_auth::AuthServices;
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use litellm_inference::provider::resolve_llm_provider;
|
||||
use litellm_llms::{
|
||||
base_llm::audio_transcription::transformation::BaseAudioTranscriptionConfig,
|
||||
bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG,
|
||||
};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
pub use prepare::prepare_audio_transcription_provider_call;
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::types::AudioTranscriptionRequest;
|
||||
|
||||
fn audio_transcription_provider(
|
||||
provider: LlmProviders,
|
||||
) -> Option<&'static dyn BaseAudioTranscriptionConfig> {
|
||||
match provider {
|
||||
LlmProviders::Bedrock => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AudioTranscriptionRoute {
|
||||
http: litellm_http::Client,
|
||||
|
|
@ -32,6 +45,10 @@ impl AudioTranscriptionRoute {
|
|||
}
|
||||
}
|
||||
|
||||
pub async fn execute(&self, request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
|
||||
self.run(request).await
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "audio_transcription",
|
||||
model = request.model,
|
||||
|
|
@ -40,13 +57,23 @@ impl AudioTranscriptionRoute {
|
|||
stream = false,
|
||||
outcome
|
||||
))]
|
||||
pub async fn execute(&self, request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
|
||||
async fn run(&self, request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
|
||||
litellm_inference::diagnostic::unary(async {
|
||||
let request =
|
||||
prepare_audio_transcription_provider_call(request, self.secrets.as_ref()).await?;
|
||||
litellm_inference::diagnostic::provider(&request.model, &request.custom_llm_provider);
|
||||
let provider = resolve_llm_provider(
|
||||
request.model,
|
||||
request.custom_llm_provider,
|
||||
"audio transcription",
|
||||
)?;
|
||||
let config = audio_transcription_provider(provider.provider).ok_or_else(|| {
|
||||
Error::InvalidProvider(<&str>::from(provider.provider).to_string())
|
||||
})?;
|
||||
litellm_inference::diagnostic::provider(
|
||||
provider.model,
|
||||
<&'static str>::from(provider.provider),
|
||||
);
|
||||
let secrets = self.secrets.resolve(&config.secret_names()).await?;
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<Value, Error>> = Box::pin(
|
||||
execute_audio_transcription_provider_call(&self.http, &self.auth, request),
|
||||
handler::execute(&self.http, &self.auth, config, provider, request, secrets),
|
||||
);
|
||||
execute.await
|
||||
})
|
||||
|
|
|
|||
|
|
@ -1,65 +0,0 @@
|
|||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use litellm_http::request::string_headers;
|
||||
use litellm_http::request::with_default_headers;
|
||||
use litellm_llms::{
|
||||
base_llm::{
|
||||
audio_transcription::transformation::BaseAudioTranscriptionConfig,
|
||||
auth::ValidatedEnvironment,
|
||||
},
|
||||
bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG,
|
||||
};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
|
||||
use super::Error;
|
||||
use crate::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
|
||||
use litellm_inference::provider::resolve_llm_provider;
|
||||
|
||||
fn provider_config(provider: LlmProviders) -> Option<&'static dyn BaseAudioTranscriptionConfig> {
|
||||
match provider {
|
||||
LlmProviders::Bedrock => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
|
||||
pub async fn prepare_audio_transcription_provider_call(
|
||||
request: AudioTranscriptionRequest<'_>,
|
||||
secrets: &dyn SecretSource,
|
||||
) -> Result<ProviderAudioTranscriptionRequest, Error> {
|
||||
let provider_info = resolve_llm_provider(
|
||||
request.model,
|
||||
request.custom_llm_provider,
|
||||
"audio transcription",
|
||||
)?;
|
||||
let model = provider_info.model.to_string();
|
||||
let config = provider_config(provider_info.provider)
|
||||
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider_info.provider).to_string()))?;
|
||||
let snapshot = secrets.resolve(&config.secret_names()).await?;
|
||||
let env_lookup = |key: &str| snapshot.get(key);
|
||||
let forwarded = string_headers("audio transcription", request.extra_headers)?;
|
||||
let validated =
|
||||
config.validate_environment(forwarded, &model, &request.optional_params, &env_lookup)?;
|
||||
let environment = ValidatedEnvironment {
|
||||
headers: with_default_headers(validated.headers, config.default_headers()),
|
||||
auth: validated.auth,
|
||||
};
|
||||
let url = config.get_complete_url(
|
||||
request.api_base,
|
||||
&model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let filtered_params = config.map_transcription_params(&request.optional_params);
|
||||
let transformed =
|
||||
config.transform_audio_transcription_request(&model, request.audio, filtered_params)?;
|
||||
Ok(ProviderAudioTranscriptionRequest {
|
||||
model,
|
||||
custom_llm_provider: <&str>::from(provider_info.provider).to_string(),
|
||||
config,
|
||||
url,
|
||||
body: transformed.body,
|
||||
environment,
|
||||
secrets: snapshot,
|
||||
timeout: request.timeout,
|
||||
})
|
||||
}
|
||||
|
|
@ -1,9 +1,5 @@
|
|||
use litellm_secrets::source::Secrets;
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_llms::base_llm::{
|
||||
audio_transcription::transformation::BaseAudioTranscriptionConfig, auth::ValidatedEnvironment,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub struct AudioTranscriptionRequest<'a> {
|
||||
|
|
@ -16,37 +12,3 @@ pub struct AudioTranscriptionRequest<'a> {
|
|||
pub optional_params: Map<String, Value>,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ProviderAudioTranscriptionRequest {
|
||||
pub model: String,
|
||||
pub custom_llm_provider: String,
|
||||
pub config: &'static dyn BaseAudioTranscriptionConfig,
|
||||
pub url: String,
|
||||
pub body: Value,
|
||||
pub environment: ValidatedEnvironment,
|
||||
pub secrets: Secrets,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl ProviderAudioTranscriptionRequest {
|
||||
pub fn model(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
|
||||
pub fn custom_llm_provider(&self) -> &str {
|
||||
&self.custom_llm_provider
|
||||
}
|
||||
|
||||
pub fn url(&self) -> &str {
|
||||
&self.url
|
||||
}
|
||||
|
||||
pub fn body(&self) -> &Value {
|
||||
&self.body
|
||||
}
|
||||
|
||||
pub fn with_body(self, body: Value) -> Self {
|
||||
Self { body, ..self }
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -195,9 +195,10 @@ impl AnthropicBeta {
|
|||
|
||||
fn rejected_by(&self) -> &'static [BetaProvider] {
|
||||
match self {
|
||||
Self::AdvancedToolUse20251120 | Self::ContextManagement20250627 => {
|
||||
&[BetaProvider::BedrockConverse]
|
||||
Self::AdvancedToolUse20251120 => {
|
||||
&[BetaProvider::BedrockConverse, BetaProvider::Databricks]
|
||||
}
|
||||
Self::ContextManagement20250627 => &[BetaProvider::BedrockConverse],
|
||||
Self::AdvisorTool20260301 | Self::Compact20260904 => &[
|
||||
BetaProvider::AzureAi,
|
||||
BetaProvider::BedrockConverse,
|
||||
|
|
@ -231,12 +232,11 @@ impl AnthropicBeta {
|
|||
BetaProvider::Bedrock,
|
||||
BetaProvider::BedrockMantle,
|
||||
BetaProvider::VertexAi,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse],
|
||||
Self::ComputerUse20250124 | Self::ComputerUse20251124 | Self::Context1m20250807 => &[],
|
||||
Self::DangerousToolUse20260903 => {
|
||||
&[BetaProvider::BedrockConverse, BetaProvider::Databricks]
|
||||
}
|
||||
Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse],
|
||||
Self::Effort20251124 => &[BetaProvider::VertexAi],
|
||||
Self::FastMode20260201 | Self::Oauth20250420 => &[
|
||||
BetaProvider::AzureAi,
|
||||
|
|
@ -244,6 +244,7 @@ impl AnthropicBeta {
|
|||
BetaProvider::Bedrock,
|
||||
BetaProvider::BedrockMantle,
|
||||
BetaProvider::VertexAi,
|
||||
BetaProvider::Databricks,
|
||||
],
|
||||
Self::FineGrainedToolStreaming20250514 => {
|
||||
&[BetaProvider::AzureAi, BetaProvider::VertexAi]
|
||||
|
|
|
|||
|
|
@ -76,7 +76,7 @@ mod tests {
|
|||
#[case::vertex_ai(BetaProvider::VertexAi, Some("tool-search-tool-2025-10-19"))]
|
||||
#[case::anthropic(BetaProvider::Anthropic, Some("advanced-tool-use-2025-11-20"))]
|
||||
#[case::azure_ai(BetaProvider::AzureAi, Some("advanced-tool-use-2025-11-20"))]
|
||||
#[case::databricks(BetaProvider::Databricks, Some("advanced-tool-use-2025-11-20"))]
|
||||
#[case::databricks(BetaProvider::Databricks, None)]
|
||||
#[case::bedrock_converse(BetaProvider::BedrockConverse, None)]
|
||||
fn filter_renames_per_host(
|
||||
#[case] provider: BetaProvider,
|
||||
|
|
|
|||
|
|
@ -74,6 +74,7 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
|
|||
headers: Headers,
|
||||
api_key: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &MessagesOptionalParams,
|
||||
_litellm_params: &LitellmParams,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
@ -109,7 +110,10 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
|
|||
}
|
||||
|
||||
fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers {
|
||||
update_headers_with_anthropic_beta(headers, request)
|
||||
merge_beta_headers(
|
||||
headers,
|
||||
provider_feature_betas(request, BetaProvider::Anthropic),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -136,12 +140,20 @@ pub(crate) fn transform_messages_request(
|
|||
messages: strip_encrypted_reasoning_blocks(messages),
|
||||
params: MessagesOptionalParams {
|
||||
context_management,
|
||||
extra: without_null_safeguards(request.params.extra),
|
||||
..request.params
|
||||
},
|
||||
..request
|
||||
})
|
||||
}
|
||||
|
||||
fn without_null_safeguards(extra: Map<String, Value>) -> Map<String, Value> {
|
||||
extra
|
||||
.into_iter()
|
||||
.filter(|(key, value)| key != "safeguards" || !value.is_null())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn update_headers_with_anthropic_beta(
|
||||
headers: Headers,
|
||||
request: &MessagesRequest,
|
||||
|
|
@ -149,12 +161,34 @@ pub(crate) fn update_headers_with_anthropic_beta(
|
|||
merge_beta_headers(headers, feature_betas(request))
|
||||
}
|
||||
|
||||
/// The betas a request needs on a host other than the first-party API: the Anthropic set
|
||||
/// filtered and renamed through the host's `BetaProvider` policy.
|
||||
pub(crate) fn update_headers_with_safeguards_beta(
|
||||
headers: Headers,
|
||||
request: &MessagesRequest,
|
||||
provider: BetaProvider,
|
||||
) -> Headers {
|
||||
merge_beta_headers(
|
||||
headers,
|
||||
feature_betas(request).union(on_provider(mapped_host_betas(&request.params), provider)),
|
||||
)
|
||||
}
|
||||
|
||||
/// The betas a request needs on a host with a column in the beta map, filtered and renamed
|
||||
/// through that column.
|
||||
pub(crate) fn provider_feature_betas(request: &MessagesRequest, provider: BetaProvider) -> BetaSet {
|
||||
feature_betas(request)
|
||||
.iter()
|
||||
.filter_map(|beta| beta.on(provider))
|
||||
on_provider(
|
||||
feature_betas(request).union(mapped_host_betas(&request.params)),
|
||||
provider,
|
||||
)
|
||||
}
|
||||
|
||||
fn on_provider(betas: BetaSet, provider: BetaProvider) -> BetaSet {
|
||||
betas.iter().filter_map(|beta| beta.on(provider)).collect()
|
||||
}
|
||||
|
||||
fn mapped_host_betas(params: &MessagesOptionalParams) -> BetaSet {
|
||||
carries_safeguards(params)
|
||||
.then_some(AnthropicBeta::DangerousToolUse20260903)
|
||||
.into_iter()
|
||||
.collect()
|
||||
}
|
||||
|
||||
|
|
@ -206,6 +240,13 @@ fn uses_structured_output(params: &MessagesOptionalParams) -> bool {
|
|||
.is_some_and(|config| config.format.is_some())
|
||||
}
|
||||
|
||||
fn carries_safeguards(params: &MessagesOptionalParams) -> bool {
|
||||
params
|
||||
.extra
|
||||
.get("safeguards")
|
||||
.is_some_and(|value| !value.is_null())
|
||||
}
|
||||
|
||||
fn messages_carry_output_config(messages: &[Message]) -> bool {
|
||||
messages
|
||||
.iter()
|
||||
|
|
@ -847,6 +888,7 @@ mod tests {
|
|||
headers(forwarded),
|
||||
api_key,
|
||||
"claude",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&env(vars),
|
||||
)
|
||||
|
|
@ -1025,10 +1067,44 @@ mod tests {
|
|||
json!({"messages": [{"role": "user", "content": "hi", "output_config": null}]}),
|
||||
&["per-turn-control-2026-07-01"]
|
||||
)]
|
||||
#[case::safeguards(
|
||||
json!({"safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}]}),
|
||||
&[]
|
||||
)]
|
||||
fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) {
|
||||
assert_eq!(feature_betas(&request(fields)), betas(expected));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic(BetaProvider::Anthropic, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])]
|
||||
#[case::vertex_ai(BetaProvider::VertexAi, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])]
|
||||
#[case::empty_list(BetaProvider::Anthropic, json!([]), &["dangerous-tool-use-2026-09-03"])]
|
||||
#[case::not_a_list(BetaProvider::Anthropic, json!(1), &["dangerous-tool-use-2026-09-03"])]
|
||||
#[case::null(BetaProvider::Anthropic, Value::Null, &[])]
|
||||
#[case::bedrock_converse_has_a_null_cell(BetaProvider::BedrockConverse, json!([{"type": "dangerous_tool_use"}]), &[])]
|
||||
fn safeguards_add_the_dangerous_tool_use_beta_where_the_beta_map_has_a_cell(
|
||||
#[case] provider: BetaProvider,
|
||||
#[case] safeguards: Value,
|
||||
#[case] expected: &[&str],
|
||||
) {
|
||||
assert_eq!(
|
||||
provider_feature_betas(&request(json!({"safeguards": safeguards})), provider),
|
||||
betas(expected)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn hosts_without_a_beta_map_column_get_safeguards_without_a_derived_beta() {
|
||||
let client = headers(&[("anthropic-beta", "claude-code-20250219")]);
|
||||
assert_eq!(
|
||||
update_headers_with_anthropic_beta(
|
||||
client.clone(),
|
||||
&request(json!({"safeguards": [{"type": "dangerous_tool_use"}]}))
|
||||
),
|
||||
client
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::no_betas(&[("x-api-key", "k"), ("anthropic-version", "2023-06-01")], json!({}), &[("x-api-key", "k"), ("anthropic-version", "2023-06-01")])]
|
||||
#[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], json!({}), &[("Anthropic-Beta", " , "), ("x-api-key", "k")])]
|
||||
|
|
@ -1067,6 +1143,11 @@ mod tests {
|
|||
json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}),
|
||||
&[("anthropic-beta", "claude-code-20250219,context-management-2025-06-27,effort-2025-11-24,interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")],
|
||||
)]
|
||||
#[case::safeguards_add_the_dangerous_tool_use_beta_to_the_client_betas(
|
||||
&[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")],
|
||||
json!({"safeguards": [{"type": "dangerous_tool_use"}]}),
|
||||
&[("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14")],
|
||||
)]
|
||||
fn request_headers_merge_the_feature_betas(
|
||||
#[case] input: &[(&str, &str)],
|
||||
#[case] fields: Value,
|
||||
|
|
@ -1127,6 +1208,7 @@ mod tests {
|
|||
Vec::new(),
|
||||
None,
|
||||
"claude",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&record,
|
||||
);
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
use litellm_auth::SecretValue;
|
||||
use litellm_http::request::{has_bearer_auth, has_header};
|
||||
use litellm_llms_types::{
|
||||
formats::messages::MessagesRequest,
|
||||
providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH},
|
||||
formats::messages::{MessagesOptionalParams, MessagesRequest},
|
||||
providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH},
|
||||
};
|
||||
use litellm_router_types::LitellmParams;
|
||||
|
||||
|
|
@ -12,7 +12,7 @@ use crate::{
|
|||
common_utils::API_KEY_PLACEMENT,
|
||||
messages::{
|
||||
handler::shape_anthropic_messages_request,
|
||||
transformation::{transform_messages_request, update_headers_with_anthropic_beta},
|
||||
transformation::{transform_messages_request, update_headers_with_safeguards_beta},
|
||||
},
|
||||
},
|
||||
azure_ai::common_utils::{
|
||||
|
|
@ -83,6 +83,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig {
|
|||
headers: Headers,
|
||||
api_key: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &MessagesOptionalParams,
|
||||
_litellm_params: &LitellmParams,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
@ -104,7 +105,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig {
|
|||
}
|
||||
|
||||
fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers {
|
||||
update_headers_with_anthropic_beta(headers, request)
|
||||
update_headers_with_safeguards_beta(headers, request, BetaProvider::AzureAi)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -237,6 +238,7 @@ mod tests {
|
|||
.collect(),
|
||||
api_key,
|
||||
"claude",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&|_| None,
|
||||
)
|
||||
|
|
@ -544,6 +546,12 @@ mod tests {
|
|||
&[("x-api-key", "k"), ("anthropic-beta", "structured-outputs-2025-11-13,web-search-2025-03-05")]
|
||||
)]
|
||||
#[case::no_feature_needs_a_beta(json!({}), &[], &[("x-api-key", "k")])]
|
||||
#[case::safeguards_need_the_dangerous_tool_use_beta(
|
||||
json!({"safeguards": [{"type": "dangerous_tool_use"}]}),
|
||||
&[("anthropic-beta", "claude-code-20250219")],
|
||||
&[("x-api-key", "k"), ("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03")]
|
||||
)]
|
||||
#[case::null_safeguards_need_no_beta(json!({"safeguards": null}), &[], &[("x-api-key", "k")])]
|
||||
fn request_headers_carry_the_anthropic_feature_betas(
|
||||
#[case] fields: serde_json::Value,
|
||||
#[case] forwarded: &[(&str, &str)],
|
||||
|
|
@ -616,6 +624,7 @@ mod tests {
|
|||
Vec::new(),
|
||||
None,
|
||||
"claude",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&record,
|
||||
);
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
|
||||
use litellm_llms_types::formats::messages::{
|
||||
MessagesOptionalParams, MessagesRequest, MessagesResponse,
|
||||
};
|
||||
use litellm_router_types::LitellmParams;
|
||||
use serde_json::Value;
|
||||
|
||||
|
|
@ -55,6 +57,7 @@ pub trait BaseMessagesConfig: Sync {
|
|||
headers: Headers,
|
||||
api_key: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &MessagesOptionalParams,
|
||||
litellm_params: &LitellmParams,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error>;
|
||||
|
|
@ -110,6 +113,7 @@ mod tests {
|
|||
headers: Headers,
|
||||
_api_key: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &MessagesOptionalParams,
|
||||
_litellm_params: &LitellmParams,
|
||||
_env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
|
|||
|
|
@ -5,8 +5,8 @@ use litellm_auth_aws::{AwsCredentialSource, resolve_aws_region};
|
|||
use litellm_core_utils::settings::resolve_non_empty;
|
||||
use litellm_http::request::{header_value, with_header, without_headers};
|
||||
use litellm_llms_types::{
|
||||
formats::messages::MessagesRequest,
|
||||
providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH},
|
||||
formats::messages::{MessagesOptionalParams, MessagesRequest},
|
||||
providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH},
|
||||
};
|
||||
use litellm_router_types::LitellmParams;
|
||||
use serde_json::Value;
|
||||
|
|
@ -15,7 +15,7 @@ use crate::{
|
|||
Error, ErrorDetail,
|
||||
anthropic::messages::{
|
||||
handler::shape_anthropic_messages_request,
|
||||
transformation::{transform_messages_request, update_headers_with_anthropic_beta},
|
||||
transformation::{transform_messages_request, update_headers_with_safeguards_beta},
|
||||
},
|
||||
base_llm::{
|
||||
auth::{AuthScheme, Headers, ValidatedEnvironment},
|
||||
|
|
@ -60,6 +60,7 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig {
|
|||
headers: Headers,
|
||||
api_key: Option<&str>,
|
||||
_model: &str,
|
||||
optional_params: &MessagesOptionalParams,
|
||||
params: &LitellmParams,
|
||||
env: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
@ -67,11 +68,8 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig {
|
|||
.filter(|value| !value.trim().is_empty())
|
||||
.or_else(|| {
|
||||
WORKSPACE_FIELDS.iter().find_map(|field| {
|
||||
params
|
||||
.extra
|
||||
.get(*field)?
|
||||
.as_str()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
workspace_value(optional_params.extra.get(*field))
|
||||
.or_else(|| workspace_value(params.extra.get(*field)))
|
||||
})
|
||||
});
|
||||
let workspace = resolve_non_empty(workspace, env, WORKSPACE_ENVS)
|
||||
|
|
@ -174,7 +172,7 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig {
|
|||
Some(workspace) => with_header(headers, WORKSPACE_HEADER, workspace.into()),
|
||||
None => headers,
|
||||
};
|
||||
update_headers_with_anthropic_beta(headers, request)
|
||||
update_headers_with_safeguards_beta(headers, request, BetaProvider::Anthropic)
|
||||
}
|
||||
|
||||
fn wire_body(&self, body: Value) -> Value {
|
||||
|
|
@ -182,6 +180,10 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig {
|
|||
}
|
||||
}
|
||||
|
||||
fn workspace_value(value: Option<&Value>) -> Option<&str> {
|
||||
value?.as_str().filter(|value| !value.trim().is_empty())
|
||||
}
|
||||
|
||||
fn platform_region(
|
||||
base: Option<&str>,
|
||||
params: &LitellmParams,
|
||||
|
|
@ -260,6 +262,7 @@ mod tests {
|
|||
vec![],
|
||||
Some("key"),
|
||||
"claude-test",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&|_| None
|
||||
)
|
||||
|
|
@ -280,6 +283,7 @@ mod tests {
|
|||
],
|
||||
Some("explicit"),
|
||||
"claude-test",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&|name| (name == "ANTHROPIC_AWS_WORKSPACE_ID").then(|| "workspace".into()),
|
||||
)
|
||||
|
|
@ -345,7 +349,14 @@ mod tests {
|
|||
.get_complete_url(None, "claude-test", ¶ms, false, &env)
|
||||
.unwrap();
|
||||
let environment = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG
|
||||
.validate_environment(vec![], None, "claude-test", ¶ms, &env)
|
||||
.validate_environment(
|
||||
vec![],
|
||||
None,
|
||||
"claude-test",
|
||||
&MessagesOptionalParams::default(),
|
||||
¶ms,
|
||||
&env,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
url,
|
||||
|
|
@ -356,6 +367,21 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::safeguards(json!([{"type": "dangerous_tool_use"}]), Some("dangerous-tool-use-2026-09-03"))]
|
||||
#[case::null_safeguards(Value::Null, None)]
|
||||
fn safeguards_need_the_dangerous_tool_use_beta(
|
||||
#[case] safeguards: Value,
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
let request = serde_json::from_value(json!({"model": "claude-test", "max_tokens": 16, "messages": [], "safeguards": safeguards})).unwrap();
|
||||
let headers = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG.request_headers(vec![], &request);
|
||||
assert_eq!(
|
||||
header_value(&headers, "anthropic-beta").map(|value| value.to_string()),
|
||||
expected.map(str::to_string)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::default(json!({}), false)]
|
||||
#[case::override_empty(json!({"claude_platform_unsupported_params": []}), true)]
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ use litellm_auth_aws::{
|
|||
resolve_bedrock_region,
|
||||
};
|
||||
use litellm_llms_types::formats::messages::{
|
||||
MessagesRequest,
|
||||
MessagesOptionalParams, MessagesRequest,
|
||||
streaming::{MessagesStreamEvent, MessagesStreamUsage},
|
||||
};
|
||||
use litellm_llms_types::providers::bedrock::{
|
||||
|
|
@ -146,6 +146,7 @@ impl BaseMessagesConfig for AmazonAnthropicClaudeMessagesConfig {
|
|||
headers: Headers,
|
||||
api_key: Option<&str>,
|
||||
model: &str,
|
||||
_optional_params: &MessagesOptionalParams,
|
||||
litellm_params: &LitellmParams,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
@ -618,6 +619,7 @@ mod tests {
|
|||
Vec::new(),
|
||||
None,
|
||||
"anthropic.claude-3",
|
||||
&MessagesOptionalParams::default(),
|
||||
&litellm_params,
|
||||
&|_| None,
|
||||
)
|
||||
|
|
@ -658,6 +660,7 @@ mod tests {
|
|||
vec![("authorization".into(), "Bearer forwarded".into())],
|
||||
api_key,
|
||||
"anthropic.claude-3",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&env,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ use litellm_auth_aws::bedrock_model_id_and_region;
|
|||
use litellm_core_utils::settings::resolve_non_empty;
|
||||
use litellm_http::request::{with_header, without_headers};
|
||||
use litellm_llms_types::{
|
||||
formats::messages::MessagesRequest,
|
||||
formats::messages::{MessagesOptionalParams, MessagesRequest},
|
||||
providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet, DEFAULT_HEADERS},
|
||||
};
|
||||
use litellm_router_types::LitellmParams;
|
||||
|
|
@ -59,6 +59,7 @@ impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig {
|
|||
headers: Headers,
|
||||
api_key: Option<&str>,
|
||||
model: &str,
|
||||
_optional_params: &MessagesOptionalParams,
|
||||
params: &LitellmParams,
|
||||
env: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
@ -317,7 +318,14 @@ mod tests {
|
|||
"https://private.test/anthropic/v1/messages"
|
||||
);
|
||||
let environment = BEDROCK_MANTLE_MESSAGES_CONFIG
|
||||
.validate_environment(vec![], None, model, ¶ms, &|_| None)
|
||||
.validate_environment(
|
||||
vec![],
|
||||
None,
|
||||
model,
|
||||
&MessagesOptionalParams::default(),
|
||||
¶ms,
|
||||
&|_| None,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(
|
||||
matches!(environment.auth, AuthScheme::AwsSigV4 { region: actual, service: "bedrock", .. } if actual == region)
|
||||
|
|
@ -347,7 +355,14 @@ mod tests {
|
|||
.get_complete_url(None, "claude-test", ¶ms, false, &|_| None)
|
||||
.unwrap();
|
||||
let environment = BEDROCK_MANTLE_MESSAGES_CONFIG
|
||||
.validate_environment(vec![], None, "claude-test", ¶ms, &|_| None)
|
||||
.validate_environment(
|
||||
vec![],
|
||||
None,
|
||||
"claude-test",
|
||||
&MessagesOptionalParams::default(),
|
||||
¶ms,
|
||||
&|_| None,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
url,
|
||||
|
|
@ -401,6 +416,7 @@ mod tests {
|
|||
],
|
||||
key,
|
||||
"claude-test",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&|name| match name {
|
||||
"BEDROCK_MANTLE_API_KEY" => Some("mantle".into()),
|
||||
|
|
|
|||
|
|
@ -104,6 +104,7 @@ impl BaseMessagesConfig for DeepSeekAnthropicMessagesConfig {
|
|||
headers: Headers,
|
||||
api_key: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &MessagesOptionalParams,
|
||||
_litellm_params: &LitellmParams,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
@ -298,6 +299,7 @@ mod tests {
|
|||
.collect(),
|
||||
api_key,
|
||||
"deepseek-v4-pro",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
env,
|
||||
)
|
||||
|
|
@ -451,4 +453,27 @@ mod tests {
|
|||
json!({"type": "enabled", "budget_tokens": 1024})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::without_client_betas(&[])]
|
||||
#[case::with_client_betas(&[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")])]
|
||||
fn safeguards_and_client_betas_reach_deepseek_as_sent(#[case] client: &[(&str, &str)]) {
|
||||
let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]);
|
||||
let body = json!({
|
||||
"model": "deepseek-v4-pro",
|
||||
"max_tokens": 100,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"safeguards": safeguards
|
||||
});
|
||||
let headers: Headers = client
|
||||
.iter()
|
||||
.map(|(name, value)| (name.to_string(), value.to_string()))
|
||||
.collect();
|
||||
assert_eq!(
|
||||
DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG
|
||||
.request_headers(headers.clone(), &request_from(body.clone())),
|
||||
headers
|
||||
);
|
||||
assert_eq!(transformed(body)["safeguards"], safeguards);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ impl BaseMessagesConfig for EdenAIAnthropicMessagesConfig {
|
|||
headers: Headers,
|
||||
api_key: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &MessagesOptionalParams,
|
||||
_params: &LitellmParams,
|
||||
env: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ impl BaseMessagesConfig for MinimaxMessagesConfig {
|
|||
headers: Headers,
|
||||
api_key: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &MessagesOptionalParams,
|
||||
_params: &LitellmParams,
|
||||
env: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
@ -183,7 +184,14 @@ mod tests {
|
|||
.flatten()
|
||||
};
|
||||
let environment = MINIMAX_MESSAGES_CONFIG
|
||||
.validate_environment(vec![], key, "test-model", &Default::default(), &env)
|
||||
.validate_environment(
|
||||
vec![],
|
||||
key,
|
||||
"test-model",
|
||||
&MessagesOptionalParams::default(),
|
||||
&Default::default(),
|
||||
&env,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
environment.auth,
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ impl BaseMessagesConfig for OpenRouterAnthropicMessagesConfig {
|
|||
headers: Headers,
|
||||
api_key: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &MessagesOptionalParams,
|
||||
_params: &LitellmParams,
|
||||
env: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ impl BaseMessagesConfig for TencentAnthropicMessagesConfig {
|
|||
headers: Headers,
|
||||
api_key: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &MessagesOptionalParams,
|
||||
_params: &LitellmParams,
|
||||
env: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
@ -229,7 +230,14 @@ mod tests {
|
|||
.flatten()
|
||||
};
|
||||
let environment = TENCENT_MESSAGES_CONFIG
|
||||
.validate_environment(vec![], key, "test-model", &Default::default(), &env)
|
||||
.validate_environment(
|
||||
vec![],
|
||||
key,
|
||||
"test-model",
|
||||
&MessagesOptionalParams::default(),
|
||||
&Default::default(),
|
||||
&env,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
environment.auth,
|
||||
|
|
|
|||
|
|
@ -114,6 +114,7 @@ impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig {
|
|||
headers: Headers,
|
||||
_api_key: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &MessagesOptionalParams,
|
||||
litellm_params: &LitellmParams,
|
||||
_env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
|
|
@ -213,23 +214,14 @@ fn has_web_search_tool(tools: Option<&[Recognized<MessagesTool>]>) -> bool {
|
|||
}
|
||||
|
||||
/// Python's Vertex `validate_anthropic_messages_environment` beta set: the shared feature
|
||||
/// betas under Vertex's policy, plus the web search and safeguards betas Vertex needs spelled
|
||||
/// out.
|
||||
/// betas under Vertex's policy, plus the web search beta Vertex needs spelled out.
|
||||
fn vertex_feature_betas(request: &MessagesRequest) -> BetaSet {
|
||||
let params = &request.params;
|
||||
provider_feature_betas(request, BetaProvider::VertexAi).union(
|
||||
[
|
||||
has_web_search_tool(params.tools.as_deref())
|
||||
.then_some(AnthropicBeta::WebSearch20250305),
|
||||
params
|
||||
.extra
|
||||
.get("safeguards")
|
||||
.is_some_and(|value| !value.is_null())
|
||||
.then_some(AnthropicBeta::DangerousToolUse20260903),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect(),
|
||||
has_web_search_tool(params.tools.as_deref())
|
||||
.then_some(AnthropicBeta::WebSearch20250305)
|
||||
.into_iter()
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -448,6 +440,7 @@ mod tests {
|
|||
.collect(),
|
||||
Some("ignored-api-key"),
|
||||
"claude-sonnet-4-5",
|
||||
&MessagesOptionalParams::default(),
|
||||
¶ms(Some("proj"), None),
|
||||
&no_env,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2,5 +2,4 @@
|
|||
|
||||
For `model_prices_and_context_window.json` validation, we should eventually:
|
||||
|
||||
- Remove any schema file like `model_prices_and_context_window.schema.json`
|
||||
- Stop skipping this crate's tests
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ fn run_chat_completions(
|
|||
) -> PyResult<Py<PyAny>> {
|
||||
let host = InferenceHost::new(
|
||||
call.view()?,
|
||||
"litellm.rust_bridge.chat_completions.route_host",
|
||||
"litellm.rust_bridge.host.routes.chat_completions",
|
||||
);
|
||||
run_inference::<ChatCompletionsRoute, _>(
|
||||
py,
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue